Add CUDA weight offloading

Lets a CUDA-backend ExecuTorch program stream model weights from a
pinned host mirror through a bounded GPU pool on demand, instead of
loading the full constant set into VRAM at program load. This lets a
process run a model whose total weight bytes exceed available GPU
memory, or share GPU across methods whose union exceeds memory.

Public surface (CudaPartitioner kwargs):

  weight_offload: bool = False
      Opt the method into offload. The runtime skips AOTI's eager
      ``update_constants_from_blob``, installs 1-byte GPU dummies in
      place of constants via
      ``update_user_managed_constant_buffer_pairs``, and routes every
      constant read through a probe op into Session::serve.

  weight_offload_pin_fqns: Optional[List[str]] = None
      FQNs to keep resident on GPU for the Session lifetime (out-of-
      pool ``cudaMalloc`` + one-time H2D + sync at create). Requires
      weight_offload=True. Currently device-0-only.

Runtime spec (LoadBackendOptionsMap, or executor_runner's
``--cuda_runtime_spec=weight_offload_budget_mb=N``):
      Total GPU bytes the offload pool + pinned constants may use.
      Defaults to floor_bytes + pinned_bytes_total when unset.

Architecture:

  Export-time pass (backends/cuda/passes/weight_offload_pass.py):
    Rewrites every parameter / buffer consumer to read through
    ``executorch_weight_offload::probe(w, probe_id)`` and serializes
    a v2 payload (schedule + floor + pin_fqns + per-FQN dtype /
    sizes / strides / storage_offset / nbytes / device) into
    NamedDataStore. ``probe_id`` is dense 0..N-1 in graph order so
    the runtime can index ``schedule[probe_id]`` directly.

  Payload parser (backends/cuda/runtime/weight_offload/payload.h):
    Single trust boundary. Validates schema version, framing bounds,
    per-FQN invariants (dtype allow-list, contiguous + offset-zero
    layout, nbytes == elementSize(dtype) * product(sizes), device
    == cuda:0), schedule <-> metadata set equality, and pin_fqns
    dedup + subset-of-schedule. Downstream code trusts the parsed
    Payload struct.

  AOTI c-shim (backends/cuda/runtime/weight_offload/probe_op.cpp):
    Looks the dummy data_ptr up in ProbeRegistry and forwards to
    Session::serve. Lookup miss with the registry empty falls back
    to identity-passthrough (preserves non-offload paths); lookup
    miss with the registry non-empty hard-fails (otherwise the
    runtime would silently read AOTI's eager constant).

  Session (backends/cuda/runtime/weight_offload/session.{h,cpp}):
    Owns the offload lifecycle. Builds the pinned host mirror by
    indexing into the ``_weights_blob`` NamedData entry, allocates a
    bounded ``cudaMemPool``, implements LRU eviction with
    event-ordered ``cudaFreeAsync`` on the compute stream, depth-1
    opportunistic prefetch on the copy stream, optional pinned-
    resident constants. ``serve()``'s miss path and
    ``opportunistic_prefetch()`` share one ``make_room_and_alloc``
    helper.

  ProbeRegistry (backends/cuda/runtime/weight_offload/probe_registry.{h,cpp}):
    Process-global ``dummy_ptr -> (callback, context)`` table in the
    shim library. ``lookup`` + ``has_any_context`` take a shared
    lock so concurrent probe dispatch across Sessions does not
    serialise.

  CudaBackend::init (backends/cuda/runtime/cuda_backend.cpp):
    Thin shim for offload methods: parse payload, fail-fast on
    cuda_graph / shared_stream / non-device-0, ``cudaSetDevice(0)``,
    walk the AOTI catalog, coverage check (catalog <-> schedule),
    AOTI <-> payload data_size cross-check (the one genuinely
    cross-source check; AOTI's compiled .so vs the .pte payload),
    fetch ``_weights_blob``, call Session::create.

Tests:

  backends/cuda/tests/test_weight_offload_probe_dispatch.py
  backends/cuda/tests/test_weight_offload_pool.py
  backends/cuda/tests/test_weight_offload_session.py
  backends/cuda/tests/test_weight_offload_transport.py
  backends/cuda/tests/test_cuda_partitioner.py
  backends/cuda/passes/tests/test_weight_offload_pass.py

  61 passed, 1 skipped on the offload surface.

Not yet wired (deferred): multi-device offload (the payload + runtime
hard-code device 0 today; the partitioner rejects target_device !=
cuda:0 when weight_offload=True). Lifting this requires teaching AOTI's
``create_with_device(..., "cuda", nullptr)`` to take a per-method
device index and propagating it through the payload, runtime, and
partitioner kwargs.

Author: mergennachin

backends/aoti/aoti_backend.py

@@ -9,7 +9,7 @@
 import typing
 from abc import ABC, abstractmethod
 from enum import Enum
-from typing import Any, Dict, List, Optional, Set
+from typing import Any, Dict, List, Optional, Set, Tuple
 
 import torch
 from executorch.backends.aoti.passes.replace_view_copy_with_view import (
@@ -129,6 +129,29 @@ def release_moved_tensors(
         """
         return
 
+    @classmethod
+    def pre_aoti_transform_and_collect_named_data(
+        cls,
+        device_edge_program: ExportedProgram,
+        compile_specs: List[CompileSpec],
+    ) -> List[Tuple[str, bytes, int, Optional[str]]]:
+        """Backend hook for graph mutation + extra NamedDataStore entries.
+
+        Called between ``run_decompositions`` and ``aot_compile``, so
+        overrides can:
+          1. Mutate ``device_edge_program`` in place (e.g. insert
+             custom ops that the backend's AOTI c-shim registry
+             handles), and/or
+          2. Return ``[(key, blob, alignment, external_tag), ...]``
+             entries to be added to the same ``NamedDataStore`` that
+             carries ``_so_blob`` / ``_weights_blob``.
+
+        Default: no-op, returns ``[]``. CudaBackend overrides this to
+        wire weight offloading when its private compile-spec opt-in
+        is set.
+        """
+        return []
+
     @classmethod
     @contextlib.contextmanager
     def collect_unsupported_fallback_kernels(cls, missing_fallback_kernels: Set[str]):
@@ -182,7 +205,7 @@ def generate_fallback_kernel_with_runtime_lookup_aot_and_collect_unsupported_ker
             )
 
     @classmethod
-    def preprocess(
+    def preprocess(  # noqa: C901
         cls,
         edge_program: ExportedProgram,
         compile_specs: List[CompileSpec],
@@ -217,6 +240,15 @@ def preprocess(
                 decomposition_table
             )
 
+        # Backend extension point: mutate ``device_edge_program`` in
+        # place (e.g. insert custom ops the AOTI shim path will pick
+        # up) and return extra NamedDataStore entries to attach
+        # alongside ``_so_blob`` / ``_weights_blob`` after compile.
+        # Default: no-op.
+        extra_named_data = cls.pre_aoti_transform_and_collect_named_data(
+            device_edge_program, compile_specs
+        )
+
         edge_program_module = device_edge_program.module()
 
         # Grab all input placeholders from the graph
@@ -289,6 +321,11 @@ def preprocess(
             method_name + "_weights_blob", blob_data, 1, external_tag
         )
 
+        for extra_key, extra_blob, extra_align, extra_tag in extra_named_data:
+            named_data_store.add_named_data(
+                extra_key, extra_blob, extra_align, extra_tag
+            )
+
         # Clean up the generated files
         os.remove(so_path)
         os.remove(blob_path)

backends/aoti/aoti_delegate_handle.h

@@ -122,6 +122,23 @@ using AOTInductorModelContainerUpdateUserManagedConstantBufferPairsFunc =
         bool use_inactive,
         bool validate_full_update);
 
+// Retrieves a constant's data size (storage bytes) by index. Required by
+// the weight-offload runtime to replicate AOTI's source-blob layout
+// without calling update_constants_from_blob.
+using AOTInductorModelContainerGetConstantDataSizeFunc = AOTIRuntimeError (*)(
+    AOTInductorModelContainerHandle container_handle,
+    size_t idx,
+    size_t* data_size);
+
+// Reports whether a constant is computed by AOTI's runtime const-folding.
+// Required by the weight-offload runtime to hard-fail at init if any
+// folded constant exists (run_const_fold reads other constants and would
+// see dummies in our pre-install model).
+using AOTInductorModelContainerGetConstantFromFoldedFunc = AOTIRuntimeError (*)(
+    AOTInductorModelContainerHandle container_handle,
+    size_t idx,
+    bool* from_folded);
+
 } // extern "C"
 
 // AOTI Delegate Handle structure
@@ -146,6 +163,12 @@ struct AOTIDelegateHandle {
   AOTInductorModelContainerExtractConstantsMapFunc extract_constants_map;
   AOTInductorModelContainerUpdateUserManagedConstantBufferPairsFunc
       update_user_managed_constant_buffer_pairs;
+
+  // Weight-offload-only function pointers. Optional for handles where
+  // offload is not enabled; cuda_backend.cpp::init dlsym's them only when
+  // the weight_offload compile spec is present, and hard-fails if missing.
+  AOTInductorModelContainerGetConstantDataSizeFunc get_constant_data_size;
+  AOTInductorModelContainerGetConstantFromFoldedFunc get_constant_from_folded;
 };
 
 } // namespace aoti

backends/cuda/CMakeLists.txt

@@ -103,8 +103,10 @@ install(
 )
 
 # CUDA-specific AOTI shim symbols (dynamically linked)
-set(_aoti_cuda_shim_sources runtime/shims/memory.cpp
-                            runtime/shims/cuda_guard.cpp
+set(_aoti_cuda_shim_sources
+    runtime/shims/memory.cpp runtime/shims/cuda_guard.cpp
+    runtime/weight_offload/probe_op.cpp
+    runtime/weight_offload/probe_registry.cpp
 )
 
 # Only build CUDA shims when CUDA language/toolchain is available.
@@ -179,7 +181,9 @@ install(
 )
 
 # CUDA backend implementation
-set(_aoti_cuda_backend_sources runtime/cuda_backend.cpp)
+set(_aoti_cuda_backend_sources runtime/cuda_backend.cpp
+                               runtime/weight_offload/session.cpp
+)
 
 # CUDA backend implementation
 add_library(aoti_cuda_backend STATIC ${_aoti_cuda_backend_sources})

backends/cuda/TARGETS

@@ -7,6 +7,7 @@ runtime.python_library(
     srcs = [
         "passes/__init__.py",
         "passes/move_cond_predicate_to_cpu.py",
+        "passes/weight_offload_pass.py",
     ],
     visibility = [
         "//executorch/backends/cuda/...",

backends/cuda/cuda_backend.py

@@ -228,6 +228,12 @@ def get_supported_fallback_kernels(cls) -> Dict[str, Any]:
             "aoti_torch_cuda_randint_low_out": None,
             "executorch_cuda::int4_plain_mm": None,
             "aoti_torch_cuda_int4_plain_mm": None,
+            # Weight-offload probe op (EXPERIMENTAL). Whitelisted so
+            # AOTI's missing-fallback-kernel check doesn't fail when a
+            # graph contains probes; the actual symbol resolution
+            # happens at .so load time against ``libaoti_cuda_shims``.
+            "executorch_weight_offload::probe": None,
+            "aoti_torch_cuda_probe": None,
         }
 
     @classmethod
@@ -262,6 +268,60 @@ def get_custom_passes(cls, compile_specs: List[CompileSpec]) -> List[typing.Any]
             passes.append(ReplaceEdgeOpWithTritonOpPass())
         return passes
 
+    @classmethod
+    def pre_aoti_transform_and_collect_named_data(
+        cls,
+        device_edge_program,
+        compile_specs: List[CompileSpec],
+    ):
+        """Wire weight offloading when the private opt-in compile spec
+        is set. Otherwise no-op.
+
+        When opted in, the pass writes a NamedData payload describing
+        the per-method probe schedule, floor budget, and constants
+        catalog. ``CudaBackend::init`` reads that payload, pre-installs
+        device dummies, skips AOTI's eager constant load, and builds a
+        ``Session`` that serves weights from a bounded GPU pool backed
+        by a host mirror (with optional pinned-resident weights).
+        """
+        # Cheap opt-in probe first — avoids importing the offload
+        # pass machinery for the (overwhelmingly common) non-offload
+        # export. The key string is duplicated against
+        # ``COMPILE_SPEC_KEY_ENABLE`` on purpose so this branch has
+        # no import dependency at all.
+        opt_in = any(
+            spec.key == "_weight_offload_internal_enable" and len(spec.value) > 0
+            for spec in compile_specs
+        )
+        if not opt_in:
+            return []
+
+        # Opt-in is set: the offload pass MUST be importable. A failing
+        # import here means a broken build, not a "skip and hope for
+        # the best" — the runtime would otherwise hard-fail at init
+        # with a "payload missing" error miles away from the actual
+        # cause.
+        from executorch.backends.cuda.passes.weight_offload_pass import (
+            _apply_weight_offload,
+            _serialize_payload,
+            named_data_key_for_method,
+            pin_fqns_from_specs,
+        )
+
+        method_name = cls.method_name_from_compile_specs(compile_specs)
+        pin_fqns = pin_fqns_from_specs(compile_specs)
+        payload = _apply_weight_offload(
+            device_edge_program,
+            method_name=method_name,
+            pin_fqns=pin_fqns,
+        )
+        blob = _serialize_payload(payload)
+        # Merge into the .pte rather than the external .ptd: the
+        # payload is metadata (a few KB at most) and the runtime
+        # parses it in ``init`` before any heavy data loading, so
+        # keeping it inline avoids dragging the .ptd open early.
+        return [(named_data_key_for_method(method_name), blob, 1, None)]
+
     @classmethod
     def get_aoti_compile_options(
         cls, compile_specs: List[CompileSpec]
@@ -314,6 +374,31 @@ def get_aoti_compile_options(
             # int4_dispatch.py not imported — op not registered, skip C shim mapping
             pass
 
+        # Weight-offload probe op (EXPERIMENTAL). The op is defined in
+        # ``backends/cuda/passes/weight_offload_pass.py``; import it
+        # lazily so a build that excludes the offload pass still
+        # compiles. ``probe`` is registered unconditionally because the
+        # AOTI wrapper only emits the call when an FX-graph node
+        # references it, and the runtime symbol always exists in
+        # ``aoti_cuda_shims``.
+        try:
+            from executorch.backends.cuda.passes import (  # noqa: F401
+                weight_offload_pass,
+            )
+
+            shims = options.setdefault("aot_inductor.custom_ops_to_c_shims", {})
+            shims.setdefault(
+                torch.ops.executorch_weight_offload.probe.default,
+                [
+                    "AOTITorchError aoti_torch_cuda_probe("
+                    "AtenTensorHandle, int64_t, AtenTensorHandle*)"
+                ],
+            )
+        except AttributeError:
+            # weight_offload_pass importable but op not registered (e.g.
+            # torch.library shim missing in the running interpreter); skip.
+            pass
+
         # Parse compile_specs to check for platform
 
         platform = "linux"