From 960ea18980c557d031f9d8d8805d8a15b994002a Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Thu, 29 Oct 2026 21:45:22 -0700 Subject: [PATCH 01/89] Refuse a CUDA graph replay when an input changes shape A CUDA graph replays the shapes it captured. Before each replay, the backend copied the new input into the captured input buffer, and it checked only that the byte count matched. An input can change shape without changing its byte count, for example from (2, 3, 4, 4) to (2, 4, 3, 4). Such a call replayed the captured shape and returned an output with the old shape and wrong values, with no error. The backend now records each input's shape at capture, and a replay with a different shape fails with InvalidArgument, the same error a different byte count already gave. test_output_stride_rearrange.py captures with (2, 3, 4, 4) inputs, then calls with (2, 4, 3, 4), for contiguous and channels-last outputs. Without this change those calls return the wrong shape. --- backends/cuda/runtime/cuda_backend.cpp | 15 +++++++----- backends/cuda/runtime/cuda_delegate_handle.h | 6 +++++ .../tests/test_output_stride_rearrange.py | 23 +++++++++++++++++++ 3 files changed, 38 insertions(+), 6 deletions(-) diff --git a/backends/cuda/runtime/cuda_backend.cpp b/backends/cuda/runtime/cuda_backend.cpp index 2dcda0004f0..32b864a705a 100644 --- a/backends/cuda/runtime/cuda_backend.cpp +++ b/backends/cuda/runtime/cuda_backend.cpp @@ -768,14 +768,14 @@ class ET_EXPERIMENTAL CudaBackend final // Copy new input data (GPU-resident) into static input buffers (D2D) for (size_t i = 0; i < n_inputs; i++) { auto* et_input = &(args[i]->toTensor()); + const auto sizes = et_input->sizes(); + const auto& captured = handle->cuda_graph_state.static_input_sizes[i]; ET_CHECK_OR_RETURN_ERROR( - et_input->nbytes() == - handle->cuda_graph_state.static_input_nbytes[i], + std::equal( + sizes.begin(), sizes.end(), captured.begin(), captured.end()), InvalidArgument, - "CUDA graph replay: input %zu size mismatch (expected %zu, got %zu)", - i, - handle->cuda_graph_state.static_input_nbytes[i], - et_input->nbytes()); + "CUDA graph replay: input %zu has a different shape than at capture", + i); ET_CUDA_CHECK_OR_RETURN_ERROR(cudaMemcpyAsync( handle->cuda_graph_state.static_input_ptrs[i], et_input->const_data_ptr(), @@ -851,6 +851,7 @@ class ET_EXPERIMENTAL CudaBackend final state_->static_output_ptrs.clear(); state_->static_input_nbytes.clear(); state_->static_output_nbytes.clear(); + state_->static_input_sizes.clear(); // Same order as ~CudaGraphState: the exec depends on the graph. if (state_->graph_exec != nullptr) { (void)cudaGraphExecDestroy(state_->graph_exec); @@ -903,6 +904,8 @@ class ET_EXPERIMENTAL CudaBackend final // through the guard instead of leaking this allocation. handle->cuda_graph_state.static_input_ptrs.push_back(static_ptr); handle->cuda_graph_state.static_input_nbytes.push_back(nbytes); + handle->cuda_graph_state.static_input_sizes.emplace_back( + et_input->sizes().begin(), et_input->sizes().end()); ET_CUDA_CHECK_OR_RETURN_ERROR(cudaMemcpy( static_ptr, diff --git a/backends/cuda/runtime/cuda_delegate_handle.h b/backends/cuda/runtime/cuda_delegate_handle.h index f42f0f2382f..605f50bf5d6 100644 --- a/backends/cuda/runtime/cuda_delegate_handle.h +++ b/backends/cuda/runtime/cuda_delegate_handle.h @@ -141,6 +141,9 @@ struct CudaGraphState { std::vector static_output_ptrs; std::vector static_input_nbytes; std::vector static_output_nbytes; + // An input with another shape of the same byte count would replay with the + // captured shape and return a wrong result with no error. + std::vector> static_input_sizes; // Allocations the graph makes and never frees -- the outputs AOTI allocates // while being captured. AutoFreeOnLaunch frees one launch's before the next @@ -227,6 +230,7 @@ struct CudaGraphState { static_output_ptrs.clear(); static_input_nbytes.clear(); static_output_nbytes.clear(); + static_input_sizes.clear(); } // Drops the captured graph so a new one is captured, for when memory the @@ -253,6 +257,7 @@ struct CudaGraphState { static_output_ptrs(std::move(other.static_output_ptrs)), static_input_nbytes(std::move(other.static_input_nbytes)), static_output_nbytes(std::move(other.static_output_nbytes)), + static_input_sizes(std::move(other.static_input_sizes)), graph_allocations(std::move(other.graph_allocations)) { other.graph = nullptr; other.graph_exec = nullptr; @@ -270,6 +275,7 @@ struct CudaGraphState { static_output_ptrs = std::move(other.static_output_ptrs); static_input_nbytes = std::move(other.static_input_nbytes); static_output_nbytes = std::move(other.static_output_nbytes); + static_input_sizes = std::move(other.static_input_sizes); graph_allocations = std::move(other.graph_allocations); other.graph = nullptr; diff --git a/backends/cuda/tests/test_output_stride_rearrange.py b/backends/cuda/tests/test_output_stride_rearrange.py index 431bb56a88f..908b58f4286 100644 --- a/backends/cuda/tests/test_output_stride_rearrange.py +++ b/backends/cuda/tests/test_output_stride_rearrange.py @@ -368,6 +368,29 @@ def test_strided_output_first_seen_at_capture(self): else: self._assert_matches_eager(module, x, out) + def test_input_shape_changes_after_capture(self): + torch.manual_seed(0) + # The last calls swap two dims, which keeps the byte count, so only the + # shape tells them apart from the captured input. + shapes = [(2, 3, 4, 4)] * 4 + [(2, 4, 3, 4)] * 2 + inputs = [torch.randn(shape) for shape in shapes] + dims = { + 1: torch.export.Dim("c", min=2, max=8), + 2: torch.export.Dim("h", min=2, max=8), + } + for memory_format in (torch.contiguous_format, torch.channels_last): + module = _Double(memory_format).eval() + outputs = self._run_with_cuda_graph( + module, (inputs[0].cuda(),), inputs, dynamic_shapes=(dims,) + ) + for i, (x, out) in enumerate(zip(inputs, outputs)): + with self.subTest(memory_format=memory_format, call=i): + if i < 4: + self._assert_matches_eager(module, x, out) + else: + # InvalidArgument, from the replay's input shape check. + self.assertIn("error 0x12", str(out)) + def test_output_shape_changes_at_capture(self): torch.manual_seed(0) # The capture call has a larger batch than the warmup calls before it. From 5f3ce38b7fe51dfe6e27151e8a4c50cca14b8831 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Thu, 29 Oct 2026 21:45:23 -0700 Subject: [PATCH 02/89] List the CUDA index shims in the wheel's shared library check The wheel check names every kernel shim that libaoti_cuda_shims must define, so a library built without some of its CUDA sources fails the check. The mask index and index_put shims were missing from that list. --- .ci/scripts/wheel/test_shared_libraries.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.ci/scripts/wheel/test_shared_libraries.py b/.ci/scripts/wheel/test_shared_libraries.py index e15df2bf074..9b0d0052a00 100644 --- a/.ci/scripts/wheel/test_shared_libraries.py +++ b/.ci/scripts/wheel/test_shared_libraries.py @@ -103,6 +103,8 @@ "aoti_torch_cuda_randint_low_out", "aoti_torch_cuda_sort_stable", "aoti_torch_cuda_guard_set_index", + "aoti_torch_cuda_index_Tensor", + "aoti_torch_index_put_out", ) _QUANTIZED_KERNEL_SYMBOLS = ( From c4cb8cf4b0c455c21ddcba48ff371313da3748a3 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Thu, 29 Oct 2026 21:45:24 -0700 Subject: [PATCH 03/89] Test aoti_torch_get_size Every getter in the shared AOTI shims has a unit test except aoti_torch_get_size, so its handling of negative dims and of dims out of range never ran. The new tests cover both, on CPU and CUDA, and the null argument tests now include it. --- .../aoti/tests/test_common_shims_slim.cpp | 35 +++++++++++++++++++ 1 file changed, 35 insertions(+) diff --git a/backends/aoti/tests/test_common_shims_slim.cpp b/backends/aoti/tests/test_common_shims_slim.cpp index 2dd8475a4d8..5446b925688 100644 --- a/backends/aoti/tests/test_common_shims_slim.cpp +++ b/backends/aoti/tests/test_common_shims_slim.cpp @@ -143,6 +143,28 @@ void runGetSizesTest(slim_c10::DeviceType device_type) { delete tensor; } +void runGetSizeTest(slim_c10::DeviceType device_type) { + std::vector sizes = {2, 3}; + std::vector strides = calculateContiguousStrides(sizes); + slim_c10::Device device(device_type, 0); + + Tensor* tensor = new Tensor(slim::empty_strided( + slim::makeArrayRef(sizes), + slim::makeArrayRef(strides), + slim_c10::ScalarType::Float, + device)); + + int64_t size = -1; + EXPECT_EQ(aoti_torch_get_size(tensor, 1, &size), Error::Ok); + EXPECT_EQ(size, 3); + EXPECT_EQ(aoti_torch_get_size(tensor, -2, &size), Error::Ok); + EXPECT_EQ(size, 2); + EXPECT_EQ(aoti_torch_get_size(tensor, 2, &size), Error::InvalidArgument); + EXPECT_EQ(aoti_torch_get_size(tensor, -3, &size), Error::InvalidArgument); + + delete tensor; +} + void runGetStridesTest(slim_c10::DeviceType device_type) { std::vector sizes = {2, 3, 4}; std::vector strides = calculateContiguousStrides(sizes); @@ -432,6 +454,10 @@ TEST_F(CommonShimsSlimTest, GetSizes_CPU) { runGetSizesTest(slim_c10::DeviceType::CPU); } +TEST_F(CommonShimsSlimTest, GetSize_CPU) { + runGetSizeTest(slim_c10::DeviceType::CPU); +} + TEST_F(CommonShimsSlimTest, GetStrides_CPU) { runGetStridesTest(slim_c10::DeviceType::CPU); } @@ -483,6 +509,13 @@ TEST_F(CommonShimsSlimTest, GetSizes_CUDA) { runGetSizesTest(slim_c10::DeviceType::CUDA); } +TEST_F(CommonShimsSlimTest, GetSize_CUDA) { + if (!isCudaAvailable()) { + GTEST_SKIP() << "CUDA not available"; + } + runGetSizeTest(slim_c10::DeviceType::CUDA); +} + TEST_F(CommonShimsSlimTest, GetStrides_CUDA) { if (!isCudaAvailable()) { GTEST_SKIP() << "CUDA not available"; @@ -555,6 +588,7 @@ TEST_F(CommonShimsSlimTest, NullTensorArgument) { EXPECT_EQ( aoti_torch_get_data_ptr(nullptr, &data_ptr), Error::InvalidArgument); EXPECT_EQ(aoti_torch_get_sizes(nullptr, &sizes), Error::InvalidArgument); + EXPECT_EQ(aoti_torch_get_size(nullptr, 0, &numel), Error::InvalidArgument); EXPECT_EQ(aoti_torch_get_strides(nullptr, &strides), Error::InvalidArgument); EXPECT_EQ(aoti_torch_get_dtype(nullptr, &dtype), Error::InvalidArgument); EXPECT_EQ(aoti_torch_get_dim(nullptr, &dim), Error::InvalidArgument); @@ -566,6 +600,7 @@ TEST_F(CommonShimsSlimTest, NullReturnPointer) { EXPECT_EQ(aoti_torch_get_data_ptr(tensor, nullptr), Error::InvalidArgument); EXPECT_EQ(aoti_torch_get_sizes(tensor, nullptr), Error::InvalidArgument); + EXPECT_EQ(aoti_torch_get_size(tensor, 0, nullptr), Error::InvalidArgument); EXPECT_EQ(aoti_torch_get_strides(tensor, nullptr), Error::InvalidArgument); EXPECT_EQ(aoti_torch_get_dtype(tensor, nullptr), Error::InvalidArgument); EXPECT_EQ(aoti_torch_get_dim(tensor, nullptr), Error::InvalidArgument); From 301e0e124cde5ec963b554c05624351952a531c4 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Thu, 29 Oct 2026 20:23:55 -0700 Subject: [PATCH 04/89] Test more CUDA index shim and CUDA graph cases The CUDA index shims run mask indexing and index_put. Their tests covered only 1-D masks and int64 accumulation, and the test of a capture call with a new output shape passed even when CUDA graphs were off. test_mask_index_and_index_put moves from test_cuda_export.py to the new test_index_shims.py. It runs programs through the Python runtime, and test_cuda_export.py is also a Buck test target that does not depend on that runtime. The new file adds a 2-D mask, and index_put with repeated indices in deterministic mode, which sends half and bfloat16 accumulation, the last-write-wins case and an int32 index to the shim. test_output_shape_changes_at_capture now ends with a call at the warmup batch size. A live graph rejects it with InvalidArgument, because a replay needs the captured input shape, so the test fails if graphs are off. --- backends/cuda/tests/test_cuda_export.py | 46 -------- backends/cuda/tests/test_index_shims.py | 111 ++++++++++++++++++ .../tests/test_output_stride_rearrange.py | 10 +- 3 files changed, 119 insertions(+), 48 deletions(-) create mode 100644 backends/cuda/tests/test_index_shims.py diff --git a/backends/cuda/tests/test_cuda_export.py b/backends/cuda/tests/test_cuda_export.py index 37f60e150a8..0380b6fe3c7 100644 --- a/backends/cuda/tests/test_cuda_export.py +++ b/backends/cuda/tests/test_cuda_export.py @@ -787,52 +787,6 @@ def forward(self, counts: torch.Tensor, idx: torch.Tensor): self._export_to_cuda_with_lower(Int64ScatterAdd(), inputs) self.assertIn("aoti_torch_cuda_scatter_reduce_two_out", str(context.exception)) - @unittest.skipIf( - torch.version.hip is not None, - "loads through the Python runtime, which the ROCm build does not include", - ) - def test_mask_index_and_index_put(self): - """A mask index makes the result size depend on the data, and Inductor - has no atomic add for int64, so these run in the index shims.""" - from executorch.runtime import Runtime - - class MaskIndex(torch.nn.Module): - def forward(self, x, mask, column, idx, counts): - out = x.clone() - out[mask] = x[mask] * 2 - # The mask and column index dims 0 and 2, which are not adjacent. - picked = x[mask, :, column].sum(0) - # A transposed input gives a result that is not row-major. - rows = x.transpose(1, 2)[mask] * 3 - ones = torch.ones_like(idx) - counts = counts.index_put((idx,), ones, accumulate=True) - return out, picked, rows, counts - - model = MaskIndex() - x = torch.randn(4, 5, 6) - column = torch.tensor([2]) - idx = torch.tensor([0, 3, 3, 1, 3]) - # Nonzero counts check that repeated writes add to the old values. - counts = torch.arange(8) * 5 - masks = ( - torch.tensor([True, False, True, True]), - torch.zeros(4, dtype=torch.bool), - torch.ones(4, dtype=torch.bool), - ) - edge = self._export_to_cuda_with_lower( - model, (x, masks[0], column, idx, counts) - ) - with tempfile.TemporaryDirectory() as directory: - path = os.path.join(directory, "model.pte") - with open(path, "wb") as f: - edge.to_executorch().write_to_file(f) - method = Runtime.get().load_program(path).load_method("forward") - for mask in masks: - inputs = (x, mask, column, idx, counts) - outputs = method.execute(list(inputs)) - for output, expected in zip(outputs, model(*inputs), strict=True): - torch.testing.assert_close(output.cpu(), expected) - @unittest.skipIf( torch.version.hip is not None, "loads through the Python runtime, which the ROCm build does not include", diff --git a/backends/cuda/tests/test_index_shims.py b/backends/cuda/tests/test_index_shims.py new file mode 100644 index 00000000000..af7a7b2fa84 --- /dev/null +++ b/backends/cuda/tests/test_index_shims.py @@ -0,0 +1,111 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Mask indexing and index_put in the CUDA backend's index shims. + +A mask index makes a result size depend on the data, so Inductor leaves it, +and an index_put with a mask, to the backend's own index shims. These tests +lower real programs, run them through the Python runtime, and compare with +eager PyTorch. +""" + +import os +import tempfile +import unittest +from typing import Sequence, Tuple + +import torch +from executorch.backends.cuda.cuda_backend import CudaBackend +from executorch.backends.cuda.cuda_partitioner import CudaPartitioner +from executorch.exir import EdgeCompileConfig, to_edge_transform_and_lower +from torch.export import export + + +@unittest.skipUnless(torch.cuda.is_available(), "requires CUDA") +@unittest.skipIf( + torch.version.hip is not None, + "loads through the Python runtime, which the ROCm build does not include", +) +class TestIndexShims(unittest.TestCase): + def _run( + self, model: torch.nn.Module, cases: Sequence[Tuple[torch.Tensor, ...]] + ) -> None: + """Lower model for cases[0], then run every case and compare with eager.""" + spec = CudaBackend.generate_method_name_compile_spec("forward") + lowered = to_edge_transform_and_lower( + export(model, cases[0]), + partitioner=[CudaPartitioner([spec])], + compile_config=EdgeCompileConfig(_check_ir_validity=False), + ) + from executorch.runtime import Runtime + + with tempfile.TemporaryDirectory() as directory: + path = os.path.join(directory, "model.pte") + with open(path, "wb") as f: + lowered.to_executorch().write_to_file(f) + method = Runtime.get().load_program(path).load_method("forward") + for inputs in cases: + outputs = method.execute(list(inputs)) + for output, expected in zip(outputs, model(*inputs), strict=True): + torch.testing.assert_close(output.cpu(), expected) + + def test_mask_index_and_index_put(self): + """Inductor also has no atomic add for int64.""" + + class MaskIndex(torch.nn.Module): + def forward(self, x, mask, column, idx, counts): + out = x.clone() + out[mask] = x[mask] * 2 + # The mask and column index dims 0 and 2, which are not adjacent. + picked = x[mask, :, column].sum(0) + # A transposed input gives a result that is not row-major. + rows = x.transpose(1, 2)[mask] * 3 + ones = torch.ones_like(idx) + counts = counts.index_put((idx,), ones, accumulate=True) + return out, picked, rows, counts + + x = torch.randn(4, 5, 6) + column = torch.tensor([2]) + idx = torch.tensor([0, 3, 3, 1, 3]) + # Nonzero counts check that repeated writes add to the old values. + counts = torch.arange(8) * 5 + masks = ( + torch.tensor([True, False, True, True]), + torch.zeros(4, dtype=torch.bool), + torch.ones(4, dtype=torch.bool), + ) + self._run(MaskIndex(), [(x, mask, column, idx, counts) for mask in masks]) + + def test_two_dim_mask(self): + class TwoDimMask(torch.nn.Module): + def forward(self, x, mask): + return (x[mask] * 2,) + + x = torch.randn(4, 5, 6) + masks = (torch.rand(4, 5) > 0.5, torch.zeros(4, 5, dtype=torch.bool)) + self._run(TwoDimMask(), [(x, mask) for mask in masks]) + + def test_index_put_with_repeated_indices(self): + """Deterministic mode sends half, bfloat16 and int32-index writes to the shim.""" + + class Put(torch.nn.Module): + def forward(self, x, idx, values): + summed = x.index_put((idx,), values, accumulate=True) + # Without accumulate, the last of the repeated writes wins. + written = x.index_put((idx.int(),), values) + return summed, written + + idx = torch.tensor([0, 3, 3, 1, 3]) + deterministic = torch.are_deterministic_algorithms_enabled() + torch.use_deterministic_algorithms(True) + try: + for dtype in (torch.float16, torch.bfloat16): + with self.subTest(dtype=dtype): + x = torch.arange(8, dtype=dtype) + values = torch.arange(1, 6, dtype=dtype) + self._run(Put(), [(x, idx, values)]) + finally: + torch.use_deterministic_algorithms(deterministic) diff --git a/backends/cuda/tests/test_output_stride_rearrange.py b/backends/cuda/tests/test_output_stride_rearrange.py index 908b58f4286..c8548d4d01d 100644 --- a/backends/cuda/tests/test_output_stride_rearrange.py +++ b/backends/cuda/tests/test_output_stride_rearrange.py @@ -394,7 +394,9 @@ def test_input_shape_changes_after_capture(self): def test_output_shape_changes_at_capture(self): torch.manual_seed(0) # The capture call has a larger batch than the warmup calls before it. - inputs = [torch.randn(batch, 3, 4, 4) for batch in (1, 1, 1, 2, 2, 2)] + # The last call goes back to the warmup batch, which only a live graph + # rejects, because a replay needs the captured input sizes. + inputs = [torch.randn(batch, 3, 4, 4) for batch in (1, 1, 1, 2, 2, 2, 1)] batch = torch.export.Dim("batch", min=1, max=8) for memory_format in (torch.contiguous_format, torch.channels_last): module = _Double(memory_format).eval() @@ -403,7 +405,11 @@ def test_output_shape_changes_at_capture(self): ) for i, (x, out) in enumerate(zip(inputs, outputs)): with self.subTest(memory_format=memory_format, call=i): - self._assert_matches_eager(module, x, out) + if i == 6: + # InvalidArgument, from the replay's input size check. + self.assertIn("error 0x12", str(out)) + else: + self._assert_matches_eager(module, x, out) if __name__ == "__main__": From 03e2c82904d79284bc78f2abdf1288a23142976b Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Thu, 29 Oct 2026 21:49:25 -0700 Subject: [PATCH 05/89] Run a CUDA graph capture call without a graph when the capture fails With enable_cuda_graph_for_method set, a method that has a mask index, such as x[mask] or x[mask] = v, fails its fourth call, the one that captures the CUDA graph. The index shim has to read the number of selected elements back to the host, which it cannot do inside a capture, so the call returns an error. Only then does the backend turn graphs off, and the later calls run without one. Now, when the 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 same retry also handles other capture failures, such as an output that is first strided at the capture call. The retry is safe because a capture records kernels without running them, so until the new graph runs, the failed attempt has not touched the inputs or any state the program updates. A failure after the graph started running is not retried and fails the call, as before. The off-graph KV cache step is still prepared once and committed once. Nothing changes in the exported program, so programs load on older runtimes as before, and programs exported before this change get the fix too. test_output_stride_rearrange.py runs a mask index model and a mask index_put model with CUDA graphs requested for six calls, and compares each call with eager PyTorch. Without this change the fourth call fails. test_strided_output_first_seen_at_capture now expects that call to match eager PyTorch too, instead of failing. --- backends/cuda/runtime/cuda_backend.cpp | 52 ++++++++++++++----- .../tests/test_output_stride_rearrange.py | 36 ++++++++++--- 2 files changed, 68 insertions(+), 20 deletions(-) diff --git a/backends/cuda/runtime/cuda_backend.cpp b/backends/cuda/runtime/cuda_backend.cpp index 32b864a705a..0d4d9f4c69f 100644 --- a/backends/cuda/runtime/cuda_backend.cpp +++ b/backends/cuda/runtime/cuda_backend.cpp @@ -808,13 +808,37 @@ class ET_EXPERIMENTAL CudaBackend final return commit_offgraph_kv_step(handle, kv_step_width); } - // --------------------------------------------------------------- - // Normal path (also used for WARMUP and CAPTURE phases) - // --------------------------------------------------------------- - bool is_capture_step = + const bool is_capture_step = (handle->cuda_graph_state.phase == CudaGraphPhase::Warmup && handle->cuda_graph_state.warmup_remaining == 0); + bool graph_ran = false; + Error error = + run_once(handle, args, n_inputs, n_outputs, is_capture_step, graph_ran); + if (error != Error::Ok && is_capture_step && !graph_ran) { + // A capture records kernels without running them, so until the new graph + // runs, the inputs and any state the program updates are untouched. + ET_LOG( + Info, + "CUDA graph: capture failed for '%s'; running this call and later " + "ones without a graph", + handle->method_name.c_str()); + error = run_once(handle, args, n_inputs, n_outputs, false, graph_ran); + } + ET_CHECK_OK_OR_RETURN_ERROR(error); + return commit_offgraph_kv_step(handle, kv_step_width); + } + private: + // Runs the program once. On the capture step it captures the CUDA graph + // instead, then launches it once, and sets graph_ran once that launch is + // queued. + Error run_once( + cuda::CudaDelegateHandle* handle, + Span args, + size_t n_inputs, + size_t n_outputs, + bool is_capture_step, + bool& graph_ran) const { // Wrrap them as SlimTensors for AOTI execution without any host<->device // copies. std::vector slim_inputs(n_inputs); @@ -841,9 +865,9 @@ class ET_EXPERIMENTAL CudaBackend final return; } // Free what this attempt pinned and disable graphs for this method, so - // a capture that cannot succeed costs one error rather than one on - // every fourth call for the life of the process. Eager execution is - // correct, just slower. + // a capture that cannot succeed is tried once rather than on every + // fourth call for the life of the process. Eager execution is correct, + // just slower. for (void* ptr : state_->static_input_ptrs) { (void)cudaFree(ptr); } @@ -1076,9 +1100,9 @@ class ET_EXPERIMENTAL CudaBackend final slim_outputs[i], &(args[i + n_inputs]->toTensor()))) { std::fill(slim_outputs.begin(), slim_outputs.end(), nullptr); ET_LOG( - Error, + Info, "CUDA graph: output %zu of '%s' is not laid out like its " - "ETensor; this call fails and later calls run without a graph", + "ETensor; running without a graph", i, handle->method_name.c_str()); return Error::NotSupported; @@ -1091,9 +1115,9 @@ class ET_EXPERIMENTAL CudaBackend final if (resize_err != Error::Ok) { std::fill(slim_outputs.begin(), slim_outputs.end(), nullptr); ET_LOG( - Error, + Info, "CUDA graph: cannot resize output %zu of '%s' to its captured " - "shape; this call fails and later calls run without a graph", + "shape; running without a graph", i, handle->method_name.c_str()); return resize_err; @@ -1135,6 +1159,7 @@ class ET_EXPERIMENTAL CudaBackend final Internal, "cudaGraphLaunch (first replay) failed: %s", cudaGetErrorString(gerr)); + graph_ran = true; // Copy capture-step outputs into the planned GPU ET buffers (D2D). for (size_t i = 0; i < n_outputs; i++) { @@ -1151,7 +1176,7 @@ class ET_EXPERIMENTAL CudaBackend final // Last failure point is behind us, so the captured state is now the state // the next call should replay from. capture_guard.disarm(); - return commit_offgraph_kv_step(handle, kv_step_width); + return Error::Ok; } // ----- Normal / WARMUP execution continues here ----- @@ -1208,9 +1233,10 @@ class ET_EXPERIMENTAL CudaBackend final slim_outputs[i] = nullptr; } - return commit_offgraph_kv_step(handle, kv_step_width); + return Error::Ok; } + public: void destroy(DelegateHandle* handle_) const override { if (handle_ == nullptr) { return; diff --git a/backends/cuda/tests/test_output_stride_rearrange.py b/backends/cuda/tests/test_output_stride_rearrange.py index c8548d4d01d..8f2530d781d 100644 --- a/backends/cuda/tests/test_output_stride_rearrange.py +++ b/backends/cuda/tests/test_output_stride_rearrange.py @@ -282,6 +282,16 @@ def forward(self, x): return self.proj(x)[:, :, :14] +class _MaskedRows(nn.Module): + def forward(self, x, mask): + return x[mask] * 2 + + +class _MaskedAdd(nn.Module): + def forward(self, x, mask): + return x.index_put((mask,), torch.ones(4), accumulate=True) + + class _Double(nn.Module): def __init__(self, memory_format): super().__init__() @@ -304,14 +314,16 @@ def setUpClass(cls): def _run_with_cuda_graph(self, module, example, inputs, dynamic_shapes=None): from executorch.backends.cuda.cuda_backend import CudaBackend from executorch.backends.cuda.cuda_partitioner import CudaPartitioner - from executorch.exir import to_edge_transform_and_lower + from executorch.exir import EdgeCompileConfig, to_edge_transform_and_lower from executorch.runtime import Runtime with torch.no_grad(): ep = torch.export.export(module, example, dynamic_shapes=dynamic_shapes) spec = CudaBackend.generate_method_name_compile_spec("forward") program = to_edge_transform_and_lower( - ep, partitioner=[CudaPartitioner([spec])] + ep, + partitioner=[CudaPartitioner([spec])], + compile_config=EdgeCompileConfig(_check_ir_validity=False), ).to_executorch() runtime = Runtime.get() option = "enable_cuda_graph_for_method" @@ -329,8 +341,9 @@ def _run_with_cuda_graph(self, module, example, inputs, dynamic_shapes=None): # The option is process wide. runtime.backend_registry.set_option("CudaBackend", {option: ""}) for x in inputs: + args = x if isinstance(x, tuple) else (x,) try: - results.append(method.execute([x])[0]) + results.append(method.execute(list(args))[0]) except RuntimeError as e: results.append(e) return results @@ -363,10 +376,7 @@ def test_strided_output_first_seen_at_capture(self): ) for i, (x, out) in enumerate(zip(inputs, outputs)): with self.subTest(call=i): - if i == 3: - self.assertIsInstance(out, RuntimeError) - else: - self._assert_matches_eager(module, x, out) + self._assert_matches_eager(module, x, out) def test_input_shape_changes_after_capture(self): torch.manual_seed(0) @@ -391,6 +401,18 @@ def test_input_shape_changes_after_capture(self): # InvalidArgument, from the replay's input shape check. self.assertIn("error 0x12", str(out)) + def test_mask_runs_without_a_graph(self): + # A mask index reads its count back to the host, which a capture + # cannot record, so the capture call has to run without a graph. + torch.manual_seed(0) + inputs = [(torch.randn(8, 4), torch.rand(8) > 0.5) for _ in range(6)] + for module in (_MaskedRows().eval(), _MaskedAdd().eval()): + outputs = self._run_with_cuda_graph(module, inputs[0], inputs) + for i, (args, out) in enumerate(zip(inputs, outputs)): + with self.subTest(module=type(module).__name__, call=i): + self.assertIsInstance(out, torch.Tensor) + torch.testing.assert_close(out, module(*args)) + def test_output_shape_changes_at_capture(self): torch.manual_seed(0) # The capture call has a larger batch than the warmup calls before it. From bf998957a250bb331937b5ab60496727ba9ecf59 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Thu, 29 Oct 2026 23:31:43 -0700 Subject: [PATCH 06/89] Say which dim differs when a CUDA graph replay refuses an input The replay check only said that an input had a different shape than at capture. It now names the dim and both sizes, or both dim counts, so the caller can see what changed. Also write down in the phase comment that a captured graph only works for the input shapes of the capture call, and that other shapes fail with InvalidArgument. --- backends/cuda/runtime/cuda_backend.cpp | 20 ++++++++++++++++---- backends/cuda/runtime/cuda_delegate_handle.h | 4 +++- 2 files changed, 19 insertions(+), 5 deletions(-) diff --git a/backends/cuda/runtime/cuda_backend.cpp b/backends/cuda/runtime/cuda_backend.cpp index 0d4d9f4c69f..9660844c430 100644 --- a/backends/cuda/runtime/cuda_backend.cpp +++ b/backends/cuda/runtime/cuda_backend.cpp @@ -771,11 +771,23 @@ class ET_EXPERIMENTAL CudaBackend final const auto sizes = et_input->sizes(); const auto& captured = handle->cuda_graph_state.static_input_sizes[i]; ET_CHECK_OR_RETURN_ERROR( - std::equal( - sizes.begin(), sizes.end(), captured.begin(), captured.end()), + sizes.size() == captured.size(), InvalidArgument, - "CUDA graph replay: input %zu has a different shape than at capture", - i); + "CUDA graph replay: input %zu has %zu dims, captured with %zu", + i, + sizes.size(), + captured.size()); + for (size_t d = 0; d < sizes.size(); d++) { + ET_CHECK_OR_RETURN_ERROR( + sizes[d] == captured[d], + InvalidArgument, + "CUDA graph replay: input %zu has size %lld in dim %zu, " + "captured with %lld", + i, + static_cast(sizes[d]), + d, + static_cast(captured[d])); + } ET_CUDA_CHECK_OR_RETURN_ERROR(cudaMemcpyAsync( handle->cuda_graph_state.static_input_ptrs[i], et_input->const_data_ptr(), diff --git a/backends/cuda/runtime/cuda_delegate_handle.h b/backends/cuda/runtime/cuda_delegate_handle.h index 605f50bf5d6..bf1b891567c 100644 --- a/backends/cuda/runtime/cuda_delegate_handle.h +++ b/backends/cuda/runtime/cuda_delegate_handle.h @@ -118,7 +118,9 @@ struct CudaWeightStorage { // - Replay: The captured `graph_exec` is launched on every execute() call. // Inputs are memcpy'd into the static input buffers, the graph is replayed, // and outputs are memcpy'd back from the static output buffers. No tensor -// setup or kernel launches happen on the host hot path. +// setup or kernel launches happen on the host hot path. The graph only +// works for the input shapes it was captured with, so a call whose input +// shapes differ from the capture call fails with InvalidArgument. enum class CudaGraphPhase { Disabled = 0, Warmup = 1, From 5202db75f73f6db6d1057d7a21b027ca4232680c Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Thu, 29 Oct 2026 23:32:02 -0700 Subject: [PATCH 07/89] Log an expected CUDA graph capture failure at Info A mask index refuses to run inside a CUDA graph capture, and the backend then runs that call without a graph, which succeeds. But the refusal and the failed run were both logged at Error, so every method with a mask index printed two ExecuTorch error lines on its fourth call although nothing failed. Log both at Info during the capture. The call that runs without a graph still logs a real failure at Error. AOTInductor's generated code still prints its own line for the refused call, which the backend does not control. The layout and resize messages no longer say "running without a graph" themselves, because the backend already says that once for every failed capture. --- backends/cuda/runtime/cuda_backend.cpp | 26 +++++++++++++++++++------- backends/cuda/runtime/shims/index.cu | 11 ++++++----- 2 files changed, 25 insertions(+), 12 deletions(-) diff --git a/backends/cuda/runtime/cuda_backend.cpp b/backends/cuda/runtime/cuda_backend.cpp index 9660844c430..3a725d68f14 100644 --- a/backends/cuda/runtime/cuda_backend.cpp +++ b/backends/cuda/runtime/cuda_backend.cpp @@ -1085,11 +1085,23 @@ class ET_EXPERIMENTAL CudaBackend final } } - ET_CHECK_OR_RETURN_ERROR( - error == Error::Ok, - Internal, - "AOTInductorModelContainerRun failed with error code %d", - error); + if (error != Error::Ok) { + // A failed capture is retried without a graph, which reports a real + // failure at Error. + if (is_capture_step) { + ET_LOG( + Info, + "AOTInductorModelContainerRun failed during CUDA graph capture " + "with error code %d", + error); + } else { + ET_LOG( + Error, + "AOTInductorModelContainerRun failed with error code %d", + error); + } + return Error::Internal; + } if (is_capture_step) { // End capture → instantiate graph @@ -1114,7 +1126,7 @@ class ET_EXPERIMENTAL CudaBackend final ET_LOG( Info, "CUDA graph: output %zu of '%s' is not laid out like its " - "ETensor; running without a graph", + "ETensor", i, handle->method_name.c_str()); return Error::NotSupported; @@ -1129,7 +1141,7 @@ class ET_EXPERIMENTAL CudaBackend final ET_LOG( Info, "CUDA graph: cannot resize output %zu of '%s' to its captured " - "shape; running without a graph", + "shape", i, handle->method_name.c_str()); return resize_err; diff --git a/backends/cuda/runtime/shims/index.cu b/backends/cuda/runtime/shims/index.cu index 97335a75ea0..88ed4f7a385 100644 --- a/backends/cuda/runtime/shims/index.cu +++ b/backends/cuda/runtime/shims/index.cu @@ -331,13 +331,14 @@ Error mask_positions( return Error::Ok; } // A graph capture cannot record the host read of the count below, so refuse - // before allocating anything. + // before allocating anything. Info, not Error: the backend then retries the + // call without a graph. cudaStreamCaptureStatus capture = cudaStreamCaptureStatusNone; ET_CUDA_CHECK_OR_RETURN_ERROR(cudaStreamIsCapturing(stream, &capture)); - ET_CHECK_OR_RETURN_ERROR( - capture == cudaStreamCaptureStatusNone, - NotSupported, - "index: a mask index cannot run inside a CUDA graph capture"); + if (capture != cudaStreamCaptureStatusNone) { + ET_LOG(Info, "index: a mask index cannot run inside a CUDA graph capture"); + return Error::NotSupported; + } // CUB counts items in an int on CUDA 12. ET_CHECK_OR_RETURN_ERROR( numel <= INT_MAX, From 1d638d8d6a471332160cd610f2bbe53c78aae227 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Thu, 29 Oct 2026 23:32:26 -0700 Subject: [PATCH 08/89] Record CUDA graph allocations only after the first launch The backend recorded the allocations a captured graph makes before it launched the graph for the first time. A graph only allocates that memory when it launches, so if that first launch failed, a later cleanup called cudaFree on memory that was never allocated. Record them after the launch succeeds. A failed capture now cleans up with CudaGraphState::release() instead of a copy of it, so it also frees the memory a successful first launch allocated before a later step failed. --- backends/cuda/runtime/cuda_backend.cpp | 20 ++----------------- backends/cuda/runtime/cuda_delegate_handle.h | 3 ++- .../cuda/runtime/test/test_cuda_kv_cache.cpp | 7 +++---- 3 files changed, 7 insertions(+), 23 deletions(-) diff --git a/backends/cuda/runtime/cuda_backend.cpp b/backends/cuda/runtime/cuda_backend.cpp index 3a725d68f14..f57425170b5 100644 --- a/backends/cuda/runtime/cuda_backend.cpp +++ b/backends/cuda/runtime/cuda_backend.cpp @@ -880,23 +880,7 @@ class ET_EXPERIMENTAL CudaBackend final // a capture that cannot succeed is tried once rather than on every // fourth call for the life of the process. Eager execution is correct, // just slower. - for (void* ptr : state_->static_input_ptrs) { - (void)cudaFree(ptr); - } - state_->static_input_ptrs.clear(); - state_->static_output_ptrs.clear(); - state_->static_input_nbytes.clear(); - state_->static_output_nbytes.clear(); - state_->static_input_sizes.clear(); - // Same order as ~CudaGraphState: the exec depends on the graph. - if (state_->graph_exec != nullptr) { - (void)cudaGraphExecDestroy(state_->graph_exec); - state_->graph_exec = nullptr; - } - if (state_->graph != nullptr) { - (void)cudaGraphDestroy(state_->graph); - state_->graph = nullptr; - } + state_->release(); state_->phase = CudaGraphPhase::Disabled; state_->warmup_remaining = 0; (void)cudaGetLastError(); @@ -1157,7 +1141,6 @@ class ET_EXPERIMENTAL CudaBackend final Internal, "cudaGraphInstantiate failed: %s", cudaGetErrorString(gerr)); - handle->cuda_graph_state.note_graph_allocations(); // Record static output pointers (stable under graph replay). Releasing // them from slim_outputs here, before the copies below, keeps the cleanup @@ -1184,6 +1167,7 @@ class ET_EXPERIMENTAL CudaBackend final "cudaGraphLaunch (first replay) failed: %s", cudaGetErrorString(gerr)); graph_ran = true; + handle->cuda_graph_state.note_graph_allocations(); // Copy capture-step outputs into the planned GPU ET buffers (D2D). for (size_t i = 0; i < n_outputs; i++) { diff --git a/backends/cuda/runtime/cuda_delegate_handle.h b/backends/cuda/runtime/cuda_delegate_handle.h index bf1b891567c..915ff4f10ac 100644 --- a/backends/cuda/runtime/cuda_delegate_handle.h +++ b/backends/cuda/runtime/cuda_delegate_handle.h @@ -168,7 +168,8 @@ struct CudaGraphState { } // Records the graph's outstanding allocations: every memory-allocation node - // without a matching free node. Call once the graph is captured. + // without a matching free node. Call once the graph has launched: only a + // launch allocates them, so freeing them before one fails. void note_graph_allocations() { graph_allocations.clear(); #if !defined(EXECUTORCH_USE_HIP) diff --git a/backends/cuda/runtime/test/test_cuda_kv_cache.cpp b/backends/cuda/runtime/test/test_cuda_kv_cache.cpp index 0c32da2079f..948a53c0c1b 100644 --- a/backends/cuda/runtime/test/test_cuda_kv_cache.cpp +++ b/backends/cuda/runtime/test/test_cuda_kv_cache.cpp @@ -1039,15 +1039,14 @@ TEST_F(CudaKVCacheTest, RetiringAGraphFreesWhatItsLastLaunchAllocated) { graph.graph, cudaGraphInstantiateFlagAutoFreeOnLaunch), cudaSuccess); + ASSERT_EQ(cudaGraphLaunch(graph.graph_exec, stream), cudaSuccess); + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + EXPECT_GE(graph_mem_in_use(), before + kBytes); graph.note_graph_allocations(); // Only the allocation the graph leaves outstanding is its to free. ASSERT_EQ(graph.graph_allocations.size(), 1); EXPECT_EQ(graph.graph_allocations[0], output); - ASSERT_EQ(cudaGraphLaunch(graph.graph_exec, stream), cudaSuccess); - ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); - EXPECT_GE(graph_mem_in_use(), before + kBytes); - graph.release(); ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); ASSERT_EQ(cudaDeviceGraphMemTrim(device), cudaSuccess); From 474e66113c5a73a83d7a449b3512ea7fd5d95f5b Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Thu, 29 Oct 2026 23:32:42 -0700 Subject: [PATCH 09/89] Check the new CUDA graph input shapes in the recapture and replay tests The KV cache recapture test now also checks that dropping the graph clears the captured input shapes, so a recapture cannot replay against the first capture's shapes. The replay tests match the execute() error, so a set_inputs() failure with the same code no longer passes them. --- backends/cuda/runtime/test/test_cuda_kv_cache.cpp | 2 ++ backends/cuda/tests/test_output_stride_rearrange.py | 4 ++-- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/backends/cuda/runtime/test/test_cuda_kv_cache.cpp b/backends/cuda/runtime/test/test_cuda_kv_cache.cpp index 948a53c0c1b..92ee8406509 100644 --- a/backends/cuda/runtime/test/test_cuda_kv_cache.cpp +++ b/backends/cuda/runtime/test/test_cuda_kv_cache.cpp @@ -932,6 +932,7 @@ TEST_F(CudaKVCacheTest, GrowthDropsACapturedGraphOnlyWhenStorageMoves) { ASSERT_EQ(cudaMalloc(&static_input, 16), cudaSuccess); graph.static_input_ptrs = {static_input}; graph.static_input_nbytes = {16}; + graph.static_input_sizes = {{4}}; ASSERT_EQ(kv.commit_step(3), Error::Ok); // Fits the current storage: the graph keeps replaying. @@ -950,6 +951,7 @@ TEST_F(CudaKVCacheTest, GrowthDropsACapturedGraphOnlyWhenStorageMoves) { EXPECT_EQ(graph.warmup_remaining, 1); EXPECT_TRUE(graph.static_input_ptrs.empty()); EXPECT_TRUE(graph.static_input_nbytes.empty()); + EXPECT_TRUE(graph.static_input_sizes.empty()); EXPECT_EQ(graph.graph_exec, nullptr); // Once recaptured, later steps that fit replay it untouched. diff --git a/backends/cuda/tests/test_output_stride_rearrange.py b/backends/cuda/tests/test_output_stride_rearrange.py index 8f2530d781d..e791dab24bc 100644 --- a/backends/cuda/tests/test_output_stride_rearrange.py +++ b/backends/cuda/tests/test_output_stride_rearrange.py @@ -399,7 +399,7 @@ def test_input_shape_changes_after_capture(self): self._assert_matches_eager(module, x, out) else: # InvalidArgument, from the replay's input shape check. - self.assertIn("error 0x12", str(out)) + self.assertIn("execute() failed with error 0x12", str(out)) def test_mask_runs_without_a_graph(self): # A mask index reads its count back to the host, which a capture @@ -429,7 +429,7 @@ def test_output_shape_changes_at_capture(self): with self.subTest(memory_format=memory_format, call=i): if i == 6: # InvalidArgument, from the replay's input size check. - self.assertIn("error 0x12", str(out)) + self.assertIn("execute() failed with error 0x12", str(out)) else: self._assert_matches_eager(module, x, out) From 69d8749c48222dfd6896d438957d186da050bdb2 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Thu, 29 Oct 2026 23:59:47 -0700 Subject: [PATCH 10/89] Delete the antialiased resize torch.compile test test_align_corners_compiles_with_dynamic_sizes could not fail. It passed with no rewrite at all, since torch.compile then calls ATen's kernel, and it passed with the min and max in the align_corners ratio replaced by a plain if, since torch.compile treats an output size of 1 as a constant and compiles again. With torch 2.14 and CI set it also failed for a reason unrelated to the resize. test_bicubic_antialias_resize already runs an exported program at an output height of 1, which is where a plain if breaks. --- .../cuda/tests/test_antialiased_resize.py | 20 ------------------- 1 file changed, 20 deletions(-) diff --git a/backends/cuda/tests/test_antialiased_resize.py b/backends/cuda/tests/test_antialiased_resize.py index 963444f80e7..be48e7fefe1 100644 --- a/backends/cuda/tests/test_antialiased_resize.py +++ b/backends/cuda/tests/test_antialiased_resize.py @@ -11,7 +11,6 @@ runtime, and compare with eager PyTorch. """ -import functools import os import tempfile import unittest @@ -152,22 +151,3 @@ def test_keeps_memory_format(self): x = torch.randn(1, 3, 48, 32, device="cuda") output = resize(x.to(memory_format=torch.channels_last), [20, 56], False) self.assertTrue(output.is_contiguous(memory_format=torch.channels_last)) - - def test_align_corners_compiles_with_dynamic_sizes(self): - from torch._inductor.compile_fx import compile_fx - - model = Resize("bicubic", align_corners=True) - compiled = torch.compile( - model, - backend=functools.partial( - compile_fx, decompositions=CudaBackend.get_decomposition_table() - ), - dynamic=True, - fullgraph=True, - ) - for in_height, out_height in ((48, 20), (10, 1)): - x = torch.randn(1, 3, in_height, 32, device="cuda") - like = torch.empty(1, 1, out_height, device="cuda") - torch.testing.assert_close( - compiled(x, like), model(x, like), atol=1e-5, rtol=1e-4 - ) From 8eafbbde7885947624bd045c4a658022c1c6f4cd Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Thu, 29 Oct 2026 23:59:48 -0700 Subject: [PATCH 11/89] Say why the antialiased resize avoids a conditional The comment on the min and max in the align_corners ratio blamed torch.compile, which handles a conditional by compiling again. The real reason is export: an exported program keeps the branch it was traced with, so a plain if divides by zero at an output size of 1 when the program was traced with a larger one. --- backends/cuda/cuda_backend.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index f2bd1a202f1..243437faa8a 100644 --- a/backends/cuda/cuda_backend.py +++ b/backends/cuda/cuda_backend.py @@ -585,8 +585,8 @@ def _antialias_taps( if align_corners or scale is None or scale <= 0: if align_corners: # ATen uses a ratio of 0 for an output size of 1. min and max give - # that without a conditional, which torch.compile cannot handle for - # dynamic sizes. + # that without a conditional: an exported program keeps the branch + # it was traced with, so a plain if would divide by zero there. num = (in_size - 1) * torch.sym_min(out_size - 1, 1) den = torch.sym_max(out_size - 1, 1) else: From ce87740d1e104d4c4b8be950cf1378edce96377b Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Fri, 30 Oct 2026 00:00:21 -0700 Subject: [PATCH 12/89] Refuse a hand-built data name in AOTI backends that keep data in the .pte Only generate_data_name_compile_spec refused a data name for a backend that keeps its data in the .pte, such as Metal. A data_name compile spec built by hand went through preprocess with no error and no name check, even for a name like "../../outside". data_name_from_compile_specs now refuses it too, and preprocess always calls it, so both paths behave the same. --- backends/aoti/aoti_backend.py | 22 ++++++++++++---------- backends/cuda/tests/test_cuda_export.py | 7 ++++++- 2 files changed, 18 insertions(+), 11 deletions(-) diff --git a/backends/aoti/aoti_backend.py b/backends/aoti/aoti_backend.py index 0d0670b7318..d6b78d66000 100644 --- a/backends/aoti/aoti_backend.py +++ b/backends/aoti/aoti_backend.py @@ -390,11 +390,8 @@ def preprocess( named_data_store.add_named_data(so_blob_key, so_data, 1, None) # Determine whether to save named data externally based on backend setting # External: save to separate .ptd file, otherwise merge with .pte file - external_tag = ( - cls.data_name_from_compile_specs(compile_specs) - if cls.save_data_externally() - else None - ) + data_name = cls.data_name_from_compile_specs(compile_specs) + external_tag = data_name if cls.save_data_externally() else None named_data_store.add_named_data(weights_blob_key, blob_data, 1, external_tag) @@ -435,11 +432,7 @@ def generate_data_name_compile_spec(cls, data_name: str) -> CompileSpec: Generate a CompileSpec that names the external data file, written as ``.ptd``, so several programs can share a folder. """ - if not cls.save_data_externally(): - raise ValueError( - f"{cls.__name__} keeps its data in the .pte file, so it takes no " - "data name" - ) + cls._check_saves_data_externally() return CompileSpec( COMPILE_SPEC_KEYS.DATA_NAME.value, _checked_data_name(data_name).encode("utf-8"), @@ -453,9 +446,18 @@ def data_name_from_compile_specs(cls, compile_specs: List[CompileSpec]) -> str: """ for spec in compile_specs: if spec.key == COMPILE_SPEC_KEYS.DATA_NAME.value: + cls._check_saves_data_externally() return _checked_data_name(spec.value.decode("utf-8")) return f"aoti_{cls.get_device_name()}_blob" + @classmethod + def _check_saves_data_externally(cls) -> None: + if not cls.save_data_externally(): + raise ValueError( + f"{cls.__name__} keeps its data in the .pte file, so it takes no " + "data name" + ) + @classmethod def method_name_from_compile_specs( cls, diff --git a/backends/cuda/tests/test_cuda_export.py b/backends/cuda/tests/test_cuda_export.py index 0380b6fe3c7..b678b2242ef 100644 --- a/backends/cuda/tests/test_cuda_export.py +++ b/backends/cuda/tests/test_cuda_export.py @@ -269,10 +269,15 @@ def test_invalid_data_name_compile_spec(self): CudaBackend.data_name_from_compile_specs( [CompileSpec(key="data_name", value=name.encode())] ) - # A backend that keeps its data in the .pte would ignore the name. + # A backend that keeps its data in the .pte refuses a name, also one + # built by hand. with patch.object(CudaBackend, "save_data_externally", return_value=False): with self.assertRaisesRegex(ValueError, "takes no data name"): CudaBackend.generate_data_name_compile_spec("encoder") + with self.assertRaisesRegex(ValueError, "takes no data name"): + CudaBackend.data_name_from_compile_specs( + [CompileSpec(key="data_name", value=b"encoder")] + ) def test_valid_data_name_compile_spec(self): self.assertEqual(CudaBackend.data_name_from_compile_specs([]), "aoti_cuda_blob") From 84d146819c53bda269505e088b65eae2d25f72f1 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Fri, 30 Oct 2026 00:00:22 -0700 Subject: [PATCH 13/89] Warn that CUDA data names differing only in case can share a file Windows and macOS usually treat Encoder.ptd and encoder.ptd as one file, so two programs exported into one folder with those data names overwrite each other's weights, and the first one then fails to load. The name check allows both, as both are valid names elsewhere, so the docs now say to pick names that differ in more than case. --- docs/source/backends/cuda/cuda-overview.md | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/docs/source/backends/cuda/cuda-overview.md b/docs/source/backends/cuda/cuda-overview.md index a78a4b831fa..56dccf5813b 100644 --- a/docs/source/backends/cuda/cuda-overview.md +++ b/docs/source/backends/cuda/cuda-overview.md @@ -92,7 +92,9 @@ partitioner = CudaPartitioner([ ``` The name can have up to 200 ASCII letters, digits, `_`, `-` and `.`. It cannot start with `.` or `-`, cannot end with `.ptd` -(the export adds it), and cannot be a Windows device name such as `NUL`. Load each program with its own data file: +(the export adds it), and cannot be a Windows device name such as `NUL`. Windows and macOS usually treat names that differ +only in case, such as `Encoder` and `encoder`, as the same file, so give programs that share a folder names that differ in +more than case. Load each program with its own data file: ```python from executorch.runtime import Runtime From f86630c91a9350943710af034a4b50c71935affb Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Fri, 30 Oct 2026 06:34:21 -0700 Subject: [PATCH 14/89] Correct comments on the CUDA index shims and set_output_data_ptr Some comments added with the CUDA mask indexing change say the wrong thing: - index.h described accumulation as out[indices] += values. With a repeated index, += keeps only one of the values, while index_put_ with accumulate adds every value, and so does the shim. - index.cu said CUB counts items in an int for the index_put sort. That is true for the mask path, but the radix sort takes a wider count. The limit there comes from the int the code passes to the sort. - index.cu said the accumulation types are ATen's. That holds for half and bfloat16, which add in float. Narrower integers add in int64 here, which gives the same wrapped result. - method.h named the output index i in the set_output_data_ptr docs, but the parameter is output_idx. Only comments change. --- backends/cuda/runtime/shims/index.cu | 5 +++-- backends/cuda/runtime/shims/index.h | 4 ++-- runtime/executor/method.h | 3 ++- 3 files changed, 7 insertions(+), 5 deletions(-) diff --git a/backends/cuda/runtime/shims/index.cu b/backends/cuda/runtime/shims/index.cu index 88ed4f7a385..4b696cb03b9 100644 --- a/backends/cuda/runtime/shims/index.cu +++ b/backends/cuda/runtime/shims/index.cu @@ -782,7 +782,7 @@ Error put( // One key per index point, so the scratch and the sort grow with the // number of indices, not with the elements they write. const int64_t points = plan.broadcast_numel; - // CUB counts items in an int on CUDA 12. + // The sort below takes the count as an int. ET_CHECK_OR_RETURN_ERROR( points <= INT_MAX, NotSupported, @@ -1018,7 +1018,8 @@ AOTITorchError aoti_torch_index_put_out( return Error::NotSupported; } } - // The accumulation types are ATen's. + // Half and bfloat16 add in float, as ATen does. Narrower integers add in + // int64, which wraps to the same result as adding in their own type. switch (out->dtype()) { case ScalarType::Float: return put( diff --git a/backends/cuda/runtime/shims/index.h b/backends/cuda/runtime/shims/index.h index 3476eaa8c88..25f015cf4c2 100644 --- a/backends/cuda/runtime/shims/index.h +++ b/backends/cuda/runtime/shims/index.h @@ -45,8 +45,8 @@ AOTI_SHIM_EXPORT AOTITorchError aoti_torch_cuda_index_Tensor( Tensor** ret0); /** - * out[indices] = values, or out[indices] += values when accumulate is set, - * with the semantics of aten::index_put_. + * Writes values into out at indices, with the semantics of aten::index_put_. + * When accumulate is set, every value is added, also when an index repeats. * * Inductor calls this fallback, with out and self the same tensor, for an * index that is a mask, for accumulation into a dtype its generated code diff --git a/runtime/executor/method.h b/runtime/executor/method.h index 501569285f9..9432e2cc891 100644 --- a/runtime/executor/method.h +++ b/runtime/executor/method.h @@ -155,7 +155,8 @@ class Method final { * @param[in] buffer The block of memory to point the specified tensor at. * * @param[in] size the length of buffer in bytes, must be >= the largest - * nbytes the output can have, `MethodMeta::output_tensor_meta(i)->nbytes()`. + * nbytes the output can have, + * `MethodMeta::output_tensor_meta(output_idx)->nbytes()`. * * @param[in] output_idx The index of the output to set the data_ptr for. Must * correspond to a tensor, and that tensor must not have had a buffer From 4b7dabd6d91c5776c20604176ad85ec7ecaffccf Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Fri, 30 Oct 2026 06:18:46 -0700 Subject: [PATCH 15/89] Seed the CUDA random shims from std::random_device The CUDA backend's random shims seeded their generator once per process from time(nullptr). Two processes that start in the same second, such as workers launched together, then drew the same random numbers. Seed from std::random_device instead, a non-deterministic source on the platforms the CUDA backend builds for (Linux and Windows). The per-call counter on the GPU is unchanged, so a captured CUDA graph still draws new numbers on every replay. No test is added: whether two processes start in the same second is not something a test can control, so a test could pass without this change. --- backends/cuda/runtime/shims/rand.cu | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/backends/cuda/runtime/shims/rand.cu b/backends/cuda/runtime/shims/rand.cu index 9916fe8fe0c..5140a59d304 100644 --- a/backends/cuda/runtime/shims/rand.cu +++ b/backends/cuda/runtime/shims/rand.cu @@ -19,9 +19,9 @@ #include #include -#include #include #include +#include #include namespace executorch::backends::cuda { @@ -76,7 +76,10 @@ bool ensure_rng_init(cudaStream_t stream) { } d_rng = static_cast(rng.get()); RngState h; - h.seed = static_cast(time(nullptr)); + // From std::random_device: a seed from the clock repeats for every + // process started in the same second. + std::random_device entropy; + h.seed = (static_cast(entropy()) << 32) | entropy(); h.counter = 0; h.base_scratch = 0; cudaMemcpyAsync( From 26d061264e17103807f0584581214d20e6b86d2e Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Fri, 30 Oct 2026 06:49:28 -0700 Subject: [PATCH 16/89] Run a CUDA method that keeps host state without a graph A CUDA graph records only the device work of the capture call. Host code runs once, during the capture, and never on replay. So for a program that keeps a constant in host memory, such as a torch.cond flag updated on every call, a graph is wrong either way: - a successful capture fixes the branch the capture call took, and replays keep returning it; - a failed capture has already run the host code, so running the call again without a graph takes the other branch. Both returned wrong results with no error. init now leaves the graph off for a method whose AOTI container has any constant in host memory, or when that cannot be checked, and logs why. Methods without host constants, such as the mask index ones, keep their graphs and the retry. The new test is a torch.cond on such a flag; with a graph it returned the captured branch on later calls. --- backends/cuda/runtime/cuda_backend.cpp | 44 +++++++++++++++++-- .../tests/test_output_stride_rearrange.py | 33 ++++++++++++++ 2 files changed, 74 insertions(+), 3 deletions(-) diff --git a/backends/cuda/runtime/cuda_backend.cpp b/backends/cuda/runtime/cuda_backend.cpp index f57425170b5..08fe78843b2 100644 --- a/backends/cuda/runtime/cuda_backend.cpp +++ b/backends/cuda/runtime/cuda_backend.cpp @@ -614,7 +614,16 @@ class ET_EXPERIMENTAL CudaBackend final // Initialize CUDA graph state if enabled for this method. Not on a device // without memory pools, where the method runs without a graph. if (should_use_cuda_graph_for_method(method_name)) { - if (handle->cuda_graph_state.start_warmup()) { + if (keeps_host_constants(handle)) { + // A graph records only the device work of the capture call. Host code, + // such as the choice of a torch.cond branch on a flag in host memory, + // would be fixed at its value from that call. + ET_LOG( + Info, + "CUDA graph requested for method '%s', but it keeps constants in " + "host memory; running without a CUDA graph", + method_name.c_str()); + } else if (handle->cuda_graph_state.start_warmup()) { ET_LOG( Info, "CUDA graph enabled for method '%s' (warmup=%d)", @@ -828,19 +837,48 @@ class ET_EXPERIMENTAL CudaBackend final run_once(handle, args, n_inputs, n_outputs, is_capture_step, graph_ran); if (error != Error::Ok && is_capture_step && !graph_ran) { // A capture records kernels without running them, so until the new graph - // runs, the inputs and any state the program updates are untouched. + // runs, the inputs and any state the program updates are untouched. A + // program that keeps state in host memory, where host code does run + // during a capture, never gets here: init leaves its graph off. ET_LOG( Info, "CUDA graph: capture failed for '%s'; running this call and later " "ones without a graph", handle->method_name.c_str()); - error = run_once(handle, args, n_inputs, n_outputs, false, graph_ran); + error = run_once( + handle, + args, + n_inputs, + n_outputs, + /*is_capture_step=*/false, + graph_ran); } ET_CHECK_OK_OR_RETURN_ERROR(error); return commit_offgraph_kv_step(handle, kv_step_width); } private: + // Whether any constant of the program is in host memory. Unknown counts as + // yes. + static bool keeps_host_constants(cuda::CudaDelegateHandle* handle) { + if (handle->extract_constants_map == nullptr) { + return true; + } + std::unordered_map constants; + if (handle->extract_constants_map( + handle->container_handle, + reinterpret_cast(&constants), + /*use_inactive=*/false) != Error::Ok) { + return true; + } + for (const auto& [fqn, tensor] : constants) { + if (reinterpret_cast(tensor)->is_cpu()) { + return true; + } + } + return false; + } + // Runs the program once. On the capture step it captures the CUDA graph // instead, then launches it once, and sets graph_ran once that launch is // queued. diff --git a/backends/cuda/tests/test_output_stride_rearrange.py b/backends/cuda/tests/test_output_stride_rearrange.py index e791dab24bc..53d4d307c29 100644 --- a/backends/cuda/tests/test_output_stride_rearrange.py +++ b/backends/cuda/tests/test_output_stride_rearrange.py @@ -292,6 +292,20 @@ def forward(self, x, mask): return x.index_put((mask,), torch.ones(4), accumulate=True) +class _HostFlag(nn.Module): + """A torch.cond on a flag the program keeps in host memory and updates on + every call.""" + + def __init__(self): + super().__init__() + self.register_buffer("flag", torch.tensor(False), persistent=False) + + def forward(self, x): + picked = torch.cond(self.flag, lambda t: t + t, lambda t: t.neg(), (x,)) + self.flag.fill_(x.shape[1] > 1) + return (picked,) + + class _Double(nn.Module): def __init__(self, memory_format): super().__init__() @@ -401,6 +415,25 @@ def test_input_shape_changes_after_capture(self): # InvalidArgument, from the replay's input shape check. self.assertIn("execute() failed with error 0x12", str(out)) + def test_host_state_runs_without_a_graph(self): + # A graph would keep the branch of the capture call, and a failed + # capture would already have updated the flag before running the call + # again, so this method must not use a graph. + torch.manual_seed(0) + module = _HostFlag().cuda().eval() + eager = _HostFlag().cuda().eval() + eager.load_state_dict(module.state_dict()) + inputs = [torch.randn(1, rows, 64) for rows in (1, 1, 1, 50, 50, 50)] + rows = torch.export.Dim("rows", min=1, max=64) + outputs = self._run_with_cuda_graph( + module, (inputs[3].cuda(),), inputs, dynamic_shapes=({1: rows},) + ) + for i, (x, out) in enumerate(zip(inputs, outputs)): + with self.subTest(call=i): + with torch.no_grad(): + expected = eager(x.cuda())[0].cpu() + torch.testing.assert_close(out, expected) + def test_mask_runs_without_a_graph(self): # A mask index reads its count back to the host, which a capture # cannot record, so the capture call has to run without a graph. From bfca1cd19f519d391d5386756e573621756225c0 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Fri, 30 Oct 2026 06:51:06 -0700 Subject: [PATCH 17/89] Check the AOTI data name before compiling A backend that keeps its data in the .pte, such as Metal, refuses a data_name compile spec. preprocess only checked it after moving the program to the device and compiling it with AOTInductor, which took about 18 seconds and left the compiled library on disk. It now checks the name first. The data name test calls the shared preprocess with such a spec, so moving the check back behind the compile, or dropping it, fails the test. --- backends/aoti/aoti_backend.py | 4 +++- backends/cuda/tests/test_cuda_export.py | 6 ++++++ 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/backends/aoti/aoti_backend.py b/backends/aoti/aoti_backend.py index d6b78d66000..87a15498de7 100644 --- a/backends/aoti/aoti_backend.py +++ b/backends/aoti/aoti_backend.py @@ -286,6 +286,9 @@ def preprocess( Preprocess the edge program and compile it using AOTInductor. Weights are always separated from the SO file. """ + # First, so that a backend keeping its data in the .pte refuses a data + # name before anything is compiled. + data_name = cls.data_name_from_compile_specs(compile_specs) device_name = cls.get_device_name() decomposition_table = cls.get_decomposition_table() options = cls.get_aoti_compile_options(compile_specs) @@ -390,7 +393,6 @@ def preprocess( named_data_store.add_named_data(so_blob_key, so_data, 1, None) # Determine whether to save named data externally based on backend setting # External: save to separate .ptd file, otherwise merge with .pte file - data_name = cls.data_name_from_compile_specs(compile_specs) external_tag = data_name if cls.save_data_externally() else None named_data_store.add_named_data(weights_blob_key, blob_data, 1, external_tag) diff --git a/backends/cuda/tests/test_cuda_export.py b/backends/cuda/tests/test_cuda_export.py index b678b2242ef..b187fb25840 100644 --- a/backends/cuda/tests/test_cuda_export.py +++ b/backends/cuda/tests/test_cuda_export.py @@ -14,6 +14,7 @@ from unittest.mock import patch import torch +from executorch.backends.aoti.aoti_backend import AotiBackend from executorch.backends.cuda.cuda_backend import CudaBackend from executorch.backends.cuda.cuda_partitioner import CudaPartitioner from executorch.examples.models.toy_model import SdpaModule @@ -278,6 +279,11 @@ def test_invalid_data_name_compile_spec(self): CudaBackend.data_name_from_compile_specs( [CompileSpec(key="data_name", value=b"encoder")] ) + # The shared preprocess, which Metal uses, refuses it before compiling. + with self.assertRaisesRegex(ValueError, "takes no data name"): + AotiBackend.preprocess( + None, [CompileSpec(key="data_name", value=b"encoder")] + ) def test_valid_data_name_compile_spec(self): self.assertEqual(CudaBackend.data_name_from_compile_specs([]), "aoti_cuda_blob") From 19a6331d21ff0e61ea51356f53c70bbea21a1da6 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Fri, 30 Oct 2026 06:54:16 -0700 Subject: [PATCH 18/89] Check delegation and restore deterministic mode in the CUDA index shim tests The index shim tests only compared results with eager, so they would still pass if the index ops ran outside the delegate. They now check that the lowered program does nothing but call the delegate. The deterministic index_put test restores warn_only along with deterministic mode, so it no longer leaves strict deterministic mode on for later tests when warnings were enabled before it. --- backends/cuda/tests/test_index_shims.py | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/backends/cuda/tests/test_index_shims.py b/backends/cuda/tests/test_index_shims.py index af7a7b2fa84..a53d1eebba7 100644 --- a/backends/cuda/tests/test_index_shims.py +++ b/backends/cuda/tests/test_index_shims.py @@ -40,6 +40,15 @@ def _run( partitioner=[CudaPartitioner([spec])], compile_config=EdgeCompileConfig(_check_ir_validity=False), ) + # Everything runs in the delegate, the index ops included. + self.assertEqual( + { + n.target.__name__ + for n in lowered.exported_program().graph.nodes + if n.op == "call_function" + }, + {"executorch_call_delegate", "getitem"}, + ) from executorch.runtime import Runtime with tempfile.TemporaryDirectory() as directory: @@ -53,7 +62,8 @@ def _run( torch.testing.assert_close(output.cpu(), expected) def test_mask_index_and_index_put(self): - """Inductor also has no atomic add for int64.""" + """index_put with accumulate goes to the shim too: Inductor has no atomic + add for int64.""" class MaskIndex(torch.nn.Module): def forward(self, x, mask, column, idx, counts): @@ -100,6 +110,7 @@ def forward(self, x, idx, values): idx = torch.tensor([0, 3, 3, 1, 3]) deterministic = torch.are_deterministic_algorithms_enabled() + warn_only = torch.is_deterministic_algorithms_warn_only_enabled() torch.use_deterministic_algorithms(True) try: for dtype in (torch.float16, torch.bfloat16): @@ -108,4 +119,4 @@ def forward(self, x, idx, values): values = torch.arange(1, 6, dtype=dtype) self._run(Put(), [(x, idx, values)]) finally: - torch.use_deterministic_algorithms(deterministic) + torch.use_deterministic_algorithms(deterministic, warn_only=warn_only) From f332a8820b8a0cab134543399b9c984e22b1b76a Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Fri, 30 Oct 2026 06:54:16 -0700 Subject: [PATCH 19/89] Fix the CUDA graph capture guard comment Without the switch to Disabled, a capture that cannot succeed would be tried again on every later call, not on every fourth one: warmup is already over. --- backends/cuda/runtime/cuda_backend.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/backends/cuda/runtime/cuda_backend.cpp b/backends/cuda/runtime/cuda_backend.cpp index 08fe78843b2..f8d656cb2f7 100644 --- a/backends/cuda/runtime/cuda_backend.cpp +++ b/backends/cuda/runtime/cuda_backend.cpp @@ -916,7 +916,7 @@ class ET_EXPERIMENTAL CudaBackend final } // Free what this attempt pinned and disable graphs for this method, so // a capture that cannot succeed is tried once rather than on every - // fourth call for the life of the process. Eager execution is correct, + // later call for the life of the process. Eager execution is correct, // just slower. state_->release(); state_->phase = CudaGraphPhase::Disabled; From 7eaf9ea9e06c34f81c6a61d8d103f5490eee3fbd Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Fri, 30 Oct 2026 07:18:05 -0700 Subject: [PATCH 20/89] Copy through a boolean mask with where in the CUDA backend `out[mask] = src[mask]` exports as `index_put(out, [mask], index(src, [mask]))`. A mask index makes a result size depend on the data, so both calls run in the backend's index shims, and each one reads the mask's count back to the host and waits for the GPU. A program that does this pays two waits per call. Transformers' SmolVLM vision embeddings do it 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 `torch.where(mask, src, out)` gives the same result, bit for bit, with no count and no wait. A new pass in the CUDA backend makes that rewrite. The two masks may be one node or two equal computations, as in SmolVLM, which writes `mask.view(batch_size, -1)` on both sides. The pass also drops the `index` and the checks on its size when nothing else reads them, since those checks hold for any mask, including one whose size is dynamic. A user check on the count, such as `torch._check(x[mask].shape[0] >= 1)`, keeps the `index`. Inductor already turns a mask `index_put` with a single value into a masked fill in the same way. --- backends/cuda/BUCK | 1 + backends/cuda/cuda_backend.py | 4 + .../passes/replace_masked_copy_with_where.py | 183 ++++++++++++++++++ .../test_replace_masked_copy_with_where.py | 95 +++++++++ 4 files changed, 283 insertions(+) create mode 100644 backends/cuda/passes/replace_masked_copy_with_where.py create mode 100644 backends/cuda/passes/tests/test_replace_masked_copy_with_where.py diff --git a/backends/cuda/BUCK b/backends/cuda/BUCK index 7f88126cd5f..8f24d828a9f 100644 --- a/backends/cuda/BUCK +++ b/backends/cuda/BUCK @@ -81,6 +81,7 @@ fbcode_target( "passes/lower_offgraph_kv.py", "passes/replace_fused_sdpa.py", "passes/replace_int64_floordiv.py", + "passes/replace_masked_copy_with_where.py", ], visibility = [ "//executorch/backends/cuda/...", diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index 243437faa8a..7be95d144d5 100644 --- a/backends/cuda/cuda_backend.py +++ b/backends/cuda/cuda_backend.py @@ -45,6 +45,9 @@ from executorch.backends.cuda.passes.replace_int64_floordiv import ( ReplaceInt64FloorDivWithFloatPass, ) +from executorch.backends.cuda.passes.replace_masked_copy_with_where import ( + ReplaceMaskedCopyWithWherePass, +) from executorch.backends.cuda.target_smem import target_smem_context from executorch.backends.cuda.triton.kernels import triton_mutation_cache from executorch.backends.cuda.triton.replacement_pass import ( @@ -1020,6 +1023,7 @@ def get_custom_passes(cls, compile_specs: List[CompileSpec]) -> List[typing.Any] MoveCondPredicateToCpuPass(), ReplaceInt64FloorDivWithFloatPass(), ReplaceFusedSdpaWithSdpaPass(), + ReplaceMaskedCopyWithWherePass(), ] cell_layout = any( spec.key == OFFGRAPH_KV_COMPILE_SPEC diff --git a/backends/cuda/passes/replace_masked_copy_with_where.py b/backends/cuda/passes/replace_masked_copy_with_where.py new file mode 100644 index 00000000000..dfcd5b5e7ff --- /dev/null +++ b/backends/cuda/passes/replace_masked_copy_with_where.py @@ -0,0 +1,183 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +""" +Graph Transformation Pass for Masked Copies. + +Rewrites a copy through a boolean mask into a select: + + out[mask] = src[mask] -> out = torch.where(mask, src, out) + +A mask index has a result size that depends on the data, so ``index.Tensor`` +and ``index_put`` with a mask run in the backend's index shims, and each call +reads the mask's count back to the host, which waits for the GPU. When the +values written are the same mask's picks from a tensor shaped like the +destination, every element goes back where it came from, so ``where`` gives +the same result, element for element, with no count and no host read. +""" + +from typing import Any, List, Optional + +import sympy +import torch +from executorch.exir.dialects._ops import ops as exir_ops +from torch.fx import GraphModule, Node +from torch.fx.experimental.symbolic_shapes import statically_known_true +from torch.fx.passes.infra.pass_base import PassBase, PassResult +from torch.utils._sympy.value_ranges import bound_sympy, ValueRanges + +_INDEX = exir_ops.edge.aten.index.Tensor +_INDEX_PUT = exir_ops.edge.aten.index_put.default +_SYM_SIZE = torch.ops.aten.sym_size.int +_ASSERT_SCALAR = torch.ops.aten._assert_scalar.default + + +def _same_value(a: Any, b: Any) -> bool: + """ + Whether a and b are the same node, or the same deterministic op applied to + the same values, as with ``x[m.view(-1)] = y[m.view(-1)]``. + """ + if isinstance(a, (list, tuple)) and isinstance(b, (list, tuple)): + return len(a) == len(b) and all(map(_same_value, a, b)) + if not isinstance(a, Node) or not isinstance(b, Node): + return not isinstance(a, Node) and not isinstance(b, Node) and a == b + if a is b: + return True + if a.op != "call_function" or a.target != b.target: + return False + op = getattr(a.target, "_op", a.target) + if torch.Tag.nondeterministic_seeded in getattr(op, "tags", ()): + return False + return ( + _same_value(a.args, b.args) + and a.kwargs.keys() == b.kwargs.keys() + and all(_same_value(a.kwargs[k], b.kwargs[k]) for k in a.kwargs) + ) + + +class ReplaceMaskedCopyWithWherePass(PassBase): + """ + Replace ``index_put(out, indices, index(src, indices))`` with + ``where(mask, src, out)`` when ``indices`` is one boolean mask (after any + number of ``None``), ``src`` and ``out`` have the same shape and dtype, and + the put does not accumulate. + + The two masks may be one node or two equal computations. The ``index`` is + then removed too if nothing else reads it but the range checks export adds + on its size, since those hold for any mask. + """ + + def call(self, graph_module: GraphModule) -> PassResult: + modified = False + for node in list(graph_module.graph.nodes): + mask = self._mask_of_copy(node) + if mask is None: + continue + self._replace(graph_module, node, mask) + modified = True + if modified: + graph_module.graph.lint() + graph_module.recompile() + return PassResult(graph_module, modified) + + @staticmethod + def _mask_of_copy(node: Node) -> Optional[Node]: + """The mask of an index_put that copies through it, else None.""" + if node.op != "call_function" or node.target != _INDEX_PUT: + return None + out, indices, values = node.args[:3] + accumulate = ( + node.args[3] if len(node.args) > 3 else node.kwargs.get("accumulate", False) + ) + if accumulate or not isinstance(values, Node) or values.target != _INDEX: + return None + src, src_indices = values.args + if not _same_value(list(src_indices), list(indices)): + return None + masks = [index for index in indices if index is not None] + if len(masks) != 1: + return None + mask_val = masks[0].meta.get("val") + out_val = out.meta.get("val") + src_val = src.meta.get("val") + if not all(isinstance(v, torch.Tensor) for v in (mask_val, out_val, src_val)): + return None + if mask_val.dtype != torch.bool or mask_val.dim() == 0: + return None + if src_val.dtype != out_val.dtype or src_val.dim() != out_val.dim(): + return None + if not all( + statically_known_true(a == b) for a, b in zip(src_val.shape, out_val.shape) + ): + return None + return masks[0] + + def _replace(self, graph_module: GraphModule, node: Node, mask: Node) -> None: + graph = graph_module.graph + out, indices, values = node.args[:3] + src = values.args[0] + # The mask covers the dims after the leading Nones; broadcasting lines + # it up from the right, so only the dims after it need a size-1 dim. + trailing = out.meta["val"].dim() - indices.index(mask) - mask.meta["val"].dim() + with graph.inserting_before(node): + condition = mask + for _ in range(trailing): + condition_val = condition.meta["val"].unsqueeze(-1) + condition = graph.call_function( + exir_ops.edge.aten.unsqueeze_copy.default, args=(condition, -1) + ) + condition.meta["val"] = condition_val + where = graph.call_function( + exir_ops.edge.aten.where.self, args=(condition, src, out) + ) + where.meta = node.meta.copy() + node.replace_all_uses_with(where) + graph.erase_node(node) + dead = self._only_range_checked(values, mask) + if dead is not None: + for user in reversed(dead): + graph.erase_node(user) + + @staticmethod + def _only_range_checked(index: Node, mask: Node) -> Optional[List[Node]]: + """ + The index and everything that reads it, in graph order, if all that + reads it is a check that its size is in [0, mask.numel()], else None. + """ + numel = mask.meta["val"].numel() + if isinstance(numel, torch.SymInt): + numel = numel.node.shape_env.bound_sympy(numel.node.expr).upper + within = ValueRanges(0, numel) + found = [index] + pending = list(index.users) + while pending: + user = pending.pop() + if user in found: + continue + if user.target == _ASSERT_SCALAR: + condition = user.args[0] + expr = ( + condition.meta.get("val").node.expr + if isinstance(condition, Node) + and isinstance(condition.meta.get("val"), torch.SymBool) + else None + ) + if expr is None or len(expr.free_symbols) != 1: + return None + (symbol,) = expr.free_symbols + bounds = bound_sympy(expr, {symbol: within}) + if not (bounds.is_singleton() and bounds.lower == sympy.true): + return None + elif user.target == _SYM_SIZE and user.args[0] is index: + pass + elif not isinstance( + user.meta.get("val"), (torch.SymInt, torch.SymBool) + ) or not any(arg in found for arg in user.all_input_nodes): + return None + found.append(user) + pending.extend(user.users) + order = {n: i for i, n in enumerate(index.graph.nodes)} + return sorted(found, key=order.__getitem__) diff --git a/backends/cuda/passes/tests/test_replace_masked_copy_with_where.py b/backends/cuda/passes/tests/test_replace_masked_copy_with_where.py new file mode 100644 index 00000000000..798253187b1 --- /dev/null +++ b/backends/cuda/passes/tests/test_replace_masked_copy_with_where.py @@ -0,0 +1,95 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import unittest + +import torch +from executorch.backends.cuda.cuda_backend import CudaBackend +from executorch.backends.cuda.passes.replace_masked_copy_with_where import ( + ReplaceMaskedCopyWithWherePass, +) +from executorch.exir import to_edge +from executorch.exir.dialects._ops import ops as exir_ops +from torch.export import export + + +def _targets(graph_module): + return [n.target for n in graph_module.graph.nodes if n.op == "call_function"] + + +class TestReplaceMaskedCopyWithWherePass(unittest.TestCase): + def test_registered_in_cuda_backend(self): + self.assertIn( + "ReplaceMaskedCopyWithWherePass", + {type(p).__name__ for p in CudaBackend.get_custom_passes([])}, + ) + + def test_masked_copy_becomes_where(self): + """Also removes the index and the checks on its size.""" + + class Copy(torch.nn.Module): + def forward(self, x, mask): + out = torch.zeros_like(x) + # A leading None and a trailing dim, and the mask computed twice, + # as Transformers' SmolVLM writes its position ids. + out[:, mask.view(-1)] = x[:, mask.view(-1)] + return out + + x = torch.randn(2, 5, 3) + masks = ( + torch.tensor([True, False, True, True, False]), + torch.zeros(5, dtype=torch.bool), + torch.ones(5, dtype=torch.bool), + ) + program = to_edge(export(Copy(), (x, masks[0]))).exported_program() + ReplaceMaskedCopyWithWherePass()(program.graph_module) + + targets = _targets(program.graph_module) + self.assertIn(exir_ops.edge.aten.where.self, targets) + for removed in ( + exir_ops.edge.aten.index.Tensor, + exir_ops.edge.aten.index_put.default, + torch.ops.aten._assert_scalar.default, + ): + self.assertNotIn(removed, targets) + for mask in masks: + torch.testing.assert_close( + program.module()(x, mask), Copy()(x, mask), rtol=0, atol=0 + ) + + def test_other_writes_are_kept(self): + class Kept(torch.nn.Module): + def forward(self, x, mask): + scaled = torch.zeros_like(x) + # The values are not the mask's picks of a tensor. + scaled[mask] = x[mask] * 2 + checked = torch.zeros_like(x) + picked = x[mask] + # A check on the count that does not hold for every mask. + torch._check(picked.shape[0] >= 1) + checked[mask] = picked + summed = torch.ones_like(x) + summed.index_put_((mask,), x[mask], accumulate=True) + other = torch.zeros_like(x) + other[mask] = x[mask.logical_not()] + return scaled, checked, summed, other + + x = torch.randn(4, 3) + mask = torch.tensor([True, False, True, False]) + program = to_edge(export(Kept(), (x, mask))).exported_program() + ReplaceMaskedCopyWithWherePass()(program.graph_module) + + targets = _targets(program.graph_module) + self.assertEqual(targets.count(exir_ops.edge.aten.index_put.default), 3) + self.assertEqual(targets.count(exir_ops.edge.aten.where.self), 1) + # The index stays for the checked count. + self.assertIn(exir_ops.edge.aten.index.Tensor, targets) + for got, expected in zip(program.module()(x, mask), Kept()(x, mask)): + torch.testing.assert_close(got, expected, rtol=0, atol=0) + + +if __name__ == "__main__": + unittest.main() From ebca4006fb43afe1a11fb22b2e9cb4f2007e070f Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Fri, 30 Oct 2026 10:38:41 -0700 Subject: [PATCH 21/89] Drop tensor checks before the CUDA backend compiles a program `torch._assert_async(cond, msg)` exports as `aten._assert_async.msg`, and the argument checks in `torch.distributions` use it. In the CUDA backend, Inductor turns each one into a device assert inside a Triton kernel. While compiling, AOTInductor runs every kernel once to pick its launch settings. It runs them on made-up buffers instead of real data: random floats, and zeros for integers and bools. When a check reads a bool that an earlier kernel computes, it sees false there and fails, even though it holds for every real input. A failed device assert breaks the CUDA context, so the export stops with "device-side assert triggered". A Gaussian policy head shows it. `MultivariateNormal(mean, scale_tril=...)` checks that `scale_tril` is lower triangular with a positive diagonal. With 14 actions, Inductor computes the triangle check in its own kernel, and the export fails. Eager PyTorch, XNNPACK and TensorRT all handle the same model. A new pass in the CUDA backend removes these checks, and the math that only fed them, before compiling. This means the CUDA delegate no longer runs them. That matches what ExecuTorch already does for the part of a program that no backend takes: `RemoveGraphAssertsPass` removes the same checks before the program is saved. Keeping them would not help much at run time anyway: a failed device assert cannot come back as an error, it breaks the CUDA context for the rest of the process. Two other ways were measured and not taken. Turning off `autotune_at_compile_time` makes Inductor tune on the real weights and zero inputs instead, which passes here but still fails any check that zero inputs break. Checking on the host would copy a value back and wait for the GPU on every check, and a CUDA graph cannot capture that. The pass only touches `aten._assert_async.msg`. Checks on sizes and other symbolic values (`_assert_scalar`) stay, since Inductor already runs those on the host. --- backends/cuda/BUCK | 1 + backends/cuda/cuda_backend.py | 4 ++ backends/cuda/passes/remove_assert_async.py | 48 +++++++++++++++++++++ backends/cuda/tests/test_cuda_export.py | 25 +++++++++++ 4 files changed, 78 insertions(+) create mode 100644 backends/cuda/passes/remove_assert_async.py diff --git a/backends/cuda/BUCK b/backends/cuda/BUCK index 8f24d828a9f..7e8d4c88675 100644 --- a/backends/cuda/BUCK +++ b/backends/cuda/BUCK @@ -79,6 +79,7 @@ fbcode_target( "passes/__init__.py", "passes/move_cond_predicate_to_cpu.py", "passes/lower_offgraph_kv.py", + "passes/remove_assert_async.py", "passes/replace_fused_sdpa.py", "passes/replace_int64_floordiv.py", "passes/replace_masked_copy_with_where.py", diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index 7be95d144d5..3e6fc1cd607 100644 --- a/backends/cuda/cuda_backend.py +++ b/backends/cuda/cuda_backend.py @@ -39,6 +39,9 @@ from executorch.backends.cuda.passes.move_cond_predicate_to_cpu import ( MoveCondPredicateToCpuPass, ) +from executorch.backends.cuda.passes.remove_assert_async import ( + RemoveAssertAsyncPass, +) from executorch.backends.cuda.passes.replace_fused_sdpa import ( ReplaceFusedSdpaWithSdpaPass, ) @@ -1020,6 +1023,7 @@ def get_custom_passes(cls, compile_specs: List[CompileSpec]) -> List[typing.Any] ) triton_kernel_mode = mode passes = [ + RemoveAssertAsyncPass(), MoveCondPredicateToCpuPass(), ReplaceInt64FloorDivWithFloatPass(), ReplaceFusedSdpaWithSdpaPass(), diff --git a/backends/cuda/passes/remove_assert_async.py b/backends/cuda/passes/remove_assert_async.py new file mode 100644 index 00000000000..b6fb062626c --- /dev/null +++ b/backends/cuda/passes/remove_assert_async.py @@ -0,0 +1,48 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import torch +from torch.fx import GraphModule +from torch.fx.passes.infra.pass_base import PassBase, PassResult + + +class RemoveAssertAsyncPass(PassBase): + """ + Removes tensor checks (``aten._assert_async.msg``) from the program the + CUDA backend compiles, along with the math that only fed them. + + ``torch._assert_async(cond, msg)`` exports as this op, and the argument + checks in ``torch.distributions`` use it. Inductor compiles each one into + a device assert inside a Triton kernel. AOTInductor then runs every kernel + once at compile time to pick its launch settings, on made-up buffers: + random floats, and zeros for integers and bools. A check that reads a bool + computed by an earlier kernel sees false there and fails, even when it + holds for every real input. A failed device assert also breaks the CUDA + context, so the export cannot go on. + + ExecuTorch already removes these checks from the rest of a program before + it is saved (see ``RemoveGraphAssertsPass``), so the delegate now does the + same. + """ + + def call(self, graph_module: GraphModule) -> PassResult: + modified = False + for module in graph_module.modules(): + if not isinstance(module, GraphModule): + continue + asserts = [ + node + for node in module.graph.nodes + if node.op == "call_function" + and node.target == torch.ops.aten._assert_async.msg + ] + for node in asserts: + module.graph.erase_node(node) + if asserts: + module.graph.eliminate_dead_code() + module.recompile() + modified = True + return PassResult(graph_module, modified) diff --git a/backends/cuda/tests/test_cuda_export.py b/backends/cuda/tests/test_cuda_export.py index b187fb25840..472781e6f58 100644 --- a/backends/cuda/tests/test_cuda_export.py +++ b/backends/cuda/tests/test_cuda_export.py @@ -447,6 +447,31 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: CheckModule().eval(), (torch.randn(4, 8),), strict=False ) + def test_distribution_argument_check(self): + """A check that holds for every input must not fail while AOTI autotunes.""" + + class GaussianHead(torch.nn.Module): + def __init__(self): + super().__init__() + # With 14 actions, Inductor checks scale_tril in a kernel of its + # own, so the assert reads a bool that autotuning makes up. + self.mean = torch.nn.Linear(16, 14) + self.log_std = torch.nn.Linear(16, 14) + + def forward(self, x: torch.Tensor, eps: torch.Tensor) -> torch.Tensor: + std = torch.clamp(torch.exp(self.log_std(x)), 1e-5, 10.0) + dist = torch.distributions.MultivariateNormal( + self.mean(x), scale_tril=torch.diag_embed(std.sqrt()) + ) + return dist.loc + (dist.scale_tril @ eps[..., None])[..., 0] + + # A strict export cannot trace torch.distributions. + self._export_to_cuda_with_lower( + GaussianHead().eval(), + (torch.randn(1, 16), torch.randn(1, 14)), + strict=False, + ) + def test_conv2d(self): """Test CUDA export for 2D convolution.""" From 9e380c80fdfab7a2be14c68cef5ddd20053ff8b0 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Fri, 30 Oct 2026 10:18:53 -0700 Subject: [PATCH 22/89] Test the CUDA masked copy pass and the capture retry end to end The masked copy pass was only tested on a CPU graph. In a real lowering it runs inside the CUDA backend, after other passes, and comparing outputs cannot tell whether it fired: a masked copy that keeps its index still gives the right values through the index shims. A new test lowers the copy through the CUDA partitioner, compares with eager for three masks, and records the graph the pass returns, so it fails if the pass leaves the index in place. The capture retry was only tested on modules without state. A new test uses a module that adds one to a device buffer on every call and reads a mask's count, which makes the capture call fail and run again without a graph. It checks that call N returns N, so an update applied twice fails. Also seed the random mask in test_two_dim_mask so a failure reproduces, and drop a load_state_dict call in the host state test that copied nothing. --- backends/cuda/tests/test_index_shims.py | 39 +++++++++++++++++++ .../tests/test_output_stride_rearrange.py | 24 +++++++++++- 2 files changed, 62 insertions(+), 1 deletion(-) diff --git a/backends/cuda/tests/test_index_shims.py b/backends/cuda/tests/test_index_shims.py index a53d1eebba7..b1899e4e794 100644 --- a/backends/cuda/tests/test_index_shims.py +++ b/backends/cuda/tests/test_index_shims.py @@ -16,11 +16,16 @@ import tempfile import unittest from typing import Sequence, Tuple +from unittest.mock import patch import torch from executorch.backends.cuda.cuda_backend import CudaBackend from executorch.backends.cuda.cuda_partitioner import CudaPartitioner +from executorch.backends.cuda.passes.replace_masked_copy_with_where import ( + ReplaceMaskedCopyWithWherePass, +) from executorch.exir import EdgeCompileConfig, to_edge_transform_and_lower +from executorch.exir.dialects._ops import ops as exir_ops from torch.export import export @@ -94,10 +99,44 @@ class TwoDimMask(torch.nn.Module): def forward(self, x, mask): return (x[mask] * 2,) + torch.manual_seed(0) x = torch.randn(4, 5, 6) masks = (torch.rand(4, 5) > 0.5, torch.zeros(4, 5, dtype=torch.bool)) self._run(TwoDimMask(), [(x, mask) for mask in masks]) + def test_masked_copy_lowers_to_where(self): + """Outputs alone cannot show the rewrite: a masked copy that keeps its + index still gives the right values through the shims.""" + + class Copy(torch.nn.Module): + def forward(self, x, mask): + out = torch.zeros_like(x) + out[:, mask.view(-1)] = x[:, mask.view(-1)] + return (out,) + + rewritten = [] + call = ReplaceMaskedCopyWithWherePass.call + + def record(self, graph_module): + result = call(self, graph_module) + rewritten.extend( + n.target + for n in result.graph_module.graph.nodes + if n.op == "call_function" + ) + return result + + x = torch.randn(2, 5, 3) + masks = ( + torch.tensor([True, False, True, True, False]), + torch.zeros(5, dtype=torch.bool), + torch.ones(5, dtype=torch.bool), + ) + with patch.object(ReplaceMaskedCopyWithWherePass, "call", record): + self._run(Copy(), [(x, mask) for mask in masks]) + self.assertIn(exir_ops.edge.aten.where.self, rewritten) + self.assertNotIn(exir_ops.edge.aten.index.Tensor, rewritten) + def test_index_put_with_repeated_indices(self): """Deterministic mode sends half, bfloat16 and int32-index writes to the shim.""" diff --git a/backends/cuda/tests/test_output_stride_rearrange.py b/backends/cuda/tests/test_output_stride_rearrange.py index 53d4d307c29..985f3e6feee 100644 --- a/backends/cuda/tests/test_output_stride_rearrange.py +++ b/backends/cuda/tests/test_output_stride_rearrange.py @@ -292,6 +292,18 @@ def forward(self, x, mask): return x.index_put((mask,), torch.ones(4), accumulate=True) +class _MaskedCounter(nn.Module): + """Counts its calls in a device buffer, and reads a mask's count.""" + + def __init__(self): + super().__init__() + self.register_buffer("calls", torch.zeros(1), persistent=False) + + def forward(self, x, mask): + self.calls.add_(1) + return self.calls.clone(), x[mask] + + class _HostFlag(nn.Module): """A torch.cond on a flag the program keeps in host memory and updates on every call.""" @@ -422,7 +434,6 @@ def test_host_state_runs_without_a_graph(self): torch.manual_seed(0) module = _HostFlag().cuda().eval() eager = _HostFlag().cuda().eval() - eager.load_state_dict(module.state_dict()) inputs = [torch.randn(1, rows, 64) for rows in (1, 1, 1, 50, 50, 50)] rows = torch.export.Dim("rows", min=1, max=64) outputs = self._run_with_cuda_graph( @@ -446,6 +457,17 @@ def test_mask_runs_without_a_graph(self): self.assertIsInstance(out, torch.Tensor) torch.testing.assert_close(out, module(*args)) + def test_mask_retry_updates_state_once(self): + # The capture call runs again without a graph, so its update to the + # count must not land twice. + torch.manual_seed(0) + inputs = [(torch.randn(8, 4), torch.rand(8) > 0.5) for _ in range(6)] + outputs = self._run_with_cuda_graph(_MaskedCounter().eval(), inputs[0], inputs) + for i, out in enumerate(outputs): + with self.subTest(call=i): + self.assertIsInstance(out, torch.Tensor) + torch.testing.assert_close(out, torch.tensor([i + 1.0])) + def test_output_shape_changes_at_capture(self): torch.manual_seed(0) # The capture call has a larger batch than the warmup calls before it. From 3c7349669a229f6674b4f261af7c50708aa0374f Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Fri, 30 Oct 2026 10:18:53 -0700 Subject: [PATCH 23/89] Remove a size check only when it tests the index's count The masked copy pass also removes the index and the checks export adds on its size, since those hold for any mask. It accepted a check on any single symbol, bounding that symbol by the mask's size. A check that reads the count but tests another size, such as count * 0 + y.size(0) <= 4, was removed too, although it does not hold for every input. The pass now requires the checked symbol to be the index's count, its size at the mask's position. A new test builds that check by hand and fails without the change. Also drop a condition that could never be false (every node it saw reads a node already found), and check the op tags with isinstance instead of getattr. --- .../passes/replace_masked_copy_with_where.py | 33 ++++++++------ .../test_replace_masked_copy_with_where.py | 43 ++++++++++++++++++- 2 files changed, 63 insertions(+), 13 deletions(-) diff --git a/backends/cuda/passes/replace_masked_copy_with_where.py b/backends/cuda/passes/replace_masked_copy_with_where.py index dfcd5b5e7ff..0f76b995f34 100644 --- a/backends/cuda/passes/replace_masked_copy_with_where.py +++ b/backends/cuda/passes/replace_masked_copy_with_where.py @@ -24,6 +24,7 @@ import sympy import torch from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.dialects.edge._ops import EdgeOpOverload from torch.fx import GraphModule, Node from torch.fx.experimental.symbolic_shapes import statically_known_true from torch.fx.passes.infra.pass_base import PassBase, PassResult @@ -48,8 +49,10 @@ def _same_value(a: Any, b: Any) -> bool: return True if a.op != "call_function" or a.target != b.target: return False - op = getattr(a.target, "_op", a.target) - if torch.Tag.nondeterministic_seeded in getattr(op, "tags", ()): + if ( + isinstance(a.target, (EdgeOpOverload, torch._ops.OpOverload)) + and torch.Tag.nondeterministic_seeded in a.target.tags + ): return False return ( _same_value(a.args, b.args) @@ -119,9 +122,10 @@ def _replace(self, graph_module: GraphModule, node: Node, mask: Node) -> None: graph = graph_module.graph out, indices, values = node.args[:3] src = values.args[0] + position = indices.index(mask) # The mask covers the dims after the leading Nones; broadcasting lines # it up from the right, so only the dims after it need a size-1 dim. - trailing = out.meta["val"].dim() - indices.index(mask) - mask.meta["val"].dim() + trailing = out.meta["val"].dim() - position - mask.meta["val"].dim() with graph.inserting_before(node): condition = mask for _ in range(trailing): @@ -136,17 +140,23 @@ def _replace(self, graph_module: GraphModule, node: Node, mask: Node) -> None: where.meta = node.meta.copy() node.replace_all_uses_with(where) graph.erase_node(node) - dead = self._only_range_checked(values, mask) + dead = self._only_range_checked(values, mask, position) if dead is not None: for user in reversed(dead): graph.erase_node(user) @staticmethod - def _only_range_checked(index: Node, mask: Node) -> Optional[List[Node]]: + def _only_range_checked( + index: Node, mask: Node, position: int + ) -> Optional[List[Node]]: """ The index and everything that reads it, in graph order, if all that - reads it is a check that its size is in [0, mask.numel()], else None. + reads it is a check that its count, its size at ``position``, is in + [0, mask.numel()], else None. """ + count = index.meta["val"].shape[position] + if not isinstance(count, torch.SymInt): + return None numel = mask.meta["val"].numel() if isinstance(numel, torch.SymInt): numel = numel.node.shape_env.bound_sympy(numel.node.expr).upper @@ -165,17 +175,16 @@ def _only_range_checked(index: Node, mask: Node) -> Optional[List[Node]]: and isinstance(condition.meta.get("val"), torch.SymBool) else None ) - if expr is None or len(expr.free_symbols) != 1: + if expr is None or expr.free_symbols != {count.node.expr}: return None - (symbol,) = expr.free_symbols - bounds = bound_sympy(expr, {symbol: within}) + bounds = bound_sympy(expr, {count.node.expr: within}) if not (bounds.is_singleton() and bounds.lower == sympy.true): return None elif user.target == _SYM_SIZE and user.args[0] is index: pass - elif not isinstance( - user.meta.get("val"), (torch.SymInt, torch.SymBool) - ) or not any(arg in found for arg in user.all_input_nodes): + # Computing a symbolic value has no effect of its own, and its + # readers are visited in turn. + elif not isinstance(user.meta.get("val"), (torch.SymInt, torch.SymBool)): return None found.append(user) pending.extend(user.users) diff --git a/backends/cuda/passes/tests/test_replace_masked_copy_with_where.py b/backends/cuda/passes/tests/test_replace_masked_copy_with_where.py index 798253187b1..6c594bcdcd0 100644 --- a/backends/cuda/passes/tests/test_replace_masked_copy_with_where.py +++ b/backends/cuda/passes/tests/test_replace_masked_copy_with_where.py @@ -4,6 +4,7 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +import operator import unittest import torch @@ -13,7 +14,7 @@ ) from executorch.exir import to_edge from executorch.exir.dialects._ops import ops as exir_ops -from torch.export import export +from torch.export import Dim, export def _targets(graph_module): @@ -90,6 +91,46 @@ def forward(self, x, mask): for got, expected in zip(program.module()(x, mask), Kept()(x, mask)): torch.testing.assert_close(got, expected, rtol=0, atol=0) + def test_check_on_another_size_keeps_the_index(self): + """A check that reads the count but tests another size does not hold + for every mask.""" + + class Copy(torch.nn.Module): + def forward(self, x, mask, y): + out = torch.zeros_like(x) + out[mask] = x[mask] + return out, y * 2 + + x = torch.randn(4, 3) + mask = torch.tensor([True, False, True, False]) + y = torch.randn(3) + program = to_edge( + export(Copy(), (x, mask, y), dynamic_shapes=(None, None, {0: Dim("n")})) + ).exported_program() + graph = program.graph_module.graph + index = next( + n for n in graph.nodes if n.target == exir_ops.edge.aten.index.Tensor + ) + y_node = next(n for n in graph.nodes if n.name == "y") + # count * 0 + y.size(0) <= 4 has the one symbol of y's size. + with graph.inserting_before(index.next): + count = graph.call_function(torch.ops.aten.sym_size.int, (index, 0)) + size = graph.call_function(torch.ops.aten.sym_size.int, (y_node, 0)) + zero = graph.call_function(operator.mul, (count, 0)) + total = graph.call_function(operator.add, (zero, size)) + check = graph.call_function(operator.le, (total, 4)) + graph.call_function(torch.ops.aten._assert_scalar.default, (check, "")) + count.meta["val"] = index.meta["val"].shape[0] + size.meta["val"] = y_node.meta["val"].shape[0] + zero.meta["val"] = count.meta["val"] * 0 + total.meta["val"] = zero.meta["val"] + size.meta["val"] + check.meta["val"] = total.meta["val"] <= 4 + ReplaceMaskedCopyWithWherePass()(program.graph_module) + + targets = _targets(program.graph_module) + self.assertIn(exir_ops.edge.aten.where.self, targets) + self.assertIn(exir_ops.edge.aten.index.Tensor, targets) + if __name__ == "__main__": unittest.main() From 423ebaf0fb6a4bacfb90f63172cd262fc6f9fc1e Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Fri, 30 Oct 2026 10:34:46 -0700 Subject: [PATCH 24/89] Say why a CUDA method keeps its graph off for host constants A method with a constant in host memory runs without a CUDA graph. The init log said the method "keeps constants in host memory", also when the constants could not be listed at all. It now names the constant it found, or says the constants could not be listed: CUDA graph requested for method 'forward', but constant 'flag' is in host memory; running without a CUDA graph The phase diagram in cuda_delegate_handle.h now shows when a method stays Disabled and when Warmup falls back to Disabled. The MoveCondPredicateToCpu pass docstring now says that a method it changes runs without a CUDA graph. After a failed capture, later calls run eagerly only because the capture guard has already set the phase to Disabled. A debug check before the retry now states that. --- .../cuda/passes/move_cond_predicate_to_cpu.py | 4 +++ backends/cuda/runtime/cuda_backend.cpp | 28 +++++++++++-------- backends/cuda/runtime/cuda_delegate_handle.h | 6 +++- 3 files changed, 26 insertions(+), 12 deletions(-) diff --git a/backends/cuda/passes/move_cond_predicate_to_cpu.py b/backends/cuda/passes/move_cond_predicate_to_cpu.py index f876a684019..ba8c36f63b7 100644 --- a/backends/cuda/passes/move_cond_predicate_to_cpu.py +++ b/backends/cuda/passes/move_cond_predicate_to_cpu.py @@ -55,6 +55,10 @@ def forward(self, hidden_states, encoder_hidden_states): In this example if we keep `self.initialized` on GPU, we will need to copy it to CPU for every forward pass. We move the predicate to CPU to avoid device to host copies. This pass is only applicable to models that use torch.cond and its predicate is a constant buffer. + + A method this pass changes keeps a constant in host memory, so the CUDA backend runs it + without a CUDA graph even when one is requested: a graph would keep the branch the capture + call took. """ requires_exported_program = True diff --git a/backends/cuda/runtime/cuda_backend.cpp b/backends/cuda/runtime/cuda_backend.cpp index f8d656cb2f7..02f5bd602d9 100644 --- a/backends/cuda/runtime/cuda_backend.cpp +++ b/backends/cuda/runtime/cuda_backend.cpp @@ -614,15 +614,17 @@ class ET_EXPERIMENTAL CudaBackend final // Initialize CUDA graph state if enabled for this method. Not on a device // without memory pools, where the method runs without a graph. if (should_use_cuda_graph_for_method(method_name)) { - if (keeps_host_constants(handle)) { + const std::string host_constants = describe_host_constants(handle); + if (!host_constants.empty()) { // A graph records only the device work of the capture call. Host code, // such as the choice of a torch.cond branch on a flag in host memory, // would be fixed at its value from that call. ET_LOG( Info, - "CUDA graph requested for method '%s', but it keeps constants in " - "host memory; running without a CUDA graph", - method_name.c_str()); + "CUDA graph requested for method '%s', but %s; running without a " + "CUDA graph", + method_name.c_str(), + host_constants.c_str()); } else if (handle->cuda_graph_state.start_warmup()) { ET_LOG( Info, @@ -845,6 +847,8 @@ class ET_EXPERIMENTAL CudaBackend final "CUDA graph: capture failed for '%s'; running this call and later " "ones without a graph", handle->method_name.c_str()); + // The capture guard turned graphs off, so later calls run eagerly too. + ET_DCHECK(handle->cuda_graph_state.phase == CudaGraphPhase::Disabled); error = run_once( handle, args, @@ -858,25 +862,27 @@ class ET_EXPERIMENTAL CudaBackend final } private: - // Whether any constant of the program is in host memory. Unknown counts as - // yes. - static bool keeps_host_constants(cuda::CudaDelegateHandle* handle) { + // Names a constant of the program in host memory, or says that the + // constants could not be listed, since unknown counts as host memory. + // Empty when every constant is on the device. + static std::string describe_host_constants(cuda::CudaDelegateHandle* handle) { + const char* unlisted = "its constants could not be listed"; if (handle->extract_constants_map == nullptr) { - return true; + return unlisted; } std::unordered_map constants; if (handle->extract_constants_map( handle->container_handle, reinterpret_cast(&constants), /*use_inactive=*/false) != Error::Ok) { - return true; + return unlisted; } for (const auto& [fqn, tensor] : constants) { if (reinterpret_cast(tensor)->is_cpu()) { - return true; + return "constant '" + fqn + "' is in host memory"; } } - return false; + return ""; } // Runs the program once. On the capture step it captures the CUDA graph diff --git a/backends/cuda/runtime/cuda_delegate_handle.h b/backends/cuda/runtime/cuda_delegate_handle.h index 915ff4f10ac..ead44960a46 100644 --- a/backends/cuda/runtime/cuda_delegate_handle.h +++ b/backends/cuda/runtime/cuda_delegate_handle.h @@ -103,8 +103,12 @@ struct CudaWeightStorage { // Phases of the CUDA graph lifecycle for a delegate handle. // // The transition flow is: -// Disabled ──(if CUDA graph is enabled for this method)──▶ Warmup +// Disabled ──(if CUDA graph is enabled for this method, the device +// has memory pools, and every constant is in device +// memory)──▶ Warmup // Warmup ──(after `warmup_remaining` execute() calls)──▶ Replay +// Warmup ──(if the capture fails, or an output is not laid out +// like its ETensor)──▶ Disabled // // - Disabled: CUDA graph is not used for this method. Every execute() runs // eagerly through the normal kernel-launch path. From ccc125923627041be31ec984f94277d4ef22b8f8 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Fri, 30 Oct 2026 10:34:47 -0700 Subject: [PATCH 25/89] Keep destroy() with the other public CUDA backend methods Splitting execute() into execute() and run_once() left destroy() in a second public section after the private helpers. Move it up next to execute(), so the public methods sit together and the private helpers join the private members. No code changes. --- backends/cuda/runtime/cuda_backend.cpp | 120 ++++++++++++------------- 1 file changed, 59 insertions(+), 61 deletions(-) diff --git a/backends/cuda/runtime/cuda_backend.cpp b/backends/cuda/runtime/cuda_backend.cpp index 02f5bd602d9..04f9ec91893 100644 --- a/backends/cuda/runtime/cuda_backend.cpp +++ b/backends/cuda/runtime/cuda_backend.cpp @@ -861,6 +861,65 @@ class ET_EXPERIMENTAL CudaBackend final return commit_offgraph_kv_step(handle, kv_step_width); } + void destroy(DelegateHandle* handle_) const override { + if (handle_ == nullptr) { + return; + } + cuda::CudaDelegateHandle* handle = (cuda::CudaDelegateHandle*)handle_; + + mutable_state_forget_handle(handle); + if (handle->kv_cache != nullptr) { + handle->kv_cache->forget_handle(handle); + } + + // NOTE: AOTInductorModelContainerDelete does not work correctly with + // multiple .so files. Deleting one container frees shared resources, + // which causes segmentation faults when attempting to delete other + // containers. As a workaround, we skip explicit container deletion and + // defer cleanup to the OS. The corresponding shared library must remain + // loaded as well: the leaked container still owns objects whose code and + // process-wide state live in that library, so dlclose/FreeLibrary can + // invalidate them and crash during multi-method teardown. + // TODO(gasoonjia): Find a proper solution for safe container deletion. + // AOTInductorModelContainerDelete(handle->container_handle); + + // Remove the temporary shared library file + if (!handle->so_path.empty()) { + std::error_code remove_error; + std::filesystem::remove(handle->so_path, remove_error); + ET_CHECK_OR_LOG_ERROR( + !remove_error, + "Failed to remove temporary shared library %s: %s", + handle->so_path.c_str(), + remove_error.message().c_str()); + } + + delete handle; + + // The allocator lets the device pool keep freed memory so that repeated + // delegate execution does not pay to map it again. Nothing is running on + // this backend once the last handle is gone, so hand that memory back + // rather than hold it for the life of the process. + if (live_handles_.fetch_sub(1, std::memory_order_acq_rel) == 1) { + // Only frees the driver has already observed can be released, so without + // this the trim gives back nothing rather than less. A device wait rather + // than a stream wait because the frees went to whichever stream was + // current when the storage was released, which is not necessarily one + // this handle recorded, and teardown is not bound to the thread that ran + // the method either. + const cudaError_t sync_err = cudaDeviceSynchronize(); + if (sync_err != cudaSuccess) { + ET_LOG( + Error, + "cudaDeviceSynchronize before releasing the pool failed: %s.", + cudaGetErrorString(sync_err)); + (void)cudaGetLastError(); + } + // This also releases unused blocks cached by other allocator users. + CudaAllocator::release_cached_memory(-1); + } + } + private: // Names a constant of the program in host memory, or says that the // constants could not be listed, since unknown counts as host memory. @@ -1288,67 +1347,6 @@ class ET_EXPERIMENTAL CudaBackend final return Error::Ok; } - public: - void destroy(DelegateHandle* handle_) const override { - if (handle_ == nullptr) { - return; - } - cuda::CudaDelegateHandle* handle = (cuda::CudaDelegateHandle*)handle_; - - mutable_state_forget_handle(handle); - if (handle->kv_cache != nullptr) { - handle->kv_cache->forget_handle(handle); - } - - // NOTE: AOTInductorModelContainerDelete does not work correctly with - // multiple .so files. Deleting one container frees shared resources, - // which causes segmentation faults when attempting to delete other - // containers. As a workaround, we skip explicit container deletion and - // defer cleanup to the OS. The corresponding shared library must remain - // loaded as well: the leaked container still owns objects whose code and - // process-wide state live in that library, so dlclose/FreeLibrary can - // invalidate them and crash during multi-method teardown. - // TODO(gasoonjia): Find a proper solution for safe container deletion. - // AOTInductorModelContainerDelete(handle->container_handle); - - // Remove the temporary shared library file - if (!handle->so_path.empty()) { - std::error_code remove_error; - std::filesystem::remove(handle->so_path, remove_error); - ET_CHECK_OR_LOG_ERROR( - !remove_error, - "Failed to remove temporary shared library %s: %s", - handle->so_path.c_str(), - remove_error.message().c_str()); - } - - delete handle; - - // The allocator lets the device pool keep freed memory so that repeated - // delegate execution does not pay to map it again. Nothing is running on - // this backend once the last handle is gone, so hand that memory back - // rather than hold it for the life of the process. - if (live_handles_.fetch_sub(1, std::memory_order_acq_rel) == 1) { - // Only frees the driver has already observed can be released, so without - // this the trim gives back nothing rather than less. A device wait rather - // than a stream wait because the frees went to whichever stream was - // current when the storage was released, which is not necessarily one - // this handle recorded, and teardown is not bound to the thread that ran - // the method either. - const cudaError_t sync_err = cudaDeviceSynchronize(); - if (sync_err != cudaSuccess) { - ET_LOG( - Error, - "cudaDeviceSynchronize before releasing the pool failed: %s.", - cudaGetErrorString(sync_err)); - (void)cudaGetLastError(); - } - // This also releases unused blocks cached by other allocator users. - CudaAllocator::release_cached_memory(-1); - } - } - - private: mutable std::mutex cuda_graph_method_mutex_; std::string cuda_graph_method_; From 984df095cb572d6671a5b45c3ba03fc8aca4d05d Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Fri, 30 Oct 2026 10:18:54 -0700 Subject: [PATCH 26/89] Draw the CUDA RNG seed before allocating its state std::random_device can throw an implementation-defined exception derived from std::exception. ensure_rng_init caught only std::bad_alloc, so that exception escaped, and because the state was allocated first, the next call allocated it again and leaked the first block. The seed is now drawn before the allocation, and any std::exception is caught and logged. --- backends/cuda/runtime/shims/rand.cu | 26 ++++++++++++++------------ 1 file changed, 14 insertions(+), 12 deletions(-) diff --git a/backends/cuda/runtime/shims/rand.cu b/backends/cuda/runtime/shims/rand.cu index 5140a59d304..31644c53a25 100644 --- a/backends/cuda/runtime/shims/rand.cu +++ b/backends/cuda/runtime/shims/rand.cu @@ -59,12 +59,20 @@ static std::once_flag g_rng_init_flag; // Initialize RNG state on the given stream. // Must be called during warmup (before graph capture). Subsequent calls -// from any thread are no-ops thanks to std::call_once. A failed allocation -// throws out of the call_once, which leaves it not done, so the next call -// tries again rather than every later launch using a null state. +// from any thread are no-ops thanks to std::call_once. A failure throws out +// of the call_once, which leaves it not done, so the next call tries again +// rather than every later launch using a null state. bool ensure_rng_init(cudaStream_t stream) { try { std::call_once(g_rng_init_flag, [&]() { + RngState h; + // From std::random_device: a seed from the clock repeats for every + // process started in the same second. Drawn before allocating, so a + // throw here leaves nothing allocated. + std::random_device entropy; + h.seed = (static_cast(entropy()) << 32) | entropy(); + h.counter = 0; + h.base_scratch = 0; // The current device: the stream callers pass from getCurrentCUDAStream // may have been created on it whatever index they looked up, and the // cudaMalloc fallback on a device without memory pools follows this @@ -75,21 +83,15 @@ bool ensure_rng_init(cudaStream_t stream) { throw std::bad_alloc(); } d_rng = static_cast(rng.get()); - RngState h; - // From std::random_device: a seed from the clock repeats for every - // process started in the same second. - std::random_device entropy; - h.seed = (static_cast(entropy()) << 32) | entropy(); - h.counter = 0; - h.base_scratch = 0; cudaMemcpyAsync( d_rng, &h, sizeof(RngState), cudaMemcpyHostToDevice, stream); // Synchronize to ensure the copy completes before we return // (the host-side RngState `h` is on the stack). cudaStreamSynchronize(stream); }); - } catch (const std::bad_alloc&) { - ET_LOG(Error, "rand: allocating the RNG state failed"); + } catch (const std::exception& e) { + // std::random_device throws an implementation-defined std::exception. + ET_LOG(Error, "rand: initializing the RNG state failed: %s", e.what()); return false; } return true; From 1bfa4de4e08b203db42e3c3f12eff63bf49469ce Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Fri, 30 Oct 2026 11:12:20 -0700 Subject: [PATCH 27/89] Format the RemoveAssertAsyncPass import with ufmt --- backends/cuda/cuda_backend.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index 3e6fc1cd607..ecbe487d45f 100644 --- a/backends/cuda/cuda_backend.py +++ b/backends/cuda/cuda_backend.py @@ -39,9 +39,7 @@ from executorch.backends.cuda.passes.move_cond_predicate_to_cpu import ( MoveCondPredicateToCpuPass, ) -from executorch.backends.cuda.passes.remove_assert_async import ( - RemoveAssertAsyncPass, -) +from executorch.backends.cuda.passes.remove_assert_async import RemoveAssertAsyncPass from executorch.backends.cuda.passes.replace_fused_sdpa import ( ReplaceFusedSdpaWithSdpaPass, ) From 76f519a025b36880b800e7d763349efa1fe5d657 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Fri, 30 Oct 2026 09:55:29 -0700 Subject: [PATCH 28/89] Load float64 weights in the CUDA backend A program with a float64 weight lowered and saved with the CUDA backend, then failed to load: Failed to load method forward, error: 0x23: Malformed CUDA FQN weight metadata; Init failed for backend CudaBackend This happens with a plain float64 buffer, and also when constant propagation folds float64 math (for example a sinusoidal time embedding) into one float64 constant. The export side writes the weight correctly, but the runtime did not know the float64 dtype: the weight cache refused it while reading the weight list, and the tensor type the runtime hands to the compiled code had no float64 either. Support float64 through the runtime: - add Double to the runtime tensor's dtype list, with its 8 byte size - export aoti_torch_dtype_float64, which compiled code calls to name the dtype of a float64 constant - accept Double in the weight cache, with its 8 byte size when checking that a weight fits its storage The same tensor type wraps a method's inputs and outputs, so a float64 input or output now works too. Before, it stopped the process at the tensor type's dtype check. Tests: a float64 buffer now lowers, loads and matches eager PyTorch. The weight cache test parses a Double entry and uses ComplexFloat for the dtype it must refuse, and the tensor tests do the same. --- backends/aoti/common_shims_slim.cpp | 4 +++ backends/aoti/common_shims_slim.h | 1 + backends/aoti/slim/c10/core/ScalarType.h | 10 +++++-- .../slim/c10/core/test/test_scalar_type.cpp | 2 ++ .../slim/core/test/test_slimtensor_basic.cpp | 6 ++-- .../aoti/slim/factory/test/test_empty.cpp | 8 ++--- .../aoti/slim/factory/test/test_from_blob.cpp | 4 +-- .../aoti/tests/test_common_shims_slim.cpp | 1 + backends/cuda/cuda_backend.py | 4 +-- backends/cuda/runtime/cuda_weight_cache.cpp | 2 ++ .../runtime/test/test_cuda_weight_cache.cpp | 12 +++++++- .../cuda/tests/test_missing_weights_blob.py | 30 +++++++++++++++++++ 12 files changed, 70 insertions(+), 14 deletions(-) diff --git a/backends/aoti/common_shims_slim.cpp b/backends/aoti/common_shims_slim.cpp index e92c9d319ae..e44c583b0b9 100644 --- a/backends/aoti/common_shims_slim.cpp +++ b/backends/aoti/common_shims_slim.cpp @@ -145,6 +145,10 @@ int32_t aoti_torch_dtype_float32() { return 6; // ScalarType::Float } +int32_t aoti_torch_dtype_float64() { + return 7; // ScalarType::Double +} + int32_t aoti_torch_dtype_bfloat16() { return 15; // ScalarType::BFloat16 } diff --git a/backends/aoti/common_shims_slim.h b/backends/aoti/common_shims_slim.h index dea91ff8e67..37cf6785c5a 100644 --- a/backends/aoti/common_shims_slim.h +++ b/backends/aoti/common_shims_slim.h @@ -83,6 +83,7 @@ aoti_torch_get_device_index(Tensor* tensor, int32_t* ret_device_index); // ============================================================ AOTI_SHIM_EXPORT int32_t aoti_torch_dtype_float32(); +AOTI_SHIM_EXPORT int32_t aoti_torch_dtype_float64(); AOTI_SHIM_EXPORT int32_t aoti_torch_dtype_bfloat16(); AOTI_SHIM_EXPORT int32_t aoti_torch_dtype_float16(); AOTI_SHIM_EXPORT int32_t aoti_torch_dtype_int64(); diff --git a/backends/aoti/slim/c10/core/ScalarType.h b/backends/aoti/slim/c10/core/ScalarType.h index 9a99aecf992..5bc06076885 100644 --- a/backends/aoti/slim/c10/core/ScalarType.h +++ b/backends/aoti/slim/c10/core/ScalarType.h @@ -30,7 +30,7 @@ enum class ScalarType : int8_t { Long = 4, // int64_t Half = 5, // float16 Float = 6, // float - // Double = 7, // double - not currently needed + Double = 7, // double // ComplexHalf = 8, // ComplexFloat = 9, // ComplexDouble = 10, @@ -50,6 +50,7 @@ constexpr ScalarType kInt = ScalarType::Int; constexpr ScalarType kLong = ScalarType::Long; constexpr ScalarType kHalf = ScalarType::Half; constexpr ScalarType kFloat = ScalarType::Float; +constexpr ScalarType kDouble = ScalarType::Double; constexpr ScalarType kBool = ScalarType::Bool; constexpr ScalarType kBFloat16 = ScalarType::BFloat16; @@ -72,6 +73,8 @@ inline size_t elementSize(ScalarType t) { return 2; // sizeof(__half) = 2 bytes case ScalarType::Float: return sizeof(float); + case ScalarType::Double: + return sizeof(double); case ScalarType::Bool: return sizeof(bool); case ScalarType::BFloat16: @@ -100,6 +103,8 @@ inline const char* toString(ScalarType t) { return "Half"; case ScalarType::Float: return "Float"; + case ScalarType::Double: + return "Double"; case ScalarType::Bool: return "Bool"; case ScalarType::BFloat16: @@ -116,7 +121,7 @@ inline const char* toString(ScalarType t) { /// @return true if the scalar type is floating point, false otherwise. inline bool isFloatingType(ScalarType t) { return t == ScalarType::Half || t == ScalarType::Float || - t == ScalarType::BFloat16; + t == ScalarType::Double || t == ScalarType::BFloat16; } /// Checks if the scalar type is an integral type (including bool optionally). @@ -157,6 +162,7 @@ inline bool isValidScalarType(ScalarType t) { case ScalarType::Long: case ScalarType::Half: case ScalarType::Float: + case ScalarType::Double: case ScalarType::Bool: case ScalarType::BFloat16: return true; diff --git a/backends/aoti/slim/c10/core/test/test_scalar_type.cpp b/backends/aoti/slim/c10/core/test/test_scalar_type.cpp index 4c06f7ef101..bc09676f2d7 100644 --- a/backends/aoti/slim/c10/core/test/test_scalar_type.cpp +++ b/backends/aoti/slim/c10/core/test/test_scalar_type.cpp @@ -38,6 +38,7 @@ const std::vector kAllScalarTypes = { {ScalarType::Long, 4, 8, "Long", false, true, true, false}, {ScalarType::Half, 5, 2, "Half", true, false, false, false}, {ScalarType::Float, 6, 4, "Float", true, false, false, false}, + {ScalarType::Double, 7, 8, "Double", true, false, false, false}, {ScalarType::Bool, 11, 1, "Bool", false, false, true, true}, {ScalarType::BFloat16, 15, 2, "BFloat16", true, false, false, false}, }; @@ -224,6 +225,7 @@ TEST_F(IsValidScalarTypeTest, AllSupportedTypesAreValid) { EXPECT_TRUE(isValidScalarType(ScalarType::Long)); EXPECT_TRUE(isValidScalarType(ScalarType::Half)); EXPECT_TRUE(isValidScalarType(ScalarType::Float)); + EXPECT_TRUE(isValidScalarType(ScalarType::Double)); EXPECT_TRUE(isValidScalarType(ScalarType::Bool)); EXPECT_TRUE(isValidScalarType(ScalarType::BFloat16)); } diff --git a/backends/aoti/slim/core/test/test_slimtensor_basic.cpp b/backends/aoti/slim/core/test/test_slimtensor_basic.cpp index 29ebe5875b4..b8e44acc22b 100644 --- a/backends/aoti/slim/core/test/test_slimtensor_basic.cpp +++ b/backends/aoti/slim/core/test/test_slimtensor_basic.cpp @@ -508,10 +508,10 @@ TEST(SlimTensorValidationTest, InvalidDtypeUndefined) { ""); } -TEST(SlimTensorValidationTest, InvalidDtypeDouble) { +TEST(SlimTensorValidationTest, InvalidDtypeComplexFloat) { std::vector sizes = {2, 3}; std::vector strides = {3, 1}; - size_t nbytes = 6 * sizeof(double); + size_t nbytes = 6 * 2 * sizeof(float); Storage storage = make_cpu_storage(nbytes); EXPECT_DEATH( @@ -519,7 +519,7 @@ TEST(SlimTensorValidationTest, InvalidDtypeDouble) { std::move(storage), makeArrayRef(sizes), makeArrayRef(strides), - static_cast(7)), // Double = 7 + static_cast(9)), // ComplexFloat = 9 ""); } diff --git a/backends/aoti/slim/factory/test/test_empty.cpp b/backends/aoti/slim/factory/test/test_empty.cpp index f4d0e31da9c..9cc3d6ef178 100644 --- a/backends/aoti/slim/factory/test/test_empty.cpp +++ b/backends/aoti/slim/factory/test/test_empty.cpp @@ -249,7 +249,7 @@ TEST(EmptyStridedTest, InvalidDtypeUndefined) { ""); } -TEST(EmptyStridedTest, InvalidDtypeDouble) { +TEST(EmptyStridedTest, InvalidDtypeComplexFloat) { std::vector sizes = {2, 3}; std::vector strides = {3, 1}; @@ -257,7 +257,7 @@ TEST(EmptyStridedTest, InvalidDtypeDouble) { empty_strided( makeArrayRef(sizes), makeArrayRef(strides), - static_cast(7)), // Double = 7 + static_cast(9)), // ComplexFloat = 9 ""); } @@ -280,9 +280,9 @@ TEST(EmptyTest, InvalidDtypeUndefined) { EXPECT_DEATH(empty({2, 3}, c10::ScalarType::Undefined), ""); } -TEST(EmptyTest, InvalidDtypeDouble) { +TEST(EmptyTest, InvalidDtypeComplexFloat) { EXPECT_DEATH( - empty({2, 3}, static_cast(7)), // Double = 7 + empty({2, 3}, static_cast(9)), // ComplexFloat = 9 ""); } diff --git a/backends/aoti/slim/factory/test/test_from_blob.cpp b/backends/aoti/slim/factory/test/test_from_blob.cpp index d066d7ba79c..40a00c0c682 100644 --- a/backends/aoti/slim/factory/test/test_from_blob.cpp +++ b/backends/aoti/slim/factory/test/test_from_blob.cpp @@ -327,12 +327,12 @@ TEST(FromBlobTest, InvalidDtypeUndefined) { from_blob(external_data, {2, 3}, c10::ScalarType::Undefined), ""); } -TEST(FromBlobTest, InvalidDtypeDouble) { +TEST(FromBlobTest, InvalidDtypeComplexFloat) { constexpr size_t kNumFloats = 6; float external_data[kNumFloats]; EXPECT_DEATH( - from_blob(external_data, {2, 3}, static_cast(7)), ""); + from_blob(external_data, {2, 3}, static_cast(9)), ""); } TEST(FromBlobTest, InvalidDeviceType) { diff --git a/backends/aoti/tests/test_common_shims_slim.cpp b/backends/aoti/tests/test_common_shims_slim.cpp index 5446b925688..952a79283f2 100644 --- a/backends/aoti/tests/test_common_shims_slim.cpp +++ b/backends/aoti/tests/test_common_shims_slim.cpp @@ -698,6 +698,7 @@ TEST_F(CommonShimsSlimTest, ConsistentPointerReturn) { TEST_F(CommonShimsSlimTest, DTypeConstants) { // Verify dtype constants match expected PyTorch ScalarType values EXPECT_EQ(aoti_torch_dtype_float32(), 6); // ScalarType::Float + EXPECT_EQ(aoti_torch_dtype_float64(), 7); // ScalarType::Double EXPECT_EQ(aoti_torch_dtype_bfloat16(), 15); // ScalarType::BFloat16 EXPECT_EQ(aoti_torch_dtype_int64(), 4); // ScalarType::Long EXPECT_EQ(aoti_torch_dtype_int32(), 3); // ScalarType::Int diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index ecbe487d45f..f45fbbbe947 100644 --- a/backends/cuda/cuda_backend.py +++ b/backends/cuda/cuda_backend.py @@ -636,8 +636,8 @@ def _antialiased_resize( each output's taps, one dim at a time. Inductor has no lowering for these resizes, so it would call ATen's kernel, which this backend does not ship. """ - # The backend's tensors have no float64, and ATen's CUDA kernel has no - # integer version. + # The rewrite computes in float32, so float64 would lose precision, and + # ATen's CUDA kernel has no integer version. if input.dtype not in (torch.float32, torch.float16, torch.bfloat16): return NotImplemented # Half types are computed in float and rounded once, so they can differ diff --git a/backends/cuda/runtime/cuda_weight_cache.cpp b/backends/cuda/runtime/cuda_weight_cache.cpp index 104b4461318..ae32c4282bd 100644 --- a/backends/cuda/runtime/cuda_weight_cache.cpp +++ b/backends/cuda/runtime/cuda_weight_cache.cpp @@ -124,6 +124,7 @@ bool is_supported_dtype(int32_t dtype) { case 4: // Long case 5: // Half case 6: // Float + case 7: // Double case 11: // Bool case 15: // BFloat16 return true; @@ -358,6 +359,7 @@ Error CudaWeightCache::validate_view(const Entry& entry) { item_size = 4; break; case aoti::slim::c10::ScalarType::Long: + case aoti::slim::c10::ScalarType::Double: item_size = 8; break; default: diff --git a/backends/cuda/runtime/test/test_cuda_weight_cache.cpp b/backends/cuda/runtime/test/test_cuda_weight_cache.cpp index daa0a41bf0f..a7b86656fc7 100644 --- a/backends/cuda/runtime/test/test_cuda_weight_cache.cpp +++ b/backends/cuda/runtime/test/test_cuda_weight_cache.cpp @@ -275,9 +275,19 @@ TEST(CudaWeightCacheTest, RejectsTruncationAndTrailingData) { Error::InvalidProgram); } +TEST(CudaWeightCacheTest, ParsesDoubleWeight) { + const std::vector bytes = serialized_metadata(7); // Double + cuda::CudaWeightCache::Metadata metadata; + ASSERT_EQ( + cuda::CudaWeightCache::parse(bytes.data(), bytes.size(), metadata), + Error::Ok); + ASSERT_EQ(metadata.entries.size(), 1u); + EXPECT_EQ(metadata.entries[0].dtype, 7); +} + TEST(CudaWeightCacheTest, RejectsUnsupportedDtype) { const std::vector bytes = - serialized_metadata(7); // Double is unsupported. + serialized_metadata(9); // ComplexFloat is unsupported. cuda::CudaWeightCache::Metadata metadata; EXPECT_EQ( cuda::CudaWeightCache::parse(bytes.data(), bytes.size(), metadata), diff --git a/backends/cuda/tests/test_missing_weights_blob.py b/backends/cuda/tests/test_missing_weights_blob.py index 6f1308b20c3..aa7667471f2 100644 --- a/backends/cuda/tests/test_missing_weights_blob.py +++ b/backends/cuda/tests/test_missing_weights_blob.py @@ -222,6 +222,36 @@ def test_programs_with_their_own_data_names_share_a_folder(self) -> None: module(value).cpu(), ) + def test_float64_weight_loads(self) -> None: + """The weight cache must accept a float64 weight the export wrote.""" + + class Float64Scale(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.register_buffer( + "scale", torch.linspace(1.0, 2.0, 64, dtype=torch.float64) + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return (x.to(torch.float64) * self.scale).to(x.dtype) + + from executorch.runtime import Runtime + + with tempfile.TemporaryDirectory() as outdir: + module = Float64Scale() + path = self._lower(outdir, module=module) + (blob,) = [f for f in os.listdir(outdir) if f.endswith(".ptd")] + method = ( + Runtime.get() + .load_program(path, data_path=os.path.join(outdir, blob)) + .load_method("forward") + ) + # The same draw _lower made for its example input. + torch.manual_seed(0) + value = torch.randn(4, 64, device="cuda") + out = method.execute([value.cpu()]) + torch.testing.assert_close(out[0].cpu(), module(value).cpu()) + def test_load_succeeds_when_the_model_has_no_constants(self) -> None: """A model with nothing to bind still loads without its weights blob. From ebcde7b3111590e3711a5be525db8ab39e9bc836 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Fri, 30 Oct 2026 09:55:40 -0700 Subject: [PATCH 29/89] Decompose aten._reshape_copy in the CUDA backend Torch-TensorRT rewrites aten.view into aten._reshape_copy. When TensorRT leaves such a view to PyTorch, the CUDA partitioner receives it as _reshape_copy, and the export failed: Method forward missing fallback kernels (1 total): - aten::_reshape_copy Inductor has no lowering for this op, and torch registers no decomposition for it, so it fell back to an ATen kernel this backend does not ship. Add it to the backend's decomposition table. ATen computes _reshape_copy as a contiguous copy of the input viewed to the new shape, and the decomposition does the same with clone and view, which Inductor lowers. Test: a module that calls _reshape_copy on a transposed input now lowers to the CUDA backend. --- backends/cuda/cuda_backend.py | 11 +++++++++++ backends/cuda/tests/test_cuda_export.py | 9 +++++++++ 2 files changed, 20 insertions(+) diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index f45fbbbe947..71ed10d9e6b 100644 --- a/backends/cuda/cuda_backend.py +++ b/backends/cuda/cuda_backend.py @@ -558,6 +558,16 @@ def _conv_transpose_as_conv( return conv(x, weight, bias, [1] * n, [0] * n, dilation, groups) +def _reshape_copy(x: torch.Tensor, size: List[int]) -> torch.Tensor: + """ + Torch-TensorRT turns a view into _reshape_copy, and a view that TensorRT + leaves to PyTorch reaches this backend in that form. Inductor has no + lowering for it and this backend no kernel. ATen computes it as a + contiguous copy viewed to the new shape. + """ + return x.clone(memory_format=torch.contiguous_format).view(size) + + def _linear_kernel(x: torch.Tensor) -> torch.Tensor: return torch.clamp(1 - x.abs(), min=0) @@ -984,6 +994,7 @@ def get_decomposition_table(cls) -> Dict[Any, Any]: # torch._check_tensor_all exports to this op; Inductor has no lowering # for it and the shim no fallback kernel. torch.ops.aten._is_all_true.default: torch.ops.aten.all.default, + torch.ops.aten._reshape_copy.default: _reshape_copy, } for resize, support, kernel in ( (torch.ops.aten._upsample_bilinear2d_aa, 1, _linear_kernel), diff --git a/backends/cuda/tests/test_cuda_export.py b/backends/cuda/tests/test_cuda_export.py index 472781e6f58..f9e30442673 100644 --- a/backends/cuda/tests/test_cuda_export.py +++ b/backends/cuda/tests/test_cuda_export.py @@ -472,6 +472,15 @@ def forward(self, x: torch.Tensor, eps: torch.Tensor) -> torch.Tensor: strict=False, ) + def test_reshape_copy(self): + """Torch-TensorRT hands back the views it leaves to PyTorch as _reshape_copy.""" + + class ReshapeCopyModule(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.ops.aten._reshape_copy(x.transpose(0, 1), [-1, 4]) * 2 + + self._export_to_cuda_with_lower(ReshapeCopyModule(), (torch.randn(8, 4, 4),)) + def test_conv2d(self): """Test CUDA export for 2D convolution.""" From 2fa510541cbf87ac646a46fa9bb83cc1bd17b017 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Fri, 30 Oct 2026 11:24:11 -0700 Subject: [PATCH 30/89] Add aoti_torch_dtype_float64 to aoti_cuda_shims.lib The CUDA backend now exports aoti_torch_dtype_float64 for compiled code that names the dtype of a float64 constant. A model lowered for Windows links its generated wrapper against this import library, so the library has to list the new name or the link fails. Regenerated per aoti_cuda_shims.lib.md with the same export list plus the new name. The archive has no member timestamps, and rebuilding it gives the same bytes. --- backends/cuda/runtime/aoti_cuda_shims.lib | Bin 47046 -> 47862 bytes 1 file changed, 0 insertions(+), 0 deletions(-) diff --git a/backends/cuda/runtime/aoti_cuda_shims.lib b/backends/cuda/runtime/aoti_cuda_shims.lib index 3294d1ec2629c60f221598224c005977bdd04da4..cfb11867977e03e4c0cc80bf24f4c6cf1f44fa51 100644 GIT binary patch delta 1100 zcmX}oZA?>V6bJD0bWkW-kTMu56EGN-DfiwW7@Pc>n{&_qdConzpW6O; zX$zG2iVLl*cfX1R0B8hA+YT^)v6c}JAQsJP**5@0POX-k10X)1r!D&vK$xRi{&@i5 zKCKmZ0*KP1ny1nS;Hd=oqD)(z1t7lSS;X3h+J-3rv55~BTiUg4oE6{j(BfN&`2#>y2eiGM5p})-{$uqvtzi~GG=;Up)d1os_lO_un(qmKXyXAz#|_Qj1t3mw zMx6dk3vjQ%ajmBrK#2!Qyl z3Ihul&|;e-F0Q(+es@!00a}jq56`WRR5BbHcD0<=Cymm zHnkjEtuX2{sXm``1+B?=IpJff`Duf~pW6(2+O;;FpVt?>H}ld}5w)`y+x zl+nIQVTs?M?hfvoJmJD;$LG-houvw2oiwS&&()SL7u`K!D!SDjR2c5I6LuL2j|U8T z+Rg9dgURGQvs~diG-)EhmAhwC>C+(hB=@}U37&Id#oSnWe3q-vdrd0q;Rzc1=i_v5 z97WH0I21A|vyao|7gDIB|7(Szu%Wk+F_1_*FRbJ*8Z@cAl#?+-pJ2Un4tWOFDBM46 z(9JwKXvYS6tguL_PkF)!-Ep(VfLo13jzG?7P<=vJkQNsYVJ QM@5z!lXI3*)kLN3E%i%gm;e9( delta 1023 zcmXZYSxA&o6bJBg#Z1j_$`MU1)K&`f`zF(ICPSi)3Te?oXd$x*)Kbf#O`CDWa9XUB zTe+o@rRL}y2Bq!6!iRj&eCRC*qaK2`hoT3q&h<9%n}@?W=l;L*>!;t?H@`Z&E5&Y& zoBL%413(c#_-}y7SIUxJ04e62vMd8YS{b6O9s`ipdX;sj0VHF$l9&b{C0i8R8UQJE zsgmaJ0!Z@**gUOl9R-kfbSk?l0HpLBWiKz5GFB@Gc~;8e(9)sz%F!nPQg(xKA`d`1 zm8j(MjFj)1;y+fLUjd|hua&kY0O>w2mAW~#)U!$H`I2an--`K^9~2uP zEaYDJFo5<)Y{{hf_Dq4F^1^8Rj7?DT>1IJKMOH!ng-wF;imd{t&RS^FksuFs(SC*d-quGla1?FDRX}pIweZHt+vrD6?k`#d{mv!n(0OI)LJOg+2VC`b8*_IQi6~6bWPil>Z|%`(!odk*9+~J*pIu;KuR(FQc=mgr z37>g_Xkvh~TOS(~KgfL}0~)!9cJS*So&I{c;vCZOsLzZqZ$3xeSB`JII=%Dpr>D^R PVLtSLPp3O0Iez~Dqkl~t From f9d29a11ef2270d884f0d3bc6b1ff7197d41e17a Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Fri, 30 Oct 2026 13:15:14 -0700 Subject: [PATCH 31/89] Retrigger CI From 89bebfa9b5b91d90acc7db961f8e9fde17237ae6 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Fri, 30 Oct 2026 12:44:13 -0700 Subject: [PATCH 32/89] Add a CUDA shim for scatter_reduce with amax, amin, prod and mean Lowering a model that calls scatter_reduce with reduce="amax" to the CUDA backend failed, and so did "amin", "prod" and "mean": Method forward missing fallback kernels (1 total): - aoti_torch_cuda_scatter_reduce_two_out Inductor writes a scatter itself only when it can add atomically, so it leaves every other reduction to an ATen fallback kernel, and this backend did not ship that kernel. The same fallback is used for a sum into int64 or bool, a sum into bfloat16 on GPUs older than sm_90, and any scatter_reduce in deterministic mode. A float32 sum was not affected. Add the shim next to the index and index_put shims, and allow it in the backend's list of fallback kernels. It follows ATen: with include_self=False an element the index names is reduced from src alone, an element it does not name keeps its value, the mean counts self when it is included and rounds integers down, a NaN wins in amax and amin, and bool supports only sum. 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. Like the index_put shim, 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 now lives in one helper that index_put uses too, and this shim sorts only the key bits an offset can use. The shim is slower than ATen's atomic kernel: 1.5 to 3.2 times at 512 output rows, 4096 source rows and 1024 columns on an A100, and up to 7 times when many values land on few elements, since one thread reduces each element. A decomposition into ops Inductor lowers natively was measured first. It has to compare every output element with every source element. It matched ATen, but it was 8 times slower than ATen's kernel at the size above, and 89 times slower at 2048, 8192 and 1024. Also add the name to aoti_cuda_shims.lib, the import library that lowering for Windows links against, regenerated with mingw binutils 2.41 as its README describes, and to the wheel's shared library symbol check. Test: test_scatter_reduce lowers amax, amin, prod and mean, with and without include_self, in float32, float16, bfloat16 (with a NaN) and int64, runs the program and compares with eager CUDA. Without this change every case fails at lowering. The index shim tests now compare with eager CUDA instead of CPU. --- .ci/scripts/wheel/test_shared_libraries.py | 1 + backends/cuda/cuda_backend.py | 1 + backends/cuda/runtime/aoti_cuda_shims.lib | Bin 47862 -> 48750 bytes backends/cuda/runtime/shims/index.cu | 427 +++++++++++++++++++-- backends/cuda/runtime/shims/index.h | 35 ++ backends/cuda/tests/test_index_shims.py | 43 ++- backends/cuda/tests/test_sort_shim.py | 1 + 7 files changed, 477 insertions(+), 31 deletions(-) diff --git a/.ci/scripts/wheel/test_shared_libraries.py b/.ci/scripts/wheel/test_shared_libraries.py index 9b0d0052a00..5a660e1ebf7 100644 --- a/.ci/scripts/wheel/test_shared_libraries.py +++ b/.ci/scripts/wheel/test_shared_libraries.py @@ -105,6 +105,7 @@ "aoti_torch_cuda_guard_set_index", "aoti_torch_cuda_index_Tensor", "aoti_torch_index_put_out", + "aoti_torch_cuda_scatter_reduce_two_out", ) _QUANTIZED_KERNEL_SYMBOLS = ( diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index 71ed10d9e6b..84c9747ca98 100644 --- a/backends/cuda/cuda_backend.py +++ b/backends/cuda/cuda_backend.py @@ -984,6 +984,7 @@ def get_supported_fallback_kernels(cls) -> Dict[str, Any]: "aoti_torch_cuda_randint_low_out": None, "aoti_torch_cuda_index_Tensor": None, "aoti_torch_index_put_out": None, + "aoti_torch_cuda_scatter_reduce_two_out": None, } @classmethod diff --git a/backends/cuda/runtime/aoti_cuda_shims.lib b/backends/cuda/runtime/aoti_cuda_shims.lib index cfb11867977e03e4c0cc80bf24f4c6cf1f44fa51..a4bc5b1414bbefceabb5695388fb9a80de008706 100644 GIT binary patch delta 1467 zcmZ|Ne@xX?6aesZU-BU25$`>G0#Q83X5+{A`}p;cXoyAT2$zTiMhfyvWDp1sVK9M2 z#en#AP=p?WKhhKw^$et(YvtC~+|c|l$J#7w%UPpU_77<1a#{bdzT0Q-?cRINJ@?%E z_;17S^M;7^v{7E;n*O7d(oK=LfWhA))VApq&E1|{u10BQ59$`(#6W!jajc>pQro{~EQ zAnlw|cJ%;AMdixw6acA=3zEv4Ik22njw{s-08;HHrJk#o8h%jr^E#=eNog$skPhc4 zAH4z~b*@mlI8N#vRr-0YG!Us249)>aZmw7wIi{Sd29Q4GYNgLOsWfI$F7g#fm$?{; zc&#+GQTcogfOLJ9{0x7tFc z8>6wt;c&z|G1M=BGHh*$s}dYe^vZcJ&rMfYF=dYjJt<2N^i1WTpqUD%pu2mI3vyR+mQk03B1gD}uzfZvtBDu5 zyDpmUxOn7Yy_Mqk#R*Jk&?sU=W@|LjKlLdBpWUyi{OwIabhnZ7%r)!kd@tk2&7m~j z^rpa{4rmn8!f{0h&G@UgKMj1CE|3msWNzirxHbz`9rB{4gS>0^5nc7!uRVa$+g=w~ z+M!crF0XsoX+e9(N{VXd&PQ}%4_J1PUA@&8FzLyPMB8G2dDXV-RE?=e};E? z@w^EinQWAOcB8<`3rn4w8}mh9UlS#s=k6t3)K$-?FZt4wF}}h1iKTNiT`}SH32(Z6 zi7Wr^D(}rNK_pP~72a@SQd95Ham|D~CM|fg`vz*0Em$(G(d^_lArdM2no}5Ju7+eZ8`bywJy4WPNV34qP30t%nV}@FeNbXE}lqftyS}nX0G_NrnTx1`>}G_ia+}OeE8URdwqBN+~@Oo-rwiB z^OtS@cUz<|w7S5`UU;e#0boBs`f~u8%d`g)0F$3rrQpuuh_*&ce0f4eeE>_A9YZY=<*(RYC ze~}i*08pxb({_CfpzMig^>RjO2(6HpXtHVh<^Yu7i1yMR0Hsa#C>?QH=njB#NCGO| zG404P0Ob`qqx3(bMP#o+Z)>Mo0hFpiNvsTKX(N9CD5H{AxgZxSZwzUd8UU0j38O@1 zuQL6qHZuvJT&*Q*sfY7yUo7IC@x2L)^Z%bSb*7~8qLhg^{v^G(kT?6bB8_cv)6V!L z{%GryNXvW{JBrH@)3~eBkMvK4$w}LNNS&3LWKXcu4e!s?SD8#5{&Kpx&EkLk@@Jqb zmHGn?4(?cwH0L+@yMPaE89-eL_j6_S8lt!rNKaII*w|H!bmDo7&(``7J$!ff z7No{nlS^CWvR`(m^5=DVNF^_r{AGu%T&r{Qxfj!!pkuqCNPHC5wp#yGe>WJg^gDa4} zX)_io;X~{VZ9=--Vd@XFLLM@b5;*73Ge}*X2AAeb;)m&OcMdzdHzL))Y;d?oa{PFB z2}P3@^L($=_4gROJuIuYdR_ef(K4jmu*Kg8r0(R=41TR|9a8);gPq4^&yhYC2P3AW z=$8~RS*eM*xPbDI&z~UVKTk@+`d1C!aaz?mxr9F)knWhL@9mBaq;dF+od0#u0p_3y{*@l>Vp*o|<$}LoAUy zqXyFLDT}{~$qjW;2NzAt?4mJaVe({6+v)EaC%a~3wV4gJ7fW|nk=L2UKU`UhRQ8U^ Y6|-{V^(%Isc{dkn`RwDI{9cXiKba+!2><{9 diff --git a/backends/cuda/runtime/shims/index.cu b/backends/cuda/runtime/shims/index.cu index 4b696cb03b9..73889e3ee5d 100644 --- a/backends/cuda/runtime/shims/index.cu +++ b/backends/cuda/runtime/shims/index.cu @@ -17,8 +17,11 @@ #include #include #include +#include #include #include +#include +#include #include #include @@ -762,6 +765,49 @@ Error gather( return Error::Ok; } +// Stably sorts count (key, value) pairs by the low end_bit bits of the key. +Error sort_pairs( + const int64_t* keys, + int64_t* sorted_keys, + const int64_t* values, + int64_t* sorted_values, + int count, + int end_bit, + Scratch& scratch, + cudaStream_t stream, + const char* op) { + size_t temp_bytes = 0; + ET_CUDA_CHECK_OR_RETURN_ERROR(cub::DeviceRadixSort::SortPairs( + nullptr, + temp_bytes, + keys, + sorted_keys, + values, + sorted_values, + count, + 0, + end_bit, + stream)); + void* temp = scratch.allocate(temp_bytes); + ET_CHECK_OR_RETURN_ERROR( + temp != nullptr, + MemoryAllocationFailed, + "%s: allocating the sort scratch failed", + op); + ET_CUDA_CHECK_OR_RETURN_ERROR(cub::DeviceRadixSort::SortPairs( + temp, + temp_bytes, + keys, + sorted_keys, + values, + sorted_values, + count, + 0, + end_bit, + stream)); + return Error::Ok; +} + template Error put( const IndexPlan& plan, @@ -801,35 +847,16 @@ Error put( index_offsets_kernel<<>>( plan, points, out_offsets, values_offsets); ET_CUDA_KERNEL_LAUNCH_CHECK_OR_RETURN_ERROR(); - const auto count = static_cast(points); - size_t temp_bytes = 0; - ET_CUDA_CHECK_OR_RETURN_ERROR(cub::DeviceRadixSort::SortPairs( - nullptr, - temp_bytes, - out_offsets, - sorted_out_offsets, - values_offsets, - sorted_values_offsets, - count, - 0, - static_cast(sizeof(int64_t) * 8), - stream)); - void* temp = scratch.allocate(temp_bytes); - ET_CHECK_OR_RETURN_ERROR( - temp != nullptr, - MemoryAllocationFailed, - "index_put: allocating the sort scratch failed"); - ET_CUDA_CHECK_OR_RETURN_ERROR(cub::DeviceRadixSort::SortPairs( - temp, - temp_bytes, + ET_CHECK_OK_OR_RETURN_ERROR(sort_pairs( out_offsets, sorted_out_offsets, values_offsets, sorted_values_offsets, - count, - 0, + static_cast(points), static_cast(sizeof(int64_t) * 8), - stream)); + scratch, + stream, + "index_put")); put_sorted_kernel <<>>( plan, @@ -843,6 +870,203 @@ Error put( return Error::Ok; } +enum class Reduction { Sum, Prod, Mean, Max, Min }; + +// A scatter visits every element of index. Dim `dim` of out is addressed by +// the index value, every other dim of out and src by the element's position. +struct ScatterPlan { + int ndim; + int dim; + int64_t sizes[kMaxDims]; // of index + int64_t index_strides[kMaxDims]; + int64_t out_strides[kMaxDims]; + int64_t src_strides[kMaxDims]; + int64_t dim_size; // of out + int64_t max_out_offset; +}; + +__global__ void scatter_offsets_kernel( + const ScatterPlan plan, + int64_t points, + const int64_t* __restrict__ index, + int64_t* __restrict__ out_offsets, + int64_t* __restrict__ src_offsets) { + for (int64_t point = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + point < points; + point += static_cast(gridDim.x) * blockDim.x) { + int64_t linear = point; + int64_t index_offset = 0; + int64_t out_offset = 0; + int64_t src_offset = 0; + for (int d = plan.ndim - 1; d >= 0; --d) { + const int64_t coordinate = linear % plan.sizes[d]; + linear /= plan.sizes[d]; + index_offset += coordinate * plan.index_strides[d]; + src_offset += coordinate * plan.src_strides[d]; + if (d != plan.dim) { + out_offset += coordinate * plan.out_strides[d]; + } + } + const int64_t position = index[index_offset]; + // As in ATen, a negative index is out of range, not counted from the end. + CUDA_KERNEL_ASSERT_MSG( + position >= 0 && position < plan.dim_size, "index out of bounds"); + out_offsets[point] = out_offset + position * plan.out_strides[plan.dim]; + src_offsets[point] = src_offset; + } +} + +// As ATen's safe_max and safe_min, a NaN on either side wins. +template +__device__ __forceinline__ AccT +combine(Reduction reduction, AccT result, AccT value) { + bool value_is_nan = false; + if constexpr (std::is_floating_point_v) { + value_is_nan = value != value; + } + switch (reduction) { + case Reduction::Prod: + return result * value; + case Reduction::Max: + return value_is_nan || result < value ? value : result; + case Reduction::Min: + return value_is_nan || value < result ? value : result; + default: + return result + value; + } +} + +// After a stable sort of the points by the element they write, the points of +// one element are a run in index order. Each run is reduced by one thread in +// that order and written once, so the result does not depend on scheduling. +// A sum adds in AccT, as ATen does in deterministic mode, so half and +// bfloat16 keep small terms. A product rounds to T at every step and a mean +// rounds its sum to T before dividing, as ATen does, so half overflows to +// infinity and integers wrap where ATen's do. +template +__global__ void scatter_reduce_sorted_kernel( + int64_t points, + const int64_t* __restrict__ out_offsets, + const int64_t* __restrict__ src_offsets, + T* __restrict__ out, + const T* __restrict__ src, + Reduction reduction, + bool include_self, + T identity) { + for (int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + i < points; + i += static_cast(gridDim.x) * blockDim.x) { + const int64_t offset = out_offsets[i]; + if (i > 0 && out_offsets[i - 1] == offset) { + continue; + } + auto result = static_cast(include_self ? out[offset] : identity); + int64_t count = include_self ? 1 : 0; + for (int64_t j = i; j < points && out_offsets[j] == offset; ++j, ++count) { + result = + combine(reduction, result, static_cast(src[src_offsets[j]])); + if (reduction == Reduction::Prod) { + result = static_cast(static_cast(result)); + } + } + if (reduction == Reduction::Mean) { + const auto sum = static_cast(static_cast(result)); + if constexpr (std::is_floating_point_v) { + result = sum / static_cast(count); + } else { + // Integers round down, as ATen's division with rounding_mode="floor". + const AccT quotient = sum / count; + result = quotient - (quotient * count != sum && sum < 0); + } + } + out[offset] = static_cast(result); + } +} + +// What ATen writes to each element the index names before it reduces into +// it, when include_self is false. +template +T reduction_identity(Reduction reduction) { + switch (reduction) { + case Reduction::Prod: + return static_cast(AccT(1)); + case Reduction::Max: + if constexpr (std::is_floating_point_v) { + return static_cast(-std::numeric_limits::infinity()); + } else { + return std::numeric_limits::lowest(); + } + case Reduction::Min: + if constexpr (std::is_floating_point_v) { + return static_cast(std::numeric_limits::infinity()); + } else { + return std::numeric_limits::max(); + } + default: + return static_cast(AccT(0)); + } +} + +template +Error scatter_reduce( + const ScatterPlan& plan, + int64_t points, + const int64_t* index, + void* out, + const void* src, + Reduction reduction, + bool include_self, + Scratch& scratch, + cudaStream_t stream) { + // The sort below takes the count as an int. + ET_CHECK_OR_RETURN_ERROR( + points <= INT_MAX, + NotSupported, + "scatter_reduce: more than %d index elements are not supported", + INT_MAX); + auto* offsets = + static_cast(scratch.allocate(4 * points * sizeof(int64_t))); + ET_CHECK_OR_RETURN_ERROR( + offsets != nullptr, + MemoryAllocationFailed, + "scatter_reduce: allocating the offsets failed"); + int64_t* out_offsets = offsets; + int64_t* src_offsets = offsets + points; + int64_t* sorted_out_offsets = offsets + 2 * points; + int64_t* sorted_src_offsets = offsets + 3 * points; + scatter_offsets_kernel<<>>( + plan, points, index, out_offsets, src_offsets); + ET_CUDA_KERNEL_LAUNCH_CHECK_OR_RETURN_ERROR(); + // Sorting only the bits an offset into out can use is several times faster. + int end_bit = 1; + while (end_bit < 63 && (plan.max_out_offset >> end_bit) != 0) { + ++end_bit; + } + ET_CHECK_OK_OR_RETURN_ERROR(sort_pairs( + out_offsets, + sorted_out_offsets, + src_offsets, + sorted_src_offsets, + static_cast(points), + end_bit, + scratch, + stream, + "scatter_reduce")); + scatter_reduce_sorted_kernel + <<>>( + points, + sorted_out_offsets, + sorted_src_offsets, + static_cast(out), + static_cast(src), + reduction, + include_self, + reduction_identity(reduction)); + ET_CUDA_KERNEL_LAUNCH_CHECK_OR_RETURN_ERROR(); + return Error::Ok; +} + } // namespace #ifdef __cplusplus @@ -1057,6 +1281,161 @@ AOTITorchError aoti_torch_index_put_out( } } +AOTITorchError aoti_torch_cuda_scatter_reduce_two_out( + Tensor* out, + Tensor* self, + int64_t dim, + Tensor* index, + Tensor* src, + const char* reduce, + int32_t include_self) { + ET_CHECK_OR_RETURN_ERROR( + out != nullptr && index != nullptr && src != nullptr && reduce != nullptr, + InvalidArgument, + "aoti_torch_cuda_scatter_reduce_two_out: null argument"); + ET_CHECK_OR_RETURN_ERROR( + out == self, + NotSupported, + "aoti_torch_cuda_scatter_reduce_two_out: out must be self, as AOTInductor calls it"); + ET_CHECK_OR_RETURN_ERROR( + out->is_cuda() && index->is_cuda() && src->is_cuda(), + NotSupported, + "aoti_torch_cuda_scatter_reduce_two_out: tensors must be on the CUDA device"); + ET_CHECK_OR_RETURN_ERROR( + src->dtype() == out->dtype(), + InvalidArgument, + "aoti_torch_cuda_scatter_reduce_two_out: src and self must have the same dtype"); + const std::string_view name(reduce); + Reduction reduction; + if (name == "sum") { + reduction = Reduction::Sum; + } else if (name == "prod") { + reduction = Reduction::Prod; + } else if (name == "mean") { + reduction = Reduction::Mean; + } else if (name == "amax") { + reduction = Reduction::Max; + } else if (name == "amin") { + reduction = Reduction::Min; + } else { + ET_LOG( + Error, + "aoti_torch_cuda_scatter_reduce_two_out: unknown reduction %s", + reduce); + return Error::InvalidArgument; + } + ET_CHECK_OR_RETURN_ERROR( + out->dtype() != ScalarType::Bool || reduction == Reduction::Sum, + NotSupported, + "aoti_torch_cuda_scatter_reduce_two_out: bool supports only sum, as in ATen"); + + // As in ATen, a 0-dim tensor counts as one of size 1. + const auto ndim = std::max(out->dim(), 1); + ET_CHECK_OR_RETURN_ERROR( + ndim <= kMaxDims, + NotSupported, + "aoti_torch_cuda_scatter_reduce_two_out: tensors with more than %d dims are not supported", + kMaxDims); + if (dim < 0) { + dim += ndim; + } + ET_CHECK_OR_RETURN_ERROR( + dim >= 0 && dim < ndim, + InvalidArgument, + "aoti_torch_cuda_scatter_reduce_two_out: dim out of range"); + // ATen checks the index only when it has elements. + if (index->numel() == 0) { + return Error::Ok; + } + ET_CHECK_OR_RETURN_ERROR( + index->dtype() == ScalarType::Long, + InvalidArgument, + "aoti_torch_cuda_scatter_reduce_two_out: index must be int64"); + ET_CHECK_OR_RETURN_ERROR( + std::max(index->dim(), 1) == ndim && + std::max(src->dim(), 1) == ndim, + InvalidArgument, + "aoti_torch_cuda_scatter_reduce_two_out: index and src must have as many dims as self"); + auto size = [](const Tensor* tensor, int64_t d) -> int64_t { + return tensor->dim() == 0 ? 1 : tensor->size(d); + }; + auto stride = [](const Tensor* tensor, int64_t d) -> int64_t { + return tensor->dim() == 0 ? 0 : tensor->stride(d); + }; + ScatterPlan plan{}; + plan.ndim = static_cast(ndim); + plan.dim = static_cast(dim); + plan.dim_size = size(out, dim); + int64_t points = 1; + for (int64_t d = 0; d < ndim; ++d) { + plan.max_out_offset += + std::max(size(out, d) - 1, 0) * stride(out, d); + const int64_t index_size = size(index, d); + ET_CHECK_OR_RETURN_ERROR( + index_size <= size(src, d) && (d == dim || index_size <= size(out, d)), + InvalidArgument, + "aoti_torch_cuda_scatter_reduce_two_out: index is larger than src, or than self, at dim %" PRId64, + d); + plan.sizes[d] = index_size; + plan.index_strides[d] = stride(index, d); + plan.out_strides[d] = stride(out, d); + plan.src_strides[d] = stride(src, d); + points *= index_size; + } + + const auto device = static_cast(out->device_index()); + auto stream_result = getCurrentCUDAStream(out->device_index()); + ET_CHECK_OR_RETURN_ERROR( + stream_result.ok(), + Internal, + "aoti_torch_cuda_scatter_reduce_two_out: failed to get CUDA stream"); + cudaStream_t stream = stream_result.get(); + Scratch scratch(device, stream); + auto run = [&](auto value, auto accumulator) { + using T = decltype(value); + using AccT = decltype(accumulator); + return scatter_reduce( + plan, + points, + static_cast(index->data_ptr()), + out->data_ptr(), + src->data_ptr(), + reduction, + include_self != 0, + scratch, + stream); + }; + // Half and bfloat16 compute in float, and integers in int64. + switch (out->dtype()) { + case ScalarType::Float: + return run(float{}, float{}); + case ScalarType::Double: + return run(double{}, double{}); + case ScalarType::Half: + return run(__half{}, float{}); + case ScalarType::BFloat16: + return run(__nv_bfloat16{}, float{}); + case ScalarType::Long: + return run(int64_t{}, int64_t{}); + case ScalarType::Int: + return run(int32_t{}, int64_t{}); + case ScalarType::Short: + return run(int16_t{}, int64_t{}); + case ScalarType::Char: + return run(int8_t{}, int64_t{}); + case ScalarType::Byte: + return run(uint8_t{}, int64_t{}); + case ScalarType::Bool: + return run(bool{}, bool{}); + default: + ET_LOG( + Error, + "aoti_torch_cuda_scatter_reduce_two_out: unsupported dtype %d", + static_cast(out->dtype())); + return Error::NotSupported; + } +} + #ifdef __cplusplus } #endif diff --git a/backends/cuda/runtime/shims/index.h b/backends/cuda/runtime/shims/index.h index 25f015cf4c2..9da350ad2e3 100644 --- a/backends/cuda/runtime/shims/index.h +++ b/backends/cuda/runtime/shims/index.h @@ -75,6 +75,41 @@ AOTI_SHIM_EXPORT AOTITorchError aoti_torch_index_put_out( Tensor* values, bool accumulate); +/** + * Reduces src into out at index along dim, with the semantics of + * aten::scatter_reduce.two. + * + * Inductor calls this fallback, with out and self the same tensor, for every + * reduction but "sum", for a sum into a dtype its generated code cannot add + * atomically, and in deterministic mode. Values that land on the same element + * are reduced in index order, so the result is deterministic. Half and + * bfloat16 sums add in float, as in ATen's deterministic mode. A product + * rounds to the dtype of out at every step and a mean rounds its sum before + * dividing, as in ATen. As in ATen, an index out of range, negative ones + * included, triggers a device-side assert. + * + * @param out Tensor to update (CUDA device, up to 16 dims) + * @param self Must be out; AOTInductor passes the tensor it updates as both + * @param dim Dimension of out that index addresses + * @param index int64 tensor with as many dims as out, no larger than src, and + * no larger than out except at dim + * @param src Values to reduce, same dtype as out. Float, double, half, + * bfloat16 and the integer dtypes are supported, and bool for "sum" only, as + * in ATen. + * @param reduce "sum", "prod", "mean", "amax" or "amin" + * @param include_self If zero, the elements index names are reduced from src + * alone; the others keep their values either way + * @return AOTITorchError + */ +AOTI_SHIM_EXPORT AOTITorchError aoti_torch_cuda_scatter_reduce_two_out( + Tensor* out, + Tensor* self, + int64_t dim, + Tensor* index, + Tensor* src, + const char* reduce, + int32_t include_self); + #ifdef __cplusplus } #endif diff --git a/backends/cuda/tests/test_index_shims.py b/backends/cuda/tests/test_index_shims.py index b1899e4e794..144fb1ebdd9 100644 --- a/backends/cuda/tests/test_index_shims.py +++ b/backends/cuda/tests/test_index_shims.py @@ -4,12 +4,13 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -"""Mask indexing and index_put in the CUDA backend's index shims. +"""Mask indexing, index_put and scatter_reduce in the CUDA backend's index shims. A mask index makes a result size depend on the data, so Inductor leaves it, -and an index_put with a mask, to the backend's own index shims. These tests -lower real programs, run them through the Python runtime, and compare with -eager PyTorch. +and an index_put with a mask, to the backend's own index shims. Inductor also +leaves them every scatter_reduce it cannot write with an atomic add. These +tests lower real programs, run them through the Python runtime, and compare +with eager PyTorch on CUDA. """ import os @@ -38,7 +39,7 @@ class TestIndexShims(unittest.TestCase): def _run( self, model: torch.nn.Module, cases: Sequence[Tuple[torch.Tensor, ...]] ) -> None: - """Lower model for cases[0], then run every case and compare with eager.""" + """Lower model for cases[0], then compare every case with eager CUDA.""" spec = CudaBackend.generate_method_name_compile_spec("forward") lowered = to_edge_transform_and_lower( export(model, cases[0]), @@ -63,8 +64,11 @@ def _run( method = Runtime.get().load_program(path).load_method("forward") for inputs in cases: outputs = method.execute(list(inputs)) - for output, expected in zip(outputs, model(*inputs), strict=True): - torch.testing.assert_close(output.cpu(), expected) + expected = model(*(t.cuda() for t in inputs)) + for output, reference in zip(outputs, expected, strict=True): + torch.testing.assert_close( + output.cpu(), reference.cpu(), equal_nan=True + ) def test_mask_index_and_index_put(self): """index_put with accumulate goes to the shim too: Inductor has no atomic @@ -159,3 +163,28 @@ def forward(self, x, idx, values): self._run(Put(), [(x, idx, values)]) finally: torch.use_deterministic_algorithms(deterministic, warn_only=warn_only) + + def test_scatter_reduce(self): + """Inductor adds atomically but has no atomic max, min or product, so + those reductions, and the mean, go to the shim.""" + + class ScatterReduce(torch.nn.Module): + def forward(self, x, index, src): + return tuple( + x.scatter_reduce(1, index, src, reduce, include_self=include) + for reduce in ("amax", "amin", "prod", "mean") + for include in (True, False) + ) + + # Smaller than src, with repeats, and column 4 is never written. + index = torch.tensor([[0, 2, 2, 5], [1, 1, 1, 0], [5, 3, 0, 5]]) + torch.manual_seed(0) + for dtype in (torch.float32, torch.float16, torch.bfloat16, torch.int64): + with self.subTest(dtype=dtype): + # Small integers keep every sum and product exact in any order. + x = torch.randint(-4, 5, (5, 6)).to(dtype) + src = torch.randint(-4, 5, (4, 5)).to(dtype) + if dtype.is_floating_point: + # Lands on x[1, 1]; every reduction keeps the NaN. + src[1, 2] = float("nan") + self._run(ScatterReduce(), [(x, index, src)]) diff --git a/backends/cuda/tests/test_sort_shim.py b/backends/cuda/tests/test_sort_shim.py index d81d9e9e187..2195ad6d585 100644 --- a/backends/cuda/tests/test_sort_shim.py +++ b/backends/cuda/tests/test_sort_shim.py @@ -43,6 +43,7 @@ "aoti_torch_cuda_randint_low_out", "aoti_torch_cuda_index_Tensor", "aoti_torch_index_put_out", + "aoti_torch_cuda_scatter_reduce_two_out", } ) From f23afa800afa1d37e49b048f4fd9eb6c9072ba2f Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Fri, 30 Oct 2026 18:58:23 -0700 Subject: [PATCH 33/89] Retrigger CI From da1db735f36e3531dacc82b37f7ca1af9d4d279b Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Fri, 30 Oct 2026 19:56:20 -0700 Subject: [PATCH 34/89] Decompose varlen_attn in the CUDA backend torch.nn.attention.varlen.varlen_attn runs attention over many sequences packed into one tensor, with cu_seqlens marking where each one starts. Vision encoders that pack the patches of several images use it. Lowering a program that calls it to the CUDA backend failed: Method forward missing fallback kernels (1 total): - torch_attn::_varlen_attn The op is a custom op that calls Flash or cuDNN attention. Inductor has no lowering for it, so it fell back to a kernel this backend does not ship. Add it to the backend's decomposition table. The decomposition attends over all the packed tokens with a mask that keeps each query inside its own sequence, using ops Inductor lowers, and computes in float32 like the math SDPA path the backend already uses. It follows Flash attention for a causal mask or a window when a query sequence and its key sequence differ in length (their ends line up), for grouped-query attention, and for a query with no key to attend to (zeros and an infinite log-sum-exp). It raises a clear error for seqused_k and block_table (paged KV), which it does not support. It holds a float32 score per head for every pair of packed tokens: 1 GiB for 16 heads and 4096 tokens, against tens of MiB for Flash attention, and it is much slower. That suits packs of a few thousand tokens, such as a few camera images. Flash attention also treats a right window of at least max_k as unlimited. That changes the result only when a query sequence is longer than every key sequence by more than the window. The decomposition keeps the documented window instead, since max_k is often a symbolic size during export. Test: test_varlen_attn lowers a module that calls varlen_attn, and checks the decomposition against the op for a causal window over query sequences shorter than their key sequences, a causal mask over query sequences longer than their key sequences, and grouped-query attention with an empty sequence. Without this change the lowering fails with the error above. --- backends/cuda/cuda_backend.py | 72 +++++++++++++++++++++++++ backends/cuda/tests/test_cuda_export.py | 31 +++++++++++ 2 files changed, 103 insertions(+) diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index 84c9747ca98..2f1fbfdb11b 100644 --- a/backends/cuda/cuda_backend.py +++ b/backends/cuda/cuda_backend.py @@ -19,6 +19,9 @@ from typing import Any, Dict, final, List, Optional import torch + +# Registers torch_attn._varlen_attn, which the decomposition table names. +import torch.nn.attention.varlen # noqa: F401 from executorch.backends.aoti.aoti_backend import AotiBackend from executorch.backends.cuda.autotune.cuda_graph_timing import cuda_graph_timing from executorch.backends.cuda.autotune.inputs import autotune_input_scenarios @@ -568,6 +571,74 @@ def _reshape_copy(x: torch.Tensor, size: List[int]) -> torch.Tensor: return x.clone(memory_format=torch.contiguous_format).view(size) +def _varlen_attn( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + cu_seq_q: torch.Tensor, + cu_seq_k: Optional[torch.Tensor], + max_q: int, + max_k: int, + is_causal: bool = False, + scale: Optional[float] = None, + window_size: Optional[List[int]] = None, + enable_gqa: bool = False, + seqused_k: Optional[torch.Tensor] = None, + block_table: Optional[torch.Tensor] = None, + num_splits: Optional[int] = None, + backend: Optional[int] = None, +) -> typing.Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """ + varlen_attn runs Flash or cuDNN attention, and this backend has a kernel + for neither. Attend over all the packed tokens instead, with a mask that + keeps each query in its own sequence. This holds a float32 score per head + for every pair of packed tokens, 1 GiB for 16 heads and 4096 tokens. As in + Flash attention, a causal mask or a window lines up the ends of each query + and key sequence, and a query with no key to attend to gets zeros and an + infinite log-sum-exp. + """ + if cu_seq_k is None or seqused_k is not None or block_table is not None: + raise NotImplementedError( + "The CUDA backend lowers varlen_attn only with cu_seq_k and without " + "seqused_k or block_table." + ) + + def sequence_ends(cu: torch.Tensor, total: int): + # Each token's index in cu of its sequence's end, and its (negative) + # distance from that end. Tokens past cu[-1] are in no sequence, and + # the clamp keeps their lookup in range. + token = torch.arange(total, dtype=cu.dtype, device=cu.device) + end = torch.searchsorted(cu, token, right=True) + return end, token - cu[end.clamp(max=cu.numel() - 1)] + + end_q, from_end_q = sequence_ends(cu_seq_q, query.size(0)) + end_k, from_end_k = sequence_ends(cu_seq_k, key.size(0)) + allowed = end_q[:, None] == end_k[None, :] + left, right = window_size if window_size is not None else (-1, -1) + if is_causal: + right = 0 + offset = from_end_k[None, :] - from_end_q[:, None] + if left >= 0: + allowed = allowed & (offset >= -left) + if right >= 0: + allowed = allowed & (offset <= right) + + if enable_gqa: + key, value = ( + t.repeat_interleave(query.size(1) // t.size(1), dim=1) for t in (key, value) + ) + if scale is None: + scale = query.size(-1) ** -0.5 + scores = torch.einsum("qhd,khd->hqk", query.float(), key.float()) * scale + scores = scores.masked_fill(~allowed, float("-inf")) + lse = scores.logsumexp(-1) + lse = lse.masked_fill(lse == float("-inf"), float("inf")) + probs = torch.exp(scores - lse[..., None]) + out = torch.einsum("hqk,khd->qhd", probs, value.float()).to(query.dtype) + rng_state = torch.zeros(2, dtype=torch.uint64, device=query.device) + return out, lse, rng_state + + def _linear_kernel(x: torch.Tensor) -> torch.Tensor: return torch.clamp(1 - x.abs(), min=0) @@ -996,6 +1067,7 @@ def get_decomposition_table(cls) -> Dict[Any, Any]: # for it and the shim no fallback kernel. torch.ops.aten._is_all_true.default: torch.ops.aten.all.default, torch.ops.aten._reshape_copy.default: _reshape_copy, + torch.ops.torch_attn._varlen_attn.default: _varlen_attn, } for resize, support, kernel in ( (torch.ops.aten._upsample_bilinear2d_aa, 1, _linear_kernel), diff --git a/backends/cuda/tests/test_cuda_export.py b/backends/cuda/tests/test_cuda_export.py index f9e30442673..9f9db0483cd 100644 --- a/backends/cuda/tests/test_cuda_export.py +++ b/backends/cuda/tests/test_cuda_export.py @@ -21,6 +21,7 @@ from executorch.exir import EdgeCompileConfig, schema, to_edge_transform_and_lower from executorch.exir.backend.compile_spec_schema import CompileSpec from torch.export import export +from torch.nn.attention.varlen import varlen_attn class TestCudaBackendCompileOptions(unittest.TestCase): @@ -481,6 +482,36 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: self._export_to_cuda_with_lower(ReshapeCopyModule(), (torch.randn(8, 4, 4),)) + def test_varlen_attn(self): + """varlen_attn calls Flash attention, which the CUDA shims lack.""" + + class VarlenAttention(torch.nn.Module): + def forward(self, q, k, v, cu_seq): + return varlen_attn(q, k, v, cu_seq, cu_seq, 16, 16) + + q, k, v = ( + torch.randn(32, 4, 64, dtype=torch.bfloat16, device="cuda") + for _ in range(3) + ) + cu_seq = torch.tensor([0, 16, 32], dtype=torch.int32, device="cuda") + self._export_to_cuda_with_lower(VarlenAttention(), (q, k, v, cu_seq)) + + # Causal masks over query sequences shorter and longer than their key + # sequences (some queries then see no key), grouped-query attention + # and an empty sequence. + op = torch.ops.torch_attn._varlen_attn.default + short = torch.tensor([0, 4, 16], dtype=torch.int32, device="cuda") + empty = torch.tensor([0, 16, 16, 32], dtype=torch.int32, device="cuda") + cases = [ + (q[:16], k, v, short, cu_seq, 12, 16, True, None, [3, -1]), + (q, k[:16], v[:16], cu_seq, short, 16, 12, True, None, [-1, 0]), + (q, k[:, :2], v[:, :2], empty, empty, 16, 16, False, 0.5, None, True), + ] + table = CudaBackend.get_decomposition_table() + for args in cases: + for actual, expected in zip(table[op](*args)[:2], op(*args)[:2]): + torch.testing.assert_close(actual, expected, atol=1e-2, rtol=1e-2) + def test_conv2d(self): """Test CUDA export for 2D convolution.""" From 8a4ebbfa3ef095c0a05c97378ab24e7396dfe337 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Fri, 30 Oct 2026 20:40:51 -0700 Subject: [PATCH 35/89] Test the int64 scatter sum in the CUDA backend instead of expecting it to fail The CUDA shims now have scatter_reduce, so an int64 scatter_add lowers and runs. test_missing_scatter_kernel_fails_at_export still expected that export to fail, so it failed on every CUDA run. Remove it, and run the sum through test_scatter_reduce with the other reductions. For int64 that sum goes to the shim, which no test ran before. --- backends/cuda/tests/test_cuda_export.py | 12 ------------ backends/cuda/tests/test_index_shims.py | 5 +++-- 2 files changed, 3 insertions(+), 14 deletions(-) diff --git a/backends/cuda/tests/test_cuda_export.py b/backends/cuda/tests/test_cuda_export.py index 9f9db0483cd..315047e0c07 100644 --- a/backends/cuda/tests/test_cuda_export.py +++ b/backends/cuda/tests/test_cuda_export.py @@ -851,18 +851,6 @@ def forward(self, q, k, v, mask): for output, reference in zip(outputs, model(*inputs), strict=True): torch.testing.assert_close(output, reference.cpu(), atol=2e-2, rtol=2e-2) - def test_missing_scatter_kernel_fails_at_export(self): - """Inductor leaves int64 scatter_add to an ATen kernel the CUDA shim lacks.""" - - class Int64ScatterAdd(torch.nn.Module): - def forward(self, counts: torch.Tensor, idx: torch.Tensor): - return counts.scatter_add(0, idx, torch.ones_like(idx)) - - inputs = (torch.zeros(8, dtype=torch.long), torch.randint(0, 8, (32,))) - with self.assertRaises(RuntimeError) as context: - self._export_to_cuda_with_lower(Int64ScatterAdd(), inputs) - self.assertIn("aoti_torch_cuda_scatter_reduce_two_out", str(context.exception)) - @unittest.skipIf( torch.version.hip is not None, "loads through the Python runtime, which the ROCm build does not include", diff --git a/backends/cuda/tests/test_index_shims.py b/backends/cuda/tests/test_index_shims.py index 144fb1ebdd9..c94c0191104 100644 --- a/backends/cuda/tests/test_index_shims.py +++ b/backends/cuda/tests/test_index_shims.py @@ -166,13 +166,14 @@ def forward(self, x, idx, values): def test_scatter_reduce(self): """Inductor adds atomically but has no atomic max, min or product, so - those reductions, and the mean, go to the shim.""" + those reductions, and the mean, go to the shim. So does an int64 sum, + which Inductor cannot add atomically.""" class ScatterReduce(torch.nn.Module): def forward(self, x, index, src): return tuple( x.scatter_reduce(1, index, src, reduce, include_self=include) - for reduce in ("amax", "amin", "prod", "mean") + for reduce in ("sum", "amax", "amin", "prod", "mean") for include in (True, False) ) From 41207618a9132259b622ee026841e1936abfeb1e Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Fri, 30 Oct 2026 20:40:51 -0700 Subject: [PATCH 36/89] Use a model PyTorch 2.14 can export in the CUDA tensor check test test_distribution_argument_check built a MultivariateNormal. Its argument check branches on a tensor, and PyTorch 2.14, the pinned version, cannot export that, so the test failed inside torch.export before the CUDA backend ran. Write the same check with torch._assert_async, the op newer PyTorch exports for it. Without RemoveAssertAsyncPass the lowering still fails with a device-side assert, and with it the lowering passes. Also check what the pass does to the graph: the check and the math that only fed it are gone, and the rest of the model stays. --- backends/cuda/tests/test_cuda_export.py | 34 +++++++++++++++---------- 1 file changed, 21 insertions(+), 13 deletions(-) diff --git a/backends/cuda/tests/test_cuda_export.py b/backends/cuda/tests/test_cuda_export.py index 315047e0c07..148c812fb91 100644 --- a/backends/cuda/tests/test_cuda_export.py +++ b/backends/cuda/tests/test_cuda_export.py @@ -17,6 +17,7 @@ from executorch.backends.aoti.aoti_backend import AotiBackend from executorch.backends.cuda.cuda_backend import CudaBackend from executorch.backends.cuda.cuda_partitioner import CudaPartitioner +from executorch.backends.cuda.passes.remove_assert_async import RemoveAssertAsyncPass from executorch.examples.models.toy_model import SdpaModule from executorch.exir import EdgeCompileConfig, schema, to_edge_transform_and_lower from executorch.exir.backend.compile_spec_schema import CompileSpec @@ -454,24 +455,31 @@ def test_distribution_argument_check(self): class GaussianHead(torch.nn.Module): def __init__(self): super().__init__() - # With 14 actions, Inductor checks scale_tril in a kernel of its - # own, so the assert reads a bool that autotuning makes up. self.mean = torch.nn.Linear(16, 14) self.log_std = torch.nn.Linear(16, 14) def forward(self, x: torch.Tensor, eps: torch.Tensor) -> torch.Tensor: std = torch.clamp(torch.exp(self.log_std(x)), 1e-5, 10.0) - dist = torch.distributions.MultivariateNormal( - self.mean(x), scale_tril=torch.diag_embed(std.sqrt()) - ) - return dist.loc + (dist.scale_tril @ eps[..., None])[..., 0] - - # A strict export cannot trace torch.distributions. - self._export_to_cuda_with_lower( - GaussianHead().eval(), - (torch.randn(1, 16), torch.randn(1, 14)), - strict=False, - ) + scale_tril = torch.diag_embed(std.sqrt()) + # The check MultivariateNormal makes on scale_tril, as newer + # PyTorch exports it. With 14 actions, Inductor computes each + # half in a kernel of its own, so the assert reads bools that + # autotuning makes up. + lower = (scale_tril.tril() == scale_tril).flatten(-2).min(-1)[0] + positive = (scale_tril.diagonal(dim1=-2, dim2=-1) > 0).min(-1)[0] + torch._assert_async((lower & positive).all(), "invalid scale_tril") + return self.mean(x) + (scale_tril @ eps[..., None])[..., 0] + + inputs = (torch.randn(1, 16), torch.randn(1, 14)) + self._export_to_cuda_with_lower(GaussianHead().eval(), inputs) + + # The pass drops the check and the math that only fed it. + graph_module = export(GaussianHead().eval(), inputs).graph_module + RemoveAssertAsyncPass()(graph_module) + targets = {n.target for n in graph_module.graph.nodes} + self.assertNotIn(torch.ops.aten._assert_async.msg, targets) + self.assertNotIn(torch.ops.aten.tril.default, targets) + self.assertIn(torch.ops.aten.linear.default, targets) def test_reshape_copy(self): """Torch-TensorRT hands back the views it leaves to PyTorch as _reshape_copy.""" From 1a0704a175ea406aa414bbedf9f754f4caae8640 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Fri, 30 Oct 2026 20:40:52 -0700 Subject: [PATCH 37/89] Check that the CUDA mask tests try a graph capture The mask tests compared values only, and eager calls give the same values, so they passed with CUDA graphs turned off. Collect the runtime log and check that the capture was tried and failed, as a mask's count read makes it. The count test also cannot pass with PyTorch 2.14: there the program does not keep a buffer update across calls, so the count reads 1 on every call, with or without a graph. Skip it before 2.15, where the count is kept. --- .../tests/test_output_stride_rearrange.py | 40 ++++++++++++++++++- 1 file changed, 38 insertions(+), 2 deletions(-) diff --git a/backends/cuda/tests/test_output_stride_rearrange.py b/backends/cuda/tests/test_output_stride_rearrange.py index 985f3e6feee..557259770ec 100644 --- a/backends/cuda/tests/test_output_stride_rearrange.py +++ b/backends/cuda/tests/test_output_stride_rearrange.py @@ -18,6 +18,7 @@ python -m pytest backends/cuda/tests/test_output_stride_rearrange.py -v """ +import contextlib import glob import os import subprocess @@ -327,6 +328,29 @@ def forward(self, x): return (x * 2).contiguous(memory_format=self.memory_format) +class _RuntimeLog: + """Collects the runtime's log lines. They go to sys.stderr while the + caller holds the GIL, and straight to file descriptor 2 otherwise.""" + + def __enter__(self): + sys.stderr.flush() + self._file = tempfile.TemporaryFile(mode="w+") + self._saved = os.dup(2) + os.dup2(self._file.fileno(), 2) + self._redirect = contextlib.redirect_stderr(self._file) + self._redirect.__enter__() + return self + + def __exit__(self, *exc): + self._redirect.__exit__(*exc) + self._file.flush() + os.dup2(self._saved, 2) + os.close(self._saved) + self._file.seek(0) + self.text = self._file.read() + self._file.close() + + class TestStridedOutputWithCudaGraph(unittest.TestCase): """A CUDA graph replay copies raw bytes, so a strided output must not reach it.""" @@ -451,18 +475,30 @@ def test_mask_runs_without_a_graph(self): torch.manual_seed(0) inputs = [(torch.randn(8, 4), torch.rand(8) > 0.5) for _ in range(6)] for module in (_MaskedRows().eval(), _MaskedAdd().eval()): - outputs = self._run_with_cuda_graph(module, inputs[0], inputs) + with _RuntimeLog() as log: + outputs = self._run_with_cuda_graph(module, inputs[0], inputs) + # Eager calls give the same values, so only the log shows the + # capture was tried. + self.assertIn("capture failed for 'forward'", log.text) for i, (args, out) in enumerate(zip(inputs, outputs)): with self.subTest(module=type(module).__name__, call=i): self.assertIsInstance(out, torch.Tensor) torch.testing.assert_close(out, module(*args)) + @unittest.skipIf( + tuple(map(int, torch.__version__.split(".")[:2])) < (2, 15), + "with torch 2.14 the program does not keep a buffer update across calls", + ) def test_mask_retry_updates_state_once(self): # The capture call runs again without a graph, so its update to the # count must not land twice. torch.manual_seed(0) inputs = [(torch.randn(8, 4), torch.rand(8) > 0.5) for _ in range(6)] - outputs = self._run_with_cuda_graph(_MaskedCounter().eval(), inputs[0], inputs) + with _RuntimeLog() as log: + outputs = self._run_with_cuda_graph( + _MaskedCounter().eval(), inputs[0], inputs + ) + self.assertIn("capture failed for 'forward'", log.text) for i, out in enumerate(outputs): with self.subTest(call=i): self.assertIsInstance(out, torch.Tensor) From 7084c6fcc0778234711a17e3cfafdf14d489d01f Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Fri, 30 Oct 2026 20:41:35 -0700 Subject: [PATCH 38/89] Pin the masked copy pass's integer index and shared picks guards Two guards in ReplaceMaskedCopyWithWherePass could be deleted without any test failing. Without the boolean mask check, out[idx] = x[idx] with an integer index becomes a where on an int64 condition and fails. Without the check that keeps an index other code reads, picked = x[mask]; out[mask] = picked; return out, picked.sum() loses the node the output reads. Add both to the cases the pass must leave working. --- .../test_replace_masked_copy_with_where.py | 26 +++++++++++++------ 1 file changed, 18 insertions(+), 8 deletions(-) diff --git a/backends/cuda/passes/tests/test_replace_masked_copy_with_where.py b/backends/cuda/passes/tests/test_replace_masked_copy_with_where.py index 6c594bcdcd0..ea0646a4b46 100644 --- a/backends/cuda/passes/tests/test_replace_masked_copy_with_where.py +++ b/backends/cuda/passes/tests/test_replace_masked_copy_with_where.py @@ -63,7 +63,7 @@ def forward(self, x, mask): def test_other_writes_are_kept(self): class Kept(torch.nn.Module): - def forward(self, x, mask): + def forward(self, x, mask, idx): scaled = torch.zeros_like(x) # The values are not the mask's picks of a tensor. scaled[mask] = x[mask] * 2 @@ -76,19 +76,29 @@ def forward(self, x, mask): summed.index_put_((mask,), x[mask], accumulate=True) other = torch.zeros_like(x) other[mask] = x[mask.logical_not()] - return scaled, checked, summed, other + # An integer index is not a mask, so where cannot take it. + rows = torch.zeros_like(x) + rows[idx] = x[idx] + # The picks are also read elsewhere, so the index stays. + shared = torch.zeros_like(x) + read = x[mask] + shared[mask] = read + return scaled, checked, summed, other, rows, shared, read.sum() x = torch.randn(4, 3) mask = torch.tensor([True, False, True, False]) - program = to_edge(export(Kept(), (x, mask))).exported_program() + idx = torch.tensor([2, 0]) + program = to_edge(export(Kept(), (x, mask, idx))).exported_program() ReplaceMaskedCopyWithWherePass()(program.graph_module) targets = _targets(program.graph_module) - self.assertEqual(targets.count(exir_ops.edge.aten.index_put.default), 3) - self.assertEqual(targets.count(exir_ops.edge.aten.where.self), 1) - # The index stays for the checked count. - self.assertIn(exir_ops.edge.aten.index.Tensor, targets) - for got, expected in zip(program.module()(x, mask), Kept()(x, mask)): + self.assertEqual(targets.count(exir_ops.edge.aten.index_put.default), 4) + self.assertEqual(targets.count(exir_ops.edge.aten.where.self), 2) + # The index stays for the checked count and for the shared picks. + self.assertEqual(targets.count(exir_ops.edge.aten.index.Tensor), 5) + for got, expected in zip( + program.module()(x, mask, idx), Kept()(x, mask, idx), strict=True + ): torch.testing.assert_close(got, expected, rtol=0, atol=0) def test_check_on_another_size_keeps_the_index(self): From 086971116bb8dc4c69e314401d7bf668a7f14e8a Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Fri, 30 Oct 2026 20:41:36 -0700 Subject: [PATCH 39/89] Compare masks without recursion in the CUDA masked copy pass _same_value compared two masks by recursing into their inputs. Two separately computed masks with long histories passed Python's recursion limit, so lowering failed with RecursionError, and an input shared by many nodes was compared again on every path that reached it, which grows exponentially. Walk a work list instead, and compare each node pair once. --- .../passes/replace_masked_copy_with_where.py | 49 ++++++++++++------- .../test_replace_masked_copy_with_where.py | 22 +++++++++ 2 files changed, 53 insertions(+), 18 deletions(-) diff --git a/backends/cuda/passes/replace_masked_copy_with_where.py b/backends/cuda/passes/replace_masked_copy_with_where.py index 0f76b995f34..9e8186639b2 100644 --- a/backends/cuda/passes/replace_masked_copy_with_where.py +++ b/backends/cuda/passes/replace_masked_copy_with_where.py @@ -41,24 +41,37 @@ def _same_value(a: Any, b: Any) -> bool: Whether a and b are the same node, or the same deterministic op applied to the same values, as with ``x[m.view(-1)] = y[m.view(-1)]``. """ - if isinstance(a, (list, tuple)) and isinstance(b, (list, tuple)): - return len(a) == len(b) and all(map(_same_value, a, b)) - if not isinstance(a, Node) or not isinstance(b, Node): - return not isinstance(a, Node) and not isinstance(b, Node) and a == b - if a is b: - return True - if a.op != "call_function" or a.target != b.target: - return False - if ( - isinstance(a.target, (EdgeOpOverload, torch._ops.OpOverload)) - and torch.Tag.nondeterministic_seeded in a.target.tags - ): - return False - return ( - _same_value(a.args, b.args) - and a.kwargs.keys() == b.kwargs.keys() - and all(_same_value(a.kwargs[k], b.kwargs[k]) for k in a.kwargs) - ) + # A work list, not recursion: two long mask computations would pass the + # recursion limit, and each node pair is compared once, since shared + # inputs would otherwise be compared again on every path that reaches them. + pending = [(a, b)] + compared = set() + while pending: + a, b = pending.pop() + if isinstance(a, (list, tuple)) and isinstance(b, (list, tuple)): + if len(a) != len(b): + return False + pending.extend(zip(a, b)) + continue + if not isinstance(a, Node) or not isinstance(b, Node): + if isinstance(a, Node) or isinstance(b, Node) or a != b: + return False + continue + if a is b or (a, b) in compared: + continue + compared.add((a, b)) + if a.op != "call_function" or a.target != b.target: + return False + if ( + isinstance(a.target, (EdgeOpOverload, torch._ops.OpOverload)) + and torch.Tag.nondeterministic_seeded in a.target.tags + ): + return False + if a.kwargs.keys() != b.kwargs.keys(): + return False + pending.append((a.args, b.args)) + pending.extend((a.kwargs[k], b.kwargs[k]) for k in a.kwargs) + return True class ReplaceMaskedCopyWithWherePass(PassBase): diff --git a/backends/cuda/passes/tests/test_replace_masked_copy_with_where.py b/backends/cuda/passes/tests/test_replace_masked_copy_with_where.py index ea0646a4b46..5062d27cfa3 100644 --- a/backends/cuda/passes/tests/test_replace_masked_copy_with_where.py +++ b/backends/cuda/passes/tests/test_replace_masked_copy_with_where.py @@ -61,6 +61,28 @@ def forward(self, x, mask): program.module()(x, mask), Copy()(x, mask), rtol=0, atol=0 ) + def test_long_equal_masks(self): + """Two long computations of one mask, with inputs shared many times.""" + + class Copy(torch.nn.Module): + def forward(self, x, mask): + def compute(m): + for _ in range(1200): + m = m.logical_not() + for _ in range(30): + m = m & m + return m + + out = torch.zeros_like(x) + out[compute(mask)] = x[compute(mask)] + return out + + x = torch.randn(4, 3) + mask = torch.tensor([True, False, True, False]) + program = to_edge(export(Copy(), (x, mask))).exported_program() + ReplaceMaskedCopyWithWherePass()(program.graph_module) + self.assertIn(exir_ops.edge.aten.where.self, _targets(program.graph_module)) + def test_other_writes_are_kept(self): class Kept(torch.nn.Module): def forward(self, x, mask, idx): From e29e302c884b9a637dad22fc1b23b1cd0d6fe685 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Fri, 30 Oct 2026 21:48:29 -0700 Subject: [PATCH 40/89] Accept float64 in the CUDA shims that still refused it The CUDA backend now runs float64 tensors, but some shims still refused them at run time: - A float64 .item() calls aoti_torch_item_float64, which did not exist. The runtime loads the library lazily, so the first execute ended the process with "undefined symbol". Add the entry point, and list it in the Windows import library. - A torch.cond on a float64 flag reads it through aoti_torch_item_int64, which refused a Double tensor. - index_put with accumulate and a mask failed with "unsupported dtype 7". - The sort shim refused Double with the same error. Add Double to each. The tests run a float64 masked accumulate and a float64 sort. The two scalar reads are not in the tests: which call the generated code makes for them depends on the PyTorch version. --- backends/cuda/runtime/aoti_cuda_shims.lib | Bin 48750 -> 49562 bytes backends/cuda/runtime/shims/index.cu | 3 +++ backends/cuda/runtime/shims/memory.cpp | 3 +++ backends/cuda/runtime/shims/memory.h | 4 ++++ backends/cuda/runtime/shims/sort.cu | 14 ++++++++++++++ backends/cuda/tests/test_index_shims.py | 13 ++++++++++++- 6 files changed, 36 insertions(+), 1 deletion(-) diff --git a/backends/cuda/runtime/aoti_cuda_shims.lib b/backends/cuda/runtime/aoti_cuda_shims.lib index a4bc5b1414bbefceabb5695388fb9a80de008706..c392bfe1184c358b3fe82a9784a40eff536808e0 100644 GIT binary patch delta 705 zcmZ|NUr5tY6bJC{T^1USjX^8_pb=CcO{a~Y|6NKDMS-M9fiQ{5e!2I4-?2>&?0Hlb z^LfjChclyNKY%;H!4p8s#?98e-mjc*@$D)}UvzGg!?GnH)TcW)LU;&p1egjzOqv)CfSPuna zv0mE9Ze@$eD!^_(6L%>+yH_LnsE+loiGlY3izmc`L4ZB%5Tmq@jqQkWs$~;H8P_p+ z9vRb=n9ckUPd@|fIi+QD)ReutDPGePu(uS38LDNeCn8Ou**m&wwv?^C#rdYfwvwCG z9lg@4M{@FWc@nsKBOHucKE0i+)fJ9JPPO;%htU1b118J)VV@p#?^Tt$(Rs*xa(|gJ z!|{tMuM)v?+s{%O=$u%r`1W&b&5>`jqAn<_Qap%$f==k0no=`gOOrS!3M nnX&t6b#B3~>m4q$HQ%80#-hixrdyS&F)tS^r7ZJxam4lqFS^32 delta 590 zcmXZZT`0q06bA6;KN3=wLh`kw#xfR}jhc@ZS}9p-QSuehibQF}e3T+-`ip3>lu{`5 zYMGHsciN2$k}G#gE?n-yIXm9^-Ml^LeV;l<>uKLfL&yCO;aHj^wS`lxfF8&fGu>3MQYCiA7Y7e zY-L%jPXKJgCpN19ww*3^XpZeh#2#hYK~2gS4u6mmrpD}KU!2YX?2Ov7b9$6r*t4&zgXG8evvRU+D_X-#@* vtXR2j?V28m)~TshuckTlncOm!A9rc`LabJ;-qW2^y*yCHdv8^a$5F!{J0Fq# diff --git a/backends/cuda/runtime/shims/index.cu b/backends/cuda/runtime/shims/index.cu index 73889e3ee5d..2c1606f753f 100644 --- a/backends/cuda/runtime/shims/index.cu +++ b/backends/cuda/runtime/shims/index.cu @@ -1248,6 +1248,9 @@ AOTITorchError aoti_torch_index_put_out( case ScalarType::Float: return put( plan, numel, may_repeat, dst, src, scratch, stream); + case ScalarType::Double: + return put( + plan, numel, may_repeat, dst, src, scratch, stream); case ScalarType::Half: return put<__half, float, true>( plan, numel, may_repeat, dst, src, scratch, stream); diff --git a/backends/cuda/runtime/shims/memory.cpp b/backends/cuda/runtime/shims/memory.cpp index 86be88e4833..755e6969ec8 100644 --- a/backends/cuda/runtime/shims/memory.cpp +++ b/backends/cuda/runtime/shims/memory.cpp @@ -80,6 +80,8 @@ AOTITorchError read_one_element(const SlimTensor* tensor, T* ret_value) { return narrow_to(tensor->item(), ret_value); case ScalarType::Float: return narrow_to(tensor->item(), ret_value); + case ScalarType::Double: + return narrow_to(tensor->item(), ret_value); case ScalarType::BFloat16: return narrow_to( static_cast(tensor->item()), ret_value); @@ -470,6 +472,7 @@ ET_CUDA_DEFINE_ITEM_SHIM(int16, int16_t) ET_CUDA_DEFINE_ITEM_SHIM(int32, int32_t) ET_CUDA_DEFINE_ITEM_SHIM(int64, int64_t) ET_CUDA_DEFINE_ITEM_SHIM(float32, float) +ET_CUDA_DEFINE_ITEM_SHIM(float64, double) ET_CUDA_DEFINE_ITEM_SHIM(bfloat16, c10::BFloat16) ET_CUDA_DEFINE_ITEM_SHIM(bool, bool) diff --git a/backends/cuda/runtime/shims/memory.h b/backends/cuda/runtime/shims/memory.h index b1b369a58d9..8e4b87517c9 100644 --- a/backends/cuda/runtime/shims/memory.h +++ b/backends/cuda/runtime/shims/memory.h @@ -256,6 +256,10 @@ aoti_torch_item_int64(SlimTensor* tensor, int64_t* ret_value); AOTI_SHIM_EXPORT AOTITorchError aoti_torch_item_float32(SlimTensor* tensor, float* ret_value); +/// See aoti_torch_item_uint8. +AOTI_SHIM_EXPORT AOTITorchError +aoti_torch_item_float64(SlimTensor* tensor, double* ret_value); + /// See aoti_torch_item_uint8. AOTI_SHIM_EXPORT AOTITorchError aoti_torch_item_bfloat16(SlimTensor* tensor, c10::BFloat16* ret_value); diff --git a/backends/cuda/runtime/shims/sort.cu b/backends/cuda/runtime/shims/sort.cu index e3bb259bbc7..30d1cf39e08 100644 --- a/backends/cuda/runtime/shims/sort.cu +++ b/backends/cuda/runtime/shims/sort.cu @@ -211,6 +211,9 @@ AOTITorchError aoti_torch_cuda_sort_stable( case c10_slim::ScalarType::Float: elem_size = sizeof(float); break; + case c10_slim::ScalarType::Double: + elem_size = sizeof(double); + break; case c10_slim::ScalarType::BFloat16: elem_size = sizeof(__nv_bfloat16); break; @@ -416,6 +419,17 @@ AOTITorchError aoti_torch_cuda_sort_stable( stream); break; } + case c10_slim::ScalarType::Double: { + sort_slice_impl( + static_cast(values_base) + offset, + idx_ptr, + sort_size, + desc, + is_stable, + device_idx, + stream); + break; + } case c10_slim::ScalarType::BFloat16: { sort_slice_impl( static_cast<__nv_bfloat16*>(values_base) + offset, diff --git a/backends/cuda/tests/test_index_shims.py b/backends/cuda/tests/test_index_shims.py index c94c0191104..937a32f8a11 100644 --- a/backends/cuda/tests/test_index_shims.py +++ b/backends/cuda/tests/test_index_shims.py @@ -84,7 +84,9 @@ def forward(self, x, mask, column, idx, counts): rows = x.transpose(1, 2)[mask] * 3 ones = torch.ones_like(idx) counts = counts.index_put((idx,), ones, accumulate=True) - return out, picked, rows, counts + wide = x.double() + wide = wide.index_put((mask,), wide[mask], accumulate=True) + return out, picked, rows, counts, wide x = torch.randn(4, 5, 6) column = torch.tensor([2]) @@ -189,3 +191,12 @@ def forward(self, x, index, src): # Lands on x[1, 1]; every reduction keeps the NaN. src[1, 2] = float("nan") self._run(ScatterReduce(), [(x, index, src)]) + + def test_float64_sort(self): + class Sort(torch.nn.Module): + def forward(self, x): + # Rows this long go to the sort shim, not to Triton. + return torch.sort(x, stable=True) + + torch.manual_seed(0) + self._run(Sort(), [(torch.randn(4, 1000, dtype=torch.float64),)]) From fe11987c945730bdac42b6fbac0179539d699f1b Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Fri, 30 Oct 2026 21:48:29 -0700 Subject: [PATCH 41/89] Count a scatter mean in the tensor's dtype, as ATen does ATen counts the terms of a scatter_reduce mean in the tensor's own dtype, then divides. So a large count rounds or wraps there too, and a count that wraps to zero becomes one. The CUDA shim counted in int64, so 128 int8 ones averaged to -1 instead of 1, and 257 bfloat16 ones to 0.996 instead of 1. Convert the count the same way, and round integer quotients down for either sign of the divisor. --- backends/cuda/runtime/shims/index.cu | 13 ++++++++++--- backends/cuda/tests/test_index_shims.py | 15 +++++++++++++++ 2 files changed, 25 insertions(+), 3 deletions(-) diff --git a/backends/cuda/runtime/shims/index.cu b/backends/cuda/runtime/shims/index.cu index 2c1606f753f..4ec710c1c35 100644 --- a/backends/cuda/runtime/shims/index.cu +++ b/backends/cuda/runtime/shims/index.cu @@ -972,12 +972,19 @@ __global__ void scatter_reduce_sorted_kernel( } if (reduction == Reduction::Mean) { const auto sum = static_cast(static_cast(result)); + // ATen counts in T too, so a large count rounds or wraps, and a count + // that wraps to zero becomes one. + auto divisor = static_cast(static_cast(count)); + if (divisor == AccT(0)) { + divisor = AccT(1); + } if constexpr (std::is_floating_point_v) { - result = sum / static_cast(count); + result = sum / divisor; } else { // Integers round down, as ATen's division with rounding_mode="floor". - const AccT quotient = sum / count; - result = quotient - (quotient * count != sum && sum < 0); + const AccT quotient = sum / divisor; + result = quotient - + (quotient * divisor != sum && (sum < 0) != (divisor < 0)); } } out[offset] = static_cast(result); diff --git a/backends/cuda/tests/test_index_shims.py b/backends/cuda/tests/test_index_shims.py index 937a32f8a11..09b8afb9379 100644 --- a/backends/cuda/tests/test_index_shims.py +++ b/backends/cuda/tests/test_index_shims.py @@ -192,6 +192,21 @@ def forward(self, x, index, src): src[1, 2] = float("nan") self._run(ScatterReduce(), [(x, index, src)]) + def test_scatter_mean_counts_in_its_dtype(self): + """ATen counts a mean's terms in the tensor's own dtype, so 128 int8 + terms count as -128.""" + + class Mean(torch.nn.Module): + def forward(self, x, index, src): + return tuple( + x.scatter_reduce(0, index, src, "mean", include_self=include) + for include in (True, False) + ) + + x = torch.zeros(2, dtype=torch.int8) + index = torch.zeros(128, dtype=torch.long) + self._run(Mean(), [(x, index, torch.ones(128, dtype=torch.int8))]) + def test_float64_sort(self): class Sort(torch.nn.Module): def forward(self, x): From 9e707a6a929c5575e4f0e8f81a5a36a23c6be466 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Fri, 30 Oct 2026 21:48:29 -0700 Subject: [PATCH 42/89] Retry the CUDA random state setup when its seed copy fails ensure_rng_init says a failed setup leaves the call_once not done, so the next call tries again. That held for the seed draw and the allocation, which throw, but the seed copy and the stream sync after it were not checked. If one of them failed, the state stayed unseeded for the rest of the process. Check both, free the allocation, and throw like the other steps. --- backends/cuda/runtime/shims/rand.cu | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/backends/cuda/runtime/shims/rand.cu b/backends/cuda/runtime/shims/rand.cu index 31644c53a25..5f53838b672 100644 --- a/backends/cuda/runtime/shims/rand.cu +++ b/backends/cuda/runtime/shims/rand.cu @@ -22,6 +22,7 @@ #include #include #include +#include #include namespace executorch::backends::cuda { @@ -82,12 +83,18 @@ bool ensure_rng_init(cudaStream_t stream) { if (!rng.ok()) { throw std::bad_alloc(); } - d_rng = static_cast(rng.get()); - cudaMemcpyAsync( - d_rng, &h, sizeof(RngState), cudaMemcpyHostToDevice, stream); // Synchronize to ensure the copy completes before we return // (the host-side RngState `h` is on the stack). - cudaStreamSynchronize(stream); + cudaError_t err = cudaMemcpyAsync( + rng.get(), &h, sizeof(RngState), cudaMemcpyHostToDevice, stream); + if (err == cudaSuccess) { + err = cudaStreamSynchronize(stream); + } + if (err != cudaSuccess) { + (void)CudaAllocator::deallocate_stream_ordered(rng.get(), -1, stream); + throw std::runtime_error(cudaGetErrorString(err)); + } + d_rng = static_cast(rng.get()); }); } catch (const std::exception& e) { // std::random_device throws an implementation-defined std::exception. From 43410221c89cd5f2c323a0a37c49929932866450 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Fri, 30 Oct 2026 21:48:30 -0700 Subject: [PATCH 43/89] List every exported CUDA shim in the Windows import library A model lowered for Windows links its generated wrapper against aoti_cuda_shims.lib. The library did not list some names the DLL exports: aoti_torch_dtype_float16, aoti_torch_zero_, aoti_torch_cuda_rand and the aoti_torch_item_ getters other than bool and float64. So a float16 model, or one that reads a scalar or clears a workspace, failed to link. Regenerated per aoti_cuda_shims.lib.md, with the names the Linux build of the same sources exports. Rebuilding it gives the same bytes. --- backends/cuda/runtime/aoti_cuda_shims.lib | Bin 49562 -> 57574 bytes 1 file changed, 0 insertions(+), 0 deletions(-) diff --git a/backends/cuda/runtime/aoti_cuda_shims.lib b/backends/cuda/runtime/aoti_cuda_shims.lib index c392bfe1184c358b3fe82a9784a40eff536808e0..c6466b68f0ef2d27f0b4f558347b4d50753f1872 100644 GIT binary patch delta 2372 zcmaKs3rv$&6vyvvDNvy0-73n!>72&duTUthFhxLTd`twzNpuB71bK?!6rELZii#6C z4i%>YJ_8Yzqd1+fnJw9L3b;*;GmBf?BUxfxvMkHyWasu4L&UAguQ%WM&N=tobISPo1+*L%<`6`Q?MYTQ!?)~ zV}kj05=RMRf}9-5qN$7t@<&Jt>5PKnuO(&94;d>XmQj#kd7@+`okFncfn-fPV}f<7 zB^y#06Kofe`qKphqArWnCqPlTIrfAx zLEBwP`w_+jrzo-D^a{yYI;G$|6(hJn#Rx99O0HHhCiske1fP=+5znsOV=ivAFed2E zmwaVrOmN#na`y+u1m8AFzMGAas*@hjqcM7vXP4wcc^k3(J}~MH8xps zY<)V1u$cx(QZ7LjD%ypWx!Dfa-8{RbiT0u3OjU4g!fXzzv>-6024KR>WDf79E7+bo zo7c|?c3t6^G|HBlVaE3KG!FOY>QFal1_#?b1q&w7tMi$I;ii`-{yjH=!@pST= zo~Oq(3n_B{Vg|on2&_) z>cM!#8OP!5Mw4uXtwsmM)uFg{{Ui>XHkO$9JcII z(ZWl#0><@y+ejc#1`4T+XQv$yqa0f4M z4uS}s27lgb;?UD#!aEV20=YBf)fx{+QX5>*lZfK@^+<4609Pd^=+|c=A z;UNoMe4CsL9)rV<=5a`FH{qFM>72*Hu%iZyXr0Bu=ag(B)OP6Muj4d*$LV&uQBDDm zgME4<-akRdYwlFAsV$8U8Tj?~X-waYD6kzp|4WCy-w_TK=j3VWYwXe+@L&hEChdZ} z4nB^w$siRlXl?HiHX{cBk;U|#BIE+Nx zD~(@OzUI)_wbWFB4JqO1Fa*+xL+=Q@m#Wip+D8IY z4rw@T*hojxxn&~FN>i;|-IH#i1BVj@j%V2DVM2;P`(YCf8x7cO^uZ5|zIZ$8DQ?Yl zV2jg6YjHd`^y5twkewm0G*ed_u4IMcs>@79*FJ%+9GxnTa^LciAo?wLzrX`WBWb3U zm!8kvj5GNg=*MGm0=MPq6j{K1gU8HR>}MuVVYEQ60-c6yxLQ{jN`)m+0vC#O>M7$& zdWi*}pY|dD@)$wgrB?(!JZ0w;WjZ}}^Z3dM3(i#d(#=nK*4rm-s(Sa71smPIR9P7% z$a1Dr(EjQefqzu!G*ZKvT$SMzUTYIH>TVabbT(F?yGHHUT5H1YdNW>2`3{Tg&tPJm zRh@V4tQ;`mvj8KVzwoZWu+OdZq9Ij?ZFI1)gpXdY)1OV8Der;_lLJlUY}v)fZ_sJF zm8*`%5M1jp;Ao&P+1lO`INM}Zcb3qieR=Ucfz7Qtm0#jktxcoFPOHGecAZ{ybLClw zMqOPy1twnNO=V@SGla5xxMIAb)At^(obL*u?Y)tL=DShQlRn zrkR0wff?7VD$-(~h7NNeUA@6i^K$?Cm3js&R5Hw0^6%jKQ&-tf;T3ihj|`5iKU&17iTuX-KWrXT4=Z-ly$L>H#jW-9tGCVgQ;;83P3;hv z^ws*s%} Date: Fri, 30 Oct 2026 21:46:21 -0700 Subject: [PATCH 44/89] Keep buffer updates across calls in AOTI backends with PyTorch 2.14 A program lowered to the CUDA backend with PyTorch 2.14, the pinned version, lost its updates to a buffer whose elements all started equal. This counter returned 1 on every call: class Counter(torch.nn.Module): def __init__(self): super().__init__() self.register_buffer("calls", torch.zeros(1), persistent=False) def forward(self, x): self.calls.add_(1) return self.calls.clone() Inductor replaces a constant whose elements are all equal with a fill. Before PyTorch 2.15 it also did that to a buffer the program writes (pytorch/pytorch#185077). So the compiled program computed 0 + 1 once, at export, and wrote that 1 back on every call. The Metal backend compiles through the same code, so it had the same problem. Turn that folding off when the program writes a buffer. The CUDA backend already turns it off for the off-graph KV cache's cell layout, whose buffers the runtime rewrites. Newer PyTorch keeps written buffers out of the folding by itself, so there this only stops the folding of the other constants of such a program. test_mask_retry_updates_state_once was skipped before PyTorch 2.15 because of this. It runs again. Test: with PyTorch 2.14, test_mask_retry_updates_state_once fails without this change (the count reads 1 on every call) and passes with it. --- backends/aoti/aoti_backend.py | 15 +++++++++++++++ .../cuda/tests/test_output_stride_rearrange.py | 4 ---- 2 files changed, 15 insertions(+), 4 deletions(-) diff --git a/backends/aoti/aoti_backend.py b/backends/aoti/aoti_backend.py index 87a15498de7..8b72f3d4598 100644 --- a/backends/aoti/aoti_backend.py +++ b/backends/aoti/aoti_backend.py @@ -53,6 +53,20 @@ def _checked_data_name(data_name: str) -> str: return data_name +def _written_buffer_options(program: ExportedProgram) -> Dict[str, Any]: + """Inductor options that keep the buffers a program writes out of constant + folding. + + Before torch 2.15, Inductor replaces a buffer whose elements are all equal, + such as a counter that starts at zero, with a constant even when the + program writes it, so every call would start from the exported value + (pytorch/pytorch#185077). + """ + if program.graph_signature.buffers_to_mutate: + return {"joint_graph_constant_folding": False} + return {} + + @experimental( "This API and all of aoti-driven backend related functionality are experimental." ) @@ -292,6 +306,7 @@ def preprocess( device_name = cls.get_device_name() decomposition_table = cls.get_decomposition_table() options = cls.get_aoti_compile_options(compile_specs) + options.update(_written_buffer_options(edge_program)) # Move the edge_program to the target device. Routed through a hook so # backends can keep large non-parameter tensors (e.g. KV-cache buffers) diff --git a/backends/cuda/tests/test_output_stride_rearrange.py b/backends/cuda/tests/test_output_stride_rearrange.py index 557259770ec..82a5ccc43ed 100644 --- a/backends/cuda/tests/test_output_stride_rearrange.py +++ b/backends/cuda/tests/test_output_stride_rearrange.py @@ -485,10 +485,6 @@ def test_mask_runs_without_a_graph(self): self.assertIsInstance(out, torch.Tensor) torch.testing.assert_close(out, module(*args)) - @unittest.skipIf( - tuple(map(int, torch.__version__.split(".")[:2])) < (2, 15), - "with torch 2.14 the program does not keep a buffer update across calls", - ) def test_mask_retry_updates_state_once(self): # The capture call runs again without a graph, so its update to the # count must not land twice. From b66b64ea3da9ba3c1d7503a26ecba75ab828de36 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Fri, 30 Oct 2026 21:47:06 -0700 Subject: [PATCH 45/89] Test that the CUDA missing-kernel check sees the scatter fallback The scatter fallback writes its kernel call itself, so the check for missing kernels reads the kernel name from the line it writes. Only test_missing_scatter_kernel_fails_at_export covered that. It was removed when the scatter_reduce shim made its int64 scatter_add lower, so the check could stop watching the scatter fallback and no test would fail. Bring the test back, and take the kernel off the backend's list for it, so it does not depend on which shims the backend has. Test: the test passes, and fails with "RuntimeError not raised" when the check stops watching the scatter fallback. --- backends/cuda/tests/test_cuda_export.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/backends/cuda/tests/test_cuda_export.py b/backends/cuda/tests/test_cuda_export.py index 148c812fb91..8f8792c97b9 100644 --- a/backends/cuda/tests/test_cuda_export.py +++ b/backends/cuda/tests/test_cuda_export.py @@ -792,6 +792,24 @@ def test_triton_kernel_mode_off(self): "SDPA kernel export with triton_kernel_mode=OFF failed", ) + def test_missing_scatter_kernel_fails_at_export(self): + """The scatter fallback writes its kernel call directly, and the check + for missing kernels still sees it.""" + + class Int64ScatterAdd(torch.nn.Module): + def forward(self, counts: torch.Tensor, idx: torch.Tensor): + return counts.scatter_add(0, idx, torch.ones_like(idx)) + + # The backend ships this kernel, so take it off the list. + kernel = "aoti_torch_cuda_scatter_reduce_two_out" + kernels = CudaBackend.get_supported_fallback_kernels() + del kernels[kernel] + inputs = (torch.zeros(8, dtype=torch.long), torch.randint(0, 8, (32,))) + with mock.patch.object( + CudaBackend, "get_supported_fallback_kernels", return_value=kernels + ), self.assertRaises(RuntimeError) as context: + self._export_to_cuda_with_lower(Int64ScatterAdd(), inputs) + self.assertIn(kernel, str(context.exception)) @unittest.skipIf( torch.version.hip is not None, "loads through the Python runtime, which the ROCm build does not include", From 951d55b1831dc9e6d4e2918580817dc3098b1667 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Fri, 30 Oct 2026 21:51:22 -0700 Subject: [PATCH 46/89] Drop a redundant branch from the masked copy pass and describe what it accepts The sym_size branch in _only_range_checked repeated the branch after it, since every size read of the index carries a symbolic value. The docstring also said only a check on the count range keeps the index removable, while the code accepts any symbolic math on the sizes and any check that holds for every count in range. Remove the branch and describe what the code does. --- backends/cuda/passes/replace_masked_copy_with_where.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/backends/cuda/passes/replace_masked_copy_with_where.py b/backends/cuda/passes/replace_masked_copy_with_where.py index 9e8186639b2..03c89dc5f55 100644 --- a/backends/cuda/passes/replace_masked_copy_with_where.py +++ b/backends/cuda/passes/replace_masked_copy_with_where.py @@ -32,7 +32,6 @@ _INDEX = exir_ops.edge.aten.index.Tensor _INDEX_PUT = exir_ops.edge.aten.index_put.default -_SYM_SIZE = torch.ops.aten.sym_size.int _ASSERT_SCALAR = torch.ops.aten._assert_scalar.default @@ -164,8 +163,8 @@ def _only_range_checked( ) -> Optional[List[Node]]: """ The index and everything that reads it, in graph order, if all that - reads it is a check that its count, its size at ``position``, is in - [0, mask.numel()], else None. + reads it is symbolic math on its sizes and checks that hold for every + count, its size at ``position``, in [0, mask.numel()], else None. """ count = index.meta["val"].shape[position] if not isinstance(count, torch.SymInt): @@ -193,8 +192,6 @@ def _only_range_checked( bounds = bound_sympy(expr, {count.node.expr: within}) if not (bounds.is_singleton() and bounds.lower == sympy.true): return None - elif user.target == _SYM_SIZE and user.args[0] is index: - pass # Computing a symbolic value has no effect of its own, and its # readers are visited in turn. elif not isinstance(user.meta.get("val"), (torch.SymInt, torch.SymBool)): From 911db99396e5bc93e058f0bfea7d7781331c45f7 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Fri, 30 Oct 2026 22:05:29 -0700 Subject: [PATCH 47/89] Read a float64 scalar in the CUDA float64 shim test The float64 test now also multiplies by a float64 .item(). With PyTorch 2.14 the generated code calls aoti_torch_item_float64 for it, which ended the process with "undefined symbol" before that shim existed. The scale is a whole number because the generated kernel takes the scalar as an integer. --- backends/cuda/tests/test_index_shims.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/backends/cuda/tests/test_index_shims.py b/backends/cuda/tests/test_index_shims.py index 09b8afb9379..970a50e8eb5 100644 --- a/backends/cuda/tests/test_index_shims.py +++ b/backends/cuda/tests/test_index_shims.py @@ -207,11 +207,15 @@ def forward(self, x, index, src): index = torch.zeros(128, dtype=torch.long) self._run(Mean(), [(x, index, torch.ones(128, dtype=torch.int8))]) - def test_float64_sort(self): - class Sort(torch.nn.Module): - def forward(self, x): + def test_float64_values(self): + class Float64Values(torch.nn.Module): + def forward(self, x, scale): # Rows this long go to the sort shim, not to Triton. - return torch.sort(x, stable=True) + values, indices = torch.sort(x, stable=True) + return values * scale.item(), indices torch.manual_seed(0) - self._run(Sort(), [(torch.randn(4, 1000, dtype=torch.float64),)]) + x = torch.randn(4, 1000, dtype=torch.float64) + # A whole number: the generated kernel takes the scalar as an integer. + scale = torch.tensor(2.0, dtype=torch.float64) + self._run(Float64Values(), [(x, scale)]) From 4a8bc18278d9daab45b70ab3947ab5db387f05c0 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Fri, 30 Oct 2026 22:22:35 -0700 Subject: [PATCH 48/89] Use a scatter the CUDA backend lacks in the missing-kernel test test_missing_scatter_kernel_fails_at_export took the scatter_reduce kernel off a patched CudaBackend.get_supported_fallback_kernels. That broke in two places: - The CUDA unit test job collects backends/cuda/passes/tests too, and one of those tests imports the backend as backends.cuda.cuda_backend. That loads a second CudaBackend class, lowering picks it by name, and the patch never reaches it. So the export did not fail. - On ROCm the list is empty, so removing the kernel raised KeyError. An int64 scatter without a reduction still goes to an ATen kernel that the backend has no shim for, aoti_torch_cuda_scatter_src_out, and the scatter fallback writes that call directly too. Test with it, with no patching, as the test did before the scatter_reduce shim. Test: the test passes alone and in a run that collects every CUDA test, with PyTorch 2.14 and 2.16. It fails when the check stops watching the scatter fallback. --- backends/cuda/tests/test_cuda_export.py | 20 +++++++------------- 1 file changed, 7 insertions(+), 13 deletions(-) diff --git a/backends/cuda/tests/test_cuda_export.py b/backends/cuda/tests/test_cuda_export.py index 8f8792c97b9..5c420fc55cd 100644 --- a/backends/cuda/tests/test_cuda_export.py +++ b/backends/cuda/tests/test_cuda_export.py @@ -793,23 +793,17 @@ def test_triton_kernel_mode_off(self): ) def test_missing_scatter_kernel_fails_at_export(self): - """The scatter fallback writes its kernel call directly, and the check - for missing kernels still sees it.""" + """Inductor leaves an int64 scatter to an ATen kernel the CUDA shims lack.""" - class Int64ScatterAdd(torch.nn.Module): + class Int64Scatter(torch.nn.Module): def forward(self, counts: torch.Tensor, idx: torch.Tensor): - return counts.scatter_add(0, idx, torch.ones_like(idx)) + return counts.scatter(0, idx, torch.ones_like(idx)) - # The backend ships this kernel, so take it off the list. - kernel = "aoti_torch_cuda_scatter_reduce_two_out" - kernels = CudaBackend.get_supported_fallback_kernels() - del kernels[kernel] inputs = (torch.zeros(8, dtype=torch.long), torch.randint(0, 8, (32,))) - with mock.patch.object( - CudaBackend, "get_supported_fallback_kernels", return_value=kernels - ), self.assertRaises(RuntimeError) as context: - self._export_to_cuda_with_lower(Int64ScatterAdd(), inputs) - self.assertIn(kernel, str(context.exception)) + with self.assertRaises(RuntimeError) as context: + self._export_to_cuda_with_lower(Int64Scatter(), inputs) + self.assertIn("aoti_torch_cuda_scatter_src_out", str(context.exception)) + @unittest.skipIf( torch.version.hip is not None, "loads through the Python runtime, which the ROCm build does not include", From 900d0ee6213bbd91e19055222a8407f59f8681ef Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sat, 31 Oct 2026 00:16:55 -0700 Subject: [PATCH 49/89] Return a contiguous output from the CUDA varlen_attn decomposition The varlen_attn decomposition in the CUDA backend returned einsum's result, a permuted view of a (heads, tokens, dim) tensor. The op allocates its output like the query, which is contiguous in the usual packed (tokens, heads, dim) layout, and the rest of the program was traced against that. So a program that reshapes the attention output, like the usual o_proj(out.reshape(T, -1)), failed to lower: ValueError: Cannot view a tensor with shape torch.Size([32, 4, 64]) and strides (64, 2048, 1) as a tensor with shape (32, 256)! Return a contiguous tensor. Test: test_varlen_attn now flattens the output it lowers. It fails with the error above without this change and passes with it. --- backends/cuda/cuda_backend.py | 4 +++- backends/cuda/tests/test_cuda_export.py | 3 ++- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index 2f1fbfdb11b..c9c77040c83 100644 --- a/backends/cuda/cuda_backend.py +++ b/backends/cuda/cuda_backend.py @@ -636,7 +636,9 @@ def sequence_ends(cu: torch.Tensor, total: int): probs = torch.exp(scores - lse[..., None]) out = torch.einsum("hqk,khd->qhd", probs, value.float()).to(query.dtype) rng_state = torch.zeros(2, dtype=torch.uint64, device=query.device) - return out, lse, rng_state + # einsum gives a permuted view, but the op's output is contiguous, and + # later views of it in the program rely on that. + return out.contiguous(), lse, rng_state def _linear_kernel(x: torch.Tensor) -> torch.Tensor: diff --git a/backends/cuda/tests/test_cuda_export.py b/backends/cuda/tests/test_cuda_export.py index 5c420fc55cd..80a69ac8186 100644 --- a/backends/cuda/tests/test_cuda_export.py +++ b/backends/cuda/tests/test_cuda_export.py @@ -495,7 +495,8 @@ def test_varlen_attn(self): class VarlenAttention(torch.nn.Module): def forward(self, q, k, v, cu_seq): - return varlen_attn(q, k, v, cu_seq, cu_seq, 16, 16) + # Flattening views the output, which needs the op's layout. + return varlen_attn(q, k, v, cu_seq, cu_seq, 16, 16).flatten(1) q, k, v = ( torch.randn(32, 4, 64, dtype=torch.bfloat16, device="cuda") From 5c5d2626fe5c7d1eb745a3e4ef04238e6c05f07b Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sat, 31 Oct 2026 00:16:56 -0700 Subject: [PATCH 50/89] Refuse varlen_attn windows on which Flash attention versions differ Flash attention 2, which PyTorch's varlen_attn runs for sliding windows by default, changes some windows based on max_k, the longest key sequence. It drops a window side of max_k or more, and it gives a window with only a left side a right side of max_k. Flash attention 3, which PyTorch can also run, keeps the window as given. The two differ only when a query sequence is longer than max_k + 1. The varlen_attn decomposition in the CUDA backend kept the window as given, so it differed from the default there. For example, window (3, -1) with 12 queries and 10 keys differed by 1.1. So lower a window other than (left, 0) only when max_q <= max_k + 1 is known at export time, where every version gives the same result, and raise a clear error otherwise. Self-attention usually passes the same value for both, also when it comes from .item(), so it lowers. No window, a causal mask and a window (left, 0) give the same result in every version, so they always lower. Test: test_varlen_attn compares a window (4, 2) in self-attention with the op, and checks that window (3, -1) raises when a query sequence can be longer than max_k + 1. --- backends/cuda/cuda_backend.py | 20 +++++++++++++++++--- backends/cuda/tests/test_cuda_export.py | 7 ++++++- 2 files changed, 23 insertions(+), 4 deletions(-) diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index c9c77040c83..1ac94221bc6 100644 --- a/backends/cuda/cuda_backend.py +++ b/backends/cuda/cuda_backend.py @@ -64,6 +64,7 @@ from torch._decomp import decomposition_table from torch._inductor.decomposition import conv1d_to_conv2d from torch._prims_common import suggest_memory_format +from torch.fx.experimental.symbolic_shapes import statically_known_true from torch.nn.attention import SDPBackend # --------------------------------------------------------------------------- @@ -602,6 +603,22 @@ def _varlen_attn( "The CUDA backend lowers varlen_attn only with cu_seq_k and without " "seqused_k or block_table." ) + left, right = window_size if window_size is not None else (-1, -1) + if is_causal: + right = 0 + # Flash attention 2, PyTorch's default, drops a window side of max_k or + # more and gives a left side alone a right side of max_k, while later + # versions keep the window. That only matters for a query sequence + # longer than max_k + 1. + if (right > 0 or (right < 0 and left >= 0)) and not statically_known_true( + max_q <= max_k + 1 + ): + raise NotImplementedError( + "The CUDA backend lowers a varlen_attn window other than (left, 0) " + "only when max_q <= max_k + 1 is known, since Flash attention " + "versions differ on it for longer query sequences. Got " + f"window_size={window_size}, max_q={max_q}, max_k={max_k}." + ) def sequence_ends(cu: torch.Tensor, total: int): # Each token's index in cu of its sequence's end, and its (negative) @@ -614,9 +631,6 @@ def sequence_ends(cu: torch.Tensor, total: int): end_q, from_end_q = sequence_ends(cu_seq_q, query.size(0)) end_k, from_end_k = sequence_ends(cu_seq_k, key.size(0)) allowed = end_q[:, None] == end_k[None, :] - left, right = window_size if window_size is not None else (-1, -1) - if is_causal: - right = 0 offset = from_end_k[None, :] - from_end_q[:, None] if left >= 0: allowed = allowed & (offset >= -left) diff --git a/backends/cuda/tests/test_cuda_export.py b/backends/cuda/tests/test_cuda_export.py index 80a69ac8186..cea9c76b1f5 100644 --- a/backends/cuda/tests/test_cuda_export.py +++ b/backends/cuda/tests/test_cuda_export.py @@ -507,7 +507,7 @@ def forward(self, q, k, v, cu_seq): # Causal masks over query sequences shorter and longer than their key # sequences (some queries then see no key), grouped-query attention - # and an empty sequence. + # with an empty sequence, and a window in self-attention. op = torch.ops.torch_attn._varlen_attn.default short = torch.tensor([0, 4, 16], dtype=torch.int32, device="cuda") empty = torch.tensor([0, 16, 16, 32], dtype=torch.int32, device="cuda") @@ -515,11 +515,16 @@ def forward(self, q, k, v, cu_seq): (q[:16], k, v, short, cu_seq, 12, 16, True, None, [3, -1]), (q, k[:16], v[:16], cu_seq, short, 16, 12, True, None, [-1, 0]), (q, k[:, :2], v[:, :2], empty, empty, 16, 16, False, 0.5, None, True), + (q, k, v, cu_seq, cu_seq, 16, 16, False, None, [4, 2]), ] table = CudaBackend.get_decomposition_table() for args in cases: for actual, expected in zip(table[op](*args)[:2], op(*args)[:2]): torch.testing.assert_close(actual, expected, atol=1e-2, rtol=1e-2) + # Flash attention versions differ on this window, since a query sequence + # can be longer than max_k + 1. + with self.assertRaisesRegex(NotImplementedError, "window"): + table[op](q, k[:16], v[:16], cu_seq, short, 16, 12, False, None, [3, -1]) def test_conv2d(self): """Test CUDA export for 2D convolution.""" From e9a232f7114a6d0e848d66aa485f89994f400252 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sat, 31 Oct 2026 00:16:56 -0700 Subject: [PATCH 51/89] Keep NaN values outside the key sequences out of varlen_attn outputs A packed key buffer can hold rows outside its sequences, before cu_seq_k[0] or past cu_seq_k[-1], for example when its size is fixed. The varlen_attn decomposition in the CUDA backend multiplies every query's attention weights with every value row, and a masked weight is exactly 0. Since NaN * 0 is NaN, a NaN or inf in one of those rows turned the outputs of all the sequences into NaN. Flash attention never reads those rows. Zero those value rows first. The antialiased resize masks the taps outside its window for the same reason. A NaN or inf in the values of one sequence still reaches the outputs of the others. Masking the full product of weights and values would keep that local too, but that product is head-dim times larger than the scores. The docstring now says so. The comment that said tokens past the last sequence are in no sequence is also fixed: the tokens before cu[0] and past cu[-1] form two more sequences. Test: one test_varlen_attn case now has NaN values before and after its key sequences. It fails without this change and passes with it. --- backends/cuda/cuda_backend.py | 12 +++++++++--- backends/cuda/tests/test_cuda_export.py | 10 +++++++--- 2 files changed, 16 insertions(+), 6 deletions(-) diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index 1ac94221bc6..c7185478c72 100644 --- a/backends/cuda/cuda_backend.py +++ b/backends/cuda/cuda_backend.py @@ -596,7 +596,8 @@ def _varlen_attn( for every pair of packed tokens, 1 GiB for 16 heads and 4096 tokens. As in Flash attention, a causal mask or a window lines up the ends of each query and key sequence, and a query with no key to attend to gets zeros and an - infinite log-sum-exp. + infinite log-sum-exp. Unlike in Flash attention, a NaN or inf in one + sequence's values also reaches the other sequences. """ if cu_seq_k is None or seqused_k is not None or block_table is not None: raise NotImplementedError( @@ -622,8 +623,9 @@ def _varlen_attn( def sequence_ends(cu: torch.Tensor, total: int): # Each token's index in cu of its sequence's end, and its (negative) - # distance from that end. Tokens past cu[-1] are in no sequence, and - # the clamp keeps their lookup in range. + # distance from that end. Tokens before cu[0] and past cu[-1] get the + # index 0 and cu.numel(), so they form two more sequences, and the + # clamp keeps the lookup in range. token = torch.arange(total, dtype=cu.dtype, device=cu.device) end = torch.searchsorted(cu, token, right=True) return end, token - cu[end.clamp(max=cu.numel() - 1)] @@ -636,6 +638,10 @@ def sequence_ends(cu: torch.Tensor, total: int): allowed = allowed & (offset >= -left) if right >= 0: allowed = allowed & (offset <= right) + # Masked keys get zero weight, but NaN * 0 is NaN. Zero the values outside + # the key sequences so that a NaN or inf there stays out of them. + outside = (end_k == 0) | (end_k == cu_seq_k.numel()) + value = value.masked_fill(outside[:, None, None], 0) if enable_gqa: key, value = ( diff --git a/backends/cuda/tests/test_cuda_export.py b/backends/cuda/tests/test_cuda_export.py index cea9c76b1f5..7a1b35d8b5b 100644 --- a/backends/cuda/tests/test_cuda_export.py +++ b/backends/cuda/tests/test_cuda_export.py @@ -506,14 +506,18 @@ def forward(self, q, k, v, cu_seq): self._export_to_cuda_with_lower(VarlenAttention(), (q, k, v, cu_seq)) # Causal masks over query sequences shorter and longer than their key - # sequences (some queries then see no key), grouped-query attention - # with an empty sequence, and a window in self-attention. + # sequences (some queries then see no key, and NaN values outside the + # key sequences must stay out), grouped-query attention with an empty + # sequence, and a window in self-attention. op = torch.ops.torch_attn._varlen_attn.default short = torch.tensor([0, 4, 16], dtype=torch.int32, device="cuda") empty = torch.tensor([0, 16, 16, 32], dtype=torch.int32, device="cuda") + shifted = torch.tensor([2, 6, 16], dtype=torch.int32, device="cuda") + nan_outside = v.clone() + nan_outside[:2] = nan_outside[16:] = float("nan") cases = [ (q[:16], k, v, short, cu_seq, 12, 16, True, None, [3, -1]), - (q, k[:16], v[:16], cu_seq, short, 16, 12, True, None, [-1, 0]), + (q, k, nan_outside, cu_seq, shifted, 16, 12, True, None, [-1, 0]), (q, k[:, :2], v[:, :2], empty, empty, 16, 16, False, 0.5, None, True), (q, k, v, cu_seq, cu_seq, 16, 16, False, None, [4, 2]), ] From 96806130f3da541273508ecd6118578fb688d8df Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Fri, 30 Oct 2026 23:24:07 -0700 Subject: [PATCH 52/89] Correct the comments on the AOTI folding fix and the scatter kernel test Inductor in PyTorch 2.14 can fold the update of a buffer only when the buffer has one element, such as a counter, and not always then: a counter updated with calls[0] += 1 keeps its updates, and so does a buffer of four zeros that the program adds to. The comment on the AOTI backend's change said Inductor replaces any buffer whose elements are all equal. Say what it can do, link the upstream fix, and note that the change can go once the pinned PyTorch is 2.15 or newer. The missing-kernel test lost the sentence that says why it exists: the scatter fallback writes its kernel call itself, and the check has to see it. Put it back, so that whoever adds the next shim sees what the test guards. The counter module of the CUDA graph tests now says that its buffer has one element, since that is the case the folding fix needs covered. Checked with an AOTInductor program on PyTorch 2.14 that adds 1 to a buffer on every call: a 1-element and a 0-dim buffer return 1, 1, 1, and a 4-element buffer returns 1, 2, 3. With joint_graph_constant_folding off, all of them return 1, 2, 3. --- backends/aoti/aoti_backend.py | 9 +++++---- backends/cuda/tests/test_cuda_export.py | 4 +++- backends/cuda/tests/test_output_stride_rearrange.py | 3 ++- 3 files changed, 10 insertions(+), 6 deletions(-) diff --git a/backends/aoti/aoti_backend.py b/backends/aoti/aoti_backend.py index 8b72f3d4598..949264459d9 100644 --- a/backends/aoti/aoti_backend.py +++ b/backends/aoti/aoti_backend.py @@ -57,10 +57,11 @@ def _written_buffer_options(program: ExportedProgram) -> Dict[str, Any]: """Inductor options that keep the buffers a program writes out of constant folding. - Before torch 2.15, Inductor replaces a buffer whose elements are all equal, - such as a counter that starts at zero, with a constant even when the - program writes it, so every call would start from the exported value - (pytorch/pytorch#185077). + Before torch 2.15, Inductor can fold the update of a one-element buffer, + such as a counter, to its value at export, so every call writes that same + value (https://github.com/pytorch/pytorch/pull/185077). + + TODO: remove once the pinned torch is 2.15 or newer. """ if program.graph_signature.buffers_to_mutate: return {"joint_graph_constant_folding": False} diff --git a/backends/cuda/tests/test_cuda_export.py b/backends/cuda/tests/test_cuda_export.py index 7a1b35d8b5b..8a126a12dbb 100644 --- a/backends/cuda/tests/test_cuda_export.py +++ b/backends/cuda/tests/test_cuda_export.py @@ -803,7 +803,9 @@ def test_triton_kernel_mode_off(self): ) def test_missing_scatter_kernel_fails_at_export(self): - """Inductor leaves an int64 scatter to an ATen kernel the CUDA shims lack.""" + """The scatter fallback writes its kernel call itself, and the check for + missing kernels still sees it. Inductor leaves an int64 scatter to that + fallback, and the CUDA shims lack its kernel.""" class Int64Scatter(torch.nn.Module): def forward(self, counts: torch.Tensor, idx: torch.Tensor): diff --git a/backends/cuda/tests/test_output_stride_rearrange.py b/backends/cuda/tests/test_output_stride_rearrange.py index 82a5ccc43ed..29a136e4902 100644 --- a/backends/cuda/tests/test_output_stride_rearrange.py +++ b/backends/cuda/tests/test_output_stride_rearrange.py @@ -294,7 +294,8 @@ def forward(self, x, mask): class _MaskedCounter(nn.Module): - """Counts its calls in a device buffer, and reads a mask's count.""" + """Counts its calls in a device buffer, and reads a mask's count. The + buffer has one element, the case torch 2.14 can fold into a constant.""" def __init__(self): super().__init__() From c3a6f91df75f9dce925bf7a3a18739027350d43b Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sat, 31 Oct 2026 01:09:07 -0700 Subject: [PATCH 53/89] Test the count check guard in the masked copy pass The checked case in test_other_writes_are_kept did not test what it said. Export attaches torch._check(picked.shape[0] >= 1) to the first x[mask] in the program, which was another write's index, so the pass rewrote the checked copy and removed its index, and deleting the guard that keeps an index for a check that does not hold for every count failed no test. Move the checked copy first. Now its index stays, and the test fails without the guard. --- .../tests/test_replace_masked_copy_with_where.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/backends/cuda/passes/tests/test_replace_masked_copy_with_where.py b/backends/cuda/passes/tests/test_replace_masked_copy_with_where.py index 5062d27cfa3..41b7fd12c01 100644 --- a/backends/cuda/passes/tests/test_replace_masked_copy_with_where.py +++ b/backends/cuda/passes/tests/test_replace_masked_copy_with_where.py @@ -86,14 +86,15 @@ def compute(m): def test_other_writes_are_kept(self): class Kept(torch.nn.Module): def forward(self, x, mask, idx): - scaled = torch.zeros_like(x) - # The values are not the mask's picks of a tensor. - scaled[mask] = x[mask] * 2 checked = torch.zeros_like(x) + # First, since export attaches the check to the first x[mask]. picked = x[mask] # A check on the count that does not hold for every mask. torch._check(picked.shape[0] >= 1) checked[mask] = picked + scaled = torch.zeros_like(x) + # The values are not the mask's picks of a tensor. + scaled[mask] = x[mask] * 2 summed = torch.ones_like(x) summed.index_put_((mask,), x[mask], accumulate=True) other = torch.zeros_like(x) @@ -116,8 +117,8 @@ def forward(self, x, mask, idx): targets = _targets(program.graph_module) self.assertEqual(targets.count(exir_ops.edge.aten.index_put.default), 4) self.assertEqual(targets.count(exir_ops.edge.aten.where.self), 2) - # The index stays for the checked count and for the shared picks. - self.assertEqual(targets.count(exir_ops.edge.aten.index.Tensor), 5) + # The checked count and the shared picks keep their index. + self.assertEqual(targets.count(exir_ops.edge.aten.index.Tensor), 6) for got, expected in zip( program.module()(x, mask, idx), Kept()(x, mask, idx), strict=True ): From a4767c3be26138316507bf9659d7d970ef6dceb2 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sat, 31 Oct 2026 01:09:08 -0700 Subject: [PATCH 54/89] Correct two comments in the CUDA index shim tests The float64 test used a whole-number scale and said the generated kernel takes the scalar as an integer, as if that were how it should work. It is an AOTInductor bug: the Triton launcher takes a float read with .item() as int64_t, so 1.5 becomes 1. Say so. The scatter_reduce docstring said only an int64 sum goes to the shim. A bfloat16 sum does too on GPUs below sm_90. --- backends/cuda/tests/test_index_shims.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/backends/cuda/tests/test_index_shims.py b/backends/cuda/tests/test_index_shims.py index 970a50e8eb5..a1f325c0520 100644 --- a/backends/cuda/tests/test_index_shims.py +++ b/backends/cuda/tests/test_index_shims.py @@ -168,8 +168,8 @@ def forward(self, x, idx, values): def test_scatter_reduce(self): """Inductor adds atomically but has no atomic max, min or product, so - those reductions, and the mean, go to the shim. So does an int64 sum, - which Inductor cannot add atomically.""" + those reductions, and the mean, go to the shim. So does a sum it cannot + add atomically: int64, or bfloat16 below sm_90.""" class ScatterReduce(torch.nn.Module): def forward(self, x, index, src): @@ -216,6 +216,7 @@ def forward(self, x, scale): torch.manual_seed(0) x = torch.randn(4, 1000, dtype=torch.float64) - # A whole number: the generated kernel takes the scalar as an integer. + # A whole number works around an AOTInductor bug: its Triton launcher + # takes a float read with .item() as int64_t, so 1.5 would become 1. scale = torch.tensor(2.0, dtype=torch.float64) self._run(Float64Values(), [(x, scale)]) From 89c4d5f6abe237b25d3ff9439f5a4f03b3e028d7 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sat, 31 Oct 2026 01:09:09 -0700 Subject: [PATCH 55/89] Test a scatter mean count that wraps to zero test_scatter_mean_counts_in_its_dtype counted 128 int8 terms, which wraps to a negative count but never to zero, so deleting the branch that turns a zero count into one failed no test. Add a second element with 255 terms: with self it counts 256, which wraps to 0 in int8. --- backends/cuda/tests/test_index_shims.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/backends/cuda/tests/test_index_shims.py b/backends/cuda/tests/test_index_shims.py index a1f325c0520..5b4c55da0d0 100644 --- a/backends/cuda/tests/test_index_shims.py +++ b/backends/cuda/tests/test_index_shims.py @@ -204,8 +204,9 @@ def forward(self, x, index, src): ) x = torch.zeros(2, dtype=torch.int8) - index = torch.zeros(128, dtype=torch.long) - self._run(Mean(), [(x, index, torch.ones(128, dtype=torch.int8))]) + # With self, the second element counts 256, which wraps to 0. + index = torch.tensor([0] * 128 + [1] * 255) + self._run(Mean(), [(x, index, torch.ones(383, dtype=torch.int8))]) def test_float64_values(self): class Float64Values(torch.nn.Module): From d6f4a34b5124fbf9c8c172cf78e60d95b03d9d54 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sat, 31 Oct 2026 01:09:09 -0700 Subject: [PATCH 56/89] Say the CUDA random state setup failed, not its allocation The rand and randint shims reported "failed to allocate the RNG state" whenever the setup failed, also when the seed copy or the stream sync was the step that failed. Say "failed to set up" instead. --- backends/cuda/runtime/shims/rand.cu | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/backends/cuda/runtime/shims/rand.cu b/backends/cuda/runtime/shims/rand.cu index 5f53838b672..85a64e1ffd3 100644 --- a/backends/cuda/runtime/shims/rand.cu +++ b/backends/cuda/runtime/shims/rand.cu @@ -245,7 +245,7 @@ AOTITorchError aoti_torch_cuda_rand( ET_CHECK_OR_RETURN_ERROR( ensure_rng_init(stream), MemoryAllocationFailed, - "aoti_torch_cuda_rand: failed to allocate the RNG state"); + "aoti_torch_cuda_rand: failed to set up the RNG state"); constexpr int kThreads = 256; int blocks = static_cast((numel + kThreads - 1) / kThreads); @@ -308,7 +308,7 @@ AOTITorchError aoti_torch_cuda_randint_low_out( ET_CHECK_OR_RETURN_ERROR( ensure_rng_init(stream), MemoryAllocationFailed, - "aoti_torch_cuda_randint_low_out: failed to allocate the RNG state"); + "aoti_torch_cuda_randint_low_out: failed to set up the RNG state"); int64_t range = high - low; int64_t* out_data = static_cast(out->data_ptr()); From 1f884d95691d86a6f03838178a6107d8f35364f5 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sat, 31 Oct 2026 01:09:09 -0700 Subject: [PATCH 57/89] Add the float16 scalar read to the CUDA shims A float16 .item() calls aoti_torch_item_float16, which did not exist. The runtime loads the library lazily, so the program lowered, then the first execute ended the process with "undefined symbol: aoti_torch_item_float16". A bfloat16 .item() already worked. Add the entry point and a Half case for reading a single element, give the SlimTensor types a Half alias next to BFloat16, and list the name in the Windows import library. The float64 values test now runs a float16 case too; with the old shims it ends the process with that error. --- backends/aoti/slim/c10/core/ScalarType.h | 4 +++- backends/cuda/runtime/aoti_cuda_shims.lib | Bin 57574 -> 58386 bytes backends/cuda/runtime/shims/memory.cpp | 4 ++++ backends/cuda/runtime/shims/memory.h | 4 ++++ backends/cuda/tests/test_index_shims.py | 17 ++++++++++------- 5 files changed, 21 insertions(+), 8 deletions(-) diff --git a/backends/aoti/slim/c10/core/ScalarType.h b/backends/aoti/slim/c10/core/ScalarType.h index 5bc06076885..0181ff9e8e8 100644 --- a/backends/aoti/slim/c10/core/ScalarType.h +++ b/backends/aoti/slim/c10/core/ScalarType.h @@ -13,12 +13,14 @@ #include #include +#include #include namespace executorch::backends::aoti::slim::c10 { -// Import BFloat16 from ExecuTorch's portable_type +// Import BFloat16 and Half from ExecuTorch's portable_type using BFloat16 = ::executorch::runtime::etensor::BFloat16; +using Half = ::executorch::runtime::etensor::Half; /// Enum representing the scalar type (dtype) of tensor elements. /// Note: Enum values must match PyTorch's c10::ScalarType for compatibility. diff --git a/backends/cuda/runtime/aoti_cuda_shims.lib b/backends/cuda/runtime/aoti_cuda_shims.lib index c6466b68f0ef2d27f0b4f558347b4d50753f1872..78888032df12ed427e0b744a3b757990fb9607d4 100644 GIT binary patch delta 851 zcmXxgZAjBm6bA6;{!J@Eox=zQnh{7TZEn6a!^|2&&>QB!2%#7Sks!uGdV@JIiYSuZ ztf&m3T0~aZ5rU!@MaQz5rZzQG5Ty?x`8E0w-Q%GPzc2Tm=Wx&Y(U%g5r?gju%FDd@ z%jZZsfGdFG50D!Z`P~59a7q;T0k$bq6vYAN>=o`NfR*kMTdD!J)g!hQ0<1bu_|igv zFAbDNnS!0K6rps0g(pQ@7r-u>8g|)C$*yh`?GAuNeu@sWQFhB@*qwj~-7^8! z^-J7;0kG(4@%SLXx=Y1#3t&B_G3y-{FU^$fmH7;NV?M*)r-}X~zy?hXi?=NE;bRq+ zB}^e3Nr==leCp~xJopwZAdTBH@VoUuW(zU@nI@f<&QB9v$>AizaTDn%JbRuRa+d!qF zocK*e^D&R2nV}nso_z2q{qVDko`nx6T`;_cf}h+7ops9cLHA_m(xZtgr8}ox olrv=-=FMbK@LQqM^|PC4Xxgt-4XvB;DyyPQpDZe(w%LI74GWD_DEWeB25=aLev&gP!ywRv>uBfGNTa$mGl_9 zRMN{R7`8&nBLxygqv3MqYq_bJtcbn^eF-A^&^?|DKe!zJ=W@;1l( z`Ey_)fKI^q0H{79YIXo@d6iiC2Vk{*V$D&2xpoV$17LN(gzr1R{6(>81Yivr(Ugt@ zO=)23WziA=SZj;eUJJ0Dl_F%CS=&porypQ@6Jq~yfQ7^2&^CZYR*0zCksbRaV)1)G z%s6G5*y%>mZnj_@WpTb2U>DAbivfULG8uNoY{{+;iyLnMmb@-*njd9drib0Niuirg zz#b;W<2Hb$mWZBz0DC$mdQ$*~&!2_`30x|UsV4v@Zxm`L@?WClo!=|IzWpr>+v&G`^(-ViQl72g%rRe1Q zdPUQjE=5}gHz?hcb<>G_Na@vs4!tyQCC89kQG4!|qSEj>6Uv)g_o!Ft$`McBcu?ui z(ZzIl!lU%~n44mGQwdHubj)rg&!k)F+XWY0i<$?kCT;Yy6i|9!(L>!{Qz@5h`mfzV qg_+Gt*G;>qVK$^xGwClrrR}qBYFrgmDnPB}dZjitem(), ret_value); case ScalarType::Double: return narrow_to(tensor->item(), ret_value); + case ScalarType::Half: + return narrow_to( + static_cast(tensor->item()), ret_value); case ScalarType::BFloat16: return narrow_to( static_cast(tensor->item()), ret_value); @@ -473,6 +476,7 @@ ET_CUDA_DEFINE_ITEM_SHIM(int32, int32_t) ET_CUDA_DEFINE_ITEM_SHIM(int64, int64_t) ET_CUDA_DEFINE_ITEM_SHIM(float32, float) ET_CUDA_DEFINE_ITEM_SHIM(float64, double) +ET_CUDA_DEFINE_ITEM_SHIM(float16, c10::Half) ET_CUDA_DEFINE_ITEM_SHIM(bfloat16, c10::BFloat16) ET_CUDA_DEFINE_ITEM_SHIM(bool, bool) diff --git a/backends/cuda/runtime/shims/memory.h b/backends/cuda/runtime/shims/memory.h index 8e4b87517c9..ef51ff5eef0 100644 --- a/backends/cuda/runtime/shims/memory.h +++ b/backends/cuda/runtime/shims/memory.h @@ -260,6 +260,10 @@ aoti_torch_item_float32(SlimTensor* tensor, float* ret_value); AOTI_SHIM_EXPORT AOTITorchError aoti_torch_item_float64(SlimTensor* tensor, double* ret_value); +/// See aoti_torch_item_uint8. +AOTI_SHIM_EXPORT AOTITorchError +aoti_torch_item_float16(SlimTensor* tensor, c10::Half* ret_value); + /// See aoti_torch_item_uint8. AOTI_SHIM_EXPORT AOTITorchError aoti_torch_item_bfloat16(SlimTensor* tensor, c10::BFloat16* ret_value); diff --git a/backends/cuda/tests/test_index_shims.py b/backends/cuda/tests/test_index_shims.py index 5b4c55da0d0..f3e8e6c4141 100644 --- a/backends/cuda/tests/test_index_shims.py +++ b/backends/cuda/tests/test_index_shims.py @@ -208,16 +208,19 @@ def forward(self, x, index, src): index = torch.tensor([0] * 128 + [1] * 255) self._run(Mean(), [(x, index, torch.ones(383, dtype=torch.int8))]) - def test_float64_values(self): - class Float64Values(torch.nn.Module): + def test_float64_and_float16_values(self): + class Values(torch.nn.Module): def forward(self, x, scale): # Rows this long go to the sort shim, not to Triton. values, indices = torch.sort(x, stable=True) return values * scale.item(), indices torch.manual_seed(0) - x = torch.randn(4, 1000, dtype=torch.float64) - # A whole number works around an AOTInductor bug: its Triton launcher - # takes a float read with .item() as int64_t, so 1.5 would become 1. - scale = torch.tensor(2.0, dtype=torch.float64) - self._run(Float64Values(), [(x, scale)]) + for dtype in (torch.float64, torch.float16): + with self.subTest(dtype=dtype): + x = torch.randn(4, 1000, dtype=dtype) + # A whole number works around an AOTInductor bug: its Triton + # launcher takes a float read with .item() as int64_t, so 1.5 + # would become 1. + scale = torch.tensor(2.0, dtype=dtype) + self._run(Values(), [(x, scale)]) From 3cd283f3001935cc470027b18dcd217aff13f640 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sat, 31 Oct 2026 01:13:25 -0700 Subject: [PATCH 58/89] Make two varlen_attn decomposition comments exact The comment on sequence_ends in the CUDA backend's varlen_attn decomposition called every distance negative. The clamp measures the tokens past cu[-1] from cu[-1] instead, so theirs are 0 or more. The comment on the contiguous output said the op's output is contiguous. The op lays it out in the query's dimension order, so it is contiguous for a query in the usual (tokens, heads, dim) order, even a slice of a packed tensor. Only comments change. --- backends/cuda/cuda_backend.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index c7185478c72..72bfa03c985 100644 --- a/backends/cuda/cuda_backend.py +++ b/backends/cuda/cuda_backend.py @@ -623,9 +623,9 @@ def _varlen_attn( def sequence_ends(cu: torch.Tensor, total: int): # Each token's index in cu of its sequence's end, and its (negative) - # distance from that end. Tokens before cu[0] and past cu[-1] get the - # index 0 and cu.numel(), so they form two more sequences, and the - # clamp keeps the lookup in range. + # distance from that end. Tokens before cu[0] get the index 0, and + # tokens past cu[-1] get cu.numel() and, through the clamp, their + # distance from cu[-1]. So each group forms one more sequence. token = torch.arange(total, dtype=cu.dtype, device=cu.device) end = torch.searchsorted(cu, token, right=True) return end, token - cu[end.clamp(max=cu.numel() - 1)] @@ -656,8 +656,10 @@ def sequence_ends(cu: torch.Tensor, total: int): probs = torch.exp(scores - lse[..., None]) out = torch.einsum("hqk,khd->qhd", probs, value.float()).to(query.dtype) rng_state = torch.zeros(2, dtype=torch.uint64, device=query.device) - # einsum gives a permuted view, but the op's output is contiguous, and - # later views of it in the program rely on that. + # einsum gives a permuted view. The op lays out its output in the + # query's dimension order, so a query in the usual (T, H, D) order, even + # a slice of a packed tensor, gives a contiguous output. Later views of + # it rely on that. return out.contiguous(), lse, rng_state From eebec99b9807a133ebe478da930763f605151559 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sat, 31 Oct 2026 08:44:50 -0700 Subject: [PATCH 59/89] Say that float accumulation in the CUDA index_put shim can still round away small terms The comment said half and bfloat16 values summed in float lose no small terms. Float still rounds, so a term far smaller than the running sum can vanish, for example many 2^-24 terms added one at a time to 1.0. Comment only. --- backends/cuda/runtime/shims/index.h | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/backends/cuda/runtime/shims/index.h b/backends/cuda/runtime/shims/index.h index 9da350ad2e3..15b7a043666 100644 --- a/backends/cuda/runtime/shims/index.h +++ b/backends/cuda/runtime/shims/index.h @@ -52,8 +52,9 @@ AOTI_SHIM_EXPORT AOTITorchError aoti_torch_cuda_index_Tensor( * index that is a mask, for accumulation into a dtype its generated code * cannot add atomically, and in deterministic mode. Writes that land on the * same element are applied in index order, and half and bfloat16 values are - * summed in float, so the result is deterministic and small terms are not - * lost. A mask makes this call wait for the stream. Indices are checked as + * summed in float, so the result is deterministic and less is lost to + * rounding, though a term far smaller than the running sum can still vanish. + * A mask makes this call wait for the stream. Indices are checked as * for aoti_torch_cuda_index_Tensor. * * @param out Tensor to write (CUDA device, up to 16 dims, as for From ba1c036da7ae9761d9ecf266012ee8585cce37be Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sat, 31 Oct 2026 08:48:41 -0700 Subject: [PATCH 60/89] Say what loading CUDA weights waits for in the stream docs The runtime.cuda docstring and the pybindings README said that loading a method's CUDA weights waits for the whole GPU. Only weights an older export keeps in one blob do that (cudaDeviceSynchronize). Current exports copy each weight with a plain cudaMemcpy, which waits for the legacy default stream and blocking streams but not for a non-blocking stream. Measured with a 2 s kernel on a non-blocking stream: load_method took 0.01 s. Docs only. --- extension/pybindings/README.md | 6 ++++-- runtime/cuda.py | 9 +++++---- 2 files changed, 9 insertions(+), 6 deletions(-) diff --git a/extension/pybindings/README.md b/extension/pybindings/README.md index f8038474c1b..c665c7657c2 100644 --- a/extension/pybindings/README.md +++ b/extension/pybindings/README.md @@ -153,8 +153,10 @@ only host inputs and outputs waits depends on where it runs: the device, for example delegates running on another thread. In both cases, the CUDA delegate copies an output whose memory layout has to -change through the legacy default stream, and loading a method's CUDA weights, -which a first call may do, waits for the whole device. +change through the legacy default stream. Loading a method's CUDA weights, which +a first call may do, copies each one with a plain `cudaMemcpy`, so it waits for +the legacy default stream and for blocking streams; weights an older export +keeps in one blob wait for the whole device. The waits below, for CUDA inputs and outputs, are for the whole device (`cudaDeviceSynchronize`), not for one stream, so they also wait for unrelated diff --git a/runtime/cuda.py b/runtime/cuda.py index 9718cc3de05..6bdc02f0837 100644 --- a/runtime/cuda.py +++ b/runtime/cuda.py @@ -27,10 +27,11 @@ blocking ``s``, such as a ``cupy.cuda.Stream()`` with default flags, also waits for work on the legacy default stream. Some steps still wait for more: the CUDA delegate copies an output whose memory layout has to change through -the legacy default stream, and loading a method's CUDA weights waits for the -whole GPU. A call with GPU inputs or outputs, and -``ExecuTorchResult.to_host()``, also wait for the whole GPU, so GPU data -written on other streams stays safe. +the legacy default stream, and loading a method's CUDA weights copies each one +with a plain ``cudaMemcpy``, which waits for the legacy default stream and for +blocking streams. Weights an older export keeps in one blob wait for the whole +GPU. A call with GPU inputs or outputs, and ``ExecuTorchResult.to_host()``, +also wait for the whole GPU, so GPU data written on other streams stays safe. Limits for now: a stream known to be on a CUDA device other than 0 is refused, so is the legacy default stream (handle 0, or 1 for ``cudaStreamLegacy``), and From f342450ff0738c48bdae5bb7186863b14a351cbd Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sat, 31 Oct 2026 08:52:01 -0700 Subject: [PATCH 61/89] Index the CUDA mask test with a negative int32 column No test passed a negative or int32 index to the CUDA index shim, so a shim that dropped the wrap or read int32 indices as int64 still passed. The mask test's column index is now -4 as int32, which eager PyTorch treats as column 2. --- backends/cuda/tests/test_index_shims.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/backends/cuda/tests/test_index_shims.py b/backends/cuda/tests/test_index_shims.py index f3e8e6c4141..6a4922cd663 100644 --- a/backends/cuda/tests/test_index_shims.py +++ b/backends/cuda/tests/test_index_shims.py @@ -89,7 +89,8 @@ def forward(self, x, mask, column, idx, counts): return out, picked, rows, counts, wide x = torch.randn(4, 5, 6) - column = torch.tensor([2]) + # Negative and int32: the shims wrap the index and read it as int32. + column = torch.tensor([-4], dtype=torch.int32) idx = torch.tensor([0, 3, 3, 1, 3]) # Nonzero counts check that repeated writes add to the old values. counts = torch.arange(8) * 5 From a9fc94a0ae227bafa1716b21609704569a7c9b95 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sat, 31 Oct 2026 10:27:34 -0700 Subject: [PATCH 62/89] Say the CUDA scatter_reduce sum loses less to rounding, not nothing Summing half and bfloat16 in float keeps more precision, but a term far smaller than the running sum can still vanish, as the index_put comment already says. Comment only. --- backends/cuda/runtime/shims/index.cu | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/backends/cuda/runtime/shims/index.cu b/backends/cuda/runtime/shims/index.cu index 4ec710c1c35..97e7fc345c2 100644 --- a/backends/cuda/runtime/shims/index.cu +++ b/backends/cuda/runtime/shims/index.cu @@ -941,8 +941,8 @@ combine(Reduction reduction, AccT result, AccT value) { // one element are a run in index order. Each run is reduced by one thread in // that order and written once, so the result does not depend on scheduling. // A sum adds in AccT, as ATen does in deterministic mode, so half and -// bfloat16 keep small terms. A product rounds to T at every step and a mean -// rounds its sum to T before dividing, as ATen does, so half overflows to +// bfloat16 lose less to rounding. A product rounds to T at every step and a +// mean rounds its sum to T before dividing, as ATen does, so half overflows to // infinity and integers wrap where ATen's do. template __global__ void scatter_reduce_sorted_kernel( From 53a0623188dcd36304c6327f92ee64e1f9e47dda Mon Sep 17 00:00:00 2001 From: shoumikhin Date: Sat, 31 Oct 2026 08:48:33 -0700 Subject: [PATCH 63/89] Drop the unused libgomp dependency from CUDA backend programs A CUDA backend program carries a shared library that AOTInductor compiles. Inductor puts OpenMP (`-fopenmp`, `-lgomp`) on every Linux link. Unless the compiler asks the linker to link as needed, as Ubuntu's GCC does by default, the linker records `libgomp.so.1` as a dependency even when nothing in the library calls OpenMP. Programs exported with GCC on Red Hat based systems, or with Clang, carry it. PyTorch ships its own libgomp, which hides this while torch is imported. A process without torch then fails to load the program on any machine that has no system libgomp, such as a minimal container: Failed to load /tmp/executorch_cuda__0.so with error: libgomp.so.1: cannot open shared object file: No such file or directory Link with `-Wl,--as-needed` while the CUDA backend compiles. The linker then records only the libraries the program references, so an unused libgomp is dropped, while a program that does call OpenMP keeps it. PyTorch already does the same for MPS builds on macOS. The flag applies only to ahead-of-time links on Linux that neither link libtorch nor cross-compile for Windows, and only in threads inside the backend's compile context. The torch-free tests in `test_pybindings_cuda_dlpack.py` no longer probe for libgomp before they run. The ones that check which libraries got loaded now also require that libgomp is not among them, so they fail on any Linux host if the dependency comes back. That check now matches library file names only. The pybindings README no longer asks users to install libgomp. Test plan, on Linux x86_64 (CentOS Stream 9) with an A100: - Exported a small MLP and CNN before and after this change. Before, the compiled library lists `libgomp.so.1` and references no OpenMP symbol. After, `libgomp` and the unused `libm` are gone, and `libcuda`, `libstdc++`, `libgcc_s` and `libc` remain. - Ran both programs from Python without torch, where the loader cannot find libgomp. Before fails with the error above. After loads, runs, and matches eager. - With the same kernels, only the link line differs, and timing showed no change beyond noise. - A graph whose CPU part Inductor parallelizes with OpenMP still records libgomp under the flag. - The torch-free tests pass in both output modes. With the backend change reverted, the ones that check loaded libraries fail at the libgomp check, and without a system libgomp all of them fail to load. --- backends/cuda/cuda_backend.py | 55 +++++++++++++++++++ extension/pybindings/README.md | 5 +- .../test/test_pybindings_cuda_dlpack.py | 30 +++------- 3 files changed, 65 insertions(+), 25 deletions(-) diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index 72bfa03c985..2a816d90e3f 100644 --- a/backends/cuda/cuda_backend.py +++ b/backends/cuda/cuda_backend.py @@ -12,6 +12,7 @@ import math import os import shutil +import sys import threading import typing import weakref @@ -166,6 +167,57 @@ def _is_big_gpu(index_or_device=0) -> bool: inductor_utils.is_big_gpu = orig_is_big_gpu +_AS_NEEDED_GUARD = threading.local() + + +@contextlib.contextmanager +def _link_libraries_as_needed(): + """Record only the shared libraries the compiled program calls into. + + Inductor puts OpenMP on every Linux link. Unless the compiler asks the + linker to link as needed, as Ubuntu's GCC does by default, the linker + records libgomp even when nothing in the program calls OpenMP. Torch ships + a libgomp, which hides this whenever torch is loaded, but a process without + torch then fails to load the program on a machine with no system libgomp. + ``--as-needed`` drops only libraries nothing references, so a program that + does call OpenMP keeps it. Links with libtorch are left alone: a library + there can be needed without being referenced, because its static + initializers register something. + + Only threads inside this context get the flag. + + TODO: remove once Inductor does this for AOTI builds that do not link + libtorch, as it already does on macOS + (https://github.com/pytorch/pytorch/pull/198232). + """ + from torch._inductor import cpp_builder + + orig_device_options = cpp_builder.get_cpp_torch_device_options + + def _device_options(device_type, aot_mode=False, *args, **kwargs): + definitions, include_dirs, cflags, ldflags, *rest = orig_device_options( + device_type, aot_mode, *args, **kwargs + ) + if ( + aot_mode + and getattr(_AS_NEEDED_GUARD, "active", False) + and sys.platform.startswith("linux") + and not torch._inductor.config.aot_inductor.link_libtorch + and torch._inductor.config.aot_inductor.cross_target_platform is None + ): + ldflags = [*ldflags, "Wl,--as-needed"] + return (definitions, include_dirs, cflags, ldflags, *rest) + + prev_active = getattr(_AS_NEEDED_GUARD, "active", False) + cpp_builder.get_cpp_torch_device_options = _device_options + _AS_NEEDED_GUARD.active = True + try: + yield + finally: + _AS_NEEDED_GUARD.active = prev_active + cpp_builder.get_cpp_torch_device_options = orig_device_options + + def _full_zeros_preserving_strides(x: torch.Tensor, device) -> torch.Tensor: """Allocate a zero-filled tensor matching ``x``'s size/stride/dtype on ``device``. @@ -1303,6 +1355,9 @@ def _combined(): stack.enter_context(torch.nn.attention.sdpa_kernel([SDPBackend.MATH])) stack.enter_context(triton_mutation_cache.enabled()) stack.enter_context(target_smem_context()) + # Keep an unused libgomp out of the program, so it loads in a + # process without torch. See `_link_libraries_as_needed`. + stack.enter_context(_link_libraries_as_needed()) # On ROCm, is_big_gpu gates on architecture, not SM count. if torch.version.hip is None: stack.enter_context(_allow_triton_templates_on_small_gpus()) diff --git a/extension/pybindings/README.md b/extension/pybindings/README.md index c665c7657c2..82c9e48754d 100644 --- a/extension/pybindings/README.md +++ b/extension/pybindings/README.md @@ -201,10 +201,7 @@ The CUDA build workflow runs `test_pybindings_cuda_dlpack.py` in both output modes on a GPU. It checks device allocations, DLPack stream ordering, channels last views, cloning, and borrowed output lifetimes using a CUDA delegate. Set `EXECUTORCH_PYBIND_TEST_REQUIRE_CUDA=1` when running this test to treat missing -GPU hardware or a missing CUDA backend as a failure. The CUDA delegate's kernels -need the OpenMP runtime (`libgomp.so.1` on Linux). PyTorch ships one; without PyTorch, -install it from the system, or the two tests that run without PyTorch fail -under this setting. +GPU hardware or a missing CUDA backend as a failure. # Link Backends diff --git a/extension/pybindings/test/test_pybindings_cuda_dlpack.py b/extension/pybindings/test/test_pybindings_cuda_dlpack.py index 0792e6c924f..b7422ada832 100644 --- a/extension/pybindings/test/test_pybindings_cuda_dlpack.py +++ b/extension/pybindings/test/test_pybindings_cuda_dlpack.py @@ -562,7 +562,7 @@ def exit_elsewhere(): self._close_after_teardown(stream) def test_torch_free_caller_shares_one_stream_across_programs(self): - self._require_libgomp_without_torch() + self._require_linux() programs = [ self._export_to_directory( f"torch_free_{index}", module, (torch.ones(4, 8),) @@ -615,9 +615,9 @@ def load(path, data): np.testing.assert_allclose(outputs[0], np.maximum(value, 0), atol=1e-6) np.testing.assert_allclose(outputs[1], np.tanh(value), atol=1e-6) assert "torch" not in sys.modules -assert re.search( - r"lib(?:torch|aten|c10)", Path("/proc/self/maps").read_text().lower() -) is None +# libgomp as well, since these programs never call OpenMP. +maps = Path("/proc/self/maps").read_text().lower() +assert re.search(r"/lib(?:torch|aten|c10|gomp)[^/]*$", maps, re.M) is None """ arguments = [item or "" for program in programs for item in program] result = subprocess.run( @@ -1264,27 +1264,16 @@ def chain(): torch.from_dlpack(final), torch.full(self.shape, expected) ) - def _require_libgomp_without_torch(self): + def _require_linux(self): # The torch-free chain checks loaded libraries through /proc. if not sys.platform.startswith("linux"): self.skipTest("needs Linux") - # The CUDA delegate's compiled kernels link libgomp, which a host may - # provide only inside the torch package. - probe = subprocess.run( - [sys.executable, "-c", "import ctypes; ctypes.CDLL('libgomp.so.1')"], - capture_output=True, - ) - if probe.returncode != 0: - message = "libgomp.so.1 is not loadable without torch" - if os.environ.get("EXECUTORCH_PYBIND_TEST_REQUIRE_CUDA") == "1": - raise AssertionError(message) - self.skipTest(message) def test_device_chain_runs_with_torch_blocked(self): # The device-resident chain path, in a process where `import torch` # fails: upload host data, run, feed the output straight into a second # run on the GPU, then read the result back. - self._require_libgomp_without_torch() + self._require_linux() script = """ import importlib.abc import re @@ -1318,9 +1307,9 @@ def find_spec(self, fullname, path=None, target=None): np.asarray(second.to_host()), np.full(shape, 18.0, dtype=np.float32) ) assert "torch" not in sys.modules -assert re.search( - r"lib(?:torch|aten|c10)", Path("/proc/self/maps").read_text().lower() -) is None +# libgomp as well, since these programs never call OpenMP. +maps = Path("/proc/self/maps").read_text().lower() +assert re.search(r"/lib(?:torch|aten|c10|gomp)[^/]*$", maps, re.M) is None print("Device chain ran without torch") """ environment = os.environ.copy() @@ -1349,7 +1338,6 @@ def find_spec(self, fullname, path=None, target=None): def test_device_input_works_beside_a_torch_without_cuda(self): # A torch built without CUDA cannot select a CUDA device. Inputs from # other producers must not depend on it. - self._require_libgomp_without_torch() script = """ import sys import types From 5a94bfa3ddc5a86de775df0d2b5e95b534f38732 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sat, 31 Oct 2026 06:52:42 -0700 Subject: [PATCH 64/89] Keep keys past the end out of the CUDA Triton SDPA softmax The CUDA backend swaps scaled_dot_product_attention for its Triton kernel, triton::sdpa, when triton_kernel_mode is on, which is the default. The 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 loads those positions as zeros, but it never gives them a score of minus infinity. Without an attention mask or a causal mask, each one enters the softmax with a score of 0 and a value of 0, and takes weight away from the real keys. With every real score equal, where the right output is 1, one key gave 0.031 and 80 keys gave 0.625. Passing kv_len without a mask has the same problem. Set the score of every position at or past kv_len to minus infinity, as the other attention kernels in this file already do. The first version of this kernel did this, and a later rewrite dropped it. Checking every tile made some shapes up to 8% slower, so the check only runs where it can matter. 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 every call that was already right compiles to the same code as before. A kv_len, or a symbolic key count, is only known at run time, so those calls check every tile. The test file called this a kernel limitation, and its basic test only used key lengths of 64 and up. Drop that note and run the basic test with a single key too. A new test runs every tile without a mask over 80 and 96 keys and over a kv_len of 80. On 96 keys only the 32-key tiles skip the check. Without the fix, these cases are off by up to 3.6. --- backends/cuda/tests/test_triton_sdpa.py | 36 ++++++++++++++++++++++--- backends/cuda/triton/kernels/sdpa.py | 28 +++++++++++++++++++ 2 files changed, 61 insertions(+), 3 deletions(-) diff --git a/backends/cuda/tests/test_triton_sdpa.py b/backends/cuda/tests/test_triton_sdpa.py index 1fdb1dedef9..76511f6cc71 100644 --- a/backends/cuda/tests/test_triton_sdpa.py +++ b/backends/cuda/tests/test_triton_sdpa.py @@ -83,10 +83,8 @@ def _max_abs_error(out, ref): HEAD_DIMS_NON_POW2 = [80, 96] # Sequence length pairs (seqlen_q, seqlen_kv) adapted from FlashAttention -# Note: Lk must be >= BLOCK_N (32) to avoid unmasked zero-padding in K -# that dilutes softmax. This is a pre-existing kernel limitation for very -# short KV sequences. SEQLEN_PAIRS_BASIC = [ + (1, 1), (1, 64), (1, 128), (1, 512), @@ -357,6 +355,38 @@ def test_every_tile_packed_or_not_matches(self): with self.subTest(Lq=Lq, tile=tile): self.assertTrue(torch.equal(unpacked, packed)) + def test_every_tile_drops_keys_past_the_end(self): + """Each standard-kernel tile without a mask: keys at or past the key + count, or past kv_len, stay out of the softmax.""" + import triton + + sdpa_module = importlib.import_module( + "executorch.backends.cuda.triton.kernels.sdpa" + ) + kernel = sdpa_module._sdpa_fwd_kernel + B, H, Lq, D = 1, 4, 16, 64 + configs = sdpa_module._sdpa_fwd_prune( + kernel.configs, None, HEAD_DIM=D, NUM_GROUPS=1 + ) + # No tile size divides 80 keys, only the 32-key tiles divide 96, and + # kv_len is only known on the device. + for Lk, kv_len in ((80, None), (96, None), (128, 80)): + torch.manual_seed(Lk) + q = torch.randn(B, H, Lq, D, dtype=torch.bfloat16, device="cuda") + k = torch.randn(B, H, Lk, D, dtype=torch.bfloat16, device="cuda") + v = torch.randn(B, H, Lk, D, dtype=torch.bfloat16, device="cuda") + valid = Lk if kv_len is None else kv_len + ref = _reference_sdpa(q, k[:, :, :valid], v[:, :, :valid]) + kwargs = {} + if kv_len is not None: + kwargs["kv_len"] = torch.tensor([kv_len], device="cuda") + for config in configs: + single = triton.autotune(configs=[config], key=kernel.keys)(kernel.fn) + with self.subTest(Lk=Lk, kv_len=kv_len, config=str(config)): + with mock.patch.object(sdpa_module, "_sdpa_fwd_kernel", single): + out = self.sdpa(q, k, v, **kwargs) + self.assertLess(_max_abs_error(out, ref), MAX_ABS_TOL) + def test_mha_never_packs(self): sdpa_module = importlib.import_module( "executorch.backends.cuda.triton.kernels.sdpa" diff --git a/backends/cuda/triton/kernels/sdpa.py b/backends/cuda/triton/kernels/sdpa.py index e01c0922cef..a642d7794d1 100644 --- a/backends/cuda/triton/kernels/sdpa.py +++ b/backends/cuda/triton/kernels/sdpa.py @@ -507,6 +507,7 @@ def _sdpa_fwd_kernel_body( # noqa: C901 NUM_GROUPS: tl.constexpr, PACK_GQA: tl.constexpr, MASK_IS_CAUSAL: tl.constexpr = False, + KV_LEN_MULTIPLE: tl.constexpr = 1, ): """ Shared kernel body for SDPA forward pass. @@ -519,6 +520,9 @@ def _sdpa_fwd_kernel_body( # noqa: C901 When False, the grid iterates over H_q heads and each program handles one Q head with simple h_kv = h_q // NUM_GROUPS remapping. + + KV_LEN_MULTIPLE: kv_len is a multiple of this. When BLOCK_N divides it, + no tile reaches past kv_len, so the kernel skips the key bound check. """ pid_m = tl.program_id(axis=0) pid_bh = tl.program_id(axis=1) @@ -625,6 +629,17 @@ def _sdpa_fwd_kernel_body( # noqa: C901 k = tl.load(k_ptrs, mask=k_mask, other=0.0).to(tl.bfloat16) qk = (tl.dot(q, tl.trans(k)).to(tl.float32) * sm_scale).to(tl.float32) + # Keys past kv_len load as zeros, which would score 0 rather than -inf. + # A mask or a causal bound already drops them. + if ( + not (HAS_MASK or IS_CAUSAL or MASK_IS_CAUSAL) + and KV_LEN_MULTIPLE % BLOCK_N != 0 + ): + qk = tl.where( + offs_n[None, :] < kv_len, + qk, + tl.full(qk.shape, -float("inf"), dtype=tl.float32), + ) if HAS_MASK: qk = tl.where( @@ -746,6 +761,7 @@ def _sdpa_prefill_prune(configs, nargs, **kwargs): for config in _SDPA_PREFILL_CONFIGS for pack in (False, True) ] +_SDPA_BLOCK_N_SIZES = {config.kwargs["BLOCK_N"] for config in _SDPA_FWD_CONFIGS} def _sdpa_fwd_prune(configs, nargs, **kwargs): @@ -813,6 +829,7 @@ def _sdpa_fwd_kernel( BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, MASK_IS_CAUSAL: tl.constexpr = False, + KV_LEN_MULTIPLE: tl.constexpr = 1, ): # Packed, the grid walks the KV heads; unpacked, the query heads. if PACK_GQA: @@ -859,6 +876,7 @@ def _sdpa_fwd_kernel( NUM_GROUPS=NUM_GROUPS, PACK_GQA=PACK_GQA, MASK_IS_CAUSAL=MASK_IS_CAUSAL, + KV_LEN_MULTIPLE=KV_LEN_MULTIPLE, ) @@ -1108,6 +1126,15 @@ def grid(meta): return (triton.cdiv(L_q * num_groups, meta["BLOCK_M"]), B * H_kv) return (triton.cdiv(L_q, meta["BLOCK_M"]), B * H_q) + # The largest tile size that divides L_kv. A kv_len or a symbolic length is + # only known at run time, so those always check bounds. A mask or a causal + # bound makes the check unneeded, so those calls keep 1 and share kernels. + kv_len_multiple = 1 + if isinstance(L_kv, int) and not (HAS_KV_LEN or HAS_MASK or is_causal): + kv_len_multiple = max( + (n for n in _SDPA_BLOCK_N_SIZES if L_kv % n == 0), default=1 + ) + # Single autotuned kernel: the config set spans BLOCK_M in {16..128}, each # with and without pack GQA, and `_sdpa_prefill_prune` keeps only # non-spilling tiles for this HEAD_DIM, so the autotuner picks a @@ -1150,6 +1177,7 @@ def grid(meta): HEAD_DIM=D, NUM_GROUPS=num_groups, MASK_IS_CAUSAL=mask_is_causal, + KV_LEN_MULTIPLE=kv_len_multiple, ) From 4dc4d5ae2780c1c7fc39599626ebcff111797d0a Mon Sep 17 00:00:00 2001 From: shoumikhin Date: Sat, 31 Oct 2026 10:55:35 -0700 Subject: [PATCH 65/89] Remove the CUDA backend's temporary library when a load fails To load a method, the CUDA backend writes the compiled library from the program to a temporary file, such as `/tmp/executorch_cuda__0.so`, and loads it from there. Only `destroy()` removes that file, and `destroy()` only runs for a handle that `init()` returned. When `init()` fails, for example because the library cannot be loaded or the weights blob is missing, it returns no handle, so the file stays behind. Every load that fails after writing the file leaves it in the temporary directory. Remove the file with a scope guard unless `init()` succeeds, so every failure path removes it. A method that loads keeps its file until `destroy()`, as before. Nothing else changes on a failure: the handle of a late failure is still not freed, and, as in `destroy()`, a library that was already loaded stays loaded. `test_missing_weights_blob.py` now also checks that a load that fails for a missing weights blob leaves nothing in the temporary directory. Test plan, on Linux x86_64 with an A100: - The extended test fails without this change, with the library left in the temporary directory, and passes with it. - A load that fails because the library itself cannot be loaded also leaves no file now. - The other tests in `test_missing_weights_blob.py` still pass. --- backends/cuda/runtime/cuda_backend.cpp | 17 +++++++++++++++++ .../cuda/tests/test_missing_weights_blob.py | 13 ++++++++++--- 2 files changed, 27 insertions(+), 3 deletions(-) diff --git a/backends/cuda/runtime/cuda_backend.cpp b/backends/cuda/runtime/cuda_backend.cpp index 04f9ec91893..f1d09d09e29 100644 --- a/backends/cuda/runtime/cuda_backend.cpp +++ b/backends/cuda/runtime/cuda_backend.cpp @@ -493,6 +493,22 @@ class ET_EXPERIMENTAL CudaBackend final to_string(so_file_counter.fetch_add(1, std::memory_order_relaxed)) + ".so"); + // A failed init returns no handle, so destroy() never runs for it. Remove + // the file here unless init succeeds. + bool keep_so_file = false; + executorch::backends::aoti::ScopeGuard remove_so_file([&]() noexcept { + if (keep_so_file) { + return; + } + std::error_code remove_error; + filesystem::remove(so_path, remove_error); + ET_CHECK_OR_LOG_ERROR( + !remove_error, + "Failed to remove temporary shared library %s: %s", + so_path.c_str(), + remove_error.message().c_str()); + }); + // Create a temporary file ofstream outfile(so_path, ios::binary); @@ -644,6 +660,7 @@ class ET_EXPERIMENTAL CudaBackend final live_handles_.fetch_add(1, std::memory_order_acq_rel); + keep_so_file = true; return (DelegateHandle*)handle; // Return the handle post-processing } diff --git a/backends/cuda/tests/test_missing_weights_blob.py b/backends/cuda/tests/test_missing_weights_blob.py index aa7667471f2..0a524bd7353 100644 --- a/backends/cuda/tests/test_missing_weights_blob.py +++ b/backends/cuda/tests/test_missing_weights_blob.py @@ -23,6 +23,7 @@ import tempfile import unittest from typing import Optional +from unittest.mock import patch import torch from executorch.backends.cuda.cuda_backend import CudaBackend @@ -167,10 +168,16 @@ def test_load_reports_not_found_when_blob_is_absent(self) -> None: # exception carries only the method name and the error number, so the # cause is asserted through the rewrite check above rather than here. program = Runtime.get().load_program(path) - with self.assertRaisesRegex( - RuntimeError, r"Failed to load method forward, error: 0x:?20" + # The failed load leaves no handle to destroy, so it has to remove the + # library it wrote itself. An empty temporary directory shows that. + with tempfile.TemporaryDirectory() as tmp, patch.dict( + os.environ, {"TMPDIR": tmp} ): - program.load_method("forward") + with self.assertRaisesRegex( + RuntimeError, r"Failed to load method forward, error: 0x:?20" + ): + program.load_method("forward") + self.assertEqual(os.listdir(tmp), []) def test_load_fails_with_another_programs_blob(self) -> None: from executorch.runtime import Runtime From 624b83ad96fe82cb731decbc14bf7fffc0c55dc5 Mon Sep 17 00:00:00 2001 From: shoumikhin Date: Sat, 31 Oct 2026 12:38:20 -0700 Subject: [PATCH 66/89] Point the CUDA backend's libgomp TODO at the PyTorch fix The CUDA backend adds `-Wl,--as-needed` to its AOTInductor links until Inductor does it itself. pytorch/pytorch#202902 makes that change in Inductor, so the TODO now names it: the hook can go once the pinned PyTorch includes it. --- backends/cuda/cuda_backend.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index 2a816d90e3f..da70b833ffe 100644 --- a/backends/cuda/cuda_backend.py +++ b/backends/cuda/cuda_backend.py @@ -186,9 +186,8 @@ def _link_libraries_as_needed(): Only threads inside this context get the flag. - TODO: remove once Inductor does this for AOTI builds that do not link - libtorch, as it already does on macOS - (https://github.com/pytorch/pytorch/pull/198232). + TODO: remove once the pinned PyTorch has the same flag in Inductor + (https://github.com/pytorch/pytorch/pull/202902). """ from torch._inductor import cpp_builder From 4f53e3de5c6abf2656eb8d5e3ee66b4b106e3ae0 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sat, 31 Oct 2026 12:12:58 -0700 Subject: [PATCH 67/89] Add a CUDA shim for scatter with a tensor source Lowering x.scatter(dim, index, src) to the CUDA backend failed for a bfloat16, int64 or bool src: Method forward missing fallback kernels (1 total): - aoti_torch_cuda_scatter_src_out 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 fallback kernel, and in deterministic mode every dtype. This backend shipped the scatter_reduce shim but not this one. A mixture-of-experts layer that writes its expert outputs back with torch.zeros_like(h).scatter(1, index, src) in bfloat16 hits it on an A100. Add the shim next to the scatter_reduce shim, sharing its checks, its sort of the writes by the element they land on, and its kernel: a scatter is one more reduction, Assign, where each write replaces the one before it. So repeated writes to an element keep the last one in index order, as ATen does in deterministic mode; ATen's default CUDA kernel keeps any one of them. As in ATen, an index out of range triggers a device-side assert. It takes every dtype the scatter_reduce shim takes, bool included. Also allow it in the backend's list of fallback kernels, add the name to aoti_cuda_shims.lib, the import library that lowering for Windows links against, regenerated with mingw binutils 2.41 as its README describes, and to the wheel's shared library symbol check. Test: test_scatter_with_tensor_source lowers a scatter with repeated indices and a column it never writes, in bfloat16, int64 and bool, runs it and compares with eager CUDA in deterministic mode. Without this change it fails at lowering. The test that a missing scatter kernel fails at export lowered an int64 scatter from a tensor, which now works, so it lowers a scatter of a scalar in deterministic mode instead, whose kernel the backend still lacks. --- .ci/scripts/wheel/test_shared_libraries.py | 1 + backends/cuda/cuda_backend.py | 1 + backends/cuda/runtime/aoti_cuda_shims.lib | Bin 58386 -> 59238 bytes backends/cuda/runtime/shims/index.cu | 304 ++++++++++++--------- backends/cuda/runtime/shims/index.h | 28 ++ backends/cuda/tests/test_cuda_export.py | 23 +- backends/cuda/tests/test_index_shims.py | 26 ++ backends/cuda/tests/test_sort_shim.py | 1 + 8 files changed, 251 insertions(+), 133 deletions(-) diff --git a/.ci/scripts/wheel/test_shared_libraries.py b/.ci/scripts/wheel/test_shared_libraries.py index 5a660e1ebf7..06d39ba3dce 100644 --- a/.ci/scripts/wheel/test_shared_libraries.py +++ b/.ci/scripts/wheel/test_shared_libraries.py @@ -106,6 +106,7 @@ "aoti_torch_cuda_index_Tensor", "aoti_torch_index_put_out", "aoti_torch_cuda_scatter_reduce_two_out", + "aoti_torch_cuda_scatter_src_out", ) _QUANTIZED_KERNEL_SYMBOLS = ( diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index da70b833ffe..7334e5d0500 100644 --- a/backends/cuda/cuda_backend.py +++ b/backends/cuda/cuda_backend.py @@ -1131,6 +1131,7 @@ def get_supported_fallback_kernels(cls) -> Dict[str, Any]: "aoti_torch_cuda_index_Tensor": None, "aoti_torch_index_put_out": None, "aoti_torch_cuda_scatter_reduce_two_out": None, + "aoti_torch_cuda_scatter_src_out": None, } @classmethod diff --git a/backends/cuda/runtime/aoti_cuda_shims.lib b/backends/cuda/runtime/aoti_cuda_shims.lib index 78888032df12ed427e0b744a3b757990fb9607d4..ad54a4fdb24e1e8c3395f0245d36a2d9881b5ef3 100644 GIT binary patch delta 1622 zcmZA0dr*{B6aesZp~Y2Lj9nkNqKG+8qV5NXfF-tp!<3MjI2M9H$iqZ|MJ-WqGZhkp zSvdp<4y)z^7#tUl8ZxQ)C?7ZmiV$WAP2=>BkDBTa#-`5Q%U|uxZ|3{Xch9}&o_n|I zhUeouo{mNK$OREI+`sg308j}qH5*`hm=fdY5j5_A!9_`!r_?L~kPe!a2F{c0jpqQQBRp6-7NeZ>0+5<-E3K^n(&^8Y_SFE=S)Nqt zcvb1-ok|zD7^#a3lCG{%z6%DBZtxfu58b!G?M@Yd)XM{To&N|L1nU|DF_tPg6oD z09OhuNHw6h-jCXoUl-V&7DCQYo^ttzQ)pjmtiY}ejY`5f5|cRr^Yljid*Kr-N%g@^ z8Hw1GrKt*nbH<}RRfjvXd?_=N({`;7r600*`jU zH&_LRXPR4KQ9Ev}G~j-3 z1NM09>GA%h0s|ZxB~|d)#shwIp)y_IlPU|vJ2;Y8Z6w!$Xn{As)Kt@#)%sIubqv>0 zr_n$bkC|0(M2C+Trq}5xr*^TxtOg79*729ehl~_auL+vlSS4_U!HfPmG+*HJ!y5hD zlPEkiK$~GAU3FRnp8d)~fk!wg;`l@=JIawACoJw^_~N7y9~pG?&T)<=H*3^#g17&$ zWeNuRj-h}iKB9=$(Ie_OWuj*-Zwb_$R!zZDmoHZMdQEGo0PdYkzRgfF)-p(E*AR*!=Zq3xm)5QuaB%8n#Z0`*1}R26axRnDbmu%!SdT zymrxqUE}p+_=XQ@Y}aV*<(G}re{q4J!7Fuwj$Gzz(EKew*{<=CY`dYUHjmsi(&FpU zg5K<|5j6f5Uz@Ra_-WvFsKB0{(F0$3cOp5z4-;5>PiybxP59q8qSa4FAKm3b()v|# zI5c3uTt9CLxzD{09Sm{bQV~nT1DxW%EBrU{SPJ;eZyMcuz}>s~vqsj3{B_xoft)`D g3*0?CpO(y7B}5#>4@C;k(6EKrHe)WihBG|=1Kn+a delta 1508 zcmXZbdrZ|;902g|LA|@+z1O%G5-Hf-8kL~G--SSk#886aG6ivqoXA7WSJHKu!Za=- z=tLNdFAv2B5J)nw5PW?gN;(o^ILS;&(-1YAnQN7QW$_hq4scQ4WwA|4^E{0HoGDr7aaea)&D&9sp7&A1s}^p`7I^ zr5;X3>g8mlO9rKH1VHNNJ<{N|32|^`0ZbU?jncJY<%enjX@uJ)-R71_cP+~30DyGA zkP1DvPs4QMbV@X~dwGUACLa`5+I-2!rW+~ql_l@er?n}fYL#tbqapnrfos!3$rQF+ z;N#V1JYclq=Wjj5@3N9{QHD;-)}#o$oH>imXMZHnzDB2$D|zLMY#(%ZTBv(noIu;h z4hq@8tF`N9Q2EBi0$;3mkh7MLEZFFe&YY>VZS#DABb#)(pTql7H(T&OPXi@77YOXi z(J8HjSHqn?X!n{*_T0GwOSU*@*>+wv=Gw4v=QMi!S(KoqTk8e&7>BI;V4pZ*Ah0ydzd@Fi)f8#<_yJkD|a~vw>zdYkXsqM!79~ zW3xB?cq~j{e9L6ez%gHHIPMVW`K3nVZG24WaVuuHE#y4GogZ$SoKn(>S+ux)k-!eO zrh4gUx8az@KtY}10{3@JmTu~_;FD86^n$qb)Gm%b&FTL{Gx5$@6FoYUAaLPn)ju(h zPM%8=7}Bk&>AKHZsO&-%zl|P#RdTWg7c98+l9@6u@@)EhCv#4?XrqumuA=UeY9#vg zS?SNqF@n;*IlzPcmS;L=;Gm$IAwF$ zEJb+40mfJZ=#5cs-rlis_cuda() && index->is_cuda() && src->is_cuda(), + NotSupported, + "%s: tensors must be on the CUDA device", + name); + ET_CHECK_OR_RETURN_ERROR( + src->dtype() == out->dtype(), + InvalidArgument, + "%s: src and self must have the same dtype", + name); + ET_CHECK_OR_RETURN_ERROR( + out->dtype() != ScalarType::Bool || reduction == Reduction::Sum || + reduction == Reduction::Assign, + NotSupported, + "%s: bool supports only sum, as in ATen", + name); + + // As in ATen, a 0-dim tensor counts as one of size 1. + const auto ndim = std::max(out->dim(), 1); + ET_CHECK_OR_RETURN_ERROR( + ndim <= kMaxDims, + NotSupported, + "%s: tensors with more than %d dims are not supported", + name, + kMaxDims); + if (dim < 0) { + dim += ndim; + } + ET_CHECK_OR_RETURN_ERROR( + dim >= 0 && dim < ndim, InvalidArgument, "%s: dim out of range", name); + // ATen checks the index only when it has elements. + if (index->numel() == 0) { + return Error::Ok; + } + ET_CHECK_OR_RETURN_ERROR( + index->dtype() == ScalarType::Long, + InvalidArgument, + "%s: index must be int64", + name); + ET_CHECK_OR_RETURN_ERROR( + std::max(index->dim(), 1) == ndim && + std::max(src->dim(), 1) == ndim, + InvalidArgument, + "%s: index and src must have as many dims as self", + name); + auto size = [](const Tensor* tensor, int64_t d) -> int64_t { + return tensor->dim() == 0 ? 1 : tensor->size(d); + }; + auto stride = [](const Tensor* tensor, int64_t d) -> int64_t { + return tensor->dim() == 0 ? 0 : tensor->stride(d); + }; + ScatterPlan plan{}; + plan.ndim = static_cast(ndim); + plan.dim = static_cast(dim); + plan.dim_size = size(out, dim); + int64_t points = 1; + for (int64_t d = 0; d < ndim; ++d) { + plan.max_out_offset += + std::max(size(out, d) - 1, 0) * stride(out, d); + const int64_t index_size = size(index, d); + ET_CHECK_OR_RETURN_ERROR( + index_size <= size(src, d) && (d == dim || index_size <= size(out, d)), + InvalidArgument, + "%s: index is larger than src, or than self, at dim %" PRId64, + name, + d); + plan.sizes[d] = index_size; + plan.index_strides[d] = stride(index, d); + plan.out_strides[d] = stride(out, d); + plan.src_strides[d] = stride(src, d); + points *= index_size; + } + + const auto device = static_cast(out->device_index()); + auto stream_result = getCurrentCUDAStream(out->device_index()); + ET_CHECK_OR_RETURN_ERROR( + stream_result.ok(), Internal, "%s: failed to get CUDA stream", name); + cudaStream_t stream = stream_result.get(); + Scratch scratch(device, stream); + auto run = [&](auto value, auto accumulator) { + using T = decltype(value); + using AccT = decltype(accumulator); + return scatter_reduce( + plan, + points, + static_cast(index->data_ptr()), + out->data_ptr(), + src->data_ptr(), + reduction, + include_self, + scratch, + stream); + }; + // Half and bfloat16 compute in float, and integers in int64. + switch (out->dtype()) { + case ScalarType::Float: + return run(float{}, float{}); + case ScalarType::Double: + return run(double{}, double{}); + case ScalarType::Half: + return run(__half{}, float{}); + case ScalarType::BFloat16: + return run(__nv_bfloat16{}, float{}); + case ScalarType::Long: + return run(int64_t{}, int64_t{}); + case ScalarType::Int: + return run(int32_t{}, int64_t{}); + case ScalarType::Short: + return run(int16_t{}, int64_t{}); + case ScalarType::Char: + return run(int8_t{}, int64_t{}); + case ScalarType::Byte: + return run(uint8_t{}, int64_t{}); + case ScalarType::Bool: + return run(bool{}, bool{}); + default: + ET_LOG( + Error, + "%s: unsupported dtype %d", + name, + static_cast(out->dtype())); + return Error::NotSupported; + } +} + } // namespace #ifdef __cplusplus @@ -1291,6 +1441,23 @@ AOTITorchError aoti_torch_index_put_out( } } +AOTITorchError aoti_torch_cuda_scatter_src_out( + Tensor* out, + Tensor* self, + int64_t dim, + Tensor* index, + Tensor* src) { + return scatter( + "aoti_torch_cuda_scatter_src_out", + out, + self, + dim, + index, + src, + Reduction::Assign, + /*include_self=*/false); +} + AOTITorchError aoti_torch_cuda_scatter_reduce_two_out( Tensor* out, Tensor* self, @@ -1300,21 +1467,9 @@ AOTITorchError aoti_torch_cuda_scatter_reduce_two_out( const char* reduce, int32_t include_self) { ET_CHECK_OR_RETURN_ERROR( - out != nullptr && index != nullptr && src != nullptr && reduce != nullptr, + reduce != nullptr, InvalidArgument, "aoti_torch_cuda_scatter_reduce_two_out: null argument"); - ET_CHECK_OR_RETURN_ERROR( - out == self, - NotSupported, - "aoti_torch_cuda_scatter_reduce_two_out: out must be self, as AOTInductor calls it"); - ET_CHECK_OR_RETURN_ERROR( - out->is_cuda() && index->is_cuda() && src->is_cuda(), - NotSupported, - "aoti_torch_cuda_scatter_reduce_two_out: tensors must be on the CUDA device"); - ET_CHECK_OR_RETURN_ERROR( - src->dtype() == out->dtype(), - InvalidArgument, - "aoti_torch_cuda_scatter_reduce_two_out: src and self must have the same dtype"); const std::string_view name(reduce); Reduction reduction; if (name == "sum") { @@ -1334,116 +1489,15 @@ AOTITorchError aoti_torch_cuda_scatter_reduce_two_out( reduce); return Error::InvalidArgument; } - ET_CHECK_OR_RETURN_ERROR( - out->dtype() != ScalarType::Bool || reduction == Reduction::Sum, - NotSupported, - "aoti_torch_cuda_scatter_reduce_two_out: bool supports only sum, as in ATen"); - - // As in ATen, a 0-dim tensor counts as one of size 1. - const auto ndim = std::max(out->dim(), 1); - ET_CHECK_OR_RETURN_ERROR( - ndim <= kMaxDims, - NotSupported, - "aoti_torch_cuda_scatter_reduce_two_out: tensors with more than %d dims are not supported", - kMaxDims); - if (dim < 0) { - dim += ndim; - } - ET_CHECK_OR_RETURN_ERROR( - dim >= 0 && dim < ndim, - InvalidArgument, - "aoti_torch_cuda_scatter_reduce_two_out: dim out of range"); - // ATen checks the index only when it has elements. - if (index->numel() == 0) { - return Error::Ok; - } - ET_CHECK_OR_RETURN_ERROR( - index->dtype() == ScalarType::Long, - InvalidArgument, - "aoti_torch_cuda_scatter_reduce_two_out: index must be int64"); - ET_CHECK_OR_RETURN_ERROR( - std::max(index->dim(), 1) == ndim && - std::max(src->dim(), 1) == ndim, - InvalidArgument, - "aoti_torch_cuda_scatter_reduce_two_out: index and src must have as many dims as self"); - auto size = [](const Tensor* tensor, int64_t d) -> int64_t { - return tensor->dim() == 0 ? 1 : tensor->size(d); - }; - auto stride = [](const Tensor* tensor, int64_t d) -> int64_t { - return tensor->dim() == 0 ? 0 : tensor->stride(d); - }; - ScatterPlan plan{}; - plan.ndim = static_cast(ndim); - plan.dim = static_cast(dim); - plan.dim_size = size(out, dim); - int64_t points = 1; - for (int64_t d = 0; d < ndim; ++d) { - plan.max_out_offset += - std::max(size(out, d) - 1, 0) * stride(out, d); - const int64_t index_size = size(index, d); - ET_CHECK_OR_RETURN_ERROR( - index_size <= size(src, d) && (d == dim || index_size <= size(out, d)), - InvalidArgument, - "aoti_torch_cuda_scatter_reduce_two_out: index is larger than src, or than self, at dim %" PRId64, - d); - plan.sizes[d] = index_size; - plan.index_strides[d] = stride(index, d); - plan.out_strides[d] = stride(out, d); - plan.src_strides[d] = stride(src, d); - points *= index_size; - } - - const auto device = static_cast(out->device_index()); - auto stream_result = getCurrentCUDAStream(out->device_index()); - ET_CHECK_OR_RETURN_ERROR( - stream_result.ok(), - Internal, - "aoti_torch_cuda_scatter_reduce_two_out: failed to get CUDA stream"); - cudaStream_t stream = stream_result.get(); - Scratch scratch(device, stream); - auto run = [&](auto value, auto accumulator) { - using T = decltype(value); - using AccT = decltype(accumulator); - return scatter_reduce( - plan, - points, - static_cast(index->data_ptr()), - out->data_ptr(), - src->data_ptr(), - reduction, - include_self != 0, - scratch, - stream); - }; - // Half and bfloat16 compute in float, and integers in int64. - switch (out->dtype()) { - case ScalarType::Float: - return run(float{}, float{}); - case ScalarType::Double: - return run(double{}, double{}); - case ScalarType::Half: - return run(__half{}, float{}); - case ScalarType::BFloat16: - return run(__nv_bfloat16{}, float{}); - case ScalarType::Long: - return run(int64_t{}, int64_t{}); - case ScalarType::Int: - return run(int32_t{}, int64_t{}); - case ScalarType::Short: - return run(int16_t{}, int64_t{}); - case ScalarType::Char: - return run(int8_t{}, int64_t{}); - case ScalarType::Byte: - return run(uint8_t{}, int64_t{}); - case ScalarType::Bool: - return run(bool{}, bool{}); - default: - ET_LOG( - Error, - "aoti_torch_cuda_scatter_reduce_two_out: unsupported dtype %d", - static_cast(out->dtype())); - return Error::NotSupported; - } + return scatter( + "aoti_torch_cuda_scatter_reduce_two_out", + out, + self, + dim, + index, + src, + reduction, + include_self != 0); } #ifdef __cplusplus diff --git a/backends/cuda/runtime/shims/index.h b/backends/cuda/runtime/shims/index.h index 15b7a043666..16488a9f52c 100644 --- a/backends/cuda/runtime/shims/index.h +++ b/backends/cuda/runtime/shims/index.h @@ -76,6 +76,34 @@ AOTI_SHIM_EXPORT AOTITorchError aoti_torch_index_put_out( Tensor* values, bool accumulate); +/** + * Writes src into out at index along dim, with the semantics of + * aten::scatter.src. + * + * Inductor calls this fallback, with out and self the same tensor, for a + * tensor source of a dtype its generated code cannot add atomically (int64, + * bool, and bfloat16 on GPUs older than sm_90), and in deterministic mode. + * When values land on the same element, the last one in index order wins, as + * in ATen's deterministic mode; ATen's default CUDA kernel keeps any one of + * them. As in ATen, an index out of range, negative ones included, triggers a + * device-side assert. + * + * @param out Tensor to update (CUDA device, up to 16 dims) + * @param self Must be out; AOTInductor passes the tensor it updates as both + * @param dim Dimension of out that index addresses + * @param index int64 tensor with as many dims as out, no larger than src, and + * no larger than out except at dim + * @param src Values to write, same dtype as out. Float, double, half, + * bfloat16, the integer dtypes and bool are supported. + * @return AOTITorchError + */ +AOTI_SHIM_EXPORT AOTITorchError aoti_torch_cuda_scatter_src_out( + Tensor* out, + Tensor* self, + int64_t dim, + Tensor* index, + Tensor* src); + /** * Reduces src into out at index along dim, with the semantics of * aten::scatter_reduce.two. diff --git a/backends/cuda/tests/test_cuda_export.py b/backends/cuda/tests/test_cuda_export.py index 8a126a12dbb..cecc59c2876 100644 --- a/backends/cuda/tests/test_cuda_export.py +++ b/backends/cuda/tests/test_cuda_export.py @@ -804,17 +804,24 @@ def test_triton_kernel_mode_off(self): def test_missing_scatter_kernel_fails_at_export(self): """The scatter fallback writes its kernel call itself, and the check for - missing kernels still sees it. Inductor leaves an int64 scatter to that - fallback, and the CUDA shims lack its kernel.""" + missing kernels still sees it. In deterministic mode Inductor leaves a + scatter of a scalar to that fallback, and the CUDA shims lack its + kernel.""" - class Int64Scatter(torch.nn.Module): + class ValueScatter(torch.nn.Module): def forward(self, counts: torch.Tensor, idx: torch.Tensor): - return counts.scatter(0, idx, torch.ones_like(idx)) + return counts.scatter(0, idx, 1.0) - inputs = (torch.zeros(8, dtype=torch.long), torch.randint(0, 8, (32,))) - with self.assertRaises(RuntimeError) as context: - self._export_to_cuda_with_lower(Int64Scatter(), inputs) - self.assertIn("aoti_torch_cuda_scatter_src_out", str(context.exception)) + inputs = (torch.zeros(8), torch.randint(0, 8, (32,))) + deterministic = torch.are_deterministic_algorithms_enabled() + warn_only = torch.is_deterministic_algorithms_warn_only_enabled() + torch.use_deterministic_algorithms(True) + try: + with self.assertRaises(RuntimeError) as context: + self._export_to_cuda_with_lower(ValueScatter(), inputs) + finally: + torch.use_deterministic_algorithms(deterministic, warn_only=warn_only) + self.assertIn("aoti_torch_cuda_scatter_value_out", str(context.exception)) @unittest.skipIf( torch.version.hip is not None, diff --git a/backends/cuda/tests/test_index_shims.py b/backends/cuda/tests/test_index_shims.py index 6a4922cd663..310d6d587f5 100644 --- a/backends/cuda/tests/test_index_shims.py +++ b/backends/cuda/tests/test_index_shims.py @@ -193,6 +193,32 @@ def forward(self, x, index, src): src[1, 2] = float("nan") self._run(ScatterReduce(), [(x, index, src)]) + def test_scatter_with_tensor_source(self): + """Inductor writes a scatter from a tensor itself only in a dtype it can + add atomically, so int64, bool and, below sm_90, bfloat16 go to the + shim. Of repeated writes the last in index order wins, as in eager's + deterministic mode, which also sends every dtype to the shim.""" + + class Scatter(torch.nn.Module): + def forward(self, x, index, src): + return (x.scatter(1, index, src),) + + # Smaller than src, with repeats, and column 4 is never written. + index = torch.tensor([[0, 2, 2, 5], [1, 1, 1, 0], [5, 3, 0, 5]]) + x = torch.full((5, 6), -1) + # Repeated writes differ from the write before them, also as bool. + src = torch.arange(20).reshape(4, 5) + deterministic = torch.are_deterministic_algorithms_enabled() + warn_only = torch.is_deterministic_algorithms_warn_only_enabled() + torch.use_deterministic_algorithms(True) + try: + for dtype in (torch.bfloat16, torch.int64, torch.bool): + with self.subTest(dtype=dtype): + values = src % 2 if dtype == torch.bool else src + self._run(Scatter(), [(x.to(dtype), index, values.to(dtype))]) + finally: + torch.use_deterministic_algorithms(deterministic, warn_only=warn_only) + def test_scatter_mean_counts_in_its_dtype(self): """ATen counts a mean's terms in the tensor's own dtype, so 128 int8 terms count as -128.""" diff --git a/backends/cuda/tests/test_sort_shim.py b/backends/cuda/tests/test_sort_shim.py index 2195ad6d585..4241405b9a0 100644 --- a/backends/cuda/tests/test_sort_shim.py +++ b/backends/cuda/tests/test_sort_shim.py @@ -44,6 +44,7 @@ "aoti_torch_cuda_index_Tensor", "aoti_torch_index_put_out", "aoti_torch_cuda_scatter_reduce_two_out", + "aoti_torch_cuda_scatter_src_out", } ) From 96e111d9434530d835c5b68218b5489fa3139707 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sat, 31 Oct 2026 13:34:07 -0700 Subject: [PATCH 68/89] Let the CUDA scatter shims read an int32 index ATen's scatter and scatter_reduce take an int32 index as well as an int64 one, and Inductor passes it to the fallback shims as it is. The shims took only int64, so a program with an int32 index exported and then failed when it ran. The shared scatter helper now reads either width, the same way the mask index shims already do. The scatter test runs its int64 case with an int32 index. --- backends/cuda/runtime/shims/index.cu | 16 ++++++++++------ backends/cuda/runtime/shims/index.h | 8 ++++---- backends/cuda/tests/test_index_shims.py | 14 ++++++++++---- 3 files changed, 24 insertions(+), 14 deletions(-) diff --git a/backends/cuda/runtime/shims/index.cu b/backends/cuda/runtime/shims/index.cu index 4c36e9c06c3..c8dede95ad2 100644 --- a/backends/cuda/runtime/shims/index.cu +++ b/backends/cuda/runtime/shims/index.cu @@ -884,12 +884,13 @@ struct ScatterPlan { int64_t src_strides[kMaxDims]; int64_t dim_size; // of out int64_t max_out_offset; + bool index_is_int32; }; __global__ void scatter_offsets_kernel( const ScatterPlan plan, int64_t points, - const int64_t* __restrict__ index, + const void* __restrict__ index, int64_t* __restrict__ out_offsets, int64_t* __restrict__ src_offsets) { for (int64_t point = @@ -909,7 +910,9 @@ __global__ void scatter_offsets_kernel( out_offset += coordinate * plan.out_strides[d]; } } - const int64_t position = index[index_offset]; + const int64_t position = plan.index_is_int32 + ? static_cast(index)[index_offset] + : static_cast(index)[index_offset]; // As in ATen, a negative index is out of range, not counted from the end. CUDA_KERNEL_ASSERT_MSG( position >= 0 && position < plan.dim_size, "index out of bounds"); @@ -1023,7 +1026,7 @@ template Error scatter_reduce( const ScatterPlan& plan, int64_t points, - const int64_t* index, + const void* index, void* out, const void* src, Reduction reduction, @@ -1134,9 +1137,9 @@ Error scatter( return Error::Ok; } ET_CHECK_OR_RETURN_ERROR( - index->dtype() == ScalarType::Long, + index->dtype() == ScalarType::Long || index->dtype() == ScalarType::Int, InvalidArgument, - "%s: index must be int64", + "%s: index must be int64 or int32", name); ET_CHECK_OR_RETURN_ERROR( std::max(index->dim(), 1) == ndim && @@ -1154,6 +1157,7 @@ Error scatter( plan.ndim = static_cast(ndim); plan.dim = static_cast(dim); plan.dim_size = size(out, dim); + plan.index_is_int32 = index->dtype() == ScalarType::Int; int64_t points = 1; for (int64_t d = 0; d < ndim; ++d) { plan.max_out_offset += @@ -1184,7 +1188,7 @@ Error scatter( return scatter_reduce( plan, points, - static_cast(index->data_ptr()), + index->data_ptr(), out->data_ptr(), src->data_ptr(), reduction, diff --git a/backends/cuda/runtime/shims/index.h b/backends/cuda/runtime/shims/index.h index 16488a9f52c..8ee1b3d24ff 100644 --- a/backends/cuda/runtime/shims/index.h +++ b/backends/cuda/runtime/shims/index.h @@ -91,8 +91,8 @@ AOTI_SHIM_EXPORT AOTITorchError aoti_torch_index_put_out( * @param out Tensor to update (CUDA device, up to 16 dims) * @param self Must be out; AOTInductor passes the tensor it updates as both * @param dim Dimension of out that index addresses - * @param index int64 tensor with as many dims as out, no larger than src, and - * no larger than out except at dim + * @param index int64 or int32 tensor with as many dims as out, no larger than + * src, and no larger than out except at dim * @param src Values to write, same dtype as out. Float, double, half, * bfloat16, the integer dtypes and bool are supported. * @return AOTITorchError @@ -120,8 +120,8 @@ AOTI_SHIM_EXPORT AOTITorchError aoti_torch_cuda_scatter_src_out( * @param out Tensor to update (CUDA device, up to 16 dims) * @param self Must be out; AOTInductor passes the tensor it updates as both * @param dim Dimension of out that index addresses - * @param index int64 tensor with as many dims as out, no larger than src, and - * no larger than out except at dim + * @param index int64 or int32 tensor with as many dims as out, no larger than + * src, and no larger than out except at dim * @param src Values to reduce, same dtype as out. Float, double, half, * bfloat16 and the integer dtypes are supported, and bool for "sum" only, as * in ATen. diff --git a/backends/cuda/tests/test_index_shims.py b/backends/cuda/tests/test_index_shims.py index 310d6d587f5..98c4a15a372 100644 --- a/backends/cuda/tests/test_index_shims.py +++ b/backends/cuda/tests/test_index_shims.py @@ -197,7 +197,8 @@ def test_scatter_with_tensor_source(self): """Inductor writes a scatter from a tensor itself only in a dtype it can add atomically, so int64, bool and, below sm_90, bfloat16 go to the shim. Of repeated writes the last in index order wins, as in eager's - deterministic mode, which also sends every dtype to the shim.""" + deterministic mode, which also sends every dtype to the shim. As in + ATen, the index can be int32.""" class Scatter(torch.nn.Module): def forward(self, x, index, src): @@ -212,10 +213,15 @@ def forward(self, x, index, src): warn_only = torch.is_deterministic_algorithms_warn_only_enabled() torch.use_deterministic_algorithms(True) try: - for dtype in (torch.bfloat16, torch.int64, torch.bool): - with self.subTest(dtype=dtype): + for dtype, index_dtype in ( + (torch.bfloat16, torch.int64), + (torch.int64, torch.int32), + (torch.bool, torch.int64), + ): + with self.subTest(dtype=dtype, index_dtype=index_dtype): values = src % 2 if dtype == torch.bool else src - self._run(Scatter(), [(x.to(dtype), index, values.to(dtype))]) + inputs = (x.to(dtype), index.to(index_dtype), values.to(dtype)) + self._run(Scatter(), [inputs]) finally: torch.use_deterministic_algorithms(deterministic, warn_only=warn_only) From f021c76cc8a19c5c5874425719035e75c2cb1907 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sat, 31 Oct 2026 13:36:00 -0700 Subject: [PATCH 69/89] Write each element once in the CUDA scatter shim A scatter with a tensor source keeps only the last write to each element. The shim found that write by walking every write to the element in one thread, so many writes to one element were slow. Now the last write of each element stores its value and the others do nothing, so the time no longer depends on how the writes are spread. --- backends/cuda/runtime/shims/index.cu | 24 ++++++++++++++---------- 1 file changed, 14 insertions(+), 10 deletions(-) diff --git a/backends/cuda/runtime/shims/index.cu b/backends/cuda/runtime/shims/index.cu index c8dede95ad2..f4c02f30282 100644 --- a/backends/cuda/runtime/shims/index.cu +++ b/backends/cuda/runtime/shims/index.cu @@ -870,7 +870,7 @@ Error put( return Error::Ok; } -// Assign is a scatter without a reduction: every write replaces the last. +// Assign is a plain scatter: each write replaces the one before it. enum class Reduction { Sum, Prod, Mean, Max, Min, Assign }; // A scatter visits every element of index. Dim `dim` of out is addressed by @@ -936,21 +936,19 @@ combine(Reduction reduction, AccT result, AccT value) { return value_is_nan || result < value ? value : result; case Reduction::Min: return value_is_nan || value < result ? value : result; - case Reduction::Assign: - return value; default: return result + value; } } // After a stable sort of the points by the element they write, the points of -// one element are a run in index order. Each run is reduced by one thread in -// that order and written once, so the result does not depend on scheduling, -// and an assignment keeps the last write, as ATen's deterministic scatter does. -// A sum adds in AccT, as ATen does in deterministic mode, so half and -// bfloat16 lose less to rounding. A product rounds to T at every step and a -// mean rounds its sum to T before dividing, as ATen does, so half overflows to -// infinity and integers wrap where ATen's do. +// one element are a run in index order, and each element is written once, so +// the result does not depend on scheduling. An assignment keeps the last point +// of its run, as ATen's deterministic scatter does. Any other reduction walks +// its run in one thread. A sum adds in AccT, as ATen does in deterministic +// mode, so half and bfloat16 lose less to rounding. A product rounds to T at +// every step and a mean rounds its sum to T before dividing, as ATen does, so +// half overflows to infinity and integers wrap where ATen's do. template __global__ void scatter_reduce_sorted_kernel( int64_t points, @@ -965,6 +963,12 @@ __global__ void scatter_reduce_sorted_kernel( i < points; i += static_cast(gridDim.x) * blockDim.x) { const int64_t offset = out_offsets[i]; + if (reduction == Reduction::Assign) { + if (i + 1 == points || out_offsets[i + 1] != offset) { + out[offset] = src[src_offsets[i]]; + } + continue; + } if (i > 0 && out_offsets[i - 1] == offset) { continue; } From f6108decbaae95869ce2fa1ab2328f2eb99bf275 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sat, 31 Oct 2026 13:36:52 -0700 Subject: [PATCH 70/89] Name the calling CUDA scatter shim in its errors Both scatter shims share one helper, and a few of its errors said scatter_reduce even when the scatter shim had failed. They now name the shim that was called. The notes on an index out of range now say the assert matches ATen's default CUDA kernel, because in deterministic mode ATen wraps a negative index instead. --- backends/cuda/runtime/shims/index.cu | 10 +++++++--- backends/cuda/runtime/shims/index.h | 8 ++++---- backends/cuda/tests/test_index_shims.py | 9 +++++---- 3 files changed, 16 insertions(+), 11 deletions(-) diff --git a/backends/cuda/runtime/shims/index.cu b/backends/cuda/runtime/shims/index.cu index f4c02f30282..919bf887df9 100644 --- a/backends/cuda/runtime/shims/index.cu +++ b/backends/cuda/runtime/shims/index.cu @@ -1028,6 +1028,7 @@ T reduction_identity(Reduction reduction) { template Error scatter_reduce( + const char* name, const ScatterPlan& plan, int64_t points, const void* index, @@ -1041,14 +1042,16 @@ Error scatter_reduce( ET_CHECK_OR_RETURN_ERROR( points <= INT_MAX, NotSupported, - "scatter_reduce: more than %d index elements are not supported", + "%s: more than %d index elements are not supported", + name, INT_MAX); auto* offsets = static_cast(scratch.allocate(4 * points * sizeof(int64_t))); ET_CHECK_OR_RETURN_ERROR( offsets != nullptr, MemoryAllocationFailed, - "scatter_reduce: allocating the offsets failed"); + "%s: allocating the offsets failed", + name); int64_t* out_offsets = offsets; int64_t* src_offsets = offsets + points; int64_t* sorted_out_offsets = offsets + 2 * points; @@ -1070,7 +1073,7 @@ Error scatter_reduce( end_bit, scratch, stream, - "scatter_reduce")); + name)); scatter_reduce_sorted_kernel <<>>( points, @@ -1190,6 +1193,7 @@ Error scatter( using T = decltype(value); using AccT = decltype(accumulator); return scatter_reduce( + name, plan, points, index->data_ptr(), diff --git a/backends/cuda/runtime/shims/index.h b/backends/cuda/runtime/shims/index.h index 8ee1b3d24ff..de4928f64f7 100644 --- a/backends/cuda/runtime/shims/index.h +++ b/backends/cuda/runtime/shims/index.h @@ -85,8 +85,8 @@ AOTI_SHIM_EXPORT AOTITorchError aoti_torch_index_put_out( * bool, and bfloat16 on GPUs older than sm_90), and in deterministic mode. * When values land on the same element, the last one in index order wins, as * in ATen's deterministic mode; ATen's default CUDA kernel keeps any one of - * them. As in ATen, an index out of range, negative ones included, triggers a - * device-side assert. + * them. As in that default kernel, an index out of range, negative ones + * included, triggers a device-side assert. * * @param out Tensor to update (CUDA device, up to 16 dims) * @param self Must be out; AOTInductor passes the tensor it updates as both @@ -114,8 +114,8 @@ AOTI_SHIM_EXPORT AOTITorchError aoti_torch_cuda_scatter_src_out( * are reduced in index order, so the result is deterministic. Half and * bfloat16 sums add in float, as in ATen's deterministic mode. A product * rounds to the dtype of out at every step and a mean rounds its sum before - * dividing, as in ATen. As in ATen, an index out of range, negative ones - * included, triggers a device-side assert. + * dividing, as in ATen. As in ATen's default CUDA kernel, an index out of + * range, negative ones included, triggers a device-side assert. * * @param out Tensor to update (CUDA device, up to 16 dims) * @param self Must be out; AOTInductor passes the tensor it updates as both diff --git a/backends/cuda/tests/test_index_shims.py b/backends/cuda/tests/test_index_shims.py index 98c4a15a372..9a243bad94d 100644 --- a/backends/cuda/tests/test_index_shims.py +++ b/backends/cuda/tests/test_index_shims.py @@ -4,13 +4,14 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -"""Mask indexing, index_put and scatter_reduce in the CUDA backend's index shims. +"""Mask indexing, index_put, scatter and scatter_reduce in the CUDA backend's +index shims. A mask index makes a result size depend on the data, so Inductor leaves it, and an index_put with a mask, to the backend's own index shims. Inductor also -leaves them every scatter_reduce it cannot write with an atomic add. These -tests lower real programs, run them through the Python runtime, and compare -with eager PyTorch on CUDA. +leaves them every scatter and scatter_reduce it cannot write with an atomic +add. These tests lower real programs, run them through the Python runtime, +and compare with eager PyTorch on CUDA. """ import os From c1baac91350f35bd8c7b29c5b6020ead3e23fa53 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sat, 31 Oct 2026 14:50:24 -0700 Subject: [PATCH 71/89] Say which ATen kernel the CUDA scatter bounds check matches The bounds check in the scatter shims said a negative index is out of range "as in ATen", but ATen's deterministic mode wraps a negative index. Like the notes in the header, the comment now names ATen's default CUDA kernel. --- backends/cuda/runtime/shims/index.cu | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/backends/cuda/runtime/shims/index.cu b/backends/cuda/runtime/shims/index.cu index 919bf887df9..42b820d2d4c 100644 --- a/backends/cuda/runtime/shims/index.cu +++ b/backends/cuda/runtime/shims/index.cu @@ -913,7 +913,8 @@ __global__ void scatter_offsets_kernel( const int64_t position = plan.index_is_int32 ? static_cast(index)[index_offset] : static_cast(index)[index_offset]; - // As in ATen, a negative index is out of range, not counted from the end. + // As in ATen's default CUDA kernel, a negative index is out of range, not + // counted from the end. CUDA_KERNEL_ASSERT_MSG( position >= 0 && position < plan.dim_size, "index out of bounds"); out_offsets[point] = out_offset + position * plan.out_strides[plan.dim]; From ab9ed1c08d25992b80354c3f9e0fbaf676f79ef2 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sat, 31 Oct 2026 17:46:16 -0700 Subject: [PATCH 72/89] Run a CUDA graph capture that reads a device value without a graph A `.item()` on a device tensor, such as a torch.cond predicate computed on the GPU, reads the value through a blocking copy to the host. A CUDA graph capture cannot record that copy, and the copy's error check aborted the process, so the capture call never reached the retry without a graph. The item shims now check whether the stream is capturing before they read, as the mask index already does, and refuse with NotSupported. The backend then runs the capture call again without a graph, and later calls run without one too. The capture query gets a HIP spelling, since ROCm builds these shims as well. The CUDA graph test that runs a mask index now also runs a module that scales by a `.item()` read. Without the check, that test aborts the process. --- backends/cuda/runtime/shims/memory.cpp | 19 +++++++++++++++++++ .../tests/test_output_stride_rearrange.py | 18 ++++++++++++++---- extension/cuda/runtime_api.h | 9 +++++++++ 3 files changed, 42 insertions(+), 4 deletions(-) diff --git a/backends/cuda/runtime/shims/memory.cpp b/backends/cuda/runtime/shims/memory.cpp index d9077db619e..a9a39a7c7b1 100644 --- a/backends/cuda/runtime/shims/memory.cpp +++ b/backends/cuda/runtime/shims/memory.cpp @@ -67,6 +67,25 @@ AOTITorchError narrow_to(From value, To* ret_value) { template AOTITorchError read_one_element(const SlimTensor* tensor, T* ret_value) { +#if defined(CUDA_AVAILABLE) + // A graph capture cannot record the copy to the host below, and the copy + // aborts the process when it fails, so refuse first. Info, not Error: the + // backend then retries the call without a graph. + if (tensor->is_cuda()) { + auto stream = getCurrentCUDAStream(tensor->device_index()); + ET_CHECK_OK_OR_RETURN_ERROR(stream.error()); + cudaStreamCaptureStatus capture = cudaStreamCaptureStatusNone; + ET_CUDA_CHECK_OR_RETURN_ERROR( + cudaStreamIsCapturing(stream.get(), &capture)); + if (capture != cudaStreamCaptureStatusNone) { + ET_LOG( + Info, + "reading a single element: a device value cannot be read inside a " + "CUDA graph capture"); + return Error::NotSupported; + } + } +#endif switch (tensor->dtype()) { case ScalarType::Byte: return narrow_to(tensor->item(), ret_value); diff --git a/backends/cuda/tests/test_output_stride_rearrange.py b/backends/cuda/tests/test_output_stride_rearrange.py index 29a136e4902..6dcc69694f0 100644 --- a/backends/cuda/tests/test_output_stride_rearrange.py +++ b/backends/cuda/tests/test_output_stride_rearrange.py @@ -293,6 +293,11 @@ def forward(self, x, mask): return x.index_put((mask,), torch.ones(4), accumulate=True) +class _ScaledByCount(nn.Module): + def forward(self, x, mask): + return x * mask.sum().item() + + class _MaskedCounter(nn.Module): """Counts its calls in a device buffer, and reads a mask's count. The buffer has one element, the case torch 2.14 can fold into a constant.""" @@ -470,12 +475,17 @@ def test_host_state_runs_without_a_graph(self): expected = eager(x.cuda())[0].cpu() torch.testing.assert_close(out, expected) - def test_mask_runs_without_a_graph(self): - # A mask index reads its count back to the host, which a capture - # cannot record, so the capture call has to run without a graph. + def test_host_read_runs_without_a_graph(self): + # A mask index reads its count back to the host, and .item() a value, + # which a capture cannot record, so the capture call has to run + # without a graph. torch.manual_seed(0) inputs = [(torch.randn(8, 4), torch.rand(8) > 0.5) for _ in range(6)] - for module in (_MaskedRows().eval(), _MaskedAdd().eval()): + for module in ( + _MaskedRows().eval(), + _MaskedAdd().eval(), + _ScaledByCount().eval(), + ): with _RuntimeLog() as log: outputs = self._run_with_cuda_graph(module, inputs[0], inputs) # Eager calls give the same values, so only the log shows the diff --git a/extension/cuda/runtime_api.h b/extension/cuda/runtime_api.h index 91a21e53689..ffb53e97db5 100644 --- a/extension/cuda/runtime_api.h +++ b/extension/cuda/runtime_api.h @@ -21,6 +21,7 @@ using cudaGraphExec_t = hipGraphExec_t; using cudaMemcpyKind = hipMemcpyKind; using cudaMemoryType = hipMemoryType; using cudaStreamCaptureMode = hipStreamCaptureMode; +using cudaStreamCaptureStatus = hipStreamCaptureStatus; using cudaStream_t = hipStream_t; inline constexpr cudaError_t cudaSuccess = hipSuccess; @@ -34,6 +35,8 @@ inline constexpr cudaMemoryType cudaMemoryTypeDevice = hipMemoryTypeDevice; inline constexpr cudaMemoryType cudaMemoryTypeManaged = hipMemoryTypeManaged; inline constexpr cudaStreamCaptureMode cudaStreamCaptureModeRelaxed = hipStreamCaptureModeRelaxed; +inline constexpr cudaStreamCaptureStatus cudaStreamCaptureStatusNone = + hipStreamCaptureStatusNone; inline constexpr unsigned long long cudaGraphInstantiateFlagAutoFreeOnLaunch = hipGraphInstantiateFlagAutoFreeOnLaunch; @@ -211,6 +214,12 @@ inline cudaError_t cudaStreamEndCapture( return hipStreamEndCapture(stream, graph); } +inline cudaError_t cudaStreamIsCapturing( + cudaStream_t stream, + cudaStreamCaptureStatus* status) { + return hipStreamIsCapturing(stream, status); +} + inline cudaError_t cudaStreamSynchronize(cudaStream_t stream) { return hipStreamSynchronize(stream); } From 10f3e2fd0c2858d0ffd15fad32c0b617be463aca Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sat, 31 Oct 2026 17:46:16 -0700 Subject: [PATCH 73/89] Rewrite masked copies inside torch.cond branches too The pass that writes `out[mask] = src[mask]` as `torch.where` only walked the top-level graph. A masked copy inside a torch.cond branch, or any other subgraph, kept its index and index_put, and with them the read of the mask count back to the host on every call. The pass now visits every graph module, as the tensor check removal pass does. --- .../passes/replace_masked_copy_with_where.py | 21 ++++++++----- .../test_replace_masked_copy_with_where.py | 30 +++++++++++++++++++ 2 files changed, 43 insertions(+), 8 deletions(-) diff --git a/backends/cuda/passes/replace_masked_copy_with_where.py b/backends/cuda/passes/replace_masked_copy_with_where.py index 03c89dc5f55..941e5f0641c 100644 --- a/backends/cuda/passes/replace_masked_copy_with_where.py +++ b/backends/cuda/passes/replace_masked_copy_with_where.py @@ -87,15 +87,20 @@ class ReplaceMaskedCopyWithWherePass(PassBase): def call(self, graph_module: GraphModule) -> PassResult: modified = False - for node in list(graph_module.graph.nodes): - mask = self._mask_of_copy(node) - if mask is None: + for module in graph_module.modules(): + if not isinstance(module, GraphModule): continue - self._replace(graph_module, node, mask) - modified = True - if modified: - graph_module.graph.lint() - graph_module.recompile() + replaced = False + for node in list(module.graph.nodes): + mask = self._mask_of_copy(node) + if mask is None: + continue + self._replace(module, node, mask) + replaced = True + if replaced: + module.graph.lint() + module.recompile() + modified = True return PassResult(graph_module, modified) @staticmethod diff --git a/backends/cuda/passes/tests/test_replace_masked_copy_with_where.py b/backends/cuda/passes/tests/test_replace_masked_copy_with_where.py index 41b7fd12c01..41bd38a9add 100644 --- a/backends/cuda/passes/tests/test_replace_masked_copy_with_where.py +++ b/backends/cuda/passes/tests/test_replace_masked_copy_with_where.py @@ -83,6 +83,36 @@ def compute(m): ReplaceMaskedCopyWithWherePass()(program.graph_module) self.assertIn(exir_ops.edge.aten.where.self, _targets(program.graph_module)) + def test_masked_copy_in_a_cond_branch(self): + class Copy(torch.nn.Module): + def forward(self, x, mask): + def copy(x, mask): + out = torch.zeros_like(x) + out[mask] = x[mask] + return out + + return torch.cond( + x.sum() > 0, copy, lambda x, mask: x.clone(), (x, mask) + ) + + x = torch.randn(4, 3) + mask = torch.tensor([True, False, True, False]) + program = to_edge(export(Copy(), (x, mask))).exported_program() + ReplaceMaskedCopyWithWherePass()(program.graph_module) + + cond = next( + n + for n in program.graph_module.graph.nodes + if n.target is torch.ops.higher_order.cond + ) + branch = getattr(program.graph_module, cond.args[1].target) + self.assertIn(exir_ops.edge.aten.where.self, _targets(branch)) + self.assertNotIn(exir_ops.edge.aten.index_put.default, _targets(branch)) + for picked in (x.abs(), -x.abs()): + torch.testing.assert_close( + program.module()(picked, mask), Copy()(picked, mask), rtol=0, atol=0 + ) + def test_other_writes_are_kept(self): class Kept(torch.nn.Module): def forward(self, x, mask, idx): From 2b350d95aa5e8a596d0f8dc523ba6d9638f21a90 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sat, 31 Oct 2026 17:46:17 -0700 Subject: [PATCH 74/89] Keep the CUDA compile patches in place while any thread lowers The CUDA backend's compile-time patches that let Triton templates run on small GPUs and that link with `--as-needed` each saved the function they replace on entry and put it back on exit. When two threads lowered at once and the first one finished first, it put back the original while the second one was still compiling, so the rest of that compile ran without the patch. For the link patch, that program then records libgomp and fails to load without torch on a machine that has none. Both patches now share a small helper that counts the threads inside: the first one in installs the patch and the last one out removes it. A thread-local depth still decides, for each thread, whether the patch applies. A new test overlaps two threads and checks that the one that leaves last still sees the patch. --- backends/cuda/cuda_backend.py | 116 ++++++++++++++---------- backends/cuda/tests/test_cuda_export.py | 25 +++++ 2 files changed, 94 insertions(+), 47 deletions(-) diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index 7334e5d0500..5eb11936296 100644 --- a/backends/cuda/cuda_backend.py +++ b/backends/cuda/cuda_backend.py @@ -8,6 +8,7 @@ import contextlib import copy import functools +import importlib import logging import math import os @@ -133,11 +134,58 @@ def _has_rmask(_self) -> bool: IndexingOptions.has_rmask = orig_has_rmask -_SMALL_GPU_GUARD = threading.local() +class _ThreadScopedPatch: + """Makes ``.`` call ``replace(original, *args, **kwargs)`` in + the threads inside ``scope()``, while every other thread keeps the + original. + Threads can lower at the same time. The first scope to enter installs the + swap and the last one to leave restores the original, so a thread that + leaves first does not take the swap away from one still inside. + """ -@contextlib.contextmanager -def _allow_triton_templates_on_small_gpus(): + def __init__( + self, module: str, name: str, replace: typing.Callable[..., Any] + ) -> None: + self._module = module + self._name = name + self._replace = replace + self._lock = threading.Lock() + self._scopes = 0 + self._original: Any = None + self._depth = threading.local() + + @contextlib.contextmanager + def scope(self) -> typing.Iterator[None]: + owner = importlib.import_module(self._module) + with self._lock: + if self._scopes == 0: + self._original = original = getattr(owner, self._name) + + def dispatch(*args: Any, **kwargs: Any) -> Any: + if getattr(self._depth, "value", 0): + return self._replace(original, *args, **kwargs) + return original(*args, **kwargs) + + setattr(owner, self._name, dispatch) + self._scopes += 1 + self._depth.value = getattr(self._depth, "value", 0) + 1 + try: + yield + finally: + self._depth.value -= 1 + with self._lock: + self._scopes -= 1 + if self._scopes == 0: + setattr(owner, self._name, self._original) + + +_SMALL_GPU_TEMPLATES = _ThreadScopedPatch( + "torch._inductor.utils", "is_big_gpu", lambda original, *args, **kwargs: True +) + + +def _allow_triton_templates_on_small_gpus() -> typing.ContextManager[None]: """Let Inductor use Triton templates on GPUs below its 68-SM floor. This backend autotunes with Triton only, so otherwise a matmul or convolution has no Triton choices left. @@ -148,30 +196,29 @@ def _allow_triton_templates_on_small_gpus(): TODO: remove once Inductor skips the floor when ATen is not an allowed backend (https://github.com/pytorch/pytorch/issues/141690). """ - from torch._inductor import utils as inductor_utils + return _SMALL_GPU_TEMPLATES.scope() - orig_is_big_gpu = inductor_utils.is_big_gpu - - def _is_big_gpu(index_or_device=0) -> bool: - return getattr(_SMALL_GPU_GUARD, "active", False) or orig_is_big_gpu( - index_or_device - ) - prev_active = getattr(_SMALL_GPU_GUARD, "active", False) - inductor_utils.is_big_gpu = _is_big_gpu - _SMALL_GPU_GUARD.active = True - try: - yield - finally: - _SMALL_GPU_GUARD.active = prev_active - inductor_utils.is_big_gpu = orig_is_big_gpu +def _with_as_needed(original, device_type, aot_mode=False, *args, **kwargs): + definitions, include_dirs, cflags, ldflags, *rest = original( + device_type, aot_mode, *args, **kwargs + ) + if ( + aot_mode + and sys.platform.startswith("linux") + and not torch._inductor.config.aot_inductor.link_libtorch + and torch._inductor.config.aot_inductor.cross_target_platform is None + ): + ldflags = [*ldflags, "Wl,--as-needed"] + return (definitions, include_dirs, cflags, ldflags, *rest) -_AS_NEEDED_GUARD = threading.local() +_AS_NEEDED_LINKS = _ThreadScopedPatch( + "torch._inductor.cpp_builder", "get_cpp_torch_device_options", _with_as_needed +) -@contextlib.contextmanager -def _link_libraries_as_needed(): +def _link_libraries_as_needed() -> typing.ContextManager[None]: """Record only the shared libraries the compiled program calls into. Inductor puts OpenMP on every Linux link. Unless the compiler asks the @@ -189,32 +236,7 @@ def _link_libraries_as_needed(): TODO: remove once the pinned PyTorch has the same flag in Inductor (https://github.com/pytorch/pytorch/pull/202902). """ - from torch._inductor import cpp_builder - - orig_device_options = cpp_builder.get_cpp_torch_device_options - - def _device_options(device_type, aot_mode=False, *args, **kwargs): - definitions, include_dirs, cflags, ldflags, *rest = orig_device_options( - device_type, aot_mode, *args, **kwargs - ) - if ( - aot_mode - and getattr(_AS_NEEDED_GUARD, "active", False) - and sys.platform.startswith("linux") - and not torch._inductor.config.aot_inductor.link_libtorch - and torch._inductor.config.aot_inductor.cross_target_platform is None - ): - ldflags = [*ldflags, "Wl,--as-needed"] - return (definitions, include_dirs, cflags, ldflags, *rest) - - prev_active = getattr(_AS_NEEDED_GUARD, "active", False) - cpp_builder.get_cpp_torch_device_options = _device_options - _AS_NEEDED_GUARD.active = True - try: - yield - finally: - _AS_NEEDED_GUARD.active = prev_active - cpp_builder.get_cpp_torch_device_options = orig_device_options + return _AS_NEEDED_LINKS.scope() def _full_zeros_preserving_strides(x: torch.Tensor, device) -> torch.Tensor: diff --git a/backends/cuda/tests/test_cuda_export.py b/backends/cuda/tests/test_cuda_export.py index cecc59c2876..c5778966d69 100644 --- a/backends/cuda/tests/test_cuda_export.py +++ b/backends/cuda/tests/test_cuda_export.py @@ -213,6 +213,31 @@ def test_small_gpu_override_answers_only_its_own_thread(self): seen["inside"] = inductor_utils.is_big_gpu(0) self.assertEqual({"inside": True, "other": False}, seen) + def test_small_gpu_override_survives_a_thread_that_leaves_first(self): + from executorch.backends.cuda.cuda_backend import ( + _allow_triton_templates_on_small_gpus, + ) + from torch._inductor import utils as inductor_utils + + entered, leave = threading.Event(), threading.Event() + + def overlap(): + with _allow_triton_templates_on_small_gpus(): + entered.set() + leave.wait(timeout=10) + + with mock.patch.object( + inductor_utils, "is_big_gpu", return_value=False + ) as is_big_gpu: + other = threading.Thread(target=overlap) + other.start() + entered.wait(timeout=10) + with _allow_triton_templates_on_small_gpus(): + leave.set() + other.join(timeout=10) + self.assertTrue(inductor_utils.is_big_gpu(0)) + self.assertIs(inductor_utils.is_big_gpu, is_big_gpu) + def test_cuda_include_ptx_compile_spec(self): with mock.patch.object( CudaBackend, "_setup_cuda_environment_for_fatbin", return_value=True From ba5d8682679a85b5d3d5b0efd0635aded85213ea Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sat, 31 Oct 2026 17:46:17 -0700 Subject: [PATCH 75/89] Pass a float read with .item() to Triton kernels as a float AOTInductor writes one C++ function per Triton kernel to launch it, and it declares every symbolic argument of that function int64_t, even when the kernel takes a float. So a float read with `.item()` reached the kernel truncated, and `x * scale.item()` with a scale of 1.5 computed `x * 1`, with no error. This happened in float32 already, and float64 and float16, whose item shims this pull request adds, ran into it too. A compile-time patch in the CUDA backend now declares such an argument with the float type the kernel's signature gives, which is what the launcher's body already converts it to. It changes nothing else, and it does nothing once PyTorch declares the type itself (https://github.com/pytorch/pytorch/pull/201882). The float64 and float16 item test now scales by 1.5 instead of a whole number. Without the patch, it fails. --- backends/cuda/cuda_backend.py | 44 ++++++++++++++++++++++--- backends/cuda/tests/test_index_shims.py | 7 ++-- 2 files changed, 42 insertions(+), 9 deletions(-) diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index 5eb11936296..78bbc6dae96 100644 --- a/backends/cuda/cuda_backend.py +++ b/backends/cuda/cuda_backend.py @@ -135,9 +135,9 @@ def _has_rmask(_self) -> bool: class _ThreadScopedPatch: - """Makes ``.`` call ``replace(original, *args, **kwargs)`` in + """Makes ``.`` call ``replace(original, *args, **kwargs)`` in the threads inside ``scope()``, while every other thread keeps the - original. + original. ``owner`` names a module, or a class in it as ``module:Class``. Threads can lower at the same time. The first scope to enter installs the swap and the last one to leave restores the original, so a thread that @@ -145,9 +145,9 @@ class _ThreadScopedPatch: """ def __init__( - self, module: str, name: str, replace: typing.Callable[..., Any] + self, owner: str, name: str, replace: typing.Callable[..., Any] ) -> None: - self._module = module + self._owner = owner self._name = name self._replace = replace self._lock = threading.Lock() @@ -157,7 +157,10 @@ def __init__( @contextlib.contextmanager def scope(self) -> typing.Iterator[None]: - owner = importlib.import_module(self._module) + module, _, attribute = self._owner.partition(":") + owner = importlib.import_module(module) + if attribute: + owner = getattr(owner, attribute) with self._lock: if self._scopes == 0: self._original = original = getattr(owner, self._name) @@ -239,6 +242,36 @@ def _link_libraries_as_needed() -> typing.ContextManager[None]: return _AS_NEEDED_LINKS.scope() +def _float_param_from_signature(original, wrapper, name, arg_type, signature=None): + from torch._inductor.codegen.cpp_wrapper_gpu import TRITON_SIGNATURE_TO_CPP + + param = original(wrapper, name, arg_type, signature) + # PyTorch 2.14 does not pass the signature on every path. + kind = (signature or (wrapper.triton_meta or {}).get("signature", {})).get(name) + if param == f"int64_t {name}" and kind in ("fp32", "fp64"): + return f"{TRITON_SIGNATURE_TO_CPP[kind]} {name}" + return param + + +_FLOAT_KERNEL_ARGS = _ThreadScopedPatch( + "torch._inductor.codegen.cpp_wrapper_gpu:DeferredTritonCallWrapper", + "_get_cpp_param_type", + _float_param_from_signature, +) + + +def _pass_float_scalars_to_triton_as_floats() -> typing.ContextManager[None]: + """Declare a float that a Triton kernel takes, such as a value read with + ``.item()``, as a float in the C++ function that launches the kernel. + AOTInductor declares every such argument int64_t, so 1.5 reaches the + kernel as 1. + + TODO: remove once the pinned PyTorch takes the type from the kernel's + signature (https://github.com/pytorch/pytorch/pull/201882). + """ + return _FLOAT_KERNEL_ARGS.scope() + + def _full_zeros_preserving_strides(x: torch.Tensor, device) -> torch.Tensor: """Allocate a zero-filled tensor matching ``x``'s size/stride/dtype on ``device``. @@ -1380,6 +1413,7 @@ def _combined(): # Keep an unused libgomp out of the program, so it loads in a # process without torch. See `_link_libraries_as_needed`. stack.enter_context(_link_libraries_as_needed()) + stack.enter_context(_pass_float_scalars_to_triton_as_floats()) # On ROCm, is_big_gpu gates on architecture, not SM count. if torch.version.hip is None: stack.enter_context(_allow_triton_templates_on_small_gpus()) diff --git a/backends/cuda/tests/test_index_shims.py b/backends/cuda/tests/test_index_shims.py index 9a243bad94d..467e3f33b02 100644 --- a/backends/cuda/tests/test_index_shims.py +++ b/backends/cuda/tests/test_index_shims.py @@ -253,8 +253,7 @@ def forward(self, x, scale): for dtype in (torch.float64, torch.float16): with self.subTest(dtype=dtype): x = torch.randn(4, 1000, dtype=dtype) - # A whole number works around an AOTInductor bug: its Triton - # launcher takes a float read with .item() as int64_t, so 1.5 - # would become 1. - scale = torch.tensor(2.0, dtype=dtype) + # A fraction, which a float passed to Triton as an integer + # loses. + scale = torch.tensor(1.5, dtype=dtype) self._run(Values(), [(x, scale)]) From d996e09730e785bb928a4715ba35da9bc3cc2d21 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sat, 31 Oct 2026 17:46:18 -0700 Subject: [PATCH 76/89] Say the AOTI folding switch covers the whole program The docstring said the option keeps the buffers a program writes out of constant folding. It turns off Inductor's joint-graph constant folding for the whole program, which is what it now says. --- backends/aoti/aoti_backend.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/backends/aoti/aoti_backend.py b/backends/aoti/aoti_backend.py index 949264459d9..a069db4c1fc 100644 --- a/backends/aoti/aoti_backend.py +++ b/backends/aoti/aoti_backend.py @@ -54,12 +54,13 @@ def _checked_data_name(data_name: str) -> str: def _written_buffer_options(program: ExportedProgram) -> Dict[str, Any]: - """Inductor options that keep the buffers a program writes out of constant - folding. + """Inductor options for a program that writes a buffer. - Before torch 2.15, Inductor can fold the update of a one-element buffer, - such as a counter, to its value at export, so every call writes that same - value (https://github.com/pytorch/pytorch/pull/185077). + They turn off Inductor's joint-graph constant folding for the whole + program, not only for the buffers it writes. Before torch 2.15, that + folding can turn the update of a one-element buffer, such as a counter, + into its value at export, so every call writes that same value + (https://github.com/pytorch/pytorch/pull/185077). TODO: remove once the pinned torch is 2.15 or newer. """ From 2a0929447ec94d5489086a6d7f0a7cdb479123e9 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sat, 31 Oct 2026 17:46:18 -0700 Subject: [PATCH 77/89] Test the CUDA scatter shims on more layouts, dtypes and edge cases The scatter and scatter_reduce tests scattered only 2-D tensors along their last dim, so nothing walked the general strides. They now also: - scatter along the middle dim of a 3-D tensor, given as a negative dim, from a strided source and from transposed operands - scatter float32 and float16 in deterministic mode - sum bools, multiply int32 and half values until they wrap and overflow, and reduce with an empty index - sum 1000 half and bfloat16 values in deterministic mode, which the shim adds in float as ATen does - replay a scatter_reduce program in a CUDA graph The deterministic tests share PyTorch's DeterministicGuard instead of saving and restoring the setting by hand. --- backends/cuda/tests/test_index_shims.py | 77 ++++++++++++++++--- .../tests/test_output_stride_rearrange.py | 20 +++++ 2 files changed, 85 insertions(+), 12 deletions(-) diff --git a/backends/cuda/tests/test_index_shims.py b/backends/cuda/tests/test_index_shims.py index 467e3f33b02..5fa60071c0a 100644 --- a/backends/cuda/tests/test_index_shims.py +++ b/backends/cuda/tests/test_index_shims.py @@ -29,6 +29,7 @@ from executorch.exir import EdgeCompileConfig, to_edge_transform_and_lower from executorch.exir.dialects._ops import ops as exir_ops from torch.export import export +from torch.testing._internal.common_utils import DeterministicGuard @unittest.skipUnless(torch.cuda.is_available(), "requires CUDA") @@ -156,17 +157,12 @@ def forward(self, x, idx, values): return summed, written idx = torch.tensor([0, 3, 3, 1, 3]) - deterministic = torch.are_deterministic_algorithms_enabled() - warn_only = torch.is_deterministic_algorithms_warn_only_enabled() - torch.use_deterministic_algorithms(True) - try: + with DeterministicGuard(True): for dtype in (torch.float16, torch.bfloat16): with self.subTest(dtype=dtype): x = torch.arange(8, dtype=dtype) values = torch.arange(1, 6, dtype=dtype) self._run(Put(), [(x, idx, values)]) - finally: - torch.use_deterministic_algorithms(deterministic, warn_only=warn_only) def test_scatter_reduce(self): """Inductor adds atomically but has no atomic max, min or product, so @@ -210,21 +206,41 @@ def forward(self, x, index, src): x = torch.full((5, 6), -1) # Repeated writes differ from the write before them, also as bool. src = torch.arange(20).reshape(4, 5) - deterministic = torch.are_deterministic_algorithms_enabled() - warn_only = torch.is_deterministic_algorithms_warn_only_enabled() - torch.use_deterministic_algorithms(True) - try: + with DeterministicGuard(True): for dtype, index_dtype in ( (torch.bfloat16, torch.int64), (torch.int64, torch.int32), (torch.bool, torch.int64), + (torch.float32, torch.int64), + (torch.float16, torch.int64), ): with self.subTest(dtype=dtype, index_dtype=index_dtype): values = src % 2 if dtype == torch.bool else src inputs = (x.to(dtype), index.to(index_dtype), values.to(dtype)) self._run(Scatter(), [inputs]) - finally: - torch.use_deterministic_algorithms(deterministic, warn_only=warn_only) + + def test_scatter_strides(self): + """Scatters along the middle dim of a 3-D tensor, given as a negative + dim, from a strided source and from transposed operands, so the shims + walk strides that differ between their tensors.""" + + class Strided(torch.nn.Module): + def forward(self, x, index, src): + picked = src[:, ::2] + return ( + x.scatter(-2, index, picked), + x.scatter_reduce(-2, index, picked, "amax"), + x.transpose(0, 2).scatter( + 1, index.transpose(0, 2), picked.transpose(0, 2) + ), + ) + + torch.manual_seed(0) + x = torch.randn(4, 6, 5) + index = torch.randint(0, 6, (3, 4, 5)) + src = torch.randn(3, 8, 5) + with DeterministicGuard(True): + self._run(Strided(), [(x, index, src)]) def test_scatter_mean_counts_in_its_dtype(self): """ATen counts a mean's terms in the tensor's own dtype, so 128 int8 @@ -242,6 +258,43 @@ def forward(self, x, index, src): index = torch.tensor([0] * 128 + [1] * 255) self._run(Mean(), [(x, index, torch.ones(383, dtype=torch.int8))]) + def test_scatter_reduce_edge_cases(self): + """A bool sum, products that overflow int32 and half, and an empty + index.""" + + class EdgeCases(torch.nn.Module): + def forward(self, flags, ints, index, empty): + halves = ints.half() + return ( + flags.scatter_reduce(0, index, flags, "sum"), + # As in ATen, int32 wraps and half overflows to infinity. + ints.scatter_reduce(0, index, ints, "prod"), + halves.scatter_reduce(0, index, halves, "prod"), + halves.scatter_reduce(0, empty, halves[:0], "mean"), + ) + + flags = torch.tensor([True, False, False, True, False]) + ints = torch.tensor([300, -7, 46341, 300, 5], dtype=torch.int32) + index = torch.tensor([0, 2, 2, 0, 1]) + empty = torch.zeros(0, dtype=torch.long) + self._run(EdgeCases(), [(flags, ints, index, empty)]) + + def test_scatter_reduce_sum_adds_in_float(self): + """Deterministic mode sends half and bfloat16 sums to the shim, which + adds in float, as ATen does there, so many small terms add up.""" + + class Sum(torch.nn.Module): + def forward(self, index, src): + return tuple( + torch.zeros(3, dtype=dtype, device=src.device).scatter_reduce( + 0, index, src.to(dtype), "sum" + ) + for dtype in (torch.float16, torch.bfloat16) + ) + + with DeterministicGuard(True): + self._run(Sum(), [(torch.arange(1000) % 3, torch.full((1000,), 0.1))]) + def test_float64_and_float16_values(self): class Values(torch.nn.Module): def forward(self, x, scale): diff --git a/backends/cuda/tests/test_output_stride_rearrange.py b/backends/cuda/tests/test_output_stride_rearrange.py index 6dcc69694f0..bb656e07919 100644 --- a/backends/cuda/tests/test_output_stride_rearrange.py +++ b/backends/cuda/tests/test_output_stride_rearrange.py @@ -298,6 +298,11 @@ def forward(self, x, mask): return x * mask.sum().item() +class _ScatterMax(nn.Module): + def forward(self, x, index, src): + return x.scatter_reduce(1, index, src, "amax") + + class _MaskedCounter(nn.Module): """Counts its calls in a device buffer, and reads a mask's count. The buffer has one element, the case torch 2.14 can fold into a constant.""" @@ -496,6 +501,21 @@ def test_host_read_runs_without_a_graph(self): self.assertIsInstance(out, torch.Tensor) torch.testing.assert_close(out, module(*args)) + def test_scatter_reduce_replays_in_a_graph(self): + # The shim sorts its writes in scratch memory, which the capture has to + # record as well. + torch.manual_seed(0) + index = torch.tensor([[0, 2, 2, 5], [1, 1, 1, 0], [5, 3, 0, 5]]) + inputs = [(torch.randn(5, 6), index, torch.randn(4, 5)) for _ in range(6)] + module = _ScatterMax().eval() + with _RuntimeLog() as log: + outputs = self._run_with_cuda_graph(module, inputs[0], inputs) + self.assertIn("captured and instantiated for 'forward'", log.text) + for i, (args, out) in enumerate(zip(inputs, outputs)): + with self.subTest(call=i): + self.assertIsInstance(out, torch.Tensor) + torch.testing.assert_close(out, module(*args)) + def test_mask_retry_updates_state_once(self): # The capture call runs again without a graph, so its update to the # count must not land twice. From cb41bb96ff3f10b687209f5da1ded3ba68b2eed3 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sat, 31 Oct 2026 17:46:19 -0700 Subject: [PATCH 78/89] Run the varlen_attn and _reshape_copy programs in the CUDA tests The tests for the varlen_attn and _reshape_copy decompositions only lowered their programs, so the code AOTInductor compiles from them never ran. They now run through the Python runtime and match eager PyTorch: - _reshape_copy on a transposed, a sliced and a contiguous input - varlen_attn with no window, a causal mask and the windows (8, 0) and (4, 2), each output flattened, with the sequence boundaries given only at run time The varlen_attn test that compares the decomposition with the op keeps its cases. --- backends/cuda/tests/test_cuda_export.py | 21 +--------- backends/cuda/tests/test_index_shims.py | 53 ++++++++++++++++++++++--- 2 files changed, 49 insertions(+), 25 deletions(-) diff --git a/backends/cuda/tests/test_cuda_export.py b/backends/cuda/tests/test_cuda_export.py index c5778966d69..4b36488ebad 100644 --- a/backends/cuda/tests/test_cuda_export.py +++ b/backends/cuda/tests/test_cuda_export.py @@ -22,7 +22,6 @@ from executorch.exir import EdgeCompileConfig, schema, to_edge_transform_and_lower from executorch.exir.backend.compile_spec_schema import CompileSpec from torch.export import export -from torch.nn.attention.varlen import varlen_attn class TestCudaBackendCompileOptions(unittest.TestCase): @@ -506,30 +505,14 @@ def forward(self, x: torch.Tensor, eps: torch.Tensor) -> torch.Tensor: self.assertNotIn(torch.ops.aten.tril.default, targets) self.assertIn(torch.ops.aten.linear.default, targets) - def test_reshape_copy(self): - """Torch-TensorRT hands back the views it leaves to PyTorch as _reshape_copy.""" - - class ReshapeCopyModule(torch.nn.Module): - def forward(self, x: torch.Tensor) -> torch.Tensor: - return torch.ops.aten._reshape_copy(x.transpose(0, 1), [-1, 4]) * 2 - - self._export_to_cuda_with_lower(ReshapeCopyModule(), (torch.randn(8, 4, 4),)) - - def test_varlen_attn(self): - """varlen_attn calls Flash attention, which the CUDA shims lack.""" - - class VarlenAttention(torch.nn.Module): - def forward(self, q, k, v, cu_seq): - # Flattening views the output, which needs the op's layout. - return varlen_attn(q, k, v, cu_seq, cu_seq, 16, 16).flatten(1) + def test_varlen_attn_decomposition(self): + """The CUDA backend decomposes varlen_attn, which calls Flash attention.""" q, k, v = ( torch.randn(32, 4, 64, dtype=torch.bfloat16, device="cuda") for _ in range(3) ) cu_seq = torch.tensor([0, 16, 32], dtype=torch.int32, device="cuda") - self._export_to_cuda_with_lower(VarlenAttention(), (q, k, v, cu_seq)) - # Causal masks over query sequences shorter and longer than their key # sequences (some queries then see no key, and NaN values outside the # key sequences must stay out), grouped-query attention with an empty diff --git a/backends/cuda/tests/test_index_shims.py b/backends/cuda/tests/test_index_shims.py index 5fa60071c0a..0516cb93c2a 100644 --- a/backends/cuda/tests/test_index_shims.py +++ b/backends/cuda/tests/test_index_shims.py @@ -4,14 +4,14 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -"""Mask indexing, index_put, scatter and scatter_reduce in the CUDA backend's -index shims. +"""Ops that Inductor does not lower by itself in the CUDA backend. A mask index makes a result size depend on the data, so Inductor leaves it, and an index_put with a mask, to the backend's own index shims. Inductor also leaves them every scatter and scatter_reduce it cannot write with an atomic -add. These tests lower real programs, run them through the Python runtime, -and compare with eager PyTorch on CUDA. +add. varlen_attn and _reshape_copy have no Inductor lowering, so the backend +decomposes them. These tests lower real programs, run them through the Python +runtime, and compare with eager PyTorch on CUDA. """ import os @@ -29,6 +29,7 @@ from executorch.exir import EdgeCompileConfig, to_edge_transform_and_lower from executorch.exir.dialects._ops import ops as exir_ops from torch.export import export +from torch.nn.attention.varlen import varlen_attn from torch.testing._internal.common_utils import DeterministicGuard @@ -39,7 +40,10 @@ ) class TestIndexShims(unittest.TestCase): def _run( - self, model: torch.nn.Module, cases: Sequence[Tuple[torch.Tensor, ...]] + self, + model: torch.nn.Module, + cases: Sequence[Tuple[torch.Tensor, ...]], + **tolerance: float, ) -> None: """Lower model for cases[0], then compare every case with eager CUDA.""" spec = CudaBackend.generate_method_name_compile_spec("forward") @@ -69,7 +73,7 @@ def _run( expected = model(*(t.cuda() for t in inputs)) for output, reference in zip(outputs, expected, strict=True): torch.testing.assert_close( - output.cpu(), reference.cpu(), equal_nan=True + output.cpu(), reference.cpu(), equal_nan=True, **tolerance ) def test_mask_index_and_index_put(self): @@ -310,3 +314,40 @@ def forward(self, x, scale): # loses. scale = torch.tensor(1.5, dtype=dtype) self._run(Values(), [(x, scale)]) + + def test_reshape_copy(self): + """Torch-TensorRT hands back the views it leaves to PyTorch as + _reshape_copy.""" + + class ReshapeCopy(torch.nn.Module): + def forward(self, x): + reshape = torch.ops.aten._reshape_copy + return ( + reshape(x.transpose(0, 1), [-1, 4]), + reshape(x[:, 1:3], [8, -1]), + reshape(x, [4, 32]), + ) + + self._run(ReshapeCopy(), [(torch.randn(8, 4, 4),)]) + + def test_varlen_attn(self): + """varlen_attn calls Flash attention, which the CUDA shims lack.""" + + class Attention(torch.nn.Module): + def forward(self, q, k, v, cu_seq): + # Flattening views each output, which needs the op's layout. + return tuple( + varlen_attn( + q, k, v, cu_seq, cu_seq, 16, 16, window_size=window + ).flatten(1) + for window in ((-1, -1), (-1, 0), (8, 0), (4, 2)) + ) + + torch.manual_seed(0) + q, k, v = (torch.randn(28, 4, 64, dtype=torch.bfloat16) for _ in range(3)) + # The sequence boundaries come only at run time. + cases = [ + (q, k, v, torch.tensor([0, 12, 28], dtype=torch.int32)), + (q, k, v, torch.tensor([0, 16, 28], dtype=torch.int32)), + ] + self._run(Attention(), cases, atol=1e-2, rtol=1e-2) From b2137cd4beafc2717e114154fca9bb5d7531504f Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sat, 31 Oct 2026 18:07:29 -0700 Subject: [PATCH 79/89] Keep PyTorch's internal test helpers out of the CUDA index shim tests The index shim tests used DeterministicGuard from torch.testing._internal.common_utils. Importing that module turns off setting torch.backends.cudnn flags for the whole process and wraps every TorchScript call, and the CUDA tests all run in one pytest process. The deterministic tests now share a small helper in the test file that saves and restores deterministic mode, as they did before. --- backends/cuda/tests/test_index_shims.py | 25 +++++++++++++++++++------ 1 file changed, 19 insertions(+), 6 deletions(-) diff --git a/backends/cuda/tests/test_index_shims.py b/backends/cuda/tests/test_index_shims.py index 0516cb93c2a..286aef0be30 100644 --- a/backends/cuda/tests/test_index_shims.py +++ b/backends/cuda/tests/test_index_shims.py @@ -14,10 +14,11 @@ runtime, and compare with eager PyTorch on CUDA. """ +import contextlib import os import tempfile import unittest -from typing import Sequence, Tuple +from typing import Iterator, Sequence, Tuple from unittest.mock import patch import torch @@ -30,7 +31,19 @@ from executorch.exir.dialects._ops import ops as exir_ops from torch.export import export from torch.nn.attention.varlen import varlen_attn -from torch.testing._internal.common_utils import DeterministicGuard + + +@contextlib.contextmanager +def _deterministic() -> Iterator[None]: + """Deterministic mode, which also sends every scatter and index_put to the + shims.""" + enabled = torch.are_deterministic_algorithms_enabled() + warn_only = torch.is_deterministic_algorithms_warn_only_enabled() + torch.use_deterministic_algorithms(True) + try: + yield + finally: + torch.use_deterministic_algorithms(enabled, warn_only=warn_only) @unittest.skipUnless(torch.cuda.is_available(), "requires CUDA") @@ -161,7 +174,7 @@ def forward(self, x, idx, values): return summed, written idx = torch.tensor([0, 3, 3, 1, 3]) - with DeterministicGuard(True): + with _deterministic(): for dtype in (torch.float16, torch.bfloat16): with self.subTest(dtype=dtype): x = torch.arange(8, dtype=dtype) @@ -210,7 +223,7 @@ def forward(self, x, index, src): x = torch.full((5, 6), -1) # Repeated writes differ from the write before them, also as bool. src = torch.arange(20).reshape(4, 5) - with DeterministicGuard(True): + with _deterministic(): for dtype, index_dtype in ( (torch.bfloat16, torch.int64), (torch.int64, torch.int32), @@ -243,7 +256,7 @@ def forward(self, x, index, src): x = torch.randn(4, 6, 5) index = torch.randint(0, 6, (3, 4, 5)) src = torch.randn(3, 8, 5) - with DeterministicGuard(True): + with _deterministic(): self._run(Strided(), [(x, index, src)]) def test_scatter_mean_counts_in_its_dtype(self): @@ -296,7 +309,7 @@ def forward(self, index, src): for dtype in (torch.float16, torch.bfloat16) ) - with DeterministicGuard(True): + with _deterministic(): self._run(Sum(), [(torch.arange(1000) % 3, torch.full((1000,), 0.1))]) def test_float64_and_float16_values(self): From ec7046c913bb00c201b442f80ad57dece485d6fe Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sat, 31 Oct 2026 19:16:00 -0700 Subject: [PATCH 80/89] Assign the CUDA compile patches explicitly The helper that keeps the CUDA backend's compile patches in place while any thread lowers found each patched function by name and swapped it with setattr. The repository's style asks for explicit state instead of dynamic setattr and getattr, and a name in a string hides the patched function from readers and tools. Each patch now has a small install function that assigns its replacement and returns the function that puts the original back, and the shared helper only counts the threads inside and tells the patched code whether the calling thread asked for it. The behavior does not change. --- backends/cuda/cuda_backend.py | 152 ++++++++++++++++++++-------------- 1 file changed, 90 insertions(+), 62 deletions(-) diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index 78bbc6dae96..7db9a4b20ea 100644 --- a/backends/cuda/cuda_backend.py +++ b/backends/cuda/cuda_backend.py @@ -8,7 +8,6 @@ import contextlib import copy import functools -import importlib import logging import math import os @@ -134,45 +133,40 @@ def _has_rmask(_self) -> bool: IndexingOptions.has_rmask = orig_has_rmask -class _ThreadScopedPatch: - """Makes ``.`` call ``replace(original, *args, **kwargs)`` in - the threads inside ``scope()``, while every other thread keeps the - original. ``owner`` names a module, or a class in it as ``module:Class``. +class _Depth(threading.local): + def __init__(self) -> None: + self.value = 0 - Threads can lower at the same time. The first scope to enter installs the - swap and the last one to leave restores the original, so a thread that - leaves first does not take the swap away from one still inside. + +class _CompilePatch: + """A patch of Inductor that changes behavior only for the threads inside + ``scope()``. + + ``install`` applies the patch and returns the function that removes it. + The patched code calls the ``active`` function it is given to decide, on + each call, whether the calling thread asked for the change. Threads can + lower at the same time, so the first scope to enter installs the patch + and the last one to leave removes it: a thread that leaves first does not + take the patch away from one still compiling. """ - def __init__( - self, owner: str, name: str, replace: typing.Callable[..., Any] - ) -> None: - self._owner = owner - self._name = name - self._replace = replace + def __init__(self, install: typing.Callable[..., typing.Callable[[], None]]): + self._install = install + self._remove: typing.Callable[[], None] = lambda: None self._lock = threading.Lock() self._scopes = 0 - self._original: Any = None - self._depth = threading.local() + self._depth = _Depth() + + def _active(self) -> bool: + return self._depth.value > 0 @contextlib.contextmanager def scope(self) -> typing.Iterator[None]: - module, _, attribute = self._owner.partition(":") - owner = importlib.import_module(module) - if attribute: - owner = getattr(owner, attribute) with self._lock: if self._scopes == 0: - self._original = original = getattr(owner, self._name) - - def dispatch(*args: Any, **kwargs: Any) -> Any: - if getattr(self._depth, "value", 0): - return self._replace(original, *args, **kwargs) - return original(*args, **kwargs) - - setattr(owner, self._name, dispatch) + self._remove = self._install(self._active) self._scopes += 1 - self._depth.value = getattr(self._depth, "value", 0) + 1 + self._depth.value += 1 try: yield finally: @@ -180,12 +174,26 @@ def dispatch(*args: Any, **kwargs: Any) -> Any: with self._lock: self._scopes -= 1 if self._scopes == 0: - setattr(owner, self._name, self._original) + self._remove() -_SMALL_GPU_TEMPLATES = _ThreadScopedPatch( - "torch._inductor.utils", "is_big_gpu", lambda original, *args, **kwargs: True -) +def _install_small_gpu_override(active): + from torch._inductor import utils as inductor_utils + + original = inductor_utils.is_big_gpu + + def is_big_gpu(index_or_device=0) -> bool: + return active() or original(index_or_device) + + inductor_utils.is_big_gpu = is_big_gpu + + def remove() -> None: + inductor_utils.is_big_gpu = original + + return remove + + +_SMALL_GPU_OVERRIDE = _CompilePatch(_install_small_gpu_override) def _allow_triton_templates_on_small_gpus() -> typing.ContextManager[None]: @@ -199,26 +207,37 @@ def _allow_triton_templates_on_small_gpus() -> typing.ContextManager[None]: TODO: remove once Inductor skips the floor when ATen is not an allowed backend (https://github.com/pytorch/pytorch/issues/141690). """ - return _SMALL_GPU_TEMPLATES.scope() + return _SMALL_GPU_OVERRIDE.scope() -def _with_as_needed(original, device_type, aot_mode=False, *args, **kwargs): - definitions, include_dirs, cflags, ldflags, *rest = original( - device_type, aot_mode, *args, **kwargs - ) - if ( - aot_mode - and sys.platform.startswith("linux") - and not torch._inductor.config.aot_inductor.link_libtorch - and torch._inductor.config.aot_inductor.cross_target_platform is None - ): - ldflags = [*ldflags, "Wl,--as-needed"] - return (definitions, include_dirs, cflags, ldflags, *rest) +def _install_as_needed_links(active): + from torch._inductor import cpp_builder + original = cpp_builder.get_cpp_torch_device_options + + def device_options(device_type, aot_mode=False, *args, **kwargs): + definitions, include_dirs, cflags, ldflags, *rest = original( + device_type, aot_mode, *args, **kwargs + ) + if ( + aot_mode + and active() + and sys.platform.startswith("linux") + and not torch._inductor.config.aot_inductor.link_libtorch + and torch._inductor.config.aot_inductor.cross_target_platform is None + ): + ldflags = [*ldflags, "Wl,--as-needed"] + return (definitions, include_dirs, cflags, ldflags, *rest) + + cpp_builder.get_cpp_torch_device_options = device_options + + def remove() -> None: + cpp_builder.get_cpp_torch_device_options = original + + return remove -_AS_NEEDED_LINKS = _ThreadScopedPatch( - "torch._inductor.cpp_builder", "get_cpp_torch_device_options", _with_as_needed -) + +_AS_NEEDED_LINKS = _CompilePatch(_install_as_needed_links) def _link_libraries_as_needed() -> typing.ContextManager[None]: @@ -242,22 +261,31 @@ def _link_libraries_as_needed() -> typing.ContextManager[None]: return _AS_NEEDED_LINKS.scope() -def _float_param_from_signature(original, wrapper, name, arg_type, signature=None): - from torch._inductor.codegen.cpp_wrapper_gpu import TRITON_SIGNATURE_TO_CPP +def _install_float_kernel_args(active): + from torch._inductor.codegen.cpp_wrapper_gpu import ( + DeferredTritonCallWrapper, + TRITON_SIGNATURE_TO_CPP, + ) + + original = DeferredTritonCallWrapper._get_cpp_param_type - param = original(wrapper, name, arg_type, signature) - # PyTorch 2.14 does not pass the signature on every path. - kind = (signature or (wrapper.triton_meta or {}).get("signature", {})).get(name) - if param == f"int64_t {name}" and kind in ("fp32", "fp64"): - return f"{TRITON_SIGNATURE_TO_CPP[kind]} {name}" - return param + def get_cpp_param_type(self, name, arg_type, signature=None): + param = original(self, name, arg_type, signature) + # PyTorch 2.14 does not pass the signature on every path. + kind = (signature or (self.triton_meta or {}).get("signature", {})).get(name) + if active() and param == f"int64_t {name}" and kind in ("fp32", "fp64"): + return f"{TRITON_SIGNATURE_TO_CPP[kind]} {name}" + return param + DeferredTritonCallWrapper._get_cpp_param_type = get_cpp_param_type -_FLOAT_KERNEL_ARGS = _ThreadScopedPatch( - "torch._inductor.codegen.cpp_wrapper_gpu:DeferredTritonCallWrapper", - "_get_cpp_param_type", - _float_param_from_signature, -) + def remove() -> None: + DeferredTritonCallWrapper._get_cpp_param_type = original + + return remove + + +_FLOAT_KERNEL_ARGS = _CompilePatch(_install_float_kernel_args) def _pass_float_scalars_to_triton_as_floats() -> typing.ContextManager[None]: From 19a80289fb201f7f9760104ddd34078258ffaa0c Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sat, 31 Oct 2026 19:16:01 -0700 Subject: [PATCH 81/89] Test that the masked copy pass keeps two random masks apart The pass treats two masks as one when they are the same deterministic computation, so `out[m] = x[m]` with `m` computed twice can become a `where`. A random mask is drawn anew each time, so two draws must not count as one; the pass already checks for that, but no test covered it. A new test writes through two draws of `torch.rand(4) > 0.5` and checks that the index_put stays. It fails without the check. --- .../tests/test_replace_masked_copy_with_where.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/backends/cuda/passes/tests/test_replace_masked_copy_with_where.py b/backends/cuda/passes/tests/test_replace_masked_copy_with_where.py index 41bd38a9add..cc1310b7a8b 100644 --- a/backends/cuda/passes/tests/test_replace_masked_copy_with_where.py +++ b/backends/cuda/passes/tests/test_replace_masked_copy_with_where.py @@ -154,6 +154,21 @@ def forward(self, x, mask, idx): ): torch.testing.assert_close(got, expected, rtol=0, atol=0) + def test_random_masks_are_kept(self): + """Two draws of a random mask are different masks.""" + + class Copy(torch.nn.Module): + def forward(self, x): + out = torch.zeros_like(x) + out[torch.rand(4) > 0.5] = x[torch.rand(4) > 0.5] + return out + + program = to_edge(export(Copy(), (torch.randn(4, 3),))).exported_program() + ReplaceMaskedCopyWithWherePass()(program.graph_module) + self.assertIn( + exir_ops.edge.aten.index_put.default, _targets(program.graph_module) + ) + def test_check_on_another_size_keeps_the_index(self): """A check that reads the count but tests another size does not hold for every mask.""" From e922a556aada8a0fd78c3b93f1763061766ba2d7 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sat, 31 Oct 2026 19:16:01 -0700 Subject: [PATCH 82/89] Check that a failed CUDA graph capture is tried only once After a capture fails, the backend turns graphs off for the method so that later calls do not try again. The CUDA graph tests only checked that the log mentions a failed capture, which still holds if every later call tried and failed. They now check that it is mentioned exactly once. --- backends/cuda/tests/test_output_stride_rearrange.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/backends/cuda/tests/test_output_stride_rearrange.py b/backends/cuda/tests/test_output_stride_rearrange.py index bb656e07919..a53156cd3f0 100644 --- a/backends/cuda/tests/test_output_stride_rearrange.py +++ b/backends/cuda/tests/test_output_stride_rearrange.py @@ -494,8 +494,8 @@ def test_host_read_runs_without_a_graph(self): with _RuntimeLog() as log: outputs = self._run_with_cuda_graph(module, inputs[0], inputs) # Eager calls give the same values, so only the log shows the - # capture was tried. - self.assertIn("capture failed for 'forward'", log.text) + # capture was tried, and only once. + self.assertEqual(log.text.count("capture failed for 'forward'"), 1) for i, (args, out) in enumerate(zip(inputs, outputs)): with self.subTest(module=type(module).__name__, call=i): self.assertIsInstance(out, torch.Tensor) @@ -525,7 +525,7 @@ def test_mask_retry_updates_state_once(self): outputs = self._run_with_cuda_graph( _MaskedCounter().eval(), inputs[0], inputs ) - self.assertIn("capture failed for 'forward'", log.text) + self.assertEqual(log.text.count("capture failed for 'forward'"), 1) for i, out in enumerate(outputs): with self.subTest(call=i): self.assertIsInstance(out, torch.Tensor) From 7e59c0a52b05319c72b50b9a4316896af66f3b0b Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sat, 31 Oct 2026 19:16:01 -0700 Subject: [PATCH 83/89] Lower the _reshape_copy and varlen_attn programs on ROCm again The runtime tests for these decompositions live in test_index_shims.py, which skips on ROCm because the ROCm build has no CUDA backend in the Python runtime. Moving the lowering checks there left nothing that lowers the two decompositions on ROCm, where the CUDA export tests still run. The lowering tests in test_cuda_export.py are back, next to the comparison of the varlen_attn decomposition with the op. --- backends/cuda/tests/test_cuda_export.py | 21 +++++++++++++++++++-- 1 file changed, 19 insertions(+), 2 deletions(-) diff --git a/backends/cuda/tests/test_cuda_export.py b/backends/cuda/tests/test_cuda_export.py index 4b36488ebad..c5778966d69 100644 --- a/backends/cuda/tests/test_cuda_export.py +++ b/backends/cuda/tests/test_cuda_export.py @@ -22,6 +22,7 @@ from executorch.exir import EdgeCompileConfig, schema, to_edge_transform_and_lower from executorch.exir.backend.compile_spec_schema import CompileSpec from torch.export import export +from torch.nn.attention.varlen import varlen_attn class TestCudaBackendCompileOptions(unittest.TestCase): @@ -505,14 +506,30 @@ def forward(self, x: torch.Tensor, eps: torch.Tensor) -> torch.Tensor: self.assertNotIn(torch.ops.aten.tril.default, targets) self.assertIn(torch.ops.aten.linear.default, targets) - def test_varlen_attn_decomposition(self): - """The CUDA backend decomposes varlen_attn, which calls Flash attention.""" + def test_reshape_copy(self): + """Torch-TensorRT hands back the views it leaves to PyTorch as _reshape_copy.""" + + class ReshapeCopyModule(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.ops.aten._reshape_copy(x.transpose(0, 1), [-1, 4]) * 2 + + self._export_to_cuda_with_lower(ReshapeCopyModule(), (torch.randn(8, 4, 4),)) + + def test_varlen_attn(self): + """varlen_attn calls Flash attention, which the CUDA shims lack.""" + + class VarlenAttention(torch.nn.Module): + def forward(self, q, k, v, cu_seq): + # Flattening views the output, which needs the op's layout. + return varlen_attn(q, k, v, cu_seq, cu_seq, 16, 16).flatten(1) q, k, v = ( torch.randn(32, 4, 64, dtype=torch.bfloat16, device="cuda") for _ in range(3) ) cu_seq = torch.tensor([0, 16, 32], dtype=torch.int32, device="cuda") + self._export_to_cuda_with_lower(VarlenAttention(), (q, k, v, cu_seq)) + # Causal masks over query sequences shorter and longer than their key # sequences (some queries then see no key, and NaN values outside the # key sequences must stay out), grouped-query attention with an empty From 8f9dc8fbd25606de980e138b377b91a13c999155 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Mon, 2 Nov 2026 08:59:39 -0800 Subject: [PATCH 84/89] Use a stable softmax in the CUDA varlen_attn decomposition At large scores, float32 rounds the log-sum-exp as coarsely as the scores, so exp(scores - logsumexp) gave weights that no longer added up to one. With four equal scores of 8e8 each key got a weight of one, and the output came out four times too large. The decomposition now uses softmax, which subtracts the row maximum first, and still gives a query with no visible key zeros and an infinite log-sum-exp. --- backends/cuda/cuda_backend.py | 7 +++++-- backends/cuda/tests/test_cuda_export.py | 3 +++ backends/cuda/tests/test_index_shims.py | 3 +++ 3 files changed, 11 insertions(+), 2 deletions(-) diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index 7db9a4b20ea..89bc1ca46f5 100644 --- a/backends/cuda/cuda_backend.py +++ b/backends/cuda/cuda_backend.py @@ -786,8 +786,11 @@ def sequence_ends(cu: torch.Tensor, total: int): scores = torch.einsum("qhd,khd->hqk", query.float(), key.float()) * scale scores = scores.masked_fill(~allowed, float("-inf")) lse = scores.logsumexp(-1) - lse = lse.masked_fill(lse == float("-inf"), float("inf")) - probs = torch.exp(scores - lse[..., None]) + empty = lse == float("-inf") + # Not exp(scores - lse): at large scores, float32 rounds lse as coarsely + # as the scores, and the weights no longer add up to one. + probs = scores.softmax(-1).masked_fill(empty[..., None], 0) + lse = lse.masked_fill(empty, float("inf")) out = torch.einsum("hqk,khd->qhd", probs, value.float()).to(query.dtype) rng_state = torch.zeros(2, dtype=torch.uint64, device=query.device) # einsum gives a permuted view. The op lays out its output in the diff --git a/backends/cuda/tests/test_cuda_export.py b/backends/cuda/tests/test_cuda_export.py index c5778966d69..1a5e5915cdf 100644 --- a/backends/cuda/tests/test_cuda_export.py +++ b/backends/cuda/tests/test_cuda_export.py @@ -540,11 +540,14 @@ def forward(self, q, k, v, cu_seq): shifted = torch.tensor([2, 6, 16], dtype=torch.int32, device="cuda") nan_outside = v.clone() nan_outside[:2] = nan_outside[16:] = float("nan") + big = torch.full_like(q, 1024) cases = [ (q[:16], k, v, short, cu_seq, 12, 16, True, None, [3, -1]), (q, k, nan_outside, cu_seq, shifted, 16, 12, True, None, [-1, 0]), (q, k[:, :2], v[:, :2], empty, empty, 16, 16, False, 0.5, None, True), (q, k, v, cu_seq, cu_seq, 16, 16, False, None, [4, 2]), + # Equal scores so large that float32 rounds their log-sum-exp. + (big, big, v, cu_seq, cu_seq, 16, 16), ] table = CudaBackend.get_decomposition_table() for args in cases: diff --git a/backends/cuda/tests/test_index_shims.py b/backends/cuda/tests/test_index_shims.py index 286aef0be30..af7275cfe21 100644 --- a/backends/cuda/tests/test_index_shims.py +++ b/backends/cuda/tests/test_index_shims.py @@ -358,9 +358,12 @@ def forward(self, q, k, v, cu_seq): torch.manual_seed(0) q, k, v = (torch.randn(28, 4, 64, dtype=torch.bfloat16) for _ in range(3)) + # Equal scores so large that float32 rounds their log-sum-exp. + big = torch.full_like(q, 1024) # The sequence boundaries come only at run time. cases = [ (q, k, v, torch.tensor([0, 12, 28], dtype=torch.int32)), (q, k, v, torch.tensor([0, 16, 28], dtype=torch.int32)), + (big, big, v, torch.tensor([0, 12, 28], dtype=torch.int32)), ] self._run(Attention(), cases, atol=1e-2, rtol=1e-2) From fa053a78abb0a4e6bb23b29e5fc98ce4c9149031 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Mon, 2 Nov 2026 08:59:39 -0800 Subject: [PATCH 85/89] Keep masked copies whose layouts differ index_put keeps the destination's strides, but where takes its layout from its inputs, the mask first. A channels_last source or mask with a contiguous destination then failed to lower when a later view needed the destination's layout. The pass now rewrites a masked copy only when where would lay out its result like the destination. --- .../passes/replace_masked_copy_with_where.py | 15 ++++++++-- .../test_replace_masked_copy_with_where.py | 29 +++++++++++++++++++ 2 files changed, 42 insertions(+), 2 deletions(-) diff --git a/backends/cuda/passes/replace_masked_copy_with_where.py b/backends/cuda/passes/replace_masked_copy_with_where.py index 941e5f0641c..2218af81341 100644 --- a/backends/cuda/passes/replace_masked_copy_with_where.py +++ b/backends/cuda/passes/replace_masked_copy_with_where.py @@ -77,8 +77,9 @@ class ReplaceMaskedCopyWithWherePass(PassBase): """ Replace ``index_put(out, indices, index(src, indices))`` with ``where(mask, src, out)`` when ``indices`` is one boolean mask (after any - number of ``None``), ``src`` and ``out`` have the same shape and dtype, and - the put does not accumulate. + number of ``None``), ``src`` and ``out`` have the same shape and dtype, + ``where`` lays out its result like ``out``, and the put does not + accumulate. The two masks may be one node or two equal computations. The ``index`` is then removed too if nothing else reads it but the range checks export adds @@ -133,6 +134,16 @@ def _mask_of_copy(node: Node) -> Optional[Node]: statically_known_true(a == b) for a, b in zip(src_val.shape, out_val.shape) ): return None + # index_put keeps out's layout, but where takes its layout from its + # inputs, the mask first. + condition = mask_val + for _ in range(out_val.dim() - indices.index(masks[0]) - mask_val.dim()): + condition = condition.unsqueeze(-1) + layout = torch.where(condition, src_val, out_val).stride() + if not all( + statically_known_true(a == b) for a, b in zip(layout, out_val.stride()) + ): + return None return masks[0] def _replace(self, graph_module: GraphModule, node: Node, mask: Node) -> None: diff --git a/backends/cuda/passes/tests/test_replace_masked_copy_with_where.py b/backends/cuda/passes/tests/test_replace_masked_copy_with_where.py index cc1310b7a8b..e2d1e8e7af8 100644 --- a/backends/cuda/passes/tests/test_replace_masked_copy_with_where.py +++ b/backends/cuda/passes/tests/test_replace_masked_copy_with_where.py @@ -154,6 +154,35 @@ def forward(self, x, mask, idx): ): torch.testing.assert_close(got, expected, rtol=0, atol=0) + def test_different_layouts_are_kept(self): + """index_put keeps the destination's layout, but where may take the + mask's or the source's, and a later view(-1) then fails.""" + + class Source(torch.nn.Module): + def forward(self, x, mask): + src = x.contiguous(memory_format=torch.channels_last) + out = torch.zeros_like(x, memory_format=torch.contiguous_format) + out[mask] = src[mask] + return out + + class Mask(torch.nn.Module): + def forward(self, x, y): + mask = y.contiguous(memory_format=torch.channels_last) > 0 + out = torch.zeros_like(x) + out[mask] = x[mask] + return out + + x = torch.randn(4, 3, 5, 6) + for model, inputs in ( + (Source(), (x, torch.tensor([True, False, True, True]))), + (Mask(), (x, torch.randn(4, 3, 5, 6))), + ): + program = to_edge(export(model, inputs)).exported_program() + ReplaceMaskedCopyWithWherePass()(program.graph_module) + targets = _targets(program.graph_module) + self.assertIn(exir_ops.edge.aten.index_put.default, targets) + self.assertNotIn(exir_ops.edge.aten.where.self, targets) + def test_random_masks_are_kept(self): """Two draws of a random mask are different masks.""" From 88f1538e810da262374e68bebcd63e993f69f138 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Mon, 2 Nov 2026 08:59:40 -0800 Subject: [PATCH 86/89] Tidy the masked copy pass and its tests Spell out the node and plain value cases in _same_value, drop two list() calls that did nothing, and stop naming a mask entry index. The copy test writes into a nonzero destination, so it fails if where loses the values the mask does not pick. --- backends/cuda/passes/replace_masked_copy_with_where.py | 10 ++++++---- .../tests/test_replace_masked_copy_with_where.py | 6 ++++-- 2 files changed, 10 insertions(+), 6 deletions(-) diff --git a/backends/cuda/passes/replace_masked_copy_with_where.py b/backends/cuda/passes/replace_masked_copy_with_where.py index 2218af81341..e00b1ae4fd7 100644 --- a/backends/cuda/passes/replace_masked_copy_with_where.py +++ b/backends/cuda/passes/replace_masked_copy_with_where.py @@ -52,8 +52,10 @@ def _same_value(a: Any, b: Any) -> bool: return False pending.extend(zip(a, b)) continue - if not isinstance(a, Node) or not isinstance(b, Node): - if isinstance(a, Node) or isinstance(b, Node) or a != b: + if isinstance(a, Node) != isinstance(b, Node): + return False + if not isinstance(a, Node): + if a != b: return False continue if a is b or (a, b) in compared: @@ -116,9 +118,9 @@ def _mask_of_copy(node: Node) -> Optional[Node]: if accumulate or not isinstance(values, Node) or values.target != _INDEX: return None src, src_indices = values.args - if not _same_value(list(src_indices), list(indices)): + if not _same_value(src_indices, indices): return None - masks = [index for index in indices if index is not None] + masks = [entry for entry in indices if entry is not None] if len(masks) != 1: return None mask_val = masks[0].meta.get("val") diff --git a/backends/cuda/passes/tests/test_replace_masked_copy_with_where.py b/backends/cuda/passes/tests/test_replace_masked_copy_with_where.py index e2d1e8e7af8..cefc8f11b80 100644 --- a/backends/cuda/passes/tests/test_replace_masked_copy_with_where.py +++ b/backends/cuda/passes/tests/test_replace_masked_copy_with_where.py @@ -33,7 +33,8 @@ def test_masked_copy_becomes_where(self): class Copy(torch.nn.Module): def forward(self, x, mask): - out = torch.zeros_like(x) + # Not zeros, so a lost destination value shows. + out = torch.full_like(x, 7.0) # A leading None and a trailing dim, and the mask computed twice, # as Transformers' SmolVLM writes its position ids. out[:, mask.view(-1)] = x[:, mask.view(-1)] @@ -219,7 +220,8 @@ def forward(self, x, mask, y): n for n in graph.nodes if n.target == exir_ops.edge.aten.index.Tensor ) y_node = next(n for n in graph.nodes if n.name == "y") - # count * 0 + y.size(0) <= 4 has the one symbol of y's size. + # count * 0 + y.size(0) <= 4 reads the count, but whether it holds + # depends only on y's size. with graph.inserting_before(index.next): count = graph.call_function(torch.ops.aten.sym_size.int, (index, 0)) size = graph.call_function(torch.ops.aten.sym_size.int, (y_node, 0)) From b508df460810a2f9ba2d055e4294225675bc64f8 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Mon, 2 Nov 2026 08:59:40 -0800 Subject: [PATCH 87/89] Check that the CUDA backend runs RemoveAssertAsyncPass The export in test_distribution_argument_check fails without the pass only when autotuning trips the assert, so the test now also checks that the backend lists the pass. The pass docstring says why it keeps the size checks. --- backends/cuda/passes/remove_assert_async.py | 4 +++- backends/cuda/tests/test_cuda_export.py | 6 ++++++ 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/backends/cuda/passes/remove_assert_async.py b/backends/cuda/passes/remove_assert_async.py index b6fb062626c..08f9964d263 100644 --- a/backends/cuda/passes/remove_assert_async.py +++ b/backends/cuda/passes/remove_assert_async.py @@ -25,7 +25,9 @@ class RemoveAssertAsyncPass(PassBase): ExecuTorch already removes these checks from the rest of a program before it is saved (see ``RemoveGraphAssertsPass``), so the delegate now does the - same. + same. Unlike that pass, this one keeps the size checks, such as + ``aten._assert_scalar``: they run on the host, so they cannot break the + CUDA context. """ def call(self, graph_module: GraphModule) -> PassResult: diff --git a/backends/cuda/tests/test_cuda_export.py b/backends/cuda/tests/test_cuda_export.py index 1a5e5915cdf..ba3ae354250 100644 --- a/backends/cuda/tests/test_cuda_export.py +++ b/backends/cuda/tests/test_cuda_export.py @@ -498,6 +498,12 @@ def forward(self, x: torch.Tensor, eps: torch.Tensor) -> torch.Tensor: inputs = (torch.randn(1, 16), torch.randn(1, 14)) self._export_to_cuda_with_lower(GaussianHead().eval(), inputs) + # The export above fails without the pass only when autotuning trips + # the assert, so check that the backend runs it. + self.assertIn( + RemoveAssertAsyncPass, + {type(p) for p in CudaBackend.get_custom_passes([])}, + ) # The pass drops the check and the math that only fed it. graph_module = export(GaussianHead().eval(), inputs).graph_module RemoveAssertAsyncPass()(graph_module) From fd5948efd286bdbc0d24068135d4bc69dbf28903 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Mon, 2 Nov 2026 08:59:41 -0800 Subject: [PATCH 88/89] Clarify CUDA backend names and comments Rename graph_ran to graph_launched, since it is set once the first launch is queued. Say that a CUDA graph starts only when every constant is known to be in device memory, state the scatter index bounds for src and out separately, narrow what deterministic mode sends to the shims, split a sentence in the pybindings README, and drop a comment that repeated its switch and a local import of sys. --- backends/cuda/cuda_backend.py | 2 -- backends/cuda/runtime/cuda_backend.cpp | 18 +++++++++--------- backends/cuda/runtime/cuda_delegate_handle.h | 4 ++-- backends/cuda/runtime/shims/index.cu | 1 - backends/cuda/runtime/shims/index.h | 8 ++++---- backends/cuda/tests/test_index_shims.py | 2 +- extension/pybindings/README.md | 4 ++-- 7 files changed, 18 insertions(+), 21 deletions(-) diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index 89bc1ca46f5..3f368cdd37e 100644 --- a/backends/cuda/cuda_backend.py +++ b/backends/cuda/cuda_backend.py @@ -426,8 +426,6 @@ def _cpu_clone_preserve_strides(x: torch.Tensor) -> torch.Tensor: # `triton_heuristics.py` (clones for benchmark — must stay on GPU for # Triton). Discriminate by caller frame so we only force CPU clones for # the buffer-lifting path. - import sys - caller = sys._getframe(1).f_code.co_name if caller == "_unlift_graph": # KV-cache buffers are emptied (storage resize_(0)) by the low-memory diff --git a/backends/cuda/runtime/cuda_backend.cpp b/backends/cuda/runtime/cuda_backend.cpp index f1d09d09e29..29c9959a6e3 100644 --- a/backends/cuda/runtime/cuda_backend.cpp +++ b/backends/cuda/runtime/cuda_backend.cpp @@ -851,10 +851,10 @@ class ET_EXPERIMENTAL CudaBackend final const bool is_capture_step = (handle->cuda_graph_state.phase == CudaGraphPhase::Warmup && handle->cuda_graph_state.warmup_remaining == 0); - bool graph_ran = false; - Error error = - run_once(handle, args, n_inputs, n_outputs, is_capture_step, graph_ran); - if (error != Error::Ok && is_capture_step && !graph_ran) { + bool graph_launched = false; + Error error = run_once( + handle, args, n_inputs, n_outputs, is_capture_step, graph_launched); + if (error != Error::Ok && is_capture_step && !graph_launched) { // A capture records kernels without running them, so until the new graph // runs, the inputs and any state the program updates are untouched. A // program that keeps state in host memory, where host code does run @@ -872,7 +872,7 @@ class ET_EXPERIMENTAL CudaBackend final n_inputs, n_outputs, /*is_capture_step=*/false, - graph_ran); + graph_launched); } ET_CHECK_OK_OR_RETURN_ERROR(error); return commit_offgraph_kv_step(handle, kv_step_width); @@ -962,15 +962,15 @@ class ET_EXPERIMENTAL CudaBackend final } // Runs the program once. On the capture step it captures the CUDA graph - // instead, then launches it once, and sets graph_ran once that launch is - // queued. + // instead, then launches it once, and sets graph_launched once that launch + // is queued. Error run_once( cuda::CudaDelegateHandle* handle, Span args, size_t n_inputs, size_t n_outputs, bool is_capture_step, - bool& graph_ran) const { + bool& graph_launched) const { // Wrrap them as SlimTensors for AOTI execution without any host<->device // copies. std::vector slim_inputs(n_inputs); @@ -1286,7 +1286,7 @@ class ET_EXPERIMENTAL CudaBackend final Internal, "cudaGraphLaunch (first replay) failed: %s", cudaGetErrorString(gerr)); - graph_ran = true; + graph_launched = true; handle->cuda_graph_state.note_graph_allocations(); // Copy capture-step outputs into the planned GPU ET buffers (D2D). diff --git a/backends/cuda/runtime/cuda_delegate_handle.h b/backends/cuda/runtime/cuda_delegate_handle.h index ead44960a46..3bf7cdae5f6 100644 --- a/backends/cuda/runtime/cuda_delegate_handle.h +++ b/backends/cuda/runtime/cuda_delegate_handle.h @@ -104,8 +104,8 @@ struct CudaWeightStorage { // // The transition flow is: // Disabled ──(if CUDA graph is enabled for this method, the device -// has memory pools, and every constant is in device -// memory)──▶ Warmup +// has memory pools, and every constant is known to be in +// device memory)──▶ Warmup // Warmup ──(after `warmup_remaining` execute() calls)──▶ Replay // Warmup ──(if the capture fails, or an output is not laid out // like its ETensor)──▶ Disabled diff --git a/backends/cuda/runtime/shims/index.cu b/backends/cuda/runtime/shims/index.cu index 42b820d2d4c..7701bb8528b 100644 --- a/backends/cuda/runtime/shims/index.cu +++ b/backends/cuda/runtime/shims/index.cu @@ -1205,7 +1205,6 @@ Error scatter( scratch, stream); }; - // Half and bfloat16 compute in float, and integers in int64. switch (out->dtype()) { case ScalarType::Float: return run(float{}, float{}); diff --git a/backends/cuda/runtime/shims/index.h b/backends/cuda/runtime/shims/index.h index de4928f64f7..7a328f57478 100644 --- a/backends/cuda/runtime/shims/index.h +++ b/backends/cuda/runtime/shims/index.h @@ -91,8 +91,8 @@ AOTI_SHIM_EXPORT AOTITorchError aoti_torch_index_put_out( * @param out Tensor to update (CUDA device, up to 16 dims) * @param self Must be out; AOTInductor passes the tensor it updates as both * @param dim Dimension of out that index addresses - * @param index int64 or int32 tensor with as many dims as out, no larger than - * src, and no larger than out except at dim + * @param index int64 or int32 tensor with as many dims as out. In every dim it + * is no larger than src, and in every dim but dim no larger than out * @param src Values to write, same dtype as out. Float, double, half, * bfloat16, the integer dtypes and bool are supported. * @return AOTITorchError @@ -120,8 +120,8 @@ AOTI_SHIM_EXPORT AOTITorchError aoti_torch_cuda_scatter_src_out( * @param out Tensor to update (CUDA device, up to 16 dims) * @param self Must be out; AOTInductor passes the tensor it updates as both * @param dim Dimension of out that index addresses - * @param index int64 or int32 tensor with as many dims as out, no larger than - * src, and no larger than out except at dim + * @param index int64 or int32 tensor with as many dims as out. In every dim it + * is no larger than src, and in every dim but dim no larger than out * @param src Values to reduce, same dtype as out. Float, double, half, * bfloat16 and the integer dtypes are supported, and bool for "sum" only, as * in ATen. diff --git a/backends/cuda/tests/test_index_shims.py b/backends/cuda/tests/test_index_shims.py index af7275cfe21..466bbd0fa80 100644 --- a/backends/cuda/tests/test_index_shims.py +++ b/backends/cuda/tests/test_index_shims.py @@ -35,7 +35,7 @@ @contextlib.contextmanager def _deterministic() -> Iterator[None]: - """Deterministic mode, which also sends every scatter and index_put to the + """Deterministic mode, which sends more scatters and index_puts to the shims.""" enabled = torch.are_deterministic_algorithms_enabled() warn_only = torch.is_deterministic_algorithms_warn_only_enabled() diff --git a/extension/pybindings/README.md b/extension/pybindings/README.md index 82c9e48754d..352c7a6f93a 100644 --- a/extension/pybindings/README.md +++ b/extension/pybindings/README.md @@ -155,8 +155,8 @@ only host inputs and outputs waits depends on where it runs: In both cases, the CUDA delegate copies an output whose memory layout has to change through the legacy default stream. Loading a method's CUDA weights, which a first call may do, copies each one with a plain `cudaMemcpy`, so it waits for -the legacy default stream and for blocking streams; weights an older export -keeps in one blob wait for the whole device. +the legacy default stream and for blocking streams. An older export keeps its +weights in one blob, and copying that blob waits for the whole device. The waits below, for CUDA inputs and outputs, are for the whole device (`cudaDeviceSynchronize`), not for one stream, so they also wait for unrelated From fab02c4b33b3f66dbbfdccdc3b873f1895b90a72 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Mon, 2 Nov 2026 10:57:35 -0800 Subject: [PATCH 89/89] Compare the varlen_attn decomposition with the op on CUDA only The decomposition follows Flash attention, which lines a causal mask or a window up with the end of each query and key sequence. On ROCm, varlen_attn runs AOTriton, which lines an explicit window up with the start of each sequence, so the two differ whenever a query sequence and its key sequence have different lengths. The lowering checks still run on ROCm. --- backends/cuda/tests/test_cuda_export.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/backends/cuda/tests/test_cuda_export.py b/backends/cuda/tests/test_cuda_export.py index ba3ae354250..5ff328015e1 100644 --- a/backends/cuda/tests/test_cuda_export.py +++ b/backends/cuda/tests/test_cuda_export.py @@ -556,9 +556,13 @@ def forward(self, q, k, v, cu_seq): (big, big, v, cu_seq, cu_seq, 16, 16), ] table = CudaBackend.get_decomposition_table() - for args in cases: - for actual, expected in zip(table[op](*args)[:2], op(*args)[:2]): - torch.testing.assert_close(actual, expected, atol=1e-2, rtol=1e-2) + # On ROCm, varlen_attn runs AOTriton, which lines an explicit window up + # with the start of each sequence. The decomposition follows Flash + # attention on CUDA, which lines it up with the end, so compare on CUDA. + if torch.version.hip is None: + for args in cases: + for actual, expected in zip(table[op](*args)[:2], op(*args)[:2]): + torch.testing.assert_close(actual, expected, atol=1e-2, rtol=1e-2) # Flash attention versions differ on this window, since a query sequence # can be longer than max_k + 1. with self.assertRaisesRegex(NotImplementedError, "window"):