diff --git a/.ci/scripts/wheel/test_shared_libraries.py b/.ci/scripts/wheel/test_shared_libraries.py
index a1d8c791b83..8bbb86a1428 100644
--- a/.ci/scripts/wheel/test_shared_libraries.py
+++ b/.ci/scripts/wheel/test_shared_libraries.py
@@ -103,6 +103,10 @@
     "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",
+    "aoti_torch_cuda_scatter_reduce_two_out",
+    "aoti_torch_cuda_scatter_src_out",
 )
 
 _QUANTIZED_KERNEL_SYMBOLS = (
diff --git a/backends/aoti/aoti_backend.py b/backends/aoti/aoti_backend.py
index 0d0670b7318..a069db4c1fc 100644
--- a/backends/aoti/aoti_backend.py
+++ b/backends/aoti/aoti_backend.py
@@ -53,6 +53,22 @@ def _checked_data_name(data_name: str) -> str:
     return data_name
 
 
+def _written_buffer_options(program: ExportedProgram) -> Dict[str, Any]:
+    """Inductor options for a program that writes a buffer.
+
+    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.
+    """
+    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."
 )
@@ -286,9 +302,13 @@ 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)
+        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)
@@ -390,11 +410,7 @@ 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
-        )
+        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 +451,7 @@ def generate_data_name_compile_spec(cls, data_name: str) -> CompileSpec:
         Generate a CompileSpec that names the external data file, written as
         ``<data_name>.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 +465,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/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..0181ff9e8e8 100644
--- a/backends/aoti/slim/c10/core/ScalarType.h
+++ b/backends/aoti/slim/c10/core/ScalarType.h
@@ -13,12 +13,14 @@
 #include <ostream>
 
 #include <executorch/runtime/core/portable_type/bfloat16.h>
+#include <executorch/runtime/core/portable_type/half.h>
 #include <executorch/runtime/platform/assert.h>
 
 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.
@@ -30,7 +32,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 +52,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 +75,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 +105,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 +123,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 +164,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<ScalarTypeTestData> 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<int64_t> sizes = {2, 3};
   std::vector<int64_t> 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<c10::ScalarType>(7)), // Double = 7
