Offload cleanup (round 6): trim error UX, module docstring, design header

Net -214 lines, three small items from the latest review pass.

Collapse the 75-line below-floor error message to a one-paragraph
form (Session::create's budget check). The two-branch "MB vs bytes"
formatter with inline rounding was a UX nicety for a single error
path; the new message keeps the substrings tests assert on
("pinned constants", "streaming pool floor", "required total",
"set via <spec>=<value>", "Set weight_offload_budget_mb >= <mb>")
and drops everything else.

Trim weight_offload_pass.py module docstring 68 -> 13 lines. The
removed prose was a list of "RESOLVED" design notes (schedule order,
probe-op preservation, payload transport) that read like a status
report. The code already embodies the decisions; the docstring now
states what the pass does and where the runtime lives.

Delete weight_offload.h. It was a 109-line README-as-header
duplicating session.h's overview after several rounds of trim, with
no code and no other includer. Folded a tightened version of the
overview into session.h's banner.

Also collapsed probe_op.h's 18-line "DISPATCH WIRED, RUNTIME BODY
TBD" banner -- the runtime body has shipped; the banner described a
state from several commits ago.

61 of 61 offload tests pass; lint clean.

Authored with Claude.

Author: mergennachin

backends/cuda/passes/weight_offload_pass.py

@@ -4,68 +4,19 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-"""Weight offloading pass for the CUDA backend.
-
-EXPERIMENTAL. The runtime now installs pre-loaded GPU dummies in
-place of AOTI's constants, skips ``update_constants_from_blob``
-entirely, sources host bytes directly from the ``_weights_blob``
-NamedData entry, and serves probe calls from a bounded pool with
-LRU + event-ordered free-on-compute-stream. Steady-state peak GPU
-= offload pool only.
-
-What's wired today:
-  * The probe op + AOTI c-shim dispatch path. The c-shim is the
-    Session::serve callback for opt-in methods; identity-passthrough
-    for non-opt-in.
-  * The ``_apply_weight_offload`` graph rewrite + v2 payload schema
-    (per-FQN dtype/sizes/strides/storage_offset/nbytes/device, plus
-    schedule + floor + pin_fqns).
-  * The public partitioner kwargs ``weight_offload=True`` and
-    ``weight_offload_pin_fqns=[...]`` on ``CudaPartitioner``
-    translate to a private compile-spec channel
-    (``_weight_offload_internal_enable``,
-    ``_weight_offload_internal_pin_fqns``) that routes the payload
-    from ``CudaBackend.preprocess`` ->
-    ``AotiBackend.preprocess``'s
-    ``pre_aoti_transform_and_collect_named_data`` hook ->
-    ``NamedDataStore`` -> ``CudaBackend::init``, where the runtime
-    parses + cross-checks against AOTI + installs dummies + builds
-    the pinned host mirror + serves probes. The internal compile
-    spec keys stay underscore-prefixed and remain accessible from
-    tests for exact-byte budget control.
-
-Pinning (``pin_fqns``): the runtime allocates each pinned weight
-once via out-of-pool ``cudaMalloc`` + a synchronous H2D, then
-serves it through a resident fast path. The pass deduplicates
-pin_fqns before serialization so the runtime never sees duplicates.
-Multi-device offload is hard-failed at init (device 0 only today).
-
-Schedule / cursor order -- RESOLVED:
-    The probe op carries an explicit ``probe_id: int`` argument assigned
-    by the pass (contiguous 0..N-1 in graph order). The runtime indexes
-    the schedule directly by ``probe_id``; no cursor is needed. This
-    closes the entire class of "graph order drifted from wrapper order"
-    silent failures the prior design left open.
-
-Probe op preservation -- RESOLVED:
-    The c-shim ``aoti_torch_cuda_probe`` is wired into the CUDA backend's
-    ``aot_inductor.custom_ops_to_c_shims`` registry. AOTI emits one
-    direct call per probe node in the FX graph; distinct ``probe_id``
-    constants make otherwise-identical probe calls syntactically distinct
-    from inductor's CSE pass. The single-consumer and multi-consumer
-    dispatch contracts are asserted by
-    ``test_weight_offload_probe_dispatch.py``.
-
-Payload transport -- RESOLVED:
-    The pass runs from ``AotiBackend.preprocess`` (via the
-    ``pre_aoti_transform_and_collect_named_data`` hook in
-    ``CudaBackend``) and the payload ships through ``NamedDataStore``
-    alongside ``_so_blob`` / ``_weights_blob``, so all AOTI-side
-    artifacts live in one mechanism. The wire format is a small
-    custom little-endian binary (see ``_serialize_payload``) — the
-    C++ parser in ``runtime/weight_offload/payload.h`` mirrors it.
-
-The runtime half lives in ``backends/cuda/runtime/weight_offload/``.
+"""EXPERIMENTAL: CUDA weight-offloading pass.
+
+Rewrites every parameter / buffer consumer to read through a
+``probe(w, probe_id)`` op, then serializes a v2 payload (schedule
++ floor + pin_fqns + per-FQN dtype / sizes / strides /
+storage_offset / nbytes / device) into ``NamedDataStore``. The
+runtime half (``backends/cuda/runtime/weight_offload/``) parses
+the payload, installs 1-byte GPU dummies in place of AOTI's
+constants, and serves probes from a pinned host mirror through a
+bounded GPU pool. Single-device only.
+
+Public surface: ``CudaPartitioner(weight_offload=True,
+weight_offload_pin_fqns=[...])``.
 """
 
 import struct

backends/cuda/runtime/weight_offload/probe_op.h

@@ -8,25 +8,10 @@
 
 #pragma once
 
-// ===========================================================================
-// EXPERIMENTAL -- DISPATCH WIRED, RUNTIME BODY TBD
-// ===========================================================================
-// State: the c-shim is registered in
-// ``cuda_backend.py::custom_ops_to_c_shims`` and implemented in
-// ``probe_op.cpp`` as an identity passthrough (returns a fresh
-// ``SlimTensor`` handle sharing the input's storage). No pass inserts
-// probe nodes in production, so this shim has no production caller yet;
-// it is exercised by
-// ``backends/cuda/tests/test_weight_offload_probe_dispatch.py`` with
-// hand-rolled probes to validate that AOTI emits one call per FX
-// probe node (including the multi-consumer case inductor CSE could
-// otherwise elide).
-//
-// Future work replaces the identity body with a Session-managed lookup:
-// ``probe_id`` selects the FQN from the schedule and ``Session::serve``
-// returns a SlimTensor backed by pool-managed GPU bytes. See
-// ``backends/cuda/runtime/weight_offload/weight_offload.h``.
-// ===========================================================================
+// EXPERIMENTAL. The c-shim AOTI dispatches to for
+// ``executorch_weight_offload::probe(w, probe_id)``. Body lives in
+// probe_op.cpp; routes through ProbeRegistry to Session::serve when
+// offload is opted in, identity-passthrough otherwise.
 
 #include <cstdint>
 

backends/cuda/runtime/weight_offload/session.cpp

@@ -367,86 +367,34 @@ ::executorch::runtime::Result<std::unique_ptr<Session>> Session::create(
 
   const uint64_t required_total = payload.floor_bytes + pinned_bytes_total;
 
-  // 3e. Unified below-budget check covers BOTH "pinned alone
-  //     exceeds total" AND "streaming budget < floor": they're
-  //     algebraically the same condition (total < pinned + floor).
-  //     Both get the same BudgetSpec-aware UX message naming the
-  //     spec the user set and the suggested fix.
+  // total < pinned + floor covers both "pinned alone exceeds total"
+  // and "streaming budget < floor". Suggest the public spec name for
+  // the fix (rounded up to whole MB so the value is loadable).
   if (session->total_budget_bytes_ < required_total) {
-    // Pick the spec name + suggested-value-to-set for the hint.
-    // Defaults are safe: caller leaves budget_spec.name == nullptr
-    // when no spec was set, in which case we suggest the PUBLIC
-    // spec since that's what users land on first.
-    const char* spec_name = (budget_spec.name != nullptr)
-        ? budget_spec.name
-        : "weight_offload_budget_mb";
-    // Round required_total UP to whole MB for the public-spec hint
-    // (so the suggested value is actually loadable). Use
-    // division/modulo (NOT ``+ (1MiB - 1)``) so the rounding can't
-    // overflow for required_total values near UINT64_MAX.
     constexpr uint64_t kMiB = 1ull << 20;
-    const uint64_t required_mb_rounded_up =
+    const uint64_t required_mb =
         (required_total / kMiB) + ((required_total % kMiB) != 0 ? 1 : 0);
-
-    if (std::strcmp(spec_name, "weight_offload_budget_mb") == 0) {
-      // Public spec — show MB-rounded values.
-      std::fprintf(
-          stderr,
-          "[ET_WEIGHT_OFFLOAD][ERROR] Weight offloading needs at "
-          "least %llu MB of GPU memory for method '%s':\n"
-          "  pinned constants:     %llu bytes\n"
-          "  streaming pool floor: %llu bytes\n"
-          "  ---\n"
-          "  required total:       %llu bytes (~%llu MB)\n"
-          "but the configured budget is %llu bytes",
-          static_cast<unsigned long long>(required_mb_rounded_up),
-          payload.method_name.c_str(),
-          static_cast<unsigned long long>(session->pinned_bytes_total_),
-          static_cast<unsigned long long>(payload.floor_bytes),
-          static_cast<unsigned long long>(required_total),
-          static_cast<unsigned long long>(required_mb_rounded_up),
-          static_cast<unsigned long long>(session->total_budget_bytes_));
-      if (budget_spec.value > 0 && budget_spec.value_is_mb) {
-        std::fprintf(
-            stderr,
-            " (set via weight_offload_budget_mb=%llu)",
-            static_cast<unsigned long long>(budget_spec.value));
-      }
-      std::fprintf(
-          stderr,
-          ". Set weight_offload_budget_mb >= %llu to load this "
-          "method.\n",
-          static_cast<unsigned long long>(required_mb_rounded_up));
-    } else {
-      // Internal spec (test path) — show exact byte values.
-      std::fprintf(
-          stderr,
-          "[ET_WEIGHT_OFFLOAD][ERROR] Weight offloading needs at "
-          "least %llu bytes of GPU memory for method '%s':\n"
-          "  pinned constants:     %llu bytes\n"
-          "  streaming pool floor: %llu bytes\n"
-          "  ---\n"
-          "  required total:       %llu bytes\n"
-          "but the configured budget is %llu bytes",
-          static_cast<unsigned long long>(required_total),
-          payload.method_name.c_str(),
-          static_cast<unsigned long long>(session->pinned_bytes_total_),
-          static_cast<unsigned long long>(payload.floor_bytes),
-          static_cast<unsigned long long>(required_total),
-          static_cast<unsigned long long>(session->total_budget_bytes_));
-      if (budget_spec.value > 0 && !budget_spec.value_is_mb) {
-        std::fprintf(
-            stderr,
-            " (set via %s=%llu)",
-            spec_name,
-            static_cast<unsigned long long>(budget_spec.value));
-      }
+    std::fprintf(
+        stderr,
+        "[ET_WEIGHT_OFFLOAD][ERROR] method '%s': budget %llu bytes < "
+        "required total %llu (pinned constants %llu + streaming pool "
+        "floor %llu)",
+        payload.method_name.c_str(),
+        static_cast<unsigned long long>(session->total_budget_bytes_),
+        static_cast<unsigned long long>(required_total),
+        static_cast<unsigned long long>(session->pinned_bytes_total_),
+        static_cast<unsigned long long>(payload.floor_bytes));
+    if (budget_spec.value > 0) {
       std::fprintf(
           stderr,
-          ". Set %s >= %llu to load this method.\n",
-          spec_name,
-          static_cast<unsigned long long>(required_total));
+          " [set via %s=%llu]",
+          budget_spec.name,
+          static_cast<unsigned long long>(budget_spec.value));
     }
+    std::fprintf(
+        stderr,
+        ". Set weight_offload_budget_mb >= %llu.\n",
+        static_cast<unsigned long long>(required_mb));
     return Error::InvalidArgument;
   }
 

backends/cuda/runtime/weight_offload/session.h

@@ -8,12 +8,23 @@
 
 #pragma once
 
-// Per-handle CUDA weight-offload Session. See weight_offload.h for
-// the cross-component design overview. Lifecycle: created in
-// CudaBackend::init after the CUDA stream is in place and the local
-// DummyInstallation is built; destroyed first in CudaBackend::destroy
-// so the ProbeRegistry unregister + stream drain + free ordering is
-// explicit. Single-device only.
+// EXPERIMENTAL: per-handle CUDA weight-offload Session.
+//
+// What it owns: a pinned host mirror of the schedule's constants, a
+// bounded cudaMemPool sized by a software byte cap on requested live
+// offload bytes, LRU eviction with event-ordered cudaFreeAsync on the
+// compute stream, depth-1 opportunistic prefetch on the copy stream,
+// optional pinned-resident constants (out-of-pool cudaMalloc), and the
+// DummyInstallation that AOTI's container uses as its "active
+// constants" so probe ops route every constant read here. See
+// session.cpp for the per-step ordering. Single-device only (the
+// payload parser hard-fails device_index != 0).
+//
+// Companion components: payload.{h} (on-wire schema parser, the single
+// trust boundary), probe_op.{h,cpp} (AOTI c-shim), probe_registry.{h,cpp}
+// (process-global dummy_ptr -> Session map). Public entry from the
+// pass side: ``CudaPartitioner(weight_offload=True,
+// weight_offload_pin_fqns=[...])``.
 
 #include <cstdint>
 #include <list>