From 25a9fce6fe12bb58cea085f895f971294adee405 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <shoumikhin@meta.com>
Date: Sat, 22 Aug 2026 18:04:39 -0700
Subject: [PATCH 01/15] Fix device placement for memory-planned buffers in
 Runtime.load_program

A .pte file records where each memory-planned buffer has to live, on the host
or on an accelerator, in the plan's non_const_buffer_device field. ProgramMemory
never read that field and allocated every planned buffer as host memory, so a
program asking for device memory received a host pointer. With the CUDA backend
this failed at execute time with "not backed by CUDA device memory", while the
same .pte loaded through _load_for_executorch worked.

Allocate each planned buffer on the device it is tagged for. Buffer indices are
plan local, so one set of arenas shared across methods by index cannot describe
a file where one method plans onto the host and another onto an accelerator. A
method with any device-tagged buffer therefore gets its own arenas, and every
other method keeps using the shared host arenas.

extension/module/module.cpp faces the same problem and answers it differently,
because it has a share_memory_arenas flag and this loader has none. Module
refuses to load a device-planned method when that flag is set, and otherwise
builds per-method arenas for every method. Runtime.load_program has no such
flag and shares arenas unconditionally today, so refusing would break loading
files that already load. It keeps the shared host arenas for host methods and
gives a device-planned method its own instead.

Those arenas are built when the method is loaded rather than when the program
is loaded, so a file containing one accelerator method still opens on a machine
without that accelerator, still lists its methods, and its host methods still
load and run.

non_const_buffer_device is optional and MethodMeta reports CPU when it is
absent, so CPU-only and older programs keep the shared arenas, the host
allocation path, and the single argument HierarchicalAllocator that leaves
MemoryManager::has_device_memory() false.

Adds test_program_loads_when_one_method_is_device_planned, which exports a two
method program where one method has a device tagged planned buffer and the other
has none, then checks that the program loads, that the host method runs, that
loading the device method reaches the device allocator and is refused there, and
that the host method still runs afterwards. It needs no GPU, and it skips itself
on a build that links the CUDA backend, since that registers a CUDA allocator
which would satisfy the request. Verified that the test fails when the change to
pybindings.cpp is reverted: the old loader put the device buffer in host memory,
so the method loaded and no error was raised.
---
 extension/pybindings/pybindings.cpp          | 171 ++++++++++++++++---
 extension/pybindings/test/BUCK               |   4 +
 extension/pybindings/test/make_test.py       |   9 +-
 extension/pybindings/test/test_pybindings.py |  43 +++++
 4 files changed, 206 insertions(+), 21 deletions(-)

diff --git a/extension/pybindings/pybindings.cpp b/extension/pybindings/pybindings.cpp
index 6905c0553f9..53438833ccb 100644
--- a/extension/pybindings/pybindings.cpp
+++ b/extension/pybindings/pybindings.cpp
@@ -7,6 +7,7 @@
  */
 
 #include <algorithm>
+#include <cinttypes>
 #include <cstdio>
 #include <iostream>
 #include <memory>
@@ -31,6 +32,7 @@
 #include <executorch/extension/threadpool/threadpool.h>
 #include <executorch/runtime/backend/interface.h>
 #include <executorch/runtime/core/data_loader.h>
+#include <executorch/runtime/core/device_memory_buffer.h>
 #include <executorch/runtime/core/exec_aten/util/scalar_type_util.h>
 #include <executorch/runtime/executor/method.h>
 #include <executorch/runtime/executor/program.h>
@@ -81,6 +83,7 @@ using ::executorch::ET_RUNTIME_NAMESPACE::get_num_registered_backends;
 using ::executorch::ET_RUNTIME_NAMESPACE::get_registered_kernels;
 using ::executorch::ET_RUNTIME_NAMESPACE::Kernel;
 using ::executorch::ET_RUNTIME_NAMESPACE::Method;
+using ::executorch::ET_RUNTIME_NAMESPACE::MethodMeta;
 using ::executorch::ET_RUNTIME_NAMESPACE::Program;
 using ::executorch::extension::BufferDataLoader;
 using ::executorch::extension::MallocMemoryAllocator;
@@ -88,6 +91,7 @@ using ::executorch::extension::MmapDataLoader;
 using ::executorch::extension::ET_BUNDLED_MODULE_NAMESPACE::BundledModule;
 using ::executorch::extension::pybindings::PyDataLoader;
 using ::executorch::runtime::DataLoader;
+using ::executorch::runtime::DeviceMemoryBuffer;
 using ::executorch::runtime::Error;
 using ::executorch::runtime::EValue;
 using ::executorch::runtime::EventTracerDebugLogLevel;