+          static_cast<c10::ScalarType>(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<int64_t> sizes = {2, 3};
   std::vector<int64_t> strides = {3, 1};
 
@@ -257,7 +257,7 @@ TEST(EmptyStridedTest, InvalidDtypeDouble) {
       empty_strided(
           makeArrayRef(sizes),
           makeArrayRef(strides),
-          static_cast<c10::ScalarType>(7)), // Double = 7
+          static_cast<c10::ScalarType>(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<c10::ScalarType>(7)), // Double = 7
+      empty({2, 3}, static_cast<c10::ScalarType>(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<c10::ScalarType>(7)), "");
+      from_blob(external_data, {2, 3}, static_cast<c10::ScalarType>(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 2dd8475a4d8..952a79283f2 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<int64_t> sizes = {2, 3};
+  std::vector<int64_t> 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<int64_t> sizes = {2, 3, 4};
   std::vector<int64_t> 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);
@@ -663,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/BUCK b/backends/cuda/BUCK
index ca8f9eee753..b48b1be48bf 100644
--- a/backends/cuda/BUCK
+++ b/backends/cuda/BUCK
@@ -79,7 +79,9 @@ fbcode_target(
         "passes/__init__.py",
         "passes/move_cond_predicate_to_cpu.py",
         "passes/lower_offgraph_kv.py",
+        "passes/remove_assert_async.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 891fb81bd06..09a8767fd6b 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
@@ -19,6 +20,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
@@ -39,9 +43,13 @@
 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_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 (
@@ -54,6 +62,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
 
 # ---------------------------------------------------------------------------
@@ -121,11 +130,70 @@ def _has_rmask(_self) -> bool:
         IndexingOptions.has_rmask = orig_has_rmask
 
 
-_SMALL_GPU_GUARD = threading.local()
+class _Depth(threading.local):
+    def __init__(self) -> None:
+        self.value = 0
 
 
-@contextlib.contextmanager
-def _allow_triton_templates_on_small_gpus():
+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, install: typing.Callable[..., typing.Callable[[], None]]):
+        self._install = install
+        self._remove: typing.Callable[[], None] = lambda: None
+        self._lock = threading.Lock()
+        self._scopes = 0
+        self._depth = _Depth()
+
+    def _active(self) -> bool:
+        return self._depth.value > 0
+
+    @contextlib.contextmanager
+    def scope(self) -> typing.Iterator[None]:
+        with self._lock:
+            if self._scopes == 0:
+                self._remove = self._install(self._active)
+            self._scopes += 1
+        self._depth.value += 1
+        try:
+            yield
+        finally:
+            self._depth.value -= 1
+            with self._lock:
+                self._scopes -= 1
+                if self._scopes == 0:
+                    self._remove()
+
+
+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]:
     """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.
@@ -136,23 +204,97 @@ 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_OVERRIDE.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
+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 = _CompilePatch(_install_as_needed_links)
 
-    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 _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
+    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 the pinned PyTorch has the same flag in Inductor
+    (https://github.com/pytorch/pytorch/pull/202902).
+    """
+    return _AS_NEEDED_LINKS.scope()
+
+
+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
+
+    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
+
+    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]:
+    """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:
@@ -551,6 +693,107 @@ 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 _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. 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(
+            "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)
+        # 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)]
+
+    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, :]
+    offset = from_end_k[None, :] - from_end_q[:, None]
+    if left >= 0:
+        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 = (
+            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)
+    # 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
+
+
 def _linear_kernel(x: torch.Tensor) -> torch.Tensor:
     return torch.clamp(1 - x.abs(), min=0)
 
@@ -582,8 +825,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:
@@ -629,8 +872,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
@@ -967,6 +1210,8 @@ 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,
+            "aoti_torch_cuda_scatter_src_out": None,
         }
 
     @classmethod
@@ -977,6 +1222,8 @@ 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,
+            torch.ops.torch_attn._varlen_attn.default: _varlen_attn,
         }
         for resize, support, kernel in (
             (torch.ops.aten._upsample_bilinear2d_aa, 1, _linear_kernel),
@@ -1013,7 +1260,12 @@ def get_custom_passes(cls, compile_specs: List[CompileSpec]) -> List[typing.Any]
                         f"Invalid triton_kernel_mode: {mode}. Expected 'ON' or 'OFF'."
                     )
                 triton_kernel_mode = mode
-        passes = [MoveCondPredicateToCpuPass(), ReplaceInt64FloorDivWithFloatPass()]
+        passes = [
+            RemoveAssertAsyncPass(),
+            MoveCondPredicateToCpuPass(),
+            ReplaceInt64FloorDivWithFloatPass(),
+            ReplaceMaskedCopyWithWherePass(),
+        ]
         cell_layout = any(
             spec.key == OFFGRAPH_KV_COMPILE_SPEC
             and parse_offgraph_kv_manifest(spec.value)["layout"] == "cell"
@@ -1182,6 +1434,10 @@ 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())
+                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/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/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/passes/replace_masked_copy_with_where.py b/backends/cuda/passes/replace_masked_copy_with_where.py
new file mode 100644
index 00000000000..941e5f0641c
--- /dev/null
+++ b/backends/cuda/passes/replace_masked_copy_with_where.py
@@ -0,0 +1,207 @@
+# 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 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
+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
+_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)]``.
+    """
+    # 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):
+    """
+    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 module in graph_module.modules():
+            if not isinstance(module, GraphModule):
+                continue
+            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
+    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]
+        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() - position - 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, position)
+        if dead is not None:
+            for user in reversed(dead):
+                graph.erase_node(user)
+
+    @staticmethod
+    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 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):
+            return 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 expr.free_symbols != {count.node.expr}:
+                    return None
+                bounds = bound_sympy(expr, {count.node.expr: within})
+                if not (bounds.is_singleton() and bounds.lower == sympy.true):
+                    return None
+            # 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)
+        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..cc1310b7a8b
--- /dev/null
+++ b/backends/cuda/passes/tests/test_replace_masked_copy_with_where.py
@@ -0,0 +1,214 @@
+# 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 operator
+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 Dim, 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_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_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):
+                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)
+                other[mask] = x[mask.logical_not()]
+                # 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])
+        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), 4)
+        self.assertEqual(targets.count(exir_ops.edge.aten.where.self), 2)
+        # 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
+        ):
+            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."""
+
+        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()
diff --git a/backends/cuda/runtime/aoti_cuda_shims.lib b/backends/cuda/runtime/aoti_cuda_shims.lib
index 3294d1ec262..ad54a4fdb24 100644
Binary files a/backends/cuda/runtime/aoti_cuda_shims.lib and b/backends/cuda/runtime/aoti_cuda_shims.lib differ
diff --git a/backends/cuda/runtime/cuda_backend.cpp b/backends/cuda/runtime/cuda_backend.cpp
index ce429a9a65a..2f545ee3e2e 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);
 
@@ -614,7 +630,18 @@ 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()) {
+      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 %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,
             "CUDA graph enabled for method '%s' (warmup=%d)",
@@ -633,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
   }
 
