[Core] Add CUDA-graph node toggling and device-side toggle support

docs/node-toggle.md

@@ -0,0 +1,175 @@
+# Runtime CUDA-graph node toggling
+
+Node toggling enables or disables individual HookPoint producers inside
+already-captured CUDA graphs, between replays, with `cudaGraphNodeSetEnabled`
+on the instantiated exec. Nothing is recaptured or re-instantiated and the
+model is unchanged. A disabled producer costs an empty-node traversal, so a
+step pays only for the hooks it enables, and the enabled set can change from
+one step to the next.
+
+The feature is opt-in and applies to the legacy ring transport. With it off
+the only code that runs is the per-step spec seam described under
+[Behaviour when off](#behaviour-when-off).
+
+## The invariant
+
+Every step, three lanes must agree on which hooks fire:
+
+1. capacity reservation (`BackendAdapter.plan_step` → `prepare_step`),
+2. metadata publication (`RingTransport.pre_push_all_metas`),
+3. the device node state (the engine's enabled set).
+
+If they disagree, the ring desynchronizes: reserved bytes differ from written
+bytes, or metadata is paired with the wrong payload. All three read
+`RingTransport.specs_for_step()`:
+
+| Step | `specs_for_step()` |
+|---|---|
+| toggle inactive | every active spec |
+| gated step (a toggled graph replays) | active ∩ enabled ∩ captured |
+| ungated step (eager, prefill, partial-model graph) | every active spec |
+
+The integration marks each step with `transport.set_gated_step(...)`. A hook that is
+enabled but has no captured node is not effective (it would push metadata
+with no payload), and the eager safety net in `HookPoint.forward` skips
+disabled hooks on a gated step that fell back to eager dispatch.
+
+## Using it
+
+```python
+from dmi import node_toggle
+
+engine = transport._ring_engine          # legacy RingEngine
+
+# 1. Open the capture window before the framework captures its graphs.
+#    exclude_capture() is asked at the start of every capture; return True
+#    for captures that must not be toggled (e.g. partial-model segment
+#    graphs, which would make the registry non-uniform).
+node_toggle.begin_capture_window(engine, exclude_capture=is_partial_capture)
+
+framework.capture_graphs()               # creates torch.cuda.CUDAGraph()s
+
+# 2. Bind every graph you want toggled, then close the window.
+for graph in framework.full_graphs():
+    node_toggle.bind_graph(engine, graph)
+node_toggle.end_capture_window(engine)
+
+# 3. Apply a static set (or None to leave the gate inactive) and arm the
+#    replay guard.
+node_toggle.activate(transport, [(HOOK_TYPE_RESID_PRE, 24)], lazy=True)
+
+# Later, per step:
+transport.set_active_hooks_lazy(new_set)
+transport.set_gated_step(step_replays_a_toggled_graph)
+```
+
+`parse_enabled_hooks()` reads the setting format `"hook_type:layer,..."`:
+`""`/unset leaves the gate inactive, `"none"` is an explicit empty set, and a
+token without a layer (`"14"`) names a global hook (layer `-1`).
+
+### Eager and lazy application
+
+* `set_active_hooks` (eager) flips every bound exec immediately. It does not
+  wait for replays in flight, so call it only at a quiescent point, such as
+  once after warmup.
+* `set_active_hooks_lazy` updates the enabled set and the host gate now; the
+  guarded replay calls `ensure_graph_current()` for the graph about to replay,
+  so a reconfiguration costs O(changed nodes × graphs actually replayed).
+
+`cudaGraphNodeSetEnabled` affects only future launches of an exec; launches
+already enqueued keep the node state they were enqueued with. The registry
+therefore does not wait for a graph's previous replay before mutating it,
+which keeps one-step-ahead schedulers asynchronous.
+`engine.set_toggle_pre_apply_sync(True)` restores the wait (an event recorded
+after every replay, synchronized before the next mutation of that exec).
+
+### The CUDA-graph patch and replay guard
+
+`begin_capture_window` installs a `torch.cuda.CUDAGraph` subclass. It forces
+`keep_graph=True` (otherwise the captured template is freed at instantiation
+and the recorded node handles dangle; keeping it costs host memory only),
+suspends node recording for excluded captures and tags them, and guards
+replays once `activate` armed the guard:
+
+* an excluded graph may replay only on an ungated step;
+* a graph that is not registered and bound (for example captured after the
+  window closed) is refused;
+* a bound graph re-instantiated since `bind_graph` is bound again to its new
+  exec (the registry then resets that graph's per-exec node state);
+* in lazy mode the graph is brought up to date before it replays.
+
+Each refusal raises: the step's metadata is already published, so replaying
+would desynchronize the ring. `bind_graph` refuses graphs created from the
+unpatched class.
+
+The guard covers only graphs that belong to the active transport's engine:
+every capture is stamped with the monitoring session it was captured in and
+with whether it captured DMI producers (the process-wide
+`producer_launch_count()` across the capture, or the nodes recorded for that
+graph). A graph that captured no producer and is not bound -- another CUDA
+graph user in the same process -- replays untouched, as does a graph captured
+for a different engine. A graph that captured producers, or whose producers
+cannot be ruled out, keeps the fail-closed checks above.
+
+The patch is process-global, so monitoring sessions own it.
+`begin_capture_window` starts the engine's session; it ends when the window
+closes with no graph bound, on `NodeToggleController.clear()`, or on
+`dmi.transport.ring.deactivate()` (binding a graph after an empty window
+renews it). When the last session ends, the original `torch.cuda.CUDAGraph`
+is restored and the replay guard disarmed: graphs created afterwards get the
+constructor arguments they ask for, and graphs created while patched replay
+as plain graphs.
+
+### Activation guards
+
+Activation fails closed when no node is registered or no graph is bound, when
+graphs registered different hook sets, when the registered graphs are not
+exactly the bound graphs, or when a capture recorded a producer the toggle
+cannot manage (a non-kernel tail dependency, unreadable capture state, or the
+chunked producer). The verdict is cached on the registry version, which every
+registry mutation bumps.
+
+## Device-side toggling
+
+Device-side toggling applies to producers paired with a per-hook device
+control row; the static and prefix producers are always host-toggled with
+`cudaGraphNodeSetEnabled`. With `begin_capture_window(..., device_toggle=True)`
+(CUDA 12.4+ toolkit and driver), such producers are captured as
+device-updatable kernel nodes, each paired with its control row; a planner
+kernel captured
+upstream of them calls `cudaGraphKernelNodeSetEnabled(node, row[0] != -1)` at
+every replay. A `-1` row already makes a producer return without a
+reservation, a task or a payload, so a disabled node is the same observation
+as an early exit and the host never toggles these nodes. An exec holding
+device-updatable nodes must be uploaded before its first launch. Because the
+patched class keeps the template (`keep_graph=True` defers instantiation),
+a capture in device mode that may hold producers -- excluded captures
+included -- is instantiated and uploaded at capture end, and a later
+re-instantiation is uploaded before it replays. Each exec is uploaded once;
+excluded captures are never enrolled in the registry.
+The mode is chosen per capture window (`clear_toggle_registry()` turns it
+off) and stays off when the device, driver or build lacks device graph
+updates.
+
+`capture_window(engine, ...)` is the context-manager form of
+`begin_capture_window` / `end_capture_window`; the window's reference to the
+engine is also dropped by `NodeToggleController.clear()` and by
+`dmi.transport.ring.deactivate()`.
+
+## Strict eager fallback
+
+`BackendAdapter(..., allow_eager_fallback=False)` refuses a step that would
+need the eager safety net: an eager-only plan raises before anything is
+reserved, and a step larger than the ring raises after `prepare_step` (which
+reserves nothing for it) but before metadata is published or a producer is
+armed. `capture_dispatch_state()` reports cumulative preparation counts.
+
+## Behaviour when off
+
+| Runs when off | Behaviour |
+|---|---|
+| `specs_for_step()` in `pre_push_all_metas` | returns `_active_specs` (the same object) |
+| `plan_step` | iterates `active_specs` unless the gate is active |
+| node recording in the producer ops | returns on a process-global flag |
+| `HookPoint` eager gate | inside the `force_eager` branch, short-circuits on `toggle_gate_active` |
+| `legacy_ring_room` | additionally saturates when the drain tail is ahead of the reservation head |

native/Makefile

@@ -27,7 +27,8 @@ HEADERS := csrc/dmx_host_utils.h \
            csrc/ring/record_sink.h \
            csrc/ring/p2p_thread.h csrc/ring/pinned_staging.h \
            csrc/ring/ring_engine.h csrc/ring/ring_engine_py.h \
-           csrc/ring/ring_torch_op.h csrc/ring/tensor_meta.h
+           csrc/ring/ring_torch_op.h csrc/ring/tensor_meta.h \
+           csrc/ring/toggle_device.h csrc/ring/toggle_registry.h
 BUILD_DIR := build
 HOST_BUILD_DIR := $(BUILD_DIR)/host
 OBJS := $(patsubst csrc/%.cpp,$(BUILD_DIR)/%.o,$(SRCS))
@@ -185,7 +186,8 @@ NVCCFLAGS := \
 
 RING_SRCS := csrc/ring/producer.cu \
              csrc/ring/ring_engine.cu \
-             csrc/ring/ring_engine_py.cu
+             csrc/ring/ring_engine_py.cu \
+             csrc/ring/toggle_registry.cu
 RING_OBJS := $(patsubst csrc/%.cu,$(BUILD_DIR)/%.o,$(RING_SRCS))
 
 OBJS += $(RING_OBJS)

native/csrc/bindings.cpp

@@ -692,7 +692,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
   py::class_<ring_py::RingEnginePy, std::shared_ptr<ring_py::RingEnginePy>>(m, "RingEngine")
       .def(py::init([](ring_py::RingConfig cfg, py::object host_engine_obj) {
              ring_py::SubmitFn submit_fn;
-             if (!host_engine_obj.is_none()) {
+             if (!host_engine_obj.is_none() &&
+                 py::isinstance<dmx_host::DMXHostEngine>(host_engine_obj)) {
                  auto host = host_engine_obj.cast<std::shared_ptr<dmx_host::DMXHostEngine>>();
                  submit_fn = [host](const std::string& model_id, int32_t shard_rank,
                                     const std::string& req_id, const std::string& act_name,
@@ -710,6 +711,38 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
                      row.emplace_back(std::move(slice));
                      host->submit_direct(std::move(row), nbytes);
                  };
+             } else if (!host_engine_obj.is_none()) {
+                 // An in-process consumer: any Python callable taking
+                 // (model_id, shard_rank, req_id, act_name, layer_no,
+                 //  start_token, end_token, tensor).  The P2P thread calls
+                 // it without the GIL, so take it for the call.  An
+                 // exception propagates to the P2P thread's submit-failure
+                 // handling (logged once, the slice is dropped).
+                 if (!PyCallable_Check(host_engine_obj.ptr())) {
+                     throw py::type_error(
+                         "RingEngine host_engine must be a DMXHostEngine, a "
+                         "callable consumer, or None");
+                 }
+                 auto fn = std::shared_ptr<py::object>(
+                     new py::object(host_engine_obj),
+                     [](py::object* held) {
+                         // The last reference can drop on the P2P thread.
+                         py::gil_scoped_acquire gil;
+                         delete held;
+                     });
+                 submit_fn = [fn](const std::string& model_id, int32_t shard_rank,
+                                  const std::string& req_id, const std::string& act_name,
+                                  int32_t layer_no, int32_t start_token, int32_t end_token,
+                                  at::Tensor slice) {
+                     py::gil_scoped_acquire gil;
+                     try {
+                         (*fn)(model_id, shard_rank, req_id, act_name, layer_no,
+                               start_token, end_token, std::move(slice));
+                     } catch (py::error_already_set& error) {
+                         std::string what = error.what();
+                         throw std::runtime_error(what);
+                     }
+                 };
              }
              return MakeRingEngine(
                  std::move(cfg), std::move(submit_fn));
@@ -918,7 +951,64 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
            py::call_guard<py::gil_scoped_release>())
       .def("notify_drain",
            &ring_py::RingEnginePy::notify_drain,
-           py::call_guard<py::gil_scoped_release>());
+           py::call_guard<py::gil_scoped_release>())
+      // ---- Runtime CUDA-graph node toggling -------------------------------
+      // Graph, node and exec handles are the integers returned by
+      // torch.cuda.CUDAGraph.raw_cuda_graph() / raw_cuda_graph_exec().
+      .def("enable_toggle_capture",
+           &ring_py::RingEnginePy::enable_toggle_capture, py::arg("enabled"))
+      .def("toggle_capture_enabled",
+           &ring_py::RingEnginePy::toggle_capture_enabled)
+      .def("capture_anomaly_count",
+           &ring_py::RingEnginePy::capture_anomaly_count)
+      .def("bind_graph_exec", &ring_py::RingEnginePy::bind_graph_exec,
+           py::arg("graph"), py::arg("exec"))
+      .def("toggle_registry_version",
+           &ring_py::RingEnginePy::toggle_registry_version)
+      .def("set_enabled_hooks", &ring_py::RingEnginePy::set_enabled_hooks,
+           py::arg("enabled"))
+      .def("apply_toggle", &ring_py::RingEnginePy::apply_toggle)
+      .def("is_hook_enabled", &ring_py::RingEnginePy::is_hook_enabled,
+           py::arg("hook_type"), py::arg("layer_no"))
+      .def("effective_enabled_mask",
+           &ring_py::RingEnginePy::effective_enabled_mask, py::arg("query"))
+      .def("ensure_graph_current",
+           &ring_py::RingEnginePy::ensure_graph_current, py::arg("graph"))
+      .def("record_replay_event",
+           &ring_py::RingEnginePy::record_replay_event, py::arg("graph"))
+      .def("set_toggle_pre_apply_sync",
+           &ring_py::RingEnginePy::set_toggle_pre_apply_sync,
+           py::arg("enabled"))
+      .def("toggle_node_count", &ring_py::RingEnginePy::toggle_node_count)
+      .def("toggle_node_count_for_graph",
+           &ring_py::RingEnginePy::toggle_node_count_for_graph,
+           py::arg("graph"))
+      .def("producer_launch_count",
+           &ring_py::RingEnginePy::producer_launch_count)
+      .def("bound_graph_count", &ring_py::RingEnginePy::bound_graph_count)
+      .def("last_apply_count", &ring_py::RingEnginePy::last_apply_count)
+      .def("toggle_registry_uniform",
+           &ring_py::RingEnginePy::toggle_registry_uniform)
+      .def("is_graph_ready", &ring_py::RingEnginePy::is_graph_ready,
+           py::arg("graph"))
+      .def("toggle_registry_complete",
+           &ring_py::RingEnginePy::toggle_registry_complete)
+      .def("clear_toggle_registry",
+           &ring_py::RingEnginePy::clear_toggle_registry)
+      .def("enable_device_toggle",
+           &ring_py::RingEnginePy::enable_device_toggle, py::arg("enabled"))
+      .def("toggle_device_mode", &ring_py::RingEnginePy::toggle_device_mode)
+      .def("toggle_device_node_count",
+           &ring_py::RingEnginePy::toggle_device_node_count)
+      .def("toggle_device_errors",
+           &ring_py::RingEnginePy::toggle_device_errors)
+      .def("upload_graph_exec", &ring_py::RingEnginePy::upload_graph_exec,
+           py::arg("graph"), py::arg("exec"),
+           py::call_guard<py::gil_scoped_release>())
+      .def("_test_force_apply_error",
+           &ring_py::RingEnginePy::_test_force_apply_error, py::arg("code"))
+      .def("_test_applied_current",
+           &ring_py::RingEnginePy::_test_applied_current, py::arg("graph"));
 
   // Register the active ring engine pointer so C++ ring_producer_impl can
   // call it during CUDA graph capture.  The raw pointer is valid as long as

native/csrc/ring/drain_thread.cpp

@@ -258,6 +258,15 @@ uint64_t DrainThread::cpu_task_tail_committed() const {
 // ---------------------------------------------------------------------------
 void DrainThread::reserve(uint64_t payload_bytes, uint32_t num_tasks) {
     std::lock_guard<std::mutex> lk(mgmt_mu_);
+    // Producers that ran without a reservation (warmup runs of captured
+    // producers) can leave the drain's tails ahead of the reservation heads.
+    // Reserve from the tails: otherwise the heads stay behind them for as
+    // many reserved tasks as the skew, the room check reads a full ring as
+    // empty, and producers publish over unread slots.
+    if (cpu_payload_head_ < cpu_payload_tail_committed_)
+        cpu_payload_head_ = cpu_payload_tail_committed_;
+    if (cpu_task_head_ < cpu_task_tail_)
+        cpu_task_head_ = cpu_task_tail_;
     cpu_payload_head_ += payload_bytes;
     cpu_task_head_    += num_tasks;
 }