@@ -1077,18 +1081,28 @@ inline std::shared_ptr<ProgramState> load_program(
 /// A wrapper/util class for executorch memory allocations/manager.
 class ProgramMemory {
  public:
-  explicit ProgramMemory(std::vector<std::vector<uint8_t>>&& non_const_buffers)
+  /// `devices` is empty when every buffer is on the host, which keeps
+  /// `MemoryManager::has_device_memory()` false for CPU-only programs.
+  /// Otherwise it holds one entry per buffer, indexed like `sizes`.
+  ProgramMemory(
+      std::vector<int64_t>&& sizes,
+      std::vector<runtime::etensor::Device>&& devices)
       : runtime_allocator_(),
-        non_const_buffers_(std::move(non_const_buffers)),
+        planned_sizes_(std::move(sizes)),
+        planned_devices_(std::move(devices)),
+        non_const_buffers_(allocate_host_buffers()),
+        device_buffers_(allocate_device_buffers()),
         non_const_spans_(create_non_const_spans()),
-        non_const_allocator_(
-            {non_const_spans_.data(), non_const_spans_.size()}),
+        non_const_allocator_(create_non_const_allocator()),
         mem_manager_(
             &const_allocator_,
             &non_const_allocator_,
             &runtime_allocator_,
             &temp_allocator_) {}
 
+  explicit ProgramMemory(std::vector<int64_t>&& sizes)
+      : ProgramMemory(std::move(sizes), {}) {}
+
   /// Returns a pointer to the internal memory manager, the Memory instance
   /// must outlive this pointer.
   MemoryManager* mem_manager() {
@@ -1105,24 +1119,126 @@ class ProgramMemory {
 
   MallocMemoryAllocator temp_allocator_{};
 
+  std::vector<int64_t> planned_sizes_;
+
+  std::vector<runtime::etensor::Device> planned_devices_;
+
+  // Backs CPU-tagged buffers; the entry is empty for a device-tagged buffer.
   std::vector<std::vector<uint8_t>> non_const_buffers_;
 
+  // Backs device-tagged buffers; the entry is empty for a CPU-tagged buffer.
+  // Parallel to non_const_buffers_ so both index by planned buffer id. Empty
+  // for an all-host program.
+  std::vector<DeviceMemoryBuffer> device_buffers_;
+
   std::vector<Span<uint8_t>> non_const_spans_;
 
   HierarchicalAllocator non_const_allocator_;
 
   MemoryManager mem_manager_;
 
+  bool is_device_buffer(size_t index) const {
+    return index < planned_devices_.size() && !planned_devices_[index].is_cpu();
+  }
+
+  std::vector<std::vector<uint8_t>> allocate_host_buffers() {
+    std::vector<std::vector<uint8_t>> result;
+    result.reserve(planned_sizes_.size());
+    for (size_t i = 0; i < planned_sizes_.size(); i++) {
+      if (is_device_buffer(i)) {
+        result.emplace_back();
+      } else {
+        result.emplace_back(planned_sizes_[i]);
+      }
+    }
+    return result;
+  }
+
+  std::vector<DeviceMemoryBuffer> allocate_device_buffers() {
+    std::vector<DeviceMemoryBuffer> result;
+    if (planned_devices_.empty()) {
+      return result;
+    }
+    result.reserve(planned_sizes_.size());
+    for (size_t i = 0; i < planned_sizes_.size(); i++) {
+      if (!is_device_buffer(i)) {
+        result.emplace_back();
+        continue;
+      }
+      auto buffer = DeviceMemoryBuffer::create(
+          planned_sizes_[i],
+          planned_devices_[i].type(),
+          planned_devices_[i].index());
+      THROW_IF_ERROR(
+          buffer.error(),
+          "Failed to allocate %" PRId64 " bytes for buffer %zu on device %d:%d",
+          planned_sizes_[i],
+          i,
+          static_cast<int>(planned_devices_[i].type()),
+          static_cast<int>(planned_devices_[i].index()));
+      result.emplace_back(std::move(buffer.get()));
+    }
+    return result;
+  }
+
   std::vector<Span<uint8_t>> create_non_const_spans() {
     std::vector<Span<uint8_t>> result;
-    for (size_t i = 0; i < non_const_buffers_.size(); i++) {
-      result.push_back(
-          {non_const_buffers_[i].data(), non_const_buffers_[i].size()});
+    result.reserve(planned_sizes_.size());
+    for (size_t i = 0; i < planned_sizes_.size(); i++) {
+      if (is_device_buffer(i)) {
+        result.push_back(device_buffers_[i].as_span());
+      } else {
+        result.push_back(
+            {non_const_buffers_[i].data(), non_const_buffers_[i].size()});
+      }
     }
     return result;
   }
+
+  HierarchicalAllocator create_non_const_allocator() {
+    Span<Span<uint8_t>> buffers(
+        non_const_spans_.data(), non_const_spans_.size());
+    return planned_devices_.empty()
+        ? HierarchicalAllocator(buffers)
+        : HierarchicalAllocator(
+              buffers, {planned_devices_.data(), planned_devices_.size()});
+  }
 };
 
+/// True if any of the method's memory-planned buffers must live off the host.
+bool has_device_buffers(const MethodMeta& method_meta) {
+  for (size_t i = 0; i < method_meta.num_memory_planned_buffers(); i++) {
+    auto device = method_meta.memory_planned_buffer_device(i);
+    THROW_IF_ERROR(
+        device.error(), "Failed to get device of planned buffer %zu", i);
+    if (!device.get().is_cpu()) {
+      return true;
+    }
+  }
+  return false;
+}
+
+/// Arenas sized and placed for a single method, used when that method's
+/// buffers cannot come from the program-wide host arenas.
+std::shared_ptr<ProgramMemory> make_method_memory(
+    const MethodMeta& method_meta) {
+  const size_t num_buffers = method_meta.num_memory_planned_buffers();
+  std::vector<int64_t> sizes;
+  std::vector<runtime::etensor::Device> devices;
+  sizes.reserve(num_buffers);
+  devices.reserve(num_buffers);
+  for (size_t i = 0; i < num_buffers; i++) {
+    auto size = method_meta.memory_planned_buffer_size(i);
+    THROW_IF_ERROR(size.error(), "Failed to get size of planned buffer %zu", i);
+    auto device = method_meta.memory_planned_buffer_device(i);
+    THROW_IF_ERROR(
+        device.error(), "Failed to get device of planned buffer %zu", i);
+    sizes.push_back(size.get());
+    devices.push_back(device.get());
+  }
+  return std::make_shared<ProgramMemory>(std::move(sizes), std::move(devices));
+}
+
 struct PyMethod final {
   explicit PyMethod(
       std::shared_ptr<ProgramMemory> memory,
@@ -1423,8 +1539,13 @@ struct PyProgram final {
     for (size_t i = 0; i < state_->program_->num_methods(); ++i) {
       auto name = state_->program_->get_method_name(i).get();
       auto method_meta = state_->program_->method_meta(name).get();
-      for (size_t j = 0; j < method_meta.num_non_const_buffers(); j++) {
-        int64_t buffer_size = method_meta.non_const_buffer_size(j).get();
+      // A device-planned method gets its own arenas in load_method, so its
+      // sizes must not inflate the shared host arenas nobody will read.
+      if (has_device_buffers(method_meta)) {
+        continue;
+      }
+      for (size_t j = 0; j < method_meta.num_memory_planned_buffers(); j++) {
+        int64_t buffer_size = method_meta.memory_planned_buffer_size(j).get();
         if (non_const_buffer_sizes.find(j) == non_const_buffer_sizes.end()) {
           non_const_buffer_sizes.insert({j, buffer_size});
         } else {
@@ -1434,16 +1555,14 @@ struct PyProgram final {
       }
     }
 
-    // Allocate the arenas. Using vector because we need to remember the size as
-    // well, so vector is easier then unique_ptr.
-    std::vector<std::vector<uint8_t>> non_const_buffers_;
-    for (std::map<size_t, int64_t>::iterator i = non_const_buffer_sizes.begin();
-         i != non_const_buffer_sizes.end();
-         i++) {
-      non_const_buffers_.push_back(std::vector<uint8_t>(i->second));
+    // Allocate the shared host arenas.
+    std::vector<int64_t> planned_sizes;
+    planned_sizes.reserve(non_const_buffer_sizes.size());
+    for (const auto& entry : non_const_buffer_sizes) {
+      planned_sizes.push_back(entry.second);
     }
 
-    memory_ = std::make_shared<ProgramMemory>(std::move(non_const_buffers_));
+    memory_ = std::make_shared<ProgramMemory>(std::move(planned_sizes));
     if (event_tracer_ && debug_buffer_size > 0) {
       // If a debug buffer was requested for the ETDump, allocate it and make
       // sure its lifetime is as long as the event_tracer.
@@ -1508,9 +1627,21 @@ struct PyProgram final {
   }
 
   std::unique_ptr<PyMethod> load_method(const std::string& method_name) {
+    Result<MethodMeta> meta =
+        state_->program_->method_meta(method_name.c_str());
+    THROW_IF_ERROR(
+        meta.error(),
+        "Failed to get method meta for method %s, error: 0x:%" PRIx32,
+        method_name.c_str(),
+        static_cast<uint32_t>(meta.error()));
+    // Device memory is claimed here rather than at program load so that one
+    // accelerator method cannot make the rest of the program unloadable.
+    auto memory = has_device_buffers(meta.get())
+        ? make_method_memory(meta.get())
+        : memory_;
     Result<Method> res = state_->program_->load_method(
         method_name.c_str(),
-        memory_->mem_manager(),
+        memory->mem_manager(),
         event_tracer_.get(),
         state_->data_map_.get());
     THROW_IF_ERROR(
@@ -1519,7 +1650,9 @@ struct PyProgram final {
         method_name.c_str(),
         static_cast<uint32_t>(res.error()));
     return std::make_unique<PyMethod>(
-        memory_, state_, std::make_unique<Method>(std::move(res.get())));
+        std::move(memory),
+        state_,
+        std::make_unique<Method>(std::move(res.get())));
   }
 
   Span<uint8_t> get_etdump_debug_buffer() {
diff --git a/extension/pybindings/test/BUCK b/extension/pybindings/test/BUCK
index a18abac2382..b99748807e4 100644
--- a/extension/pybindings/test/BUCK
+++ b/extension/pybindings/test/BUCK
@@ -26,6 +26,7 @@ fbcode_target(
         "//executorch/exir:pass_manager",
         "//executorch/exir:scalar_type",
         "//executorch/exir/_serialize:lib",
+        "//executorch/exir/backend:partitioner",
         "//executorch/exir/emit:lib",
         "//executorch/exir/passes:lib",
         "//executorch/runtime/core:core",
@@ -39,6 +40,7 @@ fbcode_target(
     preload_deps = ["//executorch/kernels/quantized:aot_lib"],
     deps = [
         ":make_test",
+        "//executorch/exir/backend/test:device_util",
         "//executorch/extension/pybindings:portable_lib",
     ],
 )
@@ -50,6 +52,7 @@ fbcode_target(
     preload_deps = ["//executorch/kernels/quantized:aot_lib"],
     deps = [
         ":make_test",
+        "//executorch/exir/backend/test:device_util",
         "//executorch/extension/pybindings:aten_lib",
         "//executorch/kernels/quantized:aot_lib",
     ],
@@ -61,6 +64,7 @@ fbcode_target(
     srcs = ["test_pybindings.py"],
     deps = [
         ":make_test",
+        "//executorch/exir/backend/test:device_util",
     ],
 )
 
diff --git a/extension/pybindings/test/make_test.py b/extension/pybindings/test/make_test.py
index a1bf4b980e0..1e47dd9b6c6 100644
--- a/extension/pybindings/test/make_test.py
+++ b/extension/pybindings/test/make_test.py
@@ -6,10 +6,11 @@
 
 # pyre-unsafe
 
-from typing import Any, Optional, Tuple
+from typing import Any, Dict, Optional, Tuple
 
 import torch
 from executorch.exir import ExecutorchBackendConfig, ExecutorchProgramManager, to_edge
+from executorch.exir.backend.partitioner import Partitioner
 from torch.export import export
 
 
@@ -151,6 +152,7 @@ def get_inputs(self):
 def create_program(
     eager_module: torch.nn.Module,
     et_config: Optional[ExecutorchBackendConfig] = None,
+    partitioner: Optional[Dict[str, Partitioner]] = None,
 ) -> Tuple[ExecutorchProgramManager, Tuple[Any, ...]]:
     """Returns an executorch program based on ModuleAdd, along with inputs."""
 
@@ -179,7 +181,10 @@ def forward(self, *args, **kwargs):
         wrapped_mod = WrapperModule(getattr(eager_module, method_name))
         exported_methods[method_name] = export(wrapped_mod, method_input, strict=True)
 
-    exec_prog = to_edge(exported_methods).to_executorch(config=et_config)
+    edge_prog = to_edge(exported_methods)
+    if partitioner is not None:
+        edge_prog = edge_prog.to_backend(partitioner)
+    exec_prog = edge_prog.to_executorch(config=et_config)
 
     # Create the ExecuTorch program from the graph.
     exec_prog.dump_executorch_program(verbose=True)
diff --git a/extension/pybindings/test/test_pybindings.py b/extension/pybindings/test/test_pybindings.py
index 19b3620cef3..06eb3e7ed88 100644
--- a/extension/pybindings/test/test_pybindings.py
+++ b/extension/pybindings/test/test_pybindings.py
@@ -13,7 +13,9 @@
 import torch
 
 from executorch.exir import ExecutorchBackendConfig, to_edge
+from executorch.exir.backend.test.device_util import DeviceAwarePartitioner
 from executorch.exir.passes import MemoryPlanningPass
+from executorch.exir.schema import DeviceType
 from executorch.extension.pybindings.test.make_test import (
     create_program,
     ModuleAdd,
@@ -786,3 +788,44 @@ def test_program_method_rejects_input_on_unrepresentable_device(self):
         message = str(caught.exception)
         self.assertIn("is on device meta", message)
         self.assertIn("only CPU and CUDA tensors", message)
+
+    def test_program_loads_when_one_method_is_device_planned(self):
+        # Linking the CUDA backend registers a CUDA allocator at static init, and the
+        # registry has no way to drop one, so there the device load would succeed with
+        # or without the fix and this test could not tell them apart.
+        if "CudaBackend" in self.runtime._get_registered_backend_names():
+            self.skipTest("a registered CUDA allocator satisfies the device load")
+
+        exported_program, inputs = create_program(
+            ModuleMulti(),
+            et_config=ExecutorchBackendConfig(enable_non_cpu_memory_planning=True),
+            partitioner={"forward2": DeviceAwarePartitioner()},
+        )
+
+        # Without this the test would quietly become a second copy of
+        # test_method_multiple_entry if the planner stopped tagging devices.
+        planned_devices = {
+            plan.name: [
+                buffer.device_type for buffer in (plan.non_const_buffer_device or [])
+            ]
+            for plan in exported_program.executorch_program.execution_plan
+        }
+        self.assertIn(DeviceType.CUDA, planned_devices["forward2"])
+        self.assertNotIn(DeviceType.CUDA, planned_devices["forward"])
+
+        program = self.load_prog_fn(exported_program.buffer)
+        self.assertEqual(program.num_methods(), 2)
+
+        method = program.load_method("forward")
+        self.assertTrue(torch.allclose(method.call(inputs)[0], torch.ones(2, 2) * 2))
+
+        # Asking for device memory at all is what this asserts. Before the fix the
+        # loader put every planned buffer in host memory, so the device-planned method
+        # never asked and simply loaded. Now it asks, and with no device allocator
+        # registered the request is refused.
+        with self.assertRaises(RuntimeError) as caught:
+            program.load_method("forward2")
+        self.assertIn("on device", str(caught.exception))
+
+        # A failed load must leave the method that already loaded usable.
+        self.assertTrue(torch.allclose(method.call(inputs)[0], torch.ones(2, 2) * 2))

From d40172709b052c437d4769c8c0951389b5b4c44c Mon Sep 17 00:00:00 2001
From: shoumikhin <shoumikhin@meta.com>
Date: Sun, 23 Aug 2026 19:49:06 -0700
Subject: [PATCH 02/15] Document how planned memory is shared between methods

Program.load_method does not treat every method the same way. A method
whose memory planned buffers are all on the host reuses one set of arenas
shared with the other host only methods of the same program, so running a
second such method can overwrite the intermediates and outputs of the
first. A method with a buffer placed on an accelerator gets its own arenas
instead.

Nothing said so. Say it in the docstring users read.
---
 runtime/__init__.py | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/runtime/__init__.py b/runtime/__init__.py
index 011a15163cb..baec5fc6e80 100644
--- a/runtime/__init__.py
+++ b/runtime/__init__.py
@@ -183,6 +183,15 @@ def method_names(self) -> Set[str]:
     def load_method(self, name: str) -> Optional[Method]:
         """Loads a method from the program.
 
+        Memory-planned buffers are allocated differently depending on where the
+        program placed them. A method whose buffers are all on the host shares
+        one set of arenas with every other host-only method of this program,
+        sized to the largest of them, so running a second host-only method may
+        overwrite the intermediate and output values of the first. A method
+        with any buffer placed on an accelerator gets its own arenas instead,
+        claimed on the first load of that method, so a missing or exhausted
+        accelerator affects only that method and not the rest of the program.
+
         Args:
             name: The name of the method to load.
 

From 6017c7eeeb0498e4577e163ed145866aa5867892 Mon Sep 17 00:00:00 2001
From: shoumikhin <shoumikhin@meta.com>
Date: Sun, 23 Aug 2026 19:49:33 -0700
Subject: [PATCH 03/15] Allocate device buffers before host arenas

Device allocation is the only step here that can fail, for example when no
allocator is registered for the requested device or when the device is out
of memory. It ran last, so every host arena of that method was already
allocated and zero filled before the failure was reported, and all of it
was then discarded.

Swap the two members so the step that can fail runs first. Member
initialization follows declaration order, so the order is what makes this
work and a reorder would silently undo it. Say that in a comment.
---
 extension/pybindings/pybindings.cpp | 14 ++++++++++----
 1 file changed, 10 insertions(+), 4 deletions(-)

diff --git a/extension/pybindings/pybindings.cpp b/extension/pybindings/pybindings.cpp
index 53438833ccb..15b101971b8 100644
--- a/extension/pybindings/pybindings.cpp
+++ b/extension/pybindings/pybindings.cpp
@@ -1084,14 +1084,20 @@ class ProgramMemory {
   /// `devices` is empty when every buffer is on the host, which keeps
   /// `MemoryManager::has_device_memory()` false for CPU-only programs.
   /// Otherwise it holds one entry per buffer, indexed like `sizes`.
+  ///
+  /// Members initialize in declaration order and each one reads the members
+  /// declared before it, so that order is load-bearing. Device buffers come
+  /// first because they hold the only allocation that can fail, and throwing
+  /// before the host arenas exist avoids zero-filling memory that is about to
+  /// be discarded.
   ProgramMemory(
       std::vector<int64_t>&& sizes,
       std::vector<runtime::etensor::Device>&& devices)
       : runtime_allocator_(),
         planned_sizes_(std::move(sizes)),
         planned_devices_(std::move(devices)),
-        non_const_buffers_(allocate_host_buffers()),
         device_buffers_(allocate_device_buffers()),
+        non_const_buffers_(allocate_host_buffers()),
         non_const_spans_(create_non_const_spans()),
         non_const_allocator_(create_non_const_allocator()),
         mem_manager_(
@@ -1123,14 +1129,14 @@ class ProgramMemory {
 
   std::vector<runtime::etensor::Device> planned_devices_;
 
-  // Backs CPU-tagged buffers; the entry is empty for a device-tagged buffer.
-  std::vector<std::vector<uint8_t>> non_const_buffers_;
-
   // Backs device-tagged buffers; the entry is empty for a CPU-tagged buffer.
   // Parallel to non_const_buffers_ so both index by planned buffer id. Empty
   // for an all-host program.
   std::vector<DeviceMemoryBuffer> device_buffers_;
 
+  // Backs CPU-tagged buffers; the entry is empty for a device-tagged buffer.
+  std::vector<std::vector<uint8_t>> non_const_buffers_;
+
   std::vector<Span<uint8_t>> non_const_spans_;
 
   HierarchicalAllocator non_const_allocator_;

From 638890f834108bc7c11f706db54683b469416a86 Mon Sep 17 00:00:00 2001
From: shoumikhin <shoumikhin@meta.com>
Date: Sun, 23 Aug 2026 19:50:05 -0700
Subject: [PATCH 04/15] Raise instead of aborting on a buffer and device count
 mismatch

ProgramMemory hands one span per planned buffer and one device tag per
planned buffer to HierarchicalAllocator, which checks that the two counts
agree. That check aborts the process. In a Python extension that means the
interpreter dies with no traceback and nothing the caller can catch.

Both constructors keep the counts equal today, so this is not reachable, but
the invariant lives only in the callers. Check it where the error can still
become a Python exception.
---
 extension/pybindings/pybindings.cpp | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/extension/pybindings/pybindings.cpp b/extension/pybindings/pybindings.cpp
index 15b101971b8..a4513345bb8 100644
--- a/extension/pybindings/pybindings.cpp
+++ b/extension/pybindings/pybindings.cpp
@@ -1165,6 +1165,16 @@ class ProgramMemory {
     if (planned_devices_.empty()) {
       return result;
     }
+    // HierarchicalAllocator aborts rather than throws if the span count and
+    // the device count disagree, so reject that here where a Python caller can
+    // still catch it.
+    THROW_IF_ERROR(
+        planned_devices_.size() == planned_sizes_.size()
+            ? Error::Ok
+            : Error::InvalidArgument,
+        "Have %zu planned buffer sizes but %zu device tags",
+        planned_sizes_.size(),
+        planned_devices_.size());
     result.reserve(planned_sizes_.size());
     for (size_t i = 0; i < planned_sizes_.size(); i++) {
       if (!is_device_buffer(i)) {

From 43f274e50fb2e99835880bad03ef5504478a6d3a Mon Sep 17 00:00:00 2001
From: shoumikhin <shoumikhin@meta.com>
Date: Sun, 23 Aug 2026 19:50:05 -0700
Subject: [PATCH 05/15] Read a method's planned buffer metadata in one pass

load_method walked every planned buffer once through has_device_buffers to
decide whether the method needs its own arenas, then walked them all again
through make_method_memory to collect the sizes and devices. Each device
lookup scans the program's sparse device list, so the second walk repeats
work the first one already did.

Have make_method_memory return nullptr when every buffer is on the host.
One pass then answers both questions. has_device_buffers stays because the
program constructor still needs only the answer, not the sizes.
---
 extension/pybindings/pybindings.cpp | 18 +++++++++++++-----
 1 file changed, 13 insertions(+), 5 deletions(-)

diff --git a/extension/pybindings/pybindings.cpp b/extension/pybindings/pybindings.cpp
index a4513345bb8..a587581f6ff 100644
--- a/extension/pybindings/pybindings.cpp
+++ b/extension/pybindings/pybindings.cpp
@@ -1235,7 +1235,9 @@ bool has_device_buffers(const MethodMeta& method_meta) {
 }
 
 /// Arenas sized and placed for a single method, used when that method's
-/// buffers cannot come from the program-wide host arenas.
+/// buffers cannot come from the program-wide host arenas. Returns nullptr when
+/// every buffer is on the host, so one pass over the metadata answers both
+/// whether the method needs its own arenas and how big they are.
 std::shared_ptr<ProgramMemory> make_method_memory(
     const MethodMeta& method_meta) {
   const size_t num_buffers = method_meta.num_memory_planned_buffers();
@@ -1243,15 +1245,20 @@ std::shared_ptr<ProgramMemory> make_method_memory(
   std::vector<runtime::etensor::Device> devices;
   sizes.reserve(num_buffers);
   devices.reserve(num_buffers);
+  bool needs_device_memory = false;
   for (size_t i = 0; i < num_buffers; i++) {
     auto size = method_meta.memory_planned_buffer_size(i);
     THROW_IF_ERROR(size.error(), "Failed to get size of planned buffer %zu", i);
     auto device = method_meta.memory_planned_buffer_device(i);
     THROW_IF_ERROR(
         device.error(), "Failed to get device of planned buffer %zu", i);
+    needs_device_memory |= !device.get().is_cpu();
     sizes.push_back(size.get());
     devices.push_back(device.get());
   }
+  if (!needs_device_memory) {
+    return nullptr;
+  }
   return std::make_shared<ProgramMemory>(std::move(sizes), std::move(devices));
 }
 
@@ -1651,10 +1658,11 @@ struct PyProgram final {
         method_name.c_str(),
         static_cast<uint32_t>(meta.error()));
     // Device memory is claimed here rather than at program load so that one
-    // accelerator method cannot make the rest of the program unloadable.
-    auto memory = has_device_buffers(meta.get())
-        ? make_method_memory(meta.get())
-        : memory_;
+    // accelerator method cannot make the rest of the program unloadable. A
+    // host-only method keeps sharing the program-wide arenas, so its planned
+    // memory is not isolated from the other host-only methods of this program.
+    auto method_memory = make_method_memory(meta.get());
+    auto memory = method_memory ? std::move(method_memory) : memory_;
     Result<Method> res = state_->program_->load_method(
         method_name.c_str(),
         memory->mem_manager(),

From 247caed2c0785f295fcfa28492e3e927fbb6d5cb Mon Sep 17 00:00:00 2001
From: shoumikhin <shoumikhin@meta.com>
Date: Sun, 23 Aug 2026 19:50:27 -0700
Subject: [PATCH 06/15] Declare the direct dependency on device_memory_buffer

pybindings.cpp includes runtime/core/device_memory_buffer.h but the
pybindings targets did not depend on it. The build worked only because the
header came in transitively through extension/module, and the rule opts out
of automatic dependency checking, so nothing would report the omission.
---
 shim_et/xplat/executorch/extension/pybindings/pybindings.bzl | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/shim_et/xplat/executorch/extension/pybindings/pybindings.bzl b/shim_et/xplat/executorch/extension/pybindings/pybindings.bzl
index 98ab3c7cf94..e6b0319ca97 100644
--- a/shim_et/xplat/executorch/extension/pybindings/pybindings.bzl
+++ b/shim_et/xplat/executorch/extension/pybindings/pybindings.bzl
@@ -10,6 +10,7 @@ MODELS_ATEN_OPS_LEAN_MODE_GENERATED_LIB = [
 PORTABLE_MODULE_DEPS = [
     "//executorch/runtime/kernel:operator_registry",
     "//executorch/runtime/executor:program",
+    "//executorch/runtime/core:device_memory_buffer",
     "//executorch/devtools/bundled_program/schema:bundled_program_schema_fbs",
     "//executorch/extension/aten_util:aten_bridge",
     "//executorch/devtools/bundled_program:runtime",
@@ -27,6 +28,7 @@ PORTABLE_MODULE_DEPS = [
 ATEN_MODULE_DEPS = [
     "//executorch/runtime/kernel:operator_registry_aten",
     "//executorch/runtime/executor:program_aten",
+    "//executorch/runtime/core:device_memory_buffer",
     "//executorch/runtime/core/exec_aten:lib_aten",
     "//executorch/devtools/bundled_program/schema:bundled_program_schema_fbs",
     "//executorch/extension/data_loader:buffer_data_loader",

From c51de12c0414b2ba30f5b8a4dd88d2efc82598a4 Mon Sep 17 00:00:00 2001
From: shoumikhin <shoumikhin@meta.com>
Date: Sun, 23 Aug 2026 22:52:02 -0700
Subject: [PATCH 07/15] Correct three comments in the pybindings memory code

Three comments describe the code inaccurately.

The first says device buffers are allocated first because they hold the
only allocation that can fail. Allocating a host arena can also fail, so
the real reason is only that a device failure should happen before the
host arenas are allocated and zero filled.

The second implies the buffer and device count guard converts a
reachable abort. Both vectors are filled in lockstep today, so the guard
only fires if a future caller breaks that.

The third says a device planned method inflates arenas nobody will read.
Every host only method reads those arenas. What that method would inflate
is their size.
---
 extension/pybindings/pybindings.cpp | 16 ++++++++--------
 1 file changed, 8 insertions(+), 8 deletions(-)

diff --git a/extension/pybindings/pybindings.cpp b/extension/pybindings/pybindings.cpp
index a587581f6ff..8169ec6c5ca 100644
--- a/extension/pybindings/pybindings.cpp
+++ b/extension/pybindings/pybindings.cpp
@@ -1087,9 +1087,8 @@ class ProgramMemory {
   ///
   /// Members initialize in declaration order and each one reads the members
   /// declared before it, so that order is load-bearing. Device buffers come
-  /// first because they hold the only allocation that can fail, and throwing
-  /// before the host arenas exist avoids zero-filling memory that is about to
-  /// be discarded.
+  /// first so that a device that is missing or out of memory throws before the
+  /// host arenas are allocated and zero-filled, rather than after.
   ProgramMemory(
       std::vector<int64_t>&& sizes,
       std::vector<runtime::etensor::Device>&& devices)
@@ -1165,9 +1164,9 @@ class ProgramMemory {
     if (planned_devices_.empty()) {
       return result;
     }
-    // HierarchicalAllocator aborts rather than throws if the span count and
-    // the device count disagree, so reject that here where a Python caller can
-    // still catch it.
+    // Both vectors are filled in lockstep today, so this only fires if a
+    // future caller breaks that. HierarchicalAllocator aborts on a mismatch,
+    // so check here instead, where a Python caller can catch it.
     THROW_IF_ERROR(
         planned_devices_.size() == planned_sizes_.size()
             ? Error::Ok
@@ -1562,8 +1561,9 @@ struct PyProgram final {
     for (size_t i = 0; i < state_->program_->num_methods(); ++i) {
       auto name = state_->program_->get_method_name(i).get();
       auto method_meta = state_->program_->method_meta(name).get();
-      // A device-planned method gets its own arenas in load_method, so its
-      // sizes must not inflate the shared host arenas nobody will read.
+      // A device-planned method gets its own arenas in load_method and never
+      // reads these, so letting its sizes in would only grow the host arenas
+      // the other methods share.
       if (has_device_buffers(method_meta)) {
         continue;
       }

From 26941767a25c181776f1871e0e37107931946195 Mon Sep 17 00:00:00 2001
From: shoumikhin <shoumikhin@meta.com>
Date: Sun, 23 Aug 2026 22:52:02 -0700
Subject: [PATCH 08/15] Check the planned buffer size before reading it

Result::get() aborts the process if the Result holds an error, so a
corrupt program file that reports a bad planned buffer size kills the
interpreter instead of raising. This line was already rewritten by an
earlier commit in this change, and the device buffer code added next to
it already checks the same call, so the two now behave the same way.
---
 extension/pybindings/pybindings.cpp | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/extension/pybindings/pybindings.cpp b/extension/pybindings/pybindings.cpp
index 8169ec6c5ca..bce8113dbb5 100644
--- a/extension/pybindings/pybindings.cpp
+++ b/extension/pybindings/pybindings.cpp
@@ -1568,7 +1568,10 @@ struct PyProgram final {
         continue;
       }
       for (size_t j = 0; j < method_meta.num_memory_planned_buffers(); j++) {
-        int64_t buffer_size = method_meta.memory_planned_buffer_size(j).get();
+        auto size = method_meta.memory_planned_buffer_size(j);
+        THROW_IF_ERROR(
+            size.error(), "Failed to get size of planned buffer %zu", j);
+        int64_t buffer_size = size.get();
         if (non_const_buffer_sizes.find(j) == non_const_buffer_sizes.end()) {
           non_const_buffer_sizes.insert({j, buffer_size});
         } else {

From 4c938f05994e6ccca025c6f7f25fb84b4cc27ac0 Mon Sep 17 00:00:00 2001
From: shoumikhin <shoumikhin@meta.com>
Date: Sun, 23 Aug 2026 22:52:25 -0700
Subject: [PATCH 09/15] Correct the load_method docstring on planned memory

The docstring said a second host only method may overwrite the outputs
of the first. It cannot: outputs are copied out on every call, so what
sharing affects is the intermediate values, not what execute returns.

It also said a failed device allocation affects only that method. There
is no unload, so whatever a device method does claim stays claimed for
every method loaded afterwards. Both points are now stated as they are.
---
 runtime/__init__.py | 13 ++++++++-----
 1 file changed, 8 insertions(+), 5 deletions(-)

diff --git a/runtime/__init__.py b/runtime/__init__.py
index baec5fc6e80..f8a0bd1d5bc 100644
--- a/runtime/__init__.py
+++ b/runtime/__init__.py
@@ -186,11 +186,14 @@ def load_method(self, name: str) -> Optional[Method]:
         Memory-planned buffers are allocated differently depending on where the
         program placed them. A method whose buffers are all on the host shares
         one set of arenas with every other host-only method of this program,
-        sized to the largest of them, so running a second host-only method may
-        overwrite the intermediate and output values of the first. A method
-        with any buffer placed on an accelerator gets its own arenas instead,
-        claimed on the first load of that method, so a missing or exhausted
-        accelerator affects only that method and not the rest of the program.
+        sized to the largest of them, so two such methods overwrite each
+        other's intermediate values. Outputs are copied out on every call, so
+        that sharing does not change what execute returns. A method with any
+        buffer placed on an accelerator gets its own arenas instead, claimed
+        when that method is first loaded and held until the program is
+        released. A missing or exhausted accelerator therefore fails that one
+        load rather than the whole program, but whatever it does claim stays
+        claimed for every method loaded after it.
 
         Args:
             name: The name of the method to load.

From 72e317ea526ae6d2b47a35af32620297619102e7 Mon Sep 17 00:00:00 2001
From: shoumikhin <shoumikhin@meta.com>
Date: Sun, 23 Aug 2026 22:52:25 -0700
Subject: [PATCH 10/15] Declare the exir schema dependency in the pybindings
 tests

The test file imports DeviceType from executorch.exir.schema and
executorch.exir.backend.test.device_util, and the three test targets
picked both up only through make_test. A direct import needs a direct
dependency, so the targets no longer break if make_test drops either
one.
---
 extension/pybindings/test/BUCK | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/extension/pybindings/test/BUCK b/extension/pybindings/test/BUCK
index b99748807e4..e90e837732d 100644
--- a/extension/pybindings/test/BUCK
+++ b/extension/pybindings/test/BUCK
@@ -40,6 +40,7 @@ fbcode_target(
     preload_deps = ["//executorch/kernels/quantized:aot_lib"],
     deps = [
         ":make_test",
+        "//executorch/exir:schema",
         "//executorch/exir/backend/test:device_util",
         "//executorch/extension/pybindings:portable_lib",
     ],
@@ -52,6 +53,7 @@ fbcode_target(
     preload_deps = ["//executorch/kernels/quantized:aot_lib"],
     deps = [
         ":make_test",
+        "//executorch/exir:schema",
         "//executorch/exir/backend/test:device_util",
         "//executorch/extension/pybindings:aten_lib",
         "//executorch/kernels/quantized:aot_lib",
@@ -64,6 +66,7 @@ fbcode_target(
     srcs = ["test_pybindings.py"],
     deps = [
         ":make_test",
+        "//executorch/exir:schema",
         "//executorch/exir/backend/test:device_util",
     ],
 )

From f30dd9d798f9d85238cdd5bb921953135b7389c3 Mon Sep 17 00:00:00 2001
From: shoumikhin <shoumikhin@meta.com>
Date: Sun, 23 Aug 2026 22:52:25 -0700
Subject: [PATCH 11/15] Test that a device planned method allocates on the
 device

The existing test covers the case where no device allocator is
registered and the load is refused. Nothing covered the case the change
exists to fix: a build that does have a device allocator, where the
arena has to come off the device.

The new test lowers one of two methods to the CUDA backend, checks that
free device memory drops by the planned buffer size when that method is
loaded, checks that both methods return the right values, and checks
that the memory comes back when the program is released. On the previous
behavior the arena stayed in host memory and the backend rejected the
pointer with an execute error, so this test fails without the fix.

A real accelerator is required, so the test skips unless one is present.
The CUDA build job now runs it, and the CUDA workflow now triggers on
changes under extension/pybindings, which is where this code lives.
---
 .ci/scripts/test-cuda-build.sh               |   5 +
 .github/workflows/cuda.yml                   |   2 +
 extension/pybindings/test/test_pybindings.py | 111 +++++++++++++++++++
 3 files changed, 118 insertions(+)

diff --git a/.ci/scripts/test-cuda-build.sh b/.ci/scripts/test-cuda-build.sh
index e717718be66..b7bf530830b 100755
--- a/.ci/scripts/test-cuda-build.sh
+++ b/.ci/scripts/test-cuda-build.sh
@@ -80,6 +80,11 @@ except Exception as e:
     exit(1)
 "
 
+    # The pybindings loader only asks for device memory on a build that has a
+    # device allocator registered, so this job is the only place that path runs.
+    python -m unittest -v \
+        executorch.extension.pybindings.test.test_pybindings.PybindingsTest.test_device_planned_method_allocates_on_the_device
+
     echo "SUCCESS: ExecuTorch CUDA ${cuda_version} build and verification completed successfully"
 }
 
diff --git a/.github/workflows/cuda.yml b/.github/workflows/cuda.yml
index b490e28c95f..c8e03d5dd6e 100644
--- a/.github/workflows/cuda.yml
+++ b/.github/workflows/cuda.yml
@@ -26,6 +26,8 @@ on:
       - .ci/scripts/export_model_artifact.sh
       - .ci/scripts/test_model_e2e.sh
       - examples/models/muse-glimmer/**
+      - extension/pybindings/**
+      - runtime/__init__.py
   workflow_dispatch:
 
 concurrency:
diff --git a/extension/pybindings/test/test_pybindings.py b/extension/pybindings/test/test_pybindings.py
index 06eb3e7ed88..7e0d3dfce25 100644
--- a/extension/pybindings/test/test_pybindings.py
+++ b/extension/pybindings/test/test_pybindings.py
@@ -6,7 +6,9 @@
 
 # pyre-unsafe
 
+import os
 import sys
+import tempfile
 import unittest
 from io import StringIO
 
@@ -829,3 +831,112 @@ def test_program_loads_when_one_method_is_device_planned(self):
 
         # A failed load must leave the method that already loaded usable.
         self.assertTrue(torch.allclose(method.call(inputs)[0], torch.ones(2, 2) * 2))
+
+    def test_device_planned_method_allocates_on_the_device(self):
+        # The other device test covers the refusal. This one covers what the
+        # refusal is protecting: on a build that does have a device allocator,
+        # the arena has to come off the device, not out of host memory. It needs
+        # a real accelerator, so it only runs where one is present.
+        if "CudaBackend" not in self.runtime._get_registered_backend_names():
+            self.skipTest("needs a build with the CUDA backend linked in")
+        if not torch.cuda.is_available():
+            self.skipTest("needs a visible CUDA device")
+
+        from executorch.backends.cuda.cuda_partitioner import CudaPartitioner
+        from executorch.exir import to_edge_transform_and_lower
+        from executorch.exir.backend.compile_spec_schema import CompileSpec
+
+        # Large enough that the arena is far bigger than the noise in
+        # mem_get_info, which moves by a few MiB as contexts are created.
+        side = 2048
+        inputs = (torch.ones(side, side), torch.ones(side, side))
+
+        class HostOnly(torch.nn.Module):
+            def forward(self, x, y):
+                return x + y
+
+        class Delegated(torch.nn.Module):
+            def forward(self, x, y):
+                return (x + y) * 2.0
+
+        edge = to_edge_transform_and_lower(
+            {
+                "forward": export(HostOnly(), inputs, strict=True),
+                "forward2": export(Delegated(), inputs, strict=True),
+            },
+            partitioner={
+                "forward2": [CudaPartitioner([CompileSpec("method_name", b"forward2")])]
+            },
+        )
+        exported_program = edge.to_executorch(
+            config=ExecutorchBackendConfig(enable_non_cpu_memory_planning=True)
+        )
+
+        plans = {
+            plan.name: plan
+            for plan in exported_program.executorch_program.execution_plan
+        }
+        device_buffers = {
+            buffer.buffer_idx
+            for buffer in (plans["forward2"].non_const_buffer_device or [])
+        }
+        self.assertTrue(device_buffers, "the planner tagged nothing for the device")
+        self.assertFalse(plans["forward"].non_const_buffer_device or [])
+        device_bytes = sum(
+            size
+            for index, size in enumerate(plans["forward2"].non_const_buffer_sizes)
+            if index in device_buffers
+        )
+
+        # The CUDA backend keeps its weights in a separate file, so the program
+        # has to be loaded from disk with that file alongside it.
+        with tempfile.TemporaryDirectory() as directory:
+            pte_path = os.path.join(directory, "program.pte")
+            with open(pte_path, "wb") as pte_file:
+                exported_program.write_to_file(pte_file)
+            data_names = sorted(exported_program._tensor_data or {})
+            exported_program.write_tensor_data_to_file(directory)
+            data_path = (
+                os.path.join(directory, data_names[0] + ".ptd") if data_names else None
+            )
+
+            torch.cuda.init()
+            program = self.runtime._load_program(pte_path, data_path=data_path)
+
+            torch.cuda.synchronize()
+            free_before, _ = torch.cuda.mem_get_info()
+
+            # The host-only method must not touch the device at all.
+            host_method = program.load_method("forward")
+            torch.cuda.synchronize()
+            free_after_host, _ = torch.cuda.mem_get_info()
+            self.assertLess(free_before - free_after_host, device_bytes // 2)
+
+            device_method = program.load_method("forward2")
+            torch.cuda.synchronize()
+            free_after_device, _ = torch.cuda.mem_get_info()
+            self.assertGreaterEqual(
+                free_after_host - free_after_device, device_bytes * 0.9
+            )
+
+            # Before the fix this ran on a host pointer and the backend rejected
+            # it, so getting the right answer back is itself part of the check.
+            expected = (inputs[0] + inputs[1]) * 2.0
+            self.assertTrue(
+                torch.allclose(device_method.call(inputs)[0].cpu(), expected)
+            )
+            self.assertTrue(
+                torch.allclose(host_method.call(inputs)[0].cpu(), inputs[0] + inputs[1])
+            )
+            # The device arenas are private, so running the host method in
+            # between must not disturb them.
+            self.assertTrue(
+                torch.allclose(device_method.call(inputs)[0].cpu(), expected)
+            )
+
+            del device_method
+            del host_method
+            del program
+            torch.cuda.synchronize()
+            free_at_end, _ = torch.cuda.mem_get_info()
+            self.assertGreaterEqual(free_at_end - free_after_device, device_bytes * 0.9)

From 704e8f7ff16803e16aa181b6a0cbb60dc52673e4 Mon Sep 17 00:00:00 2001
From: shoumikhin <shoumikhin@meta.com>
Date: Mon, 24 Aug 2026 08:23:08 -0700
Subject: [PATCH 12/15] Run the pybindings device test in the CUDA unit test
 job

The new test loads a method that the CUDA backend runs, so the process has to
open the shared library the backend produces. That library needs a newer
libstdc++ than the CI image keeps on its default search path.

Every CUDA job in this workflow that runs a model already installs
libstdcxx-ng and puts the conda library directory first on LD_LIBRARY_PATH.
The one job that does not is test-cuda-builds, because until now it only built
ExecuTorch and ran a plain torch matmul on the GPU, never a delegated method.
Putting the test there made it fail on both CUDA versions with:

  libstdc++.so.6: version `GLIBCXX_3.4.29' not found
  Init failed for backend CudaBackend: 0x22

Move the test to unittest-cuda, the job for CUDA Python unit tests, which
already has that setup. Run it right after the install and before the long
builds, since it needs nothing they produce and failing early saves GPU time.

Also add the two pybindings paths to that job's condition. They were added to
the workflow path filter, but every job carries its own condition, so without
this a change to pybindings alone would start the workflow and then skip every
job in it.

test-cuda-build.sh now matches main again, so this change no longer touches it.

Test plan: checked the workflow file parses and the job count is unchanged,
checked the shell script with bash -n and shellcheck and confirmed it is
byte-identical to main, confirmed the test class and method names resolve, and
evaluated the new job condition against sample file lists. The test itself
needs a CUDA build of the pybindings that only CI produces, so CI is the check.
---
 .ci/scripts/test-cuda-build.sh |  5 -----
 .github/workflows/cuda.yml     | 10 ++++++++++
 2 files changed, 10 insertions(+), 5 deletions(-)

diff --git a/.ci/scripts/test-cuda-build.sh b/.ci/scripts/test-cuda-build.sh
index b7bf530830b..e717718be66 100755
--- a/.ci/scripts/test-cuda-build.sh
+++ b/.ci/scripts/test-cuda-build.sh
@@ -80,11 +80,6 @@ except Exception as e:
     exit(1)
 "
 
-    # The pybindings loader only asks for device memory on a build that has a
-    # device allocator registered, so this job is the only place that path runs.
-    python -m unittest -v \
-        executorch.extension.pybindings.test.test_pybindings.PybindingsTest.test_device_planned_method_allocates_on_the_device
-
     echo "SUCCESS: ExecuTorch CUDA ${cuda_version} build and verification completed successfully"
 }
 
diff --git a/.github/workflows/cuda.yml b/.github/workflows/cuda.yml
index c8e03d5dd6e..5b67d3eecb6 100644
--- a/.github/workflows/cuda.yml
+++ b/.github/workflows/cuda.yml
@@ -162,6 +162,8 @@ jobs:
   unittest-cuda:
     name: unittest-cuda
     needs: [changed-files, run-decision]
+    # This job runs the pybindings device test, so it also has to fire on the
+    # two pybindings paths the workflow filter lists.
     if: |
       contains(needs.changed-files.outputs.changed-files, 'backends/cuda') ||
       contains(needs.changed-files.outputs.changed-files, 'backends/aoti') ||
@@ -169,6 +171,8 @@ jobs:
       contains(needs.changed-files.outputs.changed-files, '.ci/scripts/test-cuda-build.sh') ||
       contains(needs.changed-files.outputs.changed-files, '.ci/scripts/export_model_artifact.sh') ||
       contains(needs.changed-files.outputs.changed-files, '.ci/scripts/test_model_e2e.sh') ||
+      contains(needs.changed-files.outputs.changed-files, 'extension/pybindings') ||
+      contains(needs.changed-files.outputs.changed-files, 'runtime/__init__.py') ||
       needs.run-decision.outputs.is-full-run == 'true'
     uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
     permissions:
@@ -192,6 +196,12 @@ jobs:
         conda install -y -c conda-forge 'libstdcxx-ng>=12'
         export LD_LIBRARY_PATH=/opt/conda/lib:$LD_LIBRARY_PATH
 
+        # The pybindings loader only asks for device memory on a build that has
+        # a device allocator registered, so this test cannot run in the CPU
+        # jobs. It needs the install above and nothing built below, so run it
+        # here and fail before the long builds.
+        python -m unittest -v executorch.extension.pybindings.test.test_pybindings.PybindingsTest.test_device_planned_method_allocates_on_the_device
+
         # Build ExecuTorch with CUDA support
         cmake --workflow --preset llm-release-cuda
 

From e06b99383be6c2fffdda8a4ef00785f36050c86c Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <shoumikhin@users.noreply.github.com>
Date: Mon, 24 Aug 2026 08:40:51 -0700
Subject: [PATCH 13/15] Use prefix increment in the loops this change adds

The loop counters added by this change used postfix increment. Postfix
has to keep the old value around to return it, prefix does not, and for
a plain size_t the two are identical after optimization. Prefix is the
convention this file should follow, so use it consistently.

Six loops in extension/pybindings/pybindings.cpp, all of them lines this
change introduced. The three postfix loops that were already in the file
before this change are left alone.

Test plan: no behavior change. The counter type is size_t and the result
of the increment expression is discarded in every one of the six loops,
so prefix and postfix compute the same thing. Verified the file still
compiles with the same command the build uses, and verified with a word
level diff that the only tokens that moved are the six increments.
---
 extension/pybindings/pybindings.cpp | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/extension/pybindings/pybindings.cpp b/extension/pybindings/pybindings.cpp
index bce8113dbb5..27255b3b8ac 100644
--- a/extension/pybindings/pybindings.cpp
+++ b/extension/pybindings/pybindings.cpp
@@ -1149,7 +1149,7 @@ class ProgramMemory {
   std::vector<std::vector<uint8_t>> allocate_host_buffers() {
     std::vector<std::vector<uint8_t>> result;
     result.reserve(planned_sizes_.size());
-    for (size_t i = 0; i < planned_sizes_.size(); i++) {
+    for (size_t i = 0; i < planned_sizes_.size(); ++i) {
       if (is_device_buffer(i)) {
         result.emplace_back();
       } else {
@@ -1175,7 +1175,7 @@ class ProgramMemory {
         planned_sizes_.size(),
         planned_devices_.size());
     result.reserve(planned_sizes_.size());
-    for (size_t i = 0; i < planned_sizes_.size(); i++) {
+    for (size_t i = 0; i < planned_sizes_.size(); ++i) {
       if (!is_device_buffer(i)) {
         result.emplace_back();
         continue;
@@ -1199,7 +1199,7 @@ class ProgramMemory {
   std::vector<Span<uint8_t>> create_non_const_spans() {
     std::vector<Span<uint8_t>> result;
     result.reserve(planned_sizes_.size());
-    for (size_t i = 0; i < planned_sizes_.size(); i++) {
+    for (size_t i = 0; i < planned_sizes_.size(); ++i) {
       if (is_device_buffer(i)) {
         result.push_back(device_buffers_[i].as_span());
       } else {
@@ -1222,7 +1222,7 @@ class ProgramMemory {
 
 /// True if any of the method's memory-planned buffers must live off the host.
 bool has_device_buffers(const MethodMeta& method_meta) {
-  for (size_t i = 0; i < method_meta.num_memory_planned_buffers(); i++) {
+  for (size_t i = 0; i < method_meta.num_memory_planned_buffers(); ++i) {
     auto device = method_meta.memory_planned_buffer_device(i);
     THROW_IF_ERROR(
         device.error(), "Failed to get device of planned buffer %zu", i);
@@ -1245,7 +1245,7 @@ std::shared_ptr<ProgramMemory> make_method_memory(
   sizes.reserve(num_buffers);
   devices.reserve(num_buffers);
   bool needs_device_memory = false;
-  for (size_t i = 0; i < num_buffers; i++) {
+  for (size_t i = 0; i < num_buffers; ++i) {
     auto size = method_meta.memory_planned_buffer_size(i);
     THROW_IF_ERROR(size.error(), "Failed to get size of planned buffer %zu", i);
     auto device = method_meta.memory_planned_buffer_device(i);
@@ -1567,7 +1567,7 @@ struct PyProgram final {
       if (has_device_buffers(method_meta)) {
         continue;
       }
-      for (size_t j = 0; j < method_meta.num_memory_planned_buffers(); j++) {
+      for (size_t j = 0; j < method_meta.num_memory_planned_buffers(); ++j) {
         auto size = method_meta.memory_planned_buffer_size(j);
         THROW_IF_ERROR(
             size.error(), "Failed to get size of planned buffer %zu", j);

From 2452cd64823ed582911adaece05c694407d32f9a Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Mon, 24 Aug 2026 10:28:38 -0700
Subject: [PATCH 14/15] Name the right owner of device arenas in the
 load_method docstring

The docstring said the private arenas of a device-planned method are
held until the program is released. That is the wrong owner. PyMethod
keeps a shared_ptr to the ProgramMemory, so a caller who keeps the
returned method alive and drops the program keeps the accelerator
memory too. Releasing the program alone does not give it back.

Say that the method owns the arenas, and add the reason the two usually
look the same: this Program caches every method it loads, so the
methods and the program normally die together.

Documentation only, no behavior change.

Test plan:
  black --check runtime/__init__.py  ->  unchanged
  python -m py_compile runtime/__init__.py  ->  ok
---
 runtime/__init__.py | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/runtime/__init__.py b/runtime/__init__.py
index f8a0bd1d5bc..7427c1d7680 100644
--- a/runtime/__init__.py
+++ b/runtime/__init__.py
@@ -190,10 +190,12 @@ def load_method(self, name: str) -> Optional[Method]:
         other's intermediate values. Outputs are copied out on every call, so
         that sharing does not change what execute returns. A method with any
         buffer placed on an accelerator gets its own arenas instead, claimed
-        when that method is first loaded and held until the program is
-        released. A missing or exhausted accelerator therefore fails that one
-        load rather than the whole program, but whatever it does claim stays
-        claimed for every method loaded after it.
+        when that method is first loaded and owned by that method until the
+        method itself is released. This program caches every method it loads,
+        so in normal use those arenas live as long as the program does. A
+        missing or exhausted accelerator therefore fails that one load rather
+        than the whole program, but whatever it does claim stays claimed for
+        every method loaded after it.
 
         Args:
             name: The name of the method to load.

From a3f584cd8efaadb978d7b1bee970fb3fc101b33e Mon Sep 17 00:00:00 2001
From: r <r@e>
Date: Mon, 24 Aug 2026 13:05:13 -0700
Subject: [PATCH 15/15] Fail the CUDA job if the pybindings device test skips,
 and document the arena split

unittest exits 0 when every test skips, so a missing CUDA backend or GPU would
have read as a pass. Recover the status through tee with PIPESTATUS and treat a
skip as an error. The pattern is anchored so a test whose name contains
"skipped" does not trip it.

Also document why a device-planned method gets its own arenas rather than joining
the program-wide ones.
---
 .github/workflows/cuda.yml | 19 ++++++++++++++++++-
 runtime/__init__.py        |  7 +++++++
 2 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/cuda.yml b/.github/workflows/cuda.yml
index 5b67d3eecb6..25684f8b32e 100644
--- a/.github/workflows/cuda.yml
+++ b/.github/workflows/cuda.yml
@@ -200,7 +200,24 @@ jobs:
         # a device allocator registered, so this test cannot run in the CPU
         # jobs. It needs the install above and nothing built below, so run it
         # here and fail before the long builds.
-        python -m unittest -v executorch.extension.pybindings.test.test_pybindings.PybindingsTest.test_device_planned_method_allocates_on_the_device
+        #
+        # unittest exits 0 when a test skips, and this test skips when the CUDA
+        # backend is missing or no GPU is visible. Both are the prerequisites
+        # this job exists to provide, so treat a skip as a failure rather than
+        # let a packaging regression pass silently. `tee` would otherwise hide
+        # the real exit status behind its own, so read it back explicitly.
+        set +e
+        python -m unittest -v executorch.extension.pybindings.test.test_pybindings.PybindingsTest.test_device_planned_method_allocates_on_the_device 2>&1 | tee /tmp/pybindings_device_test.log
+        pybindings_device_status=${PIPESTATUS[0]}
+        set -e
+        if [ "${pybindings_device_status}" -ne 0 ]; then
+          echo "::error::the pybindings device test failed"
+          exit "${pybindings_device_status}"
+        fi
+        if grep -qE "^OK \(skipped=|skipped=[1-9]" /tmp/pybindings_device_test.log; then
+          echo "::error::the pybindings device test skipped in the CUDA job, which means the CUDA backend or a visible GPU is missing"
+          exit 1
+        fi
 
         # Build ExecuTorch with CUDA support
         cmake --workflow --preset llm-release-cuda
diff --git a/runtime/__init__.py b/runtime/__init__.py
index 7427c1d7680..ba83ff54856 100644
--- a/runtime/__init__.py
+++ b/runtime/__init__.py
@@ -197,6 +197,13 @@ def load_method(self, name: str) -> Optional[Method]:
         than the whole program, but whatever it does claim stays claimed for
         every method loaded after it.
 
+        A program exported with ``share_mutable_buffers=True`` relies on every
+        method reading its shared mutable state from one allocation. A method
+        with accelerator-placed buffers gets its own arenas instead, so it also
+        gets its own copy of that state and will not observe writes made through
+        another method. Export without ``share_mutable_buffers`` if a method
+        needs both accelerator memory and shared state.
+
         Args:
             name: The name of the method to load.