@@ -768,14 +796,26 @@ 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],
+            sizes.size() == captured.size(),
             InvalidArgument,
-            "CUDA graph replay: input %zu size mismatch (expected %zu, got %zu)",
+            "CUDA graph replay: input %zu has %zu dims, captured with %zu",
             i,
-            handle->cuda_graph_state.static_input_nbytes[i],
-            et_input->nbytes());
+            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<long long>(sizes[d]),
+              d,
+              static_cast<long long>(captured[d]));
+        }
         ET_CUDA_CHECK_OR_RETURN_ERROR(cudaMemcpyAsync(
             handle->cuda_graph_state.static_input_ptrs[i],
             et_input->const_data_ptr(),
@@ -808,13 +848,129 @@ 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. 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());
+      // 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,
+          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);
+  }
+
+  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.
+  // 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 unlisted;
+    }
+    std::unordered_map<std::string, AtenTensorHandle> constants;
+    if (handle->extract_constants_map(
+            handle->container_handle,
+            reinterpret_cast<AOTInductorConstantMapHandle>(&constants),
+            /*use_inactive=*/false) != Error::Ok) {
+      return unlisted;
+    }
+    for (const auto& [fqn, tensor] : constants) {
+      if (reinterpret_cast<SlimTensor*>(tensor)->is_cpu()) {
+        return "constant '" + fqn + "' is in host memory";
+      }
+    }
+    return "";
+  }
+
+  // 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<EValue*> 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<SlimTensor*> slim_inputs(n_inputs);
@@ -841,25 +997,10 @@ 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.
-        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();
-        // 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;
-        }
+        // a capture that cannot succeed is tried once rather than on every
+        // later call for the life of the process. Eager execution is correct,
+        // just slower.
+        state_->release();
         state_->phase = CudaGraphPhase::Disabled;
         state_->warmup_remaining = 0;
         (void)cudaGetLastError();
@@ -903,6 +1044,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,
@@ -1034,11 +1177,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
@@ -1061,9 +1216,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",
               i,
               handle->method_name.c_str());
           return Error::NotSupported;
@@ -1076,9 +1231,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",
               i,
               handle->method_name.c_str());
           return resize_err;
@@ -1094,7 +1249,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
@@ -1120,6 +1274,8 @@ class ET_EXPERIMENTAL CudaBackend final
           Internal,
           "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++) {
@@ -1136,7 +1292,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 -----
@@ -1193,69 +1349,9 @@ class ET_EXPERIMENTAL CudaBackend final
       slim_outputs[i] = nullptr;
     }
 
-    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);
-    }
+    return Error::Ok;
   }
 
- private:
   mutable std::mutex cuda_graph_method_mutex_;
   std::string cuda_graph_method_;
 
diff --git a/backends/cuda/runtime/cuda_delegate_handle.h b/backends/cuda/runtime/cuda_delegate_handle.h
index f42f0f2382f..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.
@@ -118,7 +122,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,
@@ -141,6 +147,9 @@ struct CudaGraphState {
   std::vector<void*> static_output_ptrs;
   std::vector<size_t> static_input_nbytes;
   std::vector<size_t> 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<std::vector<int64_t>> 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
@@ -163,7 +172,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)
@@ -227,6 +237,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 +264,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 +282,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/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/shims/index.cu b/backends/cuda/runtime/shims/index.cu
index 97335a75ea0..42b820d2d4c 100644
--- a/backends/cuda/runtime/shims/index.cu
+++ b/backends/cuda/runtime/shims/index.cu
@@ -17,8 +17,11 @@
 #include <cinttypes>
 #include <climits>
 #include <cstdint>
+#include <limits>
 #include <memory>
 #include <numeric>
+#include <string_view>
+#include <type_traits>
 #include <vector>
 
 #include <c10/macros/Macros.h>
@@ -331,13 +334,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,
@@ -761,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 <typename T, typename AccT, bool kAccumulate>
 Error put(
     const IndexPlan& plan,
@@ -781,7 +828,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,
@@ -800,35 +847,16 @@ Error put(
   index_offsets_kernel<<<grid_size(points), kThreads, 0, stream>>>(
       plan, points, out_offsets, values_offsets);
   ET_CUDA_KERNEL_LAUNCH_CHECK_OR_RETURN_ERROR();
-  const auto count = static_cast<int>(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<int>(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<int>(points),
       static_cast<int>(sizeof(int64_t) * 8),
-      stream));
+      scratch,
+      stream,
+      "index_put"));
   put_sorted_kernel<T, AccT, kAccumulate>
       <<<grid_size(numel), kThreads, 0, stream>>>(
           plan,
@@ -842,6 +870,373 @@ Error put(
   return Error::Ok;
 }
 
+// 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
+// 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;
+  bool index_is_int32;
+};
+
+__global__ void scatter_offsets_kernel(
+    const ScatterPlan plan,
+    int64_t points,
+    const void* __restrict__ index,
+    int64_t* __restrict__ out_offsets,
+    int64_t* __restrict__ src_offsets) {
+  for (int64_t point =
+           static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
+       point < points;
+       point += static_cast<int64_t>(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 = plan.index_is_int32
+        ? static_cast<const int32_t*>(index)[index_offset]
+        : static_cast<const int64_t*>(index)[index_offset];
+    // 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];
+    src_offsets[point] = src_offset;
+  }
+}
+
+// As ATen's safe_max and safe_min, a NaN on either side wins.
+template <typename AccT>
+__device__ __forceinline__ AccT
+combine(Reduction reduction, AccT result, AccT value) {
+  bool value_is_nan = false;
+  if constexpr (std::is_floating_point_v<AccT>) {
+    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, 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 <typename T, typename AccT>
+__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<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
+       i < points;
+       i += static_cast<int64_t>(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;
+    }
+    auto result = static_cast<AccT>(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<AccT>(src[src_offsets[j]]));
+      if (reduction == Reduction::Prod) {
+        result = static_cast<AccT>(static_cast<T>(result));
+      }
+    }
+    if (reduction == Reduction::Mean) {
+      const auto sum = static_cast<AccT>(static_cast<T>(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<AccT>(static_cast<T>(count));
+      if (divisor == AccT(0)) {
+        divisor = AccT(1);
+      }
+      if constexpr (std::is_floating_point_v<AccT>) {
+        result = sum / divisor;
+      } else {
+        // Integers round down, as ATen's division with rounding_mode="floor".
+        const AccT quotient = sum / divisor;
+        result = quotient -
+            (quotient * divisor != sum && (sum < 0) != (divisor < 0));
+      }
+    }
+    out[offset] = static_cast<T>(result);
+  }
+}
+
+// What ATen writes to each element the index names before it reduces into
+// it, when include_self is false.
+template <typename T, typename AccT>
+T reduction_identity(Reduction reduction) {
+  switch (reduction) {
+    case Reduction::Prod:
+      return static_cast<T>(AccT(1));
+    case Reduction::Max:
+      if constexpr (std::is_floating_point_v<AccT>) {
+        return static_cast<T>(-std::numeric_limits<AccT>::infinity());
+      } else {
+        return std::numeric_limits<T>::lowest();
+      }
+    case Reduction::Min:
+      if constexpr (std::is_floating_point_v<AccT>) {
+        return static_cast<T>(std::numeric_limits<AccT>::infinity());
+      } else {
+        return std::numeric_limits<T>::max();
+      }
+    default:
+      return static_cast<T>(AccT(0));
+  }
+}
+
+template <typename T, typename AccT>
+Error scatter_reduce(
+    const char* name,
+    const ScatterPlan& plan,
+    int64_t points,
+    const void* 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,
+      "%s: more than %d index elements are not supported",
+      name,
+      INT_MAX);
+  auto* offsets =
+      static_cast<int64_t*>(scratch.allocate(4 * points * sizeof(int64_t)));
+  ET_CHECK_OR_RETURN_ERROR(
+      offsets != nullptr,
+      MemoryAllocationFailed,
+      "%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;
+  int64_t* sorted_src_offsets = offsets + 3 * points;
+  scatter_offsets_kernel<<<grid_size(points), kThreads, 0, stream>>>(
+      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<int>(points),
+      end_bit,
+      scratch,
+      stream,
+      name));
+  scatter_reduce_sorted_kernel<T, AccT>
+      <<<grid_size(points), kThreads, 0, stream>>>(
+          points,
+          sorted_out_offsets,
+          sorted_src_offsets,
+          static_cast<T*>(out),
+          static_cast<const T*>(src),
+          reduction,
+          include_self,
+          reduction_identity<T, AccT>(reduction));
+  ET_CUDA_KERNEL_LAUNCH_CHECK_OR_RETURN_ERROR();
+  return Error::Ok;
+}
+
+// Runs the scatter shim called `name`, with the checks ATen applies to
+// aten::scatter and aten::scatter_reduce.
+Error scatter(
+    const char* name,
+    Tensor* out,
+    Tensor* self,
+    int64_t dim,
+    Tensor* index,
+    Tensor* src,
+    Reduction reduction,
+    bool include_self) {
+  ET_CHECK_OR_RETURN_ERROR(
+      out != nullptr && index != nullptr && src != nullptr,
+      InvalidArgument,
+      "%s: null argument",
+      name);
+  ET_CHECK_OR_RETURN_ERROR(
+      out == self,
+      NotSupported,
+      "%s: out must be self, as AOTInductor calls it",
+      name);
+  ET_CHECK_OR_RETURN_ERROR(
+      out->is_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<int64_t>(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 || index->dtype() == ScalarType::Int,
+      InvalidArgument,
+      "%s: index must be int64 or int32",
+      name);
+  ET_CHECK_OR_RETURN_ERROR(
+      std::max<int64_t>(index->dim(), 1) == ndim &&
+          std::max<int64_t>(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<int>(ndim);
+  plan.dim = static_cast<int>(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 +=
+        std::max<int64_t>(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<int32_t>(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<T, AccT>(
+        name,
+        plan,
+        points,
+        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<int>(out->dtype()));
+      return Error::NotSupported;
+  }
+}
+
 } // namespace
 
 #ifdef __cplusplus
@@ -1017,11 +1412,15 @@ 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<float, float, true>(
           plan, numel, may_repeat, dst, src, scratch, stream);
+    case ScalarType::Double:
+      return put<double, double, true>(
+          plan, numel, may_repeat, dst, src, scratch, stream);
     case ScalarType::Half:
       return put<__half, float, true>(
           plan, numel, may_repeat, dst, src, scratch, stream);
@@ -1055,6 +1454,65 @@ 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,
+    int64_t dim,
+    Tensor* index,
+    Tensor* src,
+    const char* reduce,
+    int32_t include_self) {
+  ET_CHECK_OR_RETURN_ERROR(
+      reduce != nullptr,
+      InvalidArgument,
+      "aoti_torch_cuda_scatter_reduce_two_out: null argument");
+  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;
+  }
+  return scatter(
+      "aoti_torch_cuda_scatter_reduce_two_out",
+      out,
+      self,
+      dim,
+      index,
+      src,
+      reduction,
+      include_self != 0);
+}
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/backends/cuda/runtime/shims/index.h b/backends/cuda/runtime/shims/index.h
index 3476eaa8c88..de4928f64f7 100644
--- a/backends/cuda/runtime/shims/index.h
+++ b/backends/cuda/runtime/shims/index.h
@@ -45,15 +45,16 @@ 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
  * 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
@@ -75,6 +76,69 @@ 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 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
+ * @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 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.
+ *
+ * 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'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
+ * @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 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/runtime/shims/memory.cpp b/backends/cuda/runtime/shims/memory.cpp
index 86be88e4833..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 <typename T>
 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<uint8_t>(), ret_value);
@@ -80,6 +99,11 @@ AOTITorchError read_one_element(const SlimTensor* tensor, T* ret_value) {
       return narrow_to(tensor->item<int64_t>(), ret_value);
     case ScalarType::Float:
       return narrow_to(tensor->item<float>(), ret_value);
+    case ScalarType::Double:
+      return narrow_to(tensor->item<double>(), ret_value);
+    case ScalarType::Half:
+      return narrow_to(
+          static_cast<float>(tensor->item<c10::Half>()), ret_value);
     case ScalarType::BFloat16:
       return narrow_to(
           static_cast<float>(tensor->item<c10::BFloat16>()), ret_value);
@@ -470,6 +494,8 @@ 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(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 b1b369a58d9..ef51ff5eef0 100644
--- a/backends/cuda/runtime/shims/memory.h
+++ b/backends/cuda/runtime/shims/memory.h
@@ -256,6 +256,14 @@ 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_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/runtime/shims/rand.cu b/backends/cuda/runtime/shims/rand.cu
index 9916fe8fe0c..85a64e1ffd3 100644
--- a/backends/cuda/runtime/shims/rand.cu
+++ b/backends/cuda/runtime/shims/rand.cu
@@ -19,9 +19,10 @@
 #include <curand_kernel.h>
 
 #include <cstdint>
-#include <ctime>
 #include <mutex>
 #include <new>
+#include <random>
+#include <stdexcept>
 #include <vector>
 
 namespace executorch::backends::cuda {
@@ -59,12 +60,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<unsigned long long>(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
@@ -74,19 +83,22 @@ bool ensure_rng_init(cudaStream_t stream) {
       if (!rng.ok()) {
         throw std::bad_alloc();
       }
-      d_rng = static_cast<RngState*>(rng.get());
-      RngState h;
-      h.seed = static_cast<unsigned long long>(time(nullptr));
-      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);
+      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<RngState*>(rng.get());
     });
-  } 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;
@@ -233,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<int>((numel + kThreads - 1) / kThreads);
@@ -296,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<int64_t*>(out->data_ptr());
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<double*>(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/runtime/test/test_cuda_kv_cache.cpp b/backends/cuda/runtime/test/test_cuda_kv_cache.cpp
index 0c32da2079f..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.
@@ -1039,15 +1041,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);
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<uint8_t> 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<uint8_t> 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_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
-            )
diff --git a/backends/cuda/tests/test_cuda_export.py b/backends/cuda/tests/test_cuda_export.py
index 637ff4036b1..a9357803031 100644
--- a/backends/cuda/tests/test_cuda_export.py
+++ b/backends/cuda/tests/test_cuda_export.py
@@ -6,7 +6,6 @@
 
 import contextlib
 import os
-import tempfile
 import threading
 import unittest
 from typing import Tuple
@@ -14,12 +13,15 @@
 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.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
 from torch.export import export
+from torch.nn.attention.varlen import varlen_attn
 
 
 class TestCudaBackendCompileOptions(unittest.TestCase):
@@ -210,6 +212,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
@@ -269,10 +296,20 @@ 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")]
+                )
+        # 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")
@@ -436,6 +473,87 @@ 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__()
+                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)
+                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."""
+
+        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
+        # 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, 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]),
+        ]
+        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."""
 
@@ -709,62 +827,25 @@ def test_triton_kernel_mode_off(self):
         )
 
     def test_missing_scatter_kernel_fails_at_export(self):
-        """Inductor leaves int64 scatter_add to an ATen kernel the CUDA shim lacks."""
+        """The scatter fallback writes its kernel call itself, and the check for
+        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 Int64ScatterAdd(torch.nn.Module):
+        class ValueScatter(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",
-    )
-    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)
+                return counts.scatter(0, idx, 1.0)
+
+        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))
 
     def test_device_info_propagated_to_cuda_delegate_outputs(self):
         """
diff --git a/backends/cuda/tests/test_index_shims.py b/backends/cuda/tests/test_index_shims.py
new file mode 100644
index 00000000000..286aef0be30
--- /dev/null
+++ b/backends/cuda/tests/test_index_shims.py
@@ -0,0 +1,366 @@
+# 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.
+
+"""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. 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 contextlib
+import os
+import tempfile
+import unittest
+from typing import Iterator, 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
+from torch.nn.attention.varlen import varlen_attn
+
+
+@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")
+@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, ...]],
+        **tolerance: float,
+    ) -> None:
+        """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]),
+            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:
+            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))
+                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, **tolerance
+                    )
+
+    def test_mask_index_and_index_put(self):
+        """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):
+                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)
+                wide = x.double()
+                wide = wide.index_put((mask,), wide[mask], accumulate=True)
+                return out, picked, rows, counts, wide
+
+        x = torch.randn(4, 5, 6)
+        # 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
+        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,)
+
+        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."""
+
+        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])
+        with _deterministic():
+            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)])
+
+    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 a sum it cannot
+        add atomically: int64, or bfloat16 below sm_90."""
+
+        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 ("sum", "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)])
+
+    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. As in
+        ATen, the index can be int32."""
+
+        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)
+        with _deterministic():
+            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])
+
+    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 _deterministic():
+            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
+        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)
+        # 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_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 _deterministic():
+            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):
+                # 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)
+        for dtype in (torch.float64, torch.float16):
+            with self.subTest(dtype=dtype):
+                x = torch.randn(4, 1000, 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)])
+
+    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)
diff --git a/backends/cuda/tests/test_missing_weights_blob.py b/backends/cuda/tests/test_missing_weights_blob.py
index 6f1308b20c3..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
@@ -222,6 +229,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.
 
diff --git a/backends/cuda/tests/test_output_stride_rearrange.py b/backends/cuda/tests/test_output_stride_rearrange.py
index 431bb56a88f..a53156cd3f0 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
@@ -282,6 +283,53 @@ 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 _ScaledByCount(nn.Module):
+    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."""
+
+    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."""
+
+    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__()
@@ -291,6 +339,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."""
 
@@ -304,14 +375,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 +402,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,15 +437,106 @@ 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)
+        # 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("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()
+        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_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(),
+            _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
+            # 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)
+                    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.
+        torch.manual_seed(0)
+        inputs = [(torch.randn(8, 4), torch.rand(8) > 0.5) for _ in range(6)]
+        with _RuntimeLog() as log:
+            outputs = self._run_with_cuda_graph(
+                _MaskedCounter().eval(), inputs[0], inputs
+            )
+        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)
+                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.
-        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()
@@ -380,7 +545,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("execute() failed with error 0x12", str(out))
+                    else:
+                        self._assert_matches_eager(module, x, out)
 
 
 if __name__ == "__main__":
diff --git a/backends/cuda/tests/test_sort_shim.py b/backends/cuda/tests/test_sort_shim.py
index d81d9e9e187..4241405b9a0 100644
--- a/backends/cuda/tests/test_sort_shim.py
+++ b/backends/cuda/tests/test_sort_shim.py
@@ -43,6 +43,8 @@
         "aoti_torch_cuda_randint_low_out",
         "aoti_torch_cuda_index_Tensor",
         "aoti_torch_index_put_out",
+        "aoti_torch_cuda_scatter_reduce_two_out",
+        "aoti_torch_cuda_scatter_src_out",
     }
 )
 
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,
     )
 
 
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
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);
 }
diff --git a/extension/pybindings/README.md b/extension/pybindings/README.md
index f8038474c1b..82c9e48754d 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
@@ -199,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
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
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