From ffbedf5ffec8af7b4a48642f72e8566ee4e77277 Mon Sep 17 00:00:00 2001 From: zhaoyboo Date: Tue, 3 Nov 2026 15:36:56 -0500 Subject: [PATCH] [Core] Add CUDA-graph node toggling and device-side toggle support Enable or disable individual HookPoint producers inside captured CUDA graphs between replays with cudaGraphNodeSetEnabled, without recapture or re-instantiation. Capacity reservation, metadata publication and the device node state all read RingTransport.specs_for_step(), so they agree on which hooks fire every step; ungated (eager, prefill, partial-model) steps use the full active set. Native: - ToggleRegistry (toggle_registry.{h,cu}): capture-time node registry, graph/exec bindings, the enabled set (single source of truth), diff-based eager apply, lazy per-graph apply with an optional pre-apply wait, uniform/complete/ready guards, anomaly counting, teardown that re-enables disabled nodes. - Device-side toggling (toggle_device.h): a planner kernel captured upstream of device-updatable producer nodes enables each node from its own control row (-1 = disabled), so the host never toggles those nodes. - The static and prefix producer ops record their kernel node during the capture window; the chunked producer and unreadable capture state are recorded as anomalies so activation fails closed. - Reservations start from the drain's tails when unreserved producers (warmup runs of captured producers) left the tails ahead of the heads; otherwise the room check reads a full ring as empty for as many tasks as the skew. legacy_ring_room also saturates in that state. - RingEngine accepts a Python callable as an in-process consumer. Python: - dmi.node_toggle: NodeToggleController (activation guards cached on the registry version, memoized effective-set recomputation), the enabled-hooks parser, a capture window with an exclude_capture predicate (and a capture_window context manager; clear() and deactivate() release the window's engine), a keep_graph torch.cuda.CUDAGraph subclass with a replay guard that restores node recording if a capture fails to begin, bind_graph and activate. Device-side toggling is chosen per window and each device-updatable exec is uploaded at most once. - RingTransport: effective_specs / specs_for_step and toggle delegates; metadata follows the per-step set. - BackendAdapter: reservation follows the per-step set while the gate is active; allow_eager_fallback=False rejects eager-only and oversized steps before metadata publication, with cumulative dispatch counts. - HookPoint: the eager safety net skips disabled hooks on gated steps. Lifecycle and scope of the CUDA-graph patch: - The torch.cuda.CUDAGraph replacement is owned by monitoring sessions (one per engine, started by the capture window). A session ends when its window closes with nothing bound, on NodeToggleController.clear(), or on transport deactivation; when the last one ends the original class is restored and the replay guard disarmed, so later graphs keep their constructor arguments. - Each capture is stamped with its session and with whether it captured producers (a process-wide producer-launch counter across the capture, or the nodes recorded for that graph). The replay guard covers only bound graphs and producer-carrying graphs of the active transport's engine; an unrelated graph without producers, or a graph captured for another engine, replays untouched. Unbound producer graphs still fail closed. - ToggleRegistry::bind_exec resets a graph's per-exec state when a different exec is bound (applied version, per-node actual state, replay event, upload record), so the next apply sets every host-managed node of the new exec. The replay guard rebinds a bound graph that was re-instantiated. - In device-toggle mode a capture that may hold producers, excluded captures included, is instantiated and uploaded at capture end (keep_graph=True defers instantiation to the first replay), and a re-instantiated exec is uploaded before it replays. Excluded captures are never enrolled. Supersedes #69. --- docs/node-toggle.md | 175 ++++ native/Makefile | 6 +- native/csrc/bindings.cpp | 94 +- native/csrc/ring/drain_thread.cpp | 9 + native/csrc/ring/ring_engine_py.cu | 143 +++ native/csrc/ring/ring_engine_py.h | 86 ++ native/csrc/ring/ring_torch_op.cpp | 77 ++ native/csrc/ring/ring_torch_op.h | 8 + native/csrc/ring/toggle_device.h | 83 ++ native/csrc/ring/toggle_registry.cu | 538 ++++++++++ native/csrc/ring/toggle_registry.h | 175 ++++ src/dmi/adapters/base.py | 69 +- src/dmi/hooks/point.py | 12 + src/dmi/node_toggle.py | 852 +++++++++++++++ src/dmi/transport/ring.py | 107 +- tests/native/ring/Makefile | 16 +- .../native/ring/cpu_cuda_stub/cuda_runtime.h | 71 ++ tests/native/ring/test_ring_engine.cu | 4 + tests/native/ring/test_toggle_registry.cu | 472 +++++++++ .../native/ring/test_toggle_registry_host.cpp | 184 ++++ tests/test_adapter_protocol.py | 70 ++ tests/test_node_toggle.py | 980 ++++++++++++++++++ tests/test_node_toggle_gpu.py | 503 +++++++++ tests/test_toggle_registry_host.py | 46 + 24 files changed, 4767 insertions(+), 13 deletions(-) create mode 100644 docs/node-toggle.md create mode 100644 native/csrc/ring/toggle_device.h create mode 100644 native/csrc/ring/toggle_registry.cu create mode 100644 native/csrc/ring/toggle_registry.h create mode 100644 src/dmi/node_toggle.py create mode 100644 tests/native/ring/cpu_cuda_stub/cuda_runtime.h create mode 100644 tests/native/ring/test_toggle_registry.cu create mode 100644 tests/native/ring/test_toggle_registry_host.cpp create mode 100644 tests/test_node_toggle.py create mode 100644 tests/test_node_toggle_gpu.py create mode 100644 tests/test_toggle_registry_host.py diff --git a/docs/node-toggle.md b/docs/node-toggle.md new file mode 100644 index 000000000..b25e40504 --- /dev/null +++ b/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 | diff --git a/native/Makefile b/native/Makefile index 2353eeea8..1e8c3de03 100644 --- a/native/Makefile +++ b/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) diff --git a/native/csrc/bindings.cpp b/native/csrc/bindings.cpp index 1b0439149..7f5db58ab 100644 --- a/native/csrc/bindings.cpp +++ b/native/csrc/bindings.cpp @@ -692,7 +692,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::class_>(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(host_engine_obj)) { auto host = host_engine_obj.cast>(); 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( + 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()) .def("notify_drain", &ring_py::RingEnginePy::notify_drain, - py::call_guard()); + py::call_guard()) + // ---- 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()) + .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 diff --git a/native/csrc/ring/drain_thread.cpp b/native/csrc/ring/drain_thread.cpp index ca45fa4f0..f40e4418f 100644 --- a/native/csrc/ring/drain_thread.cpp +++ b/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 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; } diff --git a/native/csrc/ring/ring_engine_py.cu b/native/csrc/ring/ring_engine_py.cu index ad27e20be..a60b8fa04 100644 --- a/native/csrc/ring/ring_engine_py.cu +++ b/native/csrc/ring/ring_engine_py.cu @@ -11,6 +11,7 @@ #include "ring/producer.cuh" #include "ring/ring_debug.h" #include "ring/record_stall_budget.h" +#include "ring/toggle_registry.h" #include // at::cuda::getCurrentCUDAStream #include #include @@ -138,7 +139,14 @@ void refuse_on_record_ring(bool record_mode, const char* what) { // on top of it without checking again. Unsaturated, cap - (head - tail) then // wraps to ~2^64, every check admits every reservation, and the producers // publish over unread slots. +// +// Saturated at the other end too: the tail can run ahead of the head when +// producers recorded into a CUDA graph also fire during capture-time warmup +// runs and the drain consumes them with no step reservation behind them. +// The ring is then empty, not full; unsaturated, head - tail wraps to ~2^64 +// and every later step takes a spurious flush. uint64_t legacy_ring_room(uint64_t cap, uint64_t head, uint64_t tail) { + if (head < tail) return cap; const uint64_t outstanding = head - tail; return outstanding >= cap ? 0 : cap - outstanding; } @@ -165,6 +173,10 @@ struct RingEnginePy::Impl { at::Tensor payload_view; bool record_mode{false}; + // Runtime CUDA-graph node toggling: all state and graph operations live + // in the registry; the engine methods delegate. + ToggleRegistry toggle; + // Step-gated drain: ring records reserved since the last allow_flush_step() // (the current step's gap). Written by the reserving thread only. uint64_t gap_records{0}; @@ -970,6 +982,137 @@ at::Tensor RingEnginePy::payload_tensor() const { // as much room as we computed. // --------------------------------------------------------------------------- +// --------------------------------------------------------------------------- +// Runtime CUDA-graph node toggling -- thin delegates to ToggleRegistry. The +// engine adds only the process-global capture flag the producer ops read +// (ring_set_toggle_capture) and the current ATen stream for replay events. +void RingEnginePy::enable_toggle_capture(bool enabled) { + if (enabled) refuse_on_record_ring(impl_->record_mode, "node-toggle capture"); + impl_->toggle.set_capture(enabled); + ring_set_toggle_capture(enabled); +} + +bool RingEnginePy::toggle_capture_enabled() const { + return impl_->toggle.capture_enabled(); +} + +void RingEnginePy::register_capture_node(uint64_t graph, int hook_type, + int layer_no, uint64_t node) { + impl_->toggle.register_node(graph, hook_type, layer_no, node); +} + +void RingEnginePy::note_capture_anomaly() { + impl_->toggle.note_anomaly(); +} + +uint64_t RingEnginePy::capture_anomaly_count() const { + return impl_->toggle.anomaly_count(); +} + +void RingEnginePy::bind_graph_exec(uint64_t graph, uint64_t exec) { + impl_->toggle.bind_exec(graph, exec); +} + +uint64_t RingEnginePy::toggle_registry_version() const { + return impl_->toggle.version(); +} + +void RingEnginePy::set_enabled_hooks( + const std::vector>& enabled) { + impl_->toggle.set_enabled(enabled); +} + +int RingEnginePy::apply_toggle() { + return impl_->toggle.apply_all(); +} + +bool RingEnginePy::is_hook_enabled(int hook_type, int layer_no) const { + return impl_->toggle.is_enabled(hook_type, layer_no); +} + +std::vector RingEnginePy::effective_enabled_mask( + const std::vector>& query) const { + return impl_->toggle.effective_mask(query); +} + +int RingEnginePy::ensure_graph_current(uint64_t graph) { + return impl_->toggle.ensure_current(graph); +} + +int RingEnginePy::record_replay_event(uint64_t graph) { + return impl_->toggle.record_replay_event( + graph, at::cuda::getCurrentCUDAStream().stream()); +} + +void RingEnginePy::set_toggle_pre_apply_sync(bool enabled) { + impl_->toggle.set_pre_apply_sync(enabled); +} + +uint64_t RingEnginePy::toggle_node_count() const { + return impl_->toggle.node_count(); +} + +uint64_t RingEnginePy::toggle_node_count_for_graph(uint64_t graph) const { + return impl_->toggle.node_count_for(graph); +} + +uint64_t RingEnginePy::producer_launch_count() const { + return ring_producer_launch_count(); +} + +uint64_t RingEnginePy::bound_graph_count() const { + return impl_->toggle.bound_count(); +} + +uint64_t RingEnginePy::last_apply_count() const { + return impl_->toggle.last_apply_count(); +} + +bool RingEnginePy::toggle_registry_uniform() const { + return impl_->toggle.uniform(); +} + +bool RingEnginePy::is_graph_ready(uint64_t graph) const { + return impl_->toggle.is_ready(graph); +} + +bool RingEnginePy::toggle_registry_complete() const { + return impl_->toggle.complete(); +} + +void RingEnginePy::clear_toggle_registry() { + impl_->toggle.clear(); + ring_set_toggle_capture(false); +} + +bool RingEnginePy::enable_device_toggle(bool enabled) { + return impl_->toggle.set_device_mode(enabled); +} + +bool RingEnginePy::toggle_device_mode() const { + return impl_->toggle.device_mode(); +} + +uint64_t RingEnginePy::toggle_device_node_count() const { + return impl_->toggle.device_node_count(); +} + +uint64_t RingEnginePy::toggle_device_errors() const { + return impl_->toggle.device_errors(); +} + +int RingEnginePy::upload_graph_exec(uint64_t graph, uint64_t exec) { + return impl_->toggle.upload_device_graph(graph, exec); +} + +void RingEnginePy::_test_force_apply_error(int code) { + impl_->toggle.force_apply_error(code); +} + +bool RingEnginePy::_test_applied_current(uint64_t graph) const { + return impl_->toggle.applied_current(graph); +} + uint64_t RingEnginePy::available_capacity() const { auto& drain = impl_->engine.drain_thread(); const uint64_t head = drain.cpu_payload_head(); diff --git a/native/csrc/ring/ring_engine_py.h b/native/csrc/ring/ring_engine_py.h index a795f4f97..915f54810 100644 --- a/native/csrc/ring/ring_engine_py.h +++ b/native/csrc/ring/ring_engine_py.h @@ -312,6 +312,92 @@ class RingEnginePy { // ensure FIFO ordering before consuming the next meta out-of-band. void flush_and_wait(); + // ---- Runtime CUDA-graph node toggling (legacy rings) ----------------- + // cudaGraph_t / cudaGraphNode_t / cudaGraphExec_t travel as opaque + // uint64_t so this header stays CUDA-free (bindings.cpp is compiled with + // g++). The enabled (hook_type, layer_no) set is the single source of + // truth: apply_toggle()/ensure_graph_current() (device node state) and + // is_hook_enabled()/effective_enabled_mask() (host metadata gate) all + // read it. State and CUDA-graph operations live in ToggleRegistry + // (toggle_registry.h); these methods delegate. + + // While enabled, the producer ops record their kernel node (through + // cudaStreamGetCaptureInfo) during stream capture. Default off. + // Refused on a record ring. + void enable_toggle_capture(bool enabled); + bool toggle_capture_enabled() const; + // Called by the producer ops during capture (C++ only, not bound). + void register_capture_node(uint64_t graph, int hook_type, int layer_no, + uint64_t node); + // Called by the producer ops when a capture-time recording cannot be + // validated (non-kernel tail dependency, unreadable capture state, or a + // producer the toggle does not manage): activation then fails closed. + void note_capture_anomaly(); + uint64_t capture_anomaly_count() const; + // Bind a captured graph to its instantiated exec (after capture). + void bind_graph_exec(uint64_t graph, uint64_t exec); + // Bumped by every registry mutation (capture flag, node recording, + // anomaly, bind, clear). Host-side guard verdicts are cached on it. + uint64_t toggle_registry_version() const; + // Replace the enabled (hook_type, layer_no) set and activate the gate. + void set_enabled_hooks(const std::vector>& enabled); + // Apply the enabled set to every bound exec (cudaGraphNodeSetEnabled, + // diff-based). Returns the first CUDA error code (0 = success). Eager: + // call only when no bound exec has a replay in flight. + int apply_toggle(); + // Host metadata gate: true iff the hook is enabled AND registered, or + // the gate is inactive. + bool is_hook_enabled(int hook_type, int layer_no) const; + // Batched is_hook_enabled, one entry per query. + std::vector effective_enabled_mask( + const std::vector>& query) const; + // Lazy per-graph apply: bring ONE graph up to the current enabled set if + // it is stale. Call just before that graph replays. Nonzero is a CUDA + // error and must be treated as fatal by the caller (the metadata for the + // new set has already been pushed). + int ensure_graph_current(uint64_t graph); + // Record an event after a graph's replay for the optional pre-apply wait + // (set_toggle_pre_apply_sync). Nonzero must be treated as fatal. + int record_replay_event(uint64_t graph); + void set_toggle_pre_apply_sync(bool enabled); + uint64_t toggle_node_count() const; + // Producer nodes recorded for one captured graph (0 when none or when + // the graph was captured outside the capture window). + uint64_t toggle_node_count_for_graph(uint64_t graph) const; + // Process-wide count of producer kernels launched by the ring producer + // ops (eager or under capture). A change across a graph capture means + // the graph captured at least one producer. + uint64_t producer_launch_count() const; + uint64_t bound_graph_count() const; // 0 => apply is a no-op + uint64_t last_apply_count() const; // SetEnabled calls in the last apply_toggle + bool toggle_registry_uniform() const; // all graphs share one hook set + // Read-only replay-time check: the graph has recorded producer nodes AND + // a bound exec. A graph captured after the capture window closed is + // neither, so its producers would run while the metadata gate filters. + bool is_graph_ready(uint64_t graph) const; + // The node registry and the exec bindings cover exactly the same graphs. + bool toggle_registry_complete() const; + // Full reset (re-enables disabled nodes of bound execs first): registry, + // bindings, enabled set, gate, capture flag and lazy state. + void clear_toggle_registry(); + + // Device-side toggling (toggle_device.h). Opt in before the capture + // window opens; returns false when the build, driver or device lacks + // device graph updates (the mode then stays off). + bool enable_device_toggle(bool enabled); + bool toggle_device_mode() const; + uint64_t toggle_device_node_count() const; + uint64_t toggle_device_errors() const; + // Upload a captured graph's exec if it holds device-updatable producers + // and was not uploaded yet (required before its first launch); a no-op + // returning 0 otherwise. Returns the CUDA error. + int upload_graph_exec(uint64_t graph, uint64_t exec); + + // Test-only: make ensure_graph_current fail with `code` (0 disables), and + // query whether a graph's applied version is current. + void _test_force_apply_error(int code); + bool _test_applied_current(uint64_t graph) const; + private: struct Impl; std::unique_ptr impl_; diff --git a/native/csrc/ring/ring_torch_op.cpp b/native/csrc/ring/ring_torch_op.cpp index 0905de54e..762c499d3 100644 --- a/native/csrc/ring/ring_torch_op.cpp +++ b/native/csrc/ring/ring_torch_op.cpp @@ -2,6 +2,7 @@ #include #include "ring_torch_op.h" #include "producer.cuh" +#include "toggle_device.h" // dmx_capture_info #include @@ -17,6 +18,15 @@ static ring_py::RingEnginePy* g_active_engine = nullptr; #define HOST_HOOK_MAX 32 static std::atomic g_host_calls[HOST_HOOK_MAX] = {}; +// Producer kernels launched by the ring producer ops, eager or captured. The +// CUDA-graph replay guard compares it across a capture to tell a graph that +// captured producers from one that did not (dmi.node_toggle). +static std::atomic g_producer_launches{0}; + +uint64_t ring_producer_launch_count() { + return g_producer_launches.load(std::memory_order_relaxed); +} + void ring_diag_reset_host_counters() { for (int i = 0; i < HOST_HOOK_MAX; ++i) g_host_calls[i].store(0); } @@ -31,8 +41,65 @@ void ring_diag_print_host_counters() { fprintf(stderr, " total=%lu\n", total); } +// Node-toggle capture flag (see ring_torch_op.h). +static bool g_toggle_capture = false; + +void ring_set_toggle_capture(bool enabled) { + g_toggle_capture = enabled; +} + void ring_set_active_engine(ring_py::RingEnginePy* e) { g_active_engine = e; + // A stale capture flag after the engine is cleared would record into + // nothing. + if (e == nullptr) g_toggle_capture = false; +} + +// Node-toggle capture: record THIS producer's kernel node (the capture's +// current tail dependency) so it can be enabled or disabled after +// instantiation with cudaGraphNodeSetEnabled. No-op outside the capture +// window. +// supported = true (static and prefix producers): each is a single kernel +// launch, so the tail dependency is its kernel node. Validated: a +// non-kernel tail (a multi-op producer, a capture event join) must not be +// registered as the wrong node. +// supported = false (chunked producer): the toggle does not manage it. +// Recorded as an anomaly so activation fails loud instead of leaving a +// producer that fires while the gate filters its metadata. +// Every state where a producer fired under capture but its node could not be +// recorded is an anomaly. The one silent case is a stream that is not +// actively capturing: an eager warmup run, not a graph node. +static void dmx_record_capture_node(cudaStream_t stream, int hook_type, + int hook_id, bool supported) { + if (!g_toggle_capture || !g_active_engine) return; + cudaStreamCaptureStatus cap_st = cudaStreamCaptureStatusNone; + unsigned long long cap_id = 0; + cudaGraph_t cap_graph = nullptr; + const cudaGraphNode_t* cap_deps = nullptr; + size_t cap_nd = 0; + cudaError_t ce = dmx_capture_info(stream, &cap_st, &cap_id, &cap_graph, + &cap_deps, &cap_nd); + if (ce != cudaSuccess) { + // A healthy non-capturing stream returns success with status None, + // so this is a genuine error: the node cannot be confirmed. + g_active_engine->note_capture_anomaly(); + return; + } + if (cap_st != cudaStreamCaptureStatusActive) return; // eager run, not a node + if (cap_nd < 1 || !supported) { + g_active_engine->note_capture_anomaly(); + return; + } + cudaGraphNode_t node = cap_deps[cap_nd - 1]; + cudaGraphNodeType ntype; + if (cudaGraphNodeGetType(node, &ntype) == cudaSuccess + && ntype == cudaGraphNodeTypeKernel) { + g_active_engine->register_capture_node( + reinterpret_cast(cap_graph), hook_type, hook_id, + reinterpret_cast(node)); + } else { + g_active_engine->note_capture_anomaly(); + } } // Three side-effect ops, one per use case: @@ -90,6 +157,9 @@ void ring_producer_impl( static_cast(tensor.nbytes()), static_cast(hook_type), reinterpret_cast(stream.stream())); + g_producer_launches.fetch_add(1, std::memory_order_relaxed); + dmx_record_capture_node(stream.stream(), static_cast(hook_type), + static_cast(hook_id), /*supported=*/true); } } @@ -114,6 +184,9 @@ void ring_producer_prefix_impl( static_cast(row_bytes), static_cast(hook_type), reinterpret_cast(stream.stream())); + g_producer_launches.fetch_add(1, std::memory_order_relaxed); + dmx_record_capture_node(stream.stream(), static_cast(hook_type), + static_cast(hook_id), /*supported=*/true); } } @@ -139,6 +212,10 @@ void ring_producer_chunked_impl( K, static_cast(hook_type), reinterpret_cast(stream.stream())); + g_producer_launches.fetch_add(1, std::memory_order_relaxed); + // Not toggle-managed: fail closed if captured inside the window. + dmx_record_capture_node(stream.stream(), static_cast(hook_type), + static_cast(hook_id), /*supported=*/false); } } diff --git a/native/csrc/ring/ring_torch_op.h b/native/csrc/ring/ring_torch_op.h index c4e8b65ab..5ea5b78fe 100644 --- a/native/csrc/ring/ring_torch_op.h +++ b/native/csrc/ring/ring_torch_op.h @@ -6,6 +6,14 @@ // During graph REPLAY, only the captured cudaLaunchKernel args are re-used -- // this pointer is never read. void ring_set_active_engine(ring_py::RingEnginePy* e); + +// Node-toggle capture flag read by the legacy producer ops. While true, +// each producer launched under stream capture records its kernel node in the +// active engine's toggle registry. Process-global like the engine pointer; +// cleared whenever the engine pointer is cleared. +void ring_set_toggle_capture(bool enabled); +// Process-wide count of producer kernels the ring producer ops launched. +uint64_t ring_producer_launch_count(); void ring_diag_reset_host_counters(); void ring_diag_print_host_counters(); diff --git a/native/csrc/ring/toggle_device.h b/native/csrc/ring/toggle_device.h new file mode 100644 index 000000000..11a2e2411 --- /dev/null +++ b/native/csrc/ring/toggle_device.h @@ -0,0 +1,83 @@ +// ring/toggle_device.h -- device-side CUDA-graph node toggling. +// +// Host-side node toggling (ToggleRegistry::apply_all / ensure_current) calls +// cudaGraphNodeSetEnabled on an instantiated exec between replays. The +// device-side mode moves that decision into the graph: a producer launched as +// a device-updatable kernel node (cudaLaunchAttributeDeviceUpdatableKernelNode) +// is paired with a one-word control row in device memory, and a planner kernel +// captured upstream of the producers calls +// cudaGraphKernelNodeSetEnabled(node, row[0] != -1) at every replay. +// +// The control-row contract is the one the producers themselves honour: a row +// whose first word is -1 makes the producer return without a reservation, a +// task or a payload. A disabled node is therefore the same observation as an +// early-exit launch, so the planner adds no host/device lockstep and the host +// never toggles these nodes. +// +// CUDA device graph updates (CUDA 12.4+, toolkit and driver): an update made +// by an upstream node of a launch applies to the downstream nodes of the same +// launch and persists to later launches until overwritten. An exec holding +// device-updatable nodes must be uploaded (cudaGraphUpload) before its first +// launch. +#pragma once +#include +#include + +// Device graph updates need CUDA >= 12.4. Older toolkits build with the +// feature compiled out; -DDMX_NO_DEVICE_GRAPH_UPDATE forces the same. +#if defined(CUDART_VERSION) && CUDART_VERSION >= 12040 && !defined(DMX_NO_DEVICE_GRAPH_UPDATE) +#define DMX_HAS_DEVICE_GRAPH_UPDATE 1 +#else +#define DMX_HAS_DEVICE_GRAPH_UPDATE 0 +#endif +#if !defined(CUDART_VERSION) || CUDART_VERSION < 12040 +typedef void* cudaGraphDeviceNode_t; // absent from pre-12.4 headers +#endif + +// cudaStreamGetCaptureInfo gained the edge-data out-parameter in CUDA 13 (it +// was cudaStreamGetCaptureInfo_v3 before); one portable call site. +static inline cudaError_t dmx_capture_info(cudaStream_t s, cudaStreamCaptureStatus* st, + unsigned long long* id, cudaGraph_t* g, + const cudaGraphNode_t** deps, size_t* nd) { +#if defined(CUDART_VERSION) && CUDART_VERSION >= 13000 + const cudaGraphEdgeData* edges = nullptr; + return cudaStreamGetCaptureInfo(s, st, id, g, deps, &edges, nd); +#else + return cudaStreamGetCaptureInfo(s, st, id, g, deps, nd); +#endif +} + +namespace ring_py { + +// One slot per captured forward ("group"), one shared append-only entry +// arena. A piecewise forward is captured as many consecutive segment graphs +// holding a few producers each, so slots are (start, n) views of the arena. +constexpr int TOGGLE_DEV_MAX_GRAPHS = 16384; +constexpr int TOGGLE_DEV_MAX_ENTRIES = 65536; + +struct ToggleDevEntry { + const int64_t* control_row; // the producer's control row (fixed at capture) + cudaGraphDeviceNode_t node; // device handle of its kernel node +}; + +struct ToggleDevSlot { + int32_t n; // entries published (0 -> planner no-op) + int32_t start; // first entry in the arena +}; + +struct ToggleDevPool { + unsigned long long errors; // device-side cudaGraphKernelNodeSetEnabled failures + int last_error, last_slot, last_entry, last_enable; // most recent failure + ToggleDevSlot slot[TOGGLE_DEV_MAX_GRAPHS]; + ToggleDevEntry entry[TOGGLE_DEV_MAX_ENTRIES]; + // Per-entry sticky flags (diagnostic, idempotent stores): 1 = a disable + // succeeded at least once, 2 = an update failed at least once. + unsigned char flags[TOGGLE_DEV_MAX_ENTRIES]; +}; + +} // namespace ring_py + +namespace ring { +// Enqueue the planner kernel (captured into the graph) for one slot. +void launch_toggle_planner(const ring_py::ToggleDevPool* pool, int slot, cudaStream_t stream); +} // namespace ring diff --git a/native/csrc/ring/toggle_registry.cu b/native/csrc/ring/toggle_registry.cu new file mode 100644 index 000000000..d4c6b15c5 --- /dev/null +++ b/native/csrc/ring/toggle_registry.cu @@ -0,0 +1,538 @@ +// ring/toggle_registry.cu -- runtime CUDA-graph node toggling (see the header). + +#include "toggle_registry.h" + +#include +#include + +namespace ring_py { + +ToggleRegistry::~ToggleRegistry() { + for (auto& kv : replay_event_) if (kv.second) cudaEventDestroy(kv.second); + if (dev_pool_) cudaFree(dev_pool_); +} + +void ToggleRegistry::set_capture(bool enabled) { + std::lock_guard lk(mu_); + toggle_capture_ = enabled; + if (enabled && device_mode_ && dev_pool_ == nullptr) { + // Outside any capture (allocation is not capturable). Managed + // memory: the host appends entries while a graph is being captured + // (no stream operation, so it is legal under any capture mode), and + // the planner's first replay migrates the pages to the device. + // Zeroed slots (n = 0) make a planner a no-op. + const size_t bytes = sizeof(ToggleDevPool); + if (cudaMallocManaged(reinterpret_cast(&dev_pool_), bytes) != cudaSuccess) { + dev_pool_ = nullptr; + ++anomaly_; + } else { + std::memset(dev_pool_, 0, bytes); + } + } + ++version_; +} + +bool ToggleRegistry::capture_enabled() const { + std::lock_guard lk(mu_); + return toggle_capture_; +} + +void ToggleRegistry::register_node(uint64_t graph, int hook_type, int layer_no, uint64_t node) { + std::lock_guard lk(mu_); + RegEntry re{hook_type, layer_no, reinterpret_cast(node)}; + re.device = dev_nodes_.count(re.node) > 0; + reg_nodes_[reinterpret_cast(graph)].push_back(re); + registered_hooks_.insert({hook_type, layer_no}); + ++version_; +} + +void ToggleRegistry::note_anomaly() { + std::lock_guard lk(mu_); + ++anomaly_; + ++version_; +} + +uint64_t ToggleRegistry::anomaly_count() const { + std::lock_guard lk(mu_); + return anomaly_; +} + +void ToggleRegistry::bind_exec(uint64_t graph, uint64_t exec) { + std::lock_guard lk(mu_); + cudaGraph_t g = reinterpret_cast(graph); + cudaGraphExec_t ex = reinterpret_cast(exec); + auto prev = reg_exec_.find(g); + if (prev != reg_exec_.end() && prev->second != ex) { + // A new exec for a retained graph (re-instantiation). Nothing + // recorded for the old exec describes the new one: its nodes start in + // whatever state instantiation gave them. Mark the graph stale and + // every node's state unknown, so the next apply_all/ensure_current + // sets each host-managed node explicitly from the enabled set (a + // cached "already disabled" would otherwise skip a node the new exec + // still runs). The replay event and upload record belonged to the + // old exec; a later exec may reuse the old handle value, so forget it. + applied_version_.erase(g); + auto nodes = reg_nodes_.find(g); + if (nodes != reg_nodes_.end()) + for (auto& e : nodes->second) e.state_known = false; + auto ev = replay_event_.find(g); + if (ev != replay_event_.end()) { + if (ev->second) cudaEventDestroy(ev->second); + replay_event_.erase(ev); + } + dev_uploaded_.erase(prev->second); + } + reg_exec_[g] = ex; + // A graph whose producers are planner-managed must be uploaded before its + // first launch (once). Binding runs after the capture, with no replay in + // flight. + if (upload_device_graph_locked(g, reinterpret_cast(exec)) + != cudaSuccess) + ++anomaly_; // fail closed at activation + ++version_; +} + +cudaError_t ToggleRegistry::upload_device_graph_locked(cudaGraph_t g, + cudaGraphExec_t exec) { + if (dev_entries_.count(g) == 0 || dev_uploaded_.count(exec) > 0) + return cudaSuccess; + cudaStream_t s = nullptr; + cudaError_t e = cudaStreamCreateWithFlags(&s, cudaStreamNonBlocking); + if (e == cudaSuccess) e = cudaGraphUpload(exec, s); + if (e == cudaSuccess) e = cudaStreamSynchronize(s); + if (s) cudaStreamDestroy(s); + if (e == cudaSuccess) dev_uploaded_.insert(exec); + return e; +} + +int ToggleRegistry::upload_device_graph(uint64_t graph, uint64_t exec) { + std::lock_guard lk(mu_); + return static_cast(upload_device_graph_locked( + reinterpret_cast(graph), + reinterpret_cast(exec))); +} + +uint64_t ToggleRegistry::version() const { + std::lock_guard lk(mu_); + return version_; +} + +void ToggleRegistry::set_enabled(const std::vector>& enabled) { + std::lock_guard lk(mu_); + enabled_hooks_.clear(); + for (const auto& p : enabled) enabled_hooks_.insert(p); + toggle_active_ = true; + ++target_version_; // marks every graph stale for the lazy path +} + +int ToggleRegistry::apply_all() { + std::lock_guard lk(mu_); + cudaError_t first = cudaSuccess; + uint64_t applied = 0; + for (auto& kv : reg_exec_) { + cudaGraphExec_t exec = kv.second; + auto it = reg_nodes_.find(kv.first); + if (it == reg_nodes_.end()) continue; + for (auto& e : it->second) { + if (e.device) continue; // planner-managed, never host-toggled + const bool on = entry_on(e); + // Diff: only touch nodes whose desired state changed since the + // last apply (last_enabled tracks the exec's actual state; DMI is + // the only toggler of its nodes), so flipping a few hooks costs + // O(changed) SetEnabled calls, not O(nodes). + if (entry_current(e, on)) continue; + cudaError_t err = cudaGraphNodeSetEnabled(exec, e.node, on ? 1u : 0u); + if (err != cudaSuccess) { if (first == cudaSuccess) first = err; continue; } + e.last_enabled = on; // update only on success + e.state_known = true; + ++applied; + } + } + // Every bound graph now holds the current set. + if (first == cudaSuccess) + for (const auto& kv : reg_exec_) applied_version_[kv.first] = target_version_; + last_apply_count_ = applied; + return static_cast(first); +} + +uint64_t ToggleRegistry::bound_count() const { + std::lock_guard lk(mu_); + return reg_exec_.size(); +} + +uint64_t ToggleRegistry::node_count() const { + std::lock_guard lk(mu_); + uint64_t n = 0; + for (const auto& kv : reg_nodes_) n += kv.second.size(); + return n; +} + +uint64_t ToggleRegistry::node_count_for(uint64_t graph) const { + std::lock_guard lk(mu_); + auto it = reg_nodes_.find(reinterpret_cast(graph)); + return it == reg_nodes_.end() ? 0 : it->second.size(); +} + +uint64_t ToggleRegistry::last_apply_count() const { + std::lock_guard lk(mu_); + return last_apply_count_; +} + +bool ToggleRegistry::uniform() const { + std::lock_guard lk(mu_); + // True iff every captured graph registered the same (hook_type, layer) + // set. The metadata gate keys on (hook_type, layer) globally: a hook + // present in one graph but absent from the graph replayed this step would + // push metadata with no payload. + bool have_ref = false; + std::set> ref; + for (const auto& kv : reg_nodes_) { + std::set> s; + for (const auto& e : kv.second) s.insert({e.hook_type, e.layer_no}); + if (!have_ref) { ref = std::move(s); have_ref = true; } + else if (s != ref) return false; + } + return true; +} + +bool ToggleRegistry::complete() const { + std::lock_guard lk(mu_); + // Exact key-set match: every captured graph bound and every bound graph + // captured. Equal sizes plus inclusion of one in the other is a + // bijection (both are keyed by unique cudaGraph_t). + if (reg_nodes_.size() != reg_exec_.size()) return false; + for (const auto& kv : reg_nodes_) + if (reg_exec_.count(kv.first) == 0) return false; + return true; +} + +bool ToggleRegistry::is_ready(uint64_t graph) const { + std::lock_guard lk(mu_); + cudaGraph_t g = reinterpret_cast(graph); + return reg_nodes_.count(g) > 0 && reg_exec_.count(g) > 0; +} + +bool ToggleRegistry::is_enabled(int hook_type, int layer_no) const { + std::lock_guard lk(mu_); + if (!toggle_active_) return true; // toggle inactive -> every hook fires + // A hook is enabled only if it is BOTH in the enabled set AND registered + // (has a captured producer node). An enabled but uncaptured hook would + // push metadata with no payload. So "metadata pushed" <=> "a producer + // fires". + const std::pair k{hook_type, layer_no}; + return enabled_hooks_.count(k) > 0 && registered_hooks_.count(k) > 0; +} + +std::vector ToggleRegistry::effective_mask( + const std::vector>& query) const { + // Batched is_enabled: one lock, one binding crossing. + std::lock_guard lk(mu_); + const bool inactive = !toggle_active_; + std::vector mask; + mask.reserve(query.size()); + for (const auto& k : query) { + mask.push_back((inactive || + (enabled_hooks_.count(k) > 0 && + registered_hooks_.count(k) > 0)) ? 1 : 0); + } + return mask; +} + +int ToggleRegistry::ensure_current(uint64_t graph) { + // Lazy: apply the current enabled set to ONE graph, only when stale. + // Called just before that graph replays. Same per-node diff as + // apply_all, scoped to one graph. + std::lock_guard lk(mu_); + cudaGraph_t g = reinterpret_cast(graph); + auto av = applied_version_.find(g); + const uint64_t cur = (av == applied_version_.end()) ? 0 : av->second; + if (cur == target_version_) return 0; // current: the common path + + auto ex = reg_exec_.find(g); + if (ex == reg_exec_.end()) return 0; // unregistered/unbound graph + auto it = reg_nodes_.find(g); + if (it == reg_nodes_.end()) return 0; + + // Test fault injection: fail before any mutation or version bump. + if (force_apply_error_ != 0) return force_apply_error_; + + // A graph whose producers are all planner-managed needs no host apply: + // the in-graph planner derives each node's state from its control row at + // every replay. + bool host_nodes = false; + for (const auto& e : it->second) host_nodes |= !e.device; + if (!host_nodes) { applied_version_[g] = target_version_; return 0; } + + // Optional pre-apply wait for this exec's previous replay. If the wait + // fails, the previous replay cannot be confirmed finished: touch nothing, + // leave the version stale and return the error (the caller treats it as + // fatal). + auto ev = pre_apply_sync_ ? replay_event_.find(g) : replay_event_.end(); + if (ev != replay_event_.end() && ev->second) { + cudaError_t se = cudaEventSynchronize(ev->second); + if (se != cudaSuccess) return static_cast(se); + } + + cudaError_t first = cudaSuccess; + for (auto& e : it->second) { + if (e.device) continue; + const bool on = entry_on(e); + if (entry_current(e, on)) continue; + cudaError_t err = cudaGraphNodeSetEnabled(ex->second, e.node, on ? 1u : 0u); + if (err != cudaSuccess) { if (first == cudaSuccess) first = err; continue; } + e.last_enabled = on; + e.state_known = true; + } + // Current only on full success: a failed node leaves the version stale so + // a later ensure retries it (nodes that did flip are already recorded in + // last_enabled, so the retry skips them). + if (first == cudaSuccess) + applied_version_[g] = target_version_; + return static_cast(first); +} + +int ToggleRegistry::record_replay_event(uint64_t graph, cudaStream_t stream) { + // Record an event right after a graph's replay so a later ensure_current() + // can wait for it before mutating that exec. Returns the CUDA error: a + // failed create or record would leave a stale or missing event, so the + // caller treats nonzero as fatal. No-op unless the pre-apply wait is on + // (the event's only reader). + std::lock_guard lk(mu_); + if (!pre_apply_sync_) return 0; + cudaGraph_t g = reinterpret_cast(graph); + if (reg_exec_.find(g) == reg_exec_.end()) return 0; // not bound + cudaEvent_t& ev = replay_event_[g]; + if (ev == nullptr) { + cudaError_t ce = cudaEventCreateWithFlags(&ev, cudaEventDisableTiming); + if (ce != cudaSuccess) { ev = nullptr; return static_cast(ce); } + } + return static_cast(cudaEventRecord(ev, stream)); +} + +void ToggleRegistry::set_pre_apply_sync(bool enabled) { + std::lock_guard lk(mu_); + pre_apply_sync_ = enabled; +} + +bool ToggleRegistry::pre_apply_sync() const { + std::lock_guard lk(mu_); + return pre_apply_sync_; +} + +void ToggleRegistry::clear() { + std::lock_guard lk(mu_); + // Restore device state before dropping the registry: re-enable every node + // we disabled, so a graph that keeps replaying after the clear does not + // run with producers off while the (now inactive) host gate pushes every + // metadata entry. Best effort, bound execs only; the call sites are + // quiescent (before a capture, at teardown). To return to all-on while + // serving, prefer an all-on enabled set over clear(). + for (auto& kv : reg_nodes_) { + auto ex = reg_exec_.find(kv.first); + if (ex == reg_exec_.end()) continue; + for (auto& e : kv.second) { + if ((e.state_known && e.last_enabled) || e.device) continue; + if (cudaGraphNodeSetEnabled(ex->second, e.node, 1u) == cudaSuccess) { + e.last_enabled = true; + e.state_known = true; + } + } + } + reg_nodes_.clear(); + reg_exec_.clear(); + registered_hooks_.clear(); + // Full reset: drop the enabled set and deactivate, so a stale gate cannot + // skip every metadata entry after a clear; close the capture window so a + // later capture does not record into a cleared registry. + enabled_hooks_.clear(); + toggle_active_ = false; + toggle_capture_ = false; + last_apply_count_ = 0; + anomaly_ = 0; + target_version_ = 0; + applied_version_.clear(); + for (auto& kv : replay_event_) if (kv.second) cudaEventDestroy(kv.second); + replay_event_.clear(); + // Device pool: forget the open group but never recycle or zero a slot. + // An exec that outlives the clear keeps its planner, which keeps gating + // each node by its own control row (always correct); zeroing a slot would + // freeze device-disabled nodes off while their rows ask them to fire. + dev_group_slot_ = -1; + dev_group_rows_.clear(); + dev_entries_.clear(); + dev_nodes_.clear(); + dev_uploaded_.clear(); + // Device-side toggling is opted into per capture window. The pool stays + // allocated for reuse. + device_mode_ = false; + ++version_; +} + +bool ToggleRegistry::set_device_mode(bool enabled) { + std::lock_guard lk(mu_); +#if DMX_HAS_DEVICE_GRAPH_UPDATE + if (enabled) { + int device = 0; + int driver = 0; + if (cudaGetDevice(&device) != cudaSuccess || + cudaDriverGetVersion(&driver) != cudaSuccess || driver < 12040) { + device_mode_ = false; + return false; + } + } + device_mode_ = enabled; + ++version_; + return device_mode_ == enabled; +#else + device_mode_ = false; + return !enabled; +#endif +} + +bool ToggleRegistry::device_mode() const { + std::lock_guard lk(mu_); + return device_mode_; +} + +bool ToggleRegistry::device_armed() const { + std::lock_guard lk(mu_); + return device_mode_ && dev_pool_ != nullptr; +} + +const ToggleDevPool* ToggleRegistry::device_pool() const { + std::lock_guard lk(mu_); + return dev_pool_; +} + +int ToggleRegistry::device_slot(uint64_t graph, const int64_t* control_row, bool* fresh) { + // One slot per captured FORWARD ("group"), not per graph: a piecewise + // forward is captured as many consecutive segment graphs holding one or + // two producers each. The group's first producer gets the planner node + // (in its graph); that planner gates every producer of the group, + // including producers in later segment graphs of the same forward, which + // replay after it on the same stream. Each producer fires once per + // forward, so a repeated control row starts the next group. + std::lock_guard lk(mu_); + (void)graph; + *fresh = false; + if (!device_mode_ || dev_pool_ == nullptr) return -1; + if (dev_group_slot_ >= 0 && dev_group_rows_.count(control_row) == 0) { + dev_group_rows_.insert(control_row); + return dev_group_slot_; + } + if (dev_slots_used_ >= TOGGLE_DEV_MAX_GRAPHS) return -1; + const int s = dev_slots_used_++; + dev_pool_->slot[s].start = dev_entries_used_; + dev_pool_->slot[s].n = 0; + dev_group_slot_ = s; + dev_group_rows_.clear(); + dev_group_rows_.insert(control_row); + *fresh = true; + ++version_; + return s; +} + +void ToggleRegistry::register_device_entry(int slot_index, uint64_t graph, uint64_t node, + const int64_t* control_row, + cudaGraphDeviceNode_t dev) { + std::lock_guard lk(mu_); + cudaGraph_t g = reinterpret_cast(graph); + // A group's entries are contiguous in the arena (captures are + // sequential): the previous append must belong to this slot. + const bool ok = dev_pool_ != nullptr && slot_index >= 0 + && slot_index < dev_slots_used_ + && dev_pool_->slot[slot_index].start + dev_pool_->slot[slot_index].n == dev_entries_used_; + if (!ok || dev == nullptr || dev_entries_used_ >= TOGGLE_DEV_MAX_ENTRIES) { + // Fail closed. The node stays enabled (the device default), which + // is still correct (it early-exits on a -1 row), but activation + // refuses an anomalous registry. + ++anomaly_; + ++version_; + return; + } + dev_entries_[g].push_back(ToggleDevEntry{control_row, dev}); + dev_nodes_.insert(reinterpret_cast(node)); + // Publish now (host write into managed memory; no graph of this group can + // replay before its capture ends). Entry before count. + ToggleDevSlot& slot = dev_pool_->slot[slot_index]; + dev_pool_->entry[dev_entries_used_++] = ToggleDevEntry{control_row, dev}; + std::atomic_thread_fence(std::memory_order_release); + slot.n += 1; + ++version_; +} + +uint64_t ToggleRegistry::device_node_count() const { + std::lock_guard lk(mu_); + uint64_t n = 0; + for (const auto& kv : dev_entries_) n += kv.second.size(); + return n; +} + +uint64_t ToggleRegistry::device_errors() const { + std::lock_guard lk(mu_); + // Managed memory: host-readable once the planner's launches complete. + return dev_pool_ ? dev_pool_->errors : 0; +} + +int ToggleRegistry::upload_exec(uint64_t exec) { + cudaStream_t s = nullptr; + cudaError_t e = cudaStreamCreateWithFlags(&s, cudaStreamNonBlocking); + if (e == cudaSuccess) e = cudaGraphUpload(reinterpret_cast(exec), s); + if (e == cudaSuccess) e = cudaStreamSynchronize(s); + if (s) cudaStreamDestroy(s); + return static_cast(e); +} + +void ToggleRegistry::force_apply_error(int code) { + std::lock_guard lk(mu_); + force_apply_error_ = code; +} + +bool ToggleRegistry::applied_current(uint64_t graph) const { + std::lock_guard lk(mu_); + cudaGraph_t g = reinterpret_cast(graph); + auto av = applied_version_.find(g); + const uint64_t cur = (av == applied_version_.end()) ? 0 : av->second; + return cur == target_version_; +} + +} // namespace ring_py + +namespace ring { + +// One thread per producer node: enable iff its control row is not -1. Runs +// upstream of every producer of the same graph launch, so the update applies +// to this launch (device graph update semantics). +__global__ void toggle_planner_kernel(const ring_py::ToggleDevPool* pool, int slot) +{ +#if DMX_HAS_DEVICE_GRAPH_UPDATE + const ring_py::ToggleDevSlot s = pool->slot[slot]; + for (int i = threadIdx.x; i < s.n; i += blockDim.x) { + const ring_py::ToggleDevEntry e = pool->entry[s.start + i]; + // A failure leaves the node in its previous state (still correct: a + // node left enabled early-exits on a -1 row). Failures are counted. + const bool on = e.control_row[0] != -1; + const cudaError_t r = cudaGraphKernelNodeSetEnabled(e.node, on); + unsigned char* fl = const_cast(&pool->flags[s.start + i]); + if (r == cudaSuccess) { if (!on && !(*fl & 1)) *fl |= 1; } + else if (!(*fl & 2)) *fl |= 2; + if (r != cudaSuccess) { + atomicAdd(const_cast(&pool->errors), 1ull); + ring_py::ToggleDevPool* p = const_cast(pool); + p->last_error = static_cast(r); + p->last_slot = slot; + p->last_entry = i; + p->last_enable = on; + } + } +#else + (void)pool; (void)slot; +#endif +} + +void launch_toggle_planner(const ring_py::ToggleDevPool* pool, int slot, cudaStream_t stream) +{ + toggle_planner_kernel<<<1, 128, 0, stream>>>(pool, slot); +} + +} // namespace ring diff --git a/native/csrc/ring/toggle_registry.h b/native/csrc/ring/toggle_registry.h new file mode 100644 index 000000000..08c3b5c3b --- /dev/null +++ b/native/csrc/ring/toggle_registry.h @@ -0,0 +1,175 @@ +// ring/toggle_registry.h -- state and CUDA-graph operations for runtime +// node toggling, kept out of the ring engine so the engine only delegates. +// +// Owns: the capture-time node registry (graph -> producer kernel nodes), the +// graph -> exec bindings, the enabled (hook_type, layer_no) set (the single +// source of truth for the host metadata gate and the device node state), the +// lazy per-graph apply state (target/applied versions and per-graph replay +// events), the capture anomaly counter, the device-side toggle pool, and the +// registry version every mutation bumps (host caches key their guard verdicts +// on it). +// +// Thread-safe: every public method takes the internal mutex. Depends only on +// the CUDA runtime (callers pass streams in). This header contains CUDA +// types, so only .cu/.cpp translation units include it, never the CUDA-free +// ring_engine_py.h. +#pragma once + +#include +#include "toggle_device.h" + +#include +#include +#include +#include +#include +#include + +namespace ring_py { + +class ToggleRegistry { +public: + ~ToggleRegistry(); + + // -- capture window ---------------------------------------------------- + void set_capture(bool enabled); + bool capture_enabled() const; + // Called from the producer op during capture (single kernel just enqueued). + void register_node(uint64_t graph, int hook_type, int layer_no, uint64_t node); + // Called when a capture-time recording could not be validated (non-kernel + // tail, unreadable capture state, unsupported producer) -> fail closed. + void note_anomaly(); + uint64_t anomaly_count() const; + + // -- binding / introspection -------------------------------------------- + // Bind (or rebind) a graph to an instantiated exec. Binding a different + // exec than the one bound before resets every per-exec record for the + // graph: the applied version, each node's cached state (unknown until the + // next apply sets every host-managed node explicitly), the replay event + // and the previous exec's upload record. Rebinding the same exec is a + // no-op apart from the version bump. + void bind_exec(uint64_t graph, uint64_t exec); + uint64_t version() const; // bumped by every registry mutation + uint64_t bound_count() const; + uint64_t node_count() const; + uint64_t node_count_for(uint64_t graph) const; // nodes recorded for one graph + uint64_t last_apply_count() const; + bool uniform() const; // all graphs registered the same hook set + bool complete() const; // registered graphs == bound graphs + bool is_ready(uint64_t graph) const; // registered AND bound + + // -- enabled set + host apply --------------------------------------------- + void set_enabled(const std::vector>& enabled); + int apply_all(); // eager: every bound exec, diff-based + bool is_enabled(int hook_type, int layer_no) const; + std::vector effective_mask( + const std::vector>& query) const; + int ensure_current(uint64_t graph); // lazy, one graph + int record_replay_event(uint64_t graph, cudaStream_t stream); + // Wait for a graph's previous replay before mutating its exec (see + // pre_apply_sync_). Off by default. + void set_pre_apply_sync(bool enabled); + bool pre_apply_sync() const; + void clear(); + + // -- device-side toggle (toggle_device.h) -------------------------------- + // Opt in before the capture window opens. Ignored (returns false) when + // the build or the runtime lacks device graph updates. clear() turns the + // mode off again. + bool set_device_mode(bool enabled); + bool device_mode() const; + // The pool exists (allocated when the first capture window opens in + // device mode). + bool device_armed() const; + // Slot of the forward ("group") being captured (-1 = pool exhausted). + // *fresh = true for the group's first producer: the caller then enqueues + // the planner node right before it. + int device_slot(uint64_t graph, const int64_t* control_row, bool* fresh); + const ToggleDevPool* device_pool() const; + // Called right after a device-updatable producer was captured. + void register_device_entry(int slot, uint64_t graph, uint64_t node, + const int64_t* control_row, + cudaGraphDeviceNode_t dev); + uint64_t device_node_count() const; + uint64_t device_errors() const; + // cudaGraphUpload + synchronize: required before the first launch of an + // exec holding device-updatable nodes. + int upload_exec(uint64_t exec); + // Upload `exec` iff `graph` captured device-updatable producers and the + // exec was not uploaded yet: each such exec is uploaded at most once, + // whether it is bound (bind_exec) or not (a capture outside the window). + // Returns the CUDA error (0 also when nothing needed uploading). + int upload_device_graph(uint64_t graph, uint64_t exec); + + // -- test-only fault injection / introspection --------------------------- + void force_apply_error(int code); + bool applied_current(uint64_t graph) const; + +private: + struct RegEntry { + int hook_type; + int layer_no; + cudaGraphNode_t node; + bool last_enabled{true}; // device default after instantiate + // False when the bound exec changed since last_enabled was written: + // the next apply sets the node unconditionally. + bool state_known{true}; + bool device{false}; // toggled by the in-graph planner, never by the host + }; + cudaError_t upload_device_graph_locked(cudaGraph_t g, cudaGraphExec_t exec); + // The apply loops skip a node only when its exec state is known to match. + static bool entry_current(const RegEntry& e, bool on) { + return e.state_known && on == e.last_enabled; + } + bool entry_on(const RegEntry& e) const { + return enabled_hooks_.count({e.hook_type, e.layer_no}) > 0; + } + std::unordered_map> reg_nodes_; // per captured graph + std::unordered_map reg_exec_; // graph -> exec + std::set> registered_hooks_; // (ht, layer) captured + std::set> enabled_hooks_; // single source + bool toggle_active_{false}; + bool toggle_capture_{false}; + uint64_t last_apply_count_{0}; + // Lazy per-graph apply: target_version_ bumps on every enabled-set change; + // each graph tracks the version it was last applied to. ensure_current() + // applies the current enabled set to ONE graph only when stale (just + // before that graph replays), so a reconfigure costs + // O(changed x graphs actually replayed). + // + // cudaGraphNodeSetEnabled "only affects future launches of hGraphExec. + // Already enqueued or running launches of hGraphExec are not affected", + // so mutating an exec whose previous launch is still queued is defined + // and the in-flight launch keeps the node state it was enqueued with. + // The pre-apply wait (an event recorded after each replay and + // synchronized before the first mutation of a stale exec) is therefore + // off by default: it serializes the host with the GPU on every set + // change and defeats schedulers that run one step ahead of the GPU. + // set_pre_apply_sync(true) restores it; with it off, + // record_replay_event() is a no-op. + bool pre_apply_sync_{false}; + uint64_t target_version_{0}; + std::unordered_map applied_version_; + std::unordered_map replay_event_; + // Test-only fault injection: when nonzero, ensure_current returns this + // code as if a device apply failed (no SetEnabled, no version bump). + int force_apply_error_{0}; + // Capture-time recordings that could not be validated. Nonzero means the + // registry may be misaligned, so activation fails closed. + uint64_t anomaly_{0}; + uint64_t version_{0}; + // Device-side toggle state. The pool is allocated at the first + // set_capture(true) in device mode (outside any capture). + bool device_mode_{false}; + ToggleDevPool* dev_pool_{nullptr}; + int dev_slots_used_{0}; + int dev_entries_used_{0}; + int dev_group_slot_{-1}; + std::set dev_group_rows_; + std::unordered_map> dev_entries_; + std::set dev_nodes_; + std::set dev_uploaded_; + mutable std::mutex mu_; +}; + +} // namespace ring_py diff --git a/src/dmi/adapters/base.py b/src/dmi/adapters/base.py index 333ea5d8d..619c0b071 100644 --- a/src/dmi/adapters/base.py +++ b/src/dmi/adapters/base.py @@ -11,8 +11,12 @@ -> plan_step (base; walks active_specs once for total_bytes + hook_count + needs_eager) -> commit_step (base) + -> reject an eager-only plan before reservation when eager fallback + is disabled -> prepare_step (skipped when hook_count == 0) on result == 2: + -> reject before metadata publication when eager fallback is + disabled -> adapt_for_cpu_direct (subclass; default no-op) -> on_capacity_exceeded (subclass) -> _warn_once_capacity (base stub; subclass overrides) @@ -21,7 +25,9 @@ -> pre_push_all_metas The unit-test gate (``tests/test_adapter_protocol.py``) exercises the driver -ordering with mocks. +ordering and the strict rejection with mocks. Native ``prepare_step`` may +synchronize and drain earlier work before returning OVERSIZED; a strict +rejection does not undo that, and that native branch reserves nothing. """ from __future__ import annotations @@ -127,7 +133,10 @@ class BackendAdapter(abc.ABC): ``commit_step`` (reservation + context/metadata publication). """ - def __init__(self, engine: "MonitoringEngine", model_id: str) -> None: + def __init__(self, engine: "MonitoringEngine", model_id: str, + *, allow_eager_fallback: bool = True) -> None: + if type(allow_eager_fallback) is not bool: + raise ValueError("allow_eager_fallback must be a bool") self.engine = engine self.model_id = model_id self.transport = engine_ring_transport(engine) @@ -138,6 +147,27 @@ def __init__(self, engine: "MonitoringEngine", model_id: str) -> None: # Steps seen by this adapter since attach. Feeds the capture # schedule's stride/warmup gate in before_forward. self._step_counter: int = 0 + # False: a step that would need the eager safety net (a hook that + # needs eager dispatch, or a step larger than the ring) raises + # instead, before anything is reserved or published. For + # deployments where a silent switch to eager dispatch (and CPU-direct + # copies) is worse than a refused step. + self.allow_eager_fallback = allow_eager_fallback + self._dispatch_counts = dict( + prepared_steps=0, ring_flushed_steps=0, oversized_steps=0, + eager_required_steps=0, forced_eager_steps=0, + rejected_eager_steps=0) + + def capture_dispatch_state(self) -> dict: + """Cumulative step-preparation counts (not completed executions).""" + return {"allow_eager_fallback": self.allow_eager_fallback, + **self._dispatch_counts} + + def _reject_eager(self, reason: str, plan: "StepPlan") -> None: + self._dispatch_counts["rejected_eager_steps"] += 1 + raise RuntimeError( + f"capture eager fallback forbidden: {reason}; " + f"planned_bytes={plan.total_bytes}, hooks={plan.hook_count}") @property def active_hook_specs(self) -> tuple[HookSpec, ...]: @@ -395,6 +425,15 @@ def commit_step( # their internal loops produce nothing to push when the # active-spec list is empty or every shape was empty. if plan.hook_count > 0: + counts = getattr(self, "_dispatch_counts", None) + strict = not getattr(self, "allow_eager_fallback", True) + if plan.needs_eager: + if counts is not None: + counts["eager_required_steps"] += 1 + if strict: + # The plan alone establishes that graph dispatch is + # impossible: reject before reserving ring space. + self._reject_eager("hook requires eager execution", plan) # A needs_eager step runs every hook through the eager safety # net, which reserves each hook's entry itself (reserve_one), so # prepare_step only makes room for it. Reserving the step as @@ -406,10 +445,22 @@ def commit_step( reserve=not plan.needs_eager, ) ) + if counts is not None: + counts["prepared_steps"] += 1 + counts["ring_flushed_steps"] += int( + reservation is StepReservation.FLUSHED) + counts["oversized_steps"] += int( + reservation is StepReservation.OVERSIZED) + if reservation is StepReservation.OVERSIZED and strict: + # OVERSIZED flushed earlier work but reserved nothing: publish + # no metadata and arm no producer for this step. + self._reject_eager("step exceeds ring capacity", plan) self.transport.force_eager = ( reservation is StepReservation.OVERSIZED or plan.needs_eager ) + if counts is not None: + counts["forced_eager_steps"] += int(self.transport.force_eager) if reservation is StepReservation.OVERSIZED: ctx = self.adapt_for_cpu_direct(ctx) self.on_capacity_exceeded(ctx) @@ -439,7 +490,8 @@ def close(self) -> None: def plan_step(self, ctx: StepContext) -> StepPlan: """Build a side-effect-free reservation plan for one step. - Single walk over ``active_specs``: + Single walk over ``active_specs`` (``transport.specs_for_step()`` + while the node toggle is active): - ``total_bytes`` and ``hook_count`` feed ``prepare_step``. ``hook_count`` counts only specs whose ``compute_hook_shape`` returns a @@ -456,7 +508,16 @@ def plan_step(self, ctx: StepContext) -> StepPlan: total = 0 n = 0 needs_eager = False - for spec in self.active_specs: + # While the node toggle gates steps, reserve for the per-step spec + # set the transport also publishes metadata for (the enabled subset + # on a gated step, the full active set otherwise). Reserving for a + # different set than the one that fires drifts the ring's head and + # tail. With the toggle inactive this is active_specs, as before. + specs = self.active_specs + transport = getattr(self, "transport", None) + if getattr(transport, "toggle_gate_active", False) is True: + specs = transport.specs_for_step() + for spec in specs: # When the adapter has populated actual_q_len AND the spec is # prefix-strip-eligible, size the byte budget for the unpadded # data the producer will actually write. Other specs and the diff --git a/src/dmi/hooks/point.py b/src/dmi/hooks/point.py index 52b7a50e5..3f202eb70 100644 --- a/src/dmi/hooks/point.py +++ b/src/dmi/hooks/point.py @@ -290,6 +290,18 @@ def forward(self, x: Tensor) -> Tensor: from ..transport import ring as _rt transport = _rt._active_transport if transport is not None and transport.force_eager: + # Node toggle: the eager safety net fires producers outside any + # graph, where cudaGraphNodeSetEnabled cannot gate them. On a + # gated step that fell back to eager (a step larger than the + # ring), a disabled hook must be skipped here, or it fires while + # metadata and reservation counted only the enabled subset. + # This branch is never captured or compiled, so a host check is + # safe; toggle_gate_active keeps the stock eager path untouched. + if (getattr(transport, "toggle_gate_active", False) + and getattr(transport, "_gated_step", True) + and not transport._ring_engine.is_hook_enabled( + self._ring_hook_type, self._ring_hook_id)): + return x_cont engine = transport._ring_engine if engine is not None: nbytes = x_cont.nbytes diff --git a/src/dmi/node_toggle.py b/src/dmi/node_toggle.py new file mode 100644 index 000000000..43fc5e39a --- /dev/null +++ b/src/dmi/node_toggle.py @@ -0,0 +1,852 @@ +"""Runtime CUDA-graph node toggling for the legacy ring transport. + +A monitored step that replays a captured CUDA graph runs every installed +HookPoint's producer as a kernel node of that graph. Node toggling enables or +disables individual producer nodes between replays with +``cudaGraphNodeSetEnabled`` on the instantiated exec -- no recapture, no +re-instantiation, no model change. A disabled node costs an empty-node +traversal, so a step pays only for the hooks it enables. + +One invariant drives the design. Every step, three lanes must agree on which +hooks fire: + +1. capacity reservation (``BackendAdapter.plan_step``), +2. metadata publication (``RingTransport.pre_push_all_metas``), +3. the device node state (the engine's enabled set). + +Otherwise the ring desynchronizes (reserved bytes != written bytes, or +metadata != payloads). All three read one source of truth, +``RingTransport.specs_for_step()``, which this controller feeds. + +The toggle governs only producers inside a replayed graph. An eager or +prefill step runs producers ungated, so on such a step the integration marks +the transport ungated (``transport.set_gated_step(False)``) and all three lanes +use the full active set. + +Two application modes: + +* eager (``set_active_hooks``): flip every bound exec now. Only at a + quiescent point (no bound exec has a replay in flight), e.g. once after + warmup. +* lazy (``set_active_hooks_lazy``): update the enabled set and the host gate + now; each graph is brought up to date by ``ensure_graph_current`` just + before it replays. For per-step reconfiguration. + +The second half of this module holds framework-neutral CUDA-graph helpers: a +capture window that scopes node recording, a ``torch.cuda.CUDAGraph`` subclass +that keeps the captured template alive and guards replays, and the binding of +captured graphs to the engine's registry. A serving integration decides which +captures belong to the toggle (for example only its full-model decode graphs) +through the ``exclude_capture`` predicate, and calls ``bind_graph`` for each +graph it wants toggled. + +The ``torch.cuda.CUDAGraph`` replacement is process-global, so monitoring +sessions own it: opening a capture window starts (or renews) the session of +that engine. The session ends when the window closes with no graph bound, +when the controller is cleared, or when the transport is deactivated; once the +last session ends the original class is restored and the replay guard +disarmed. Each capture is stamped with the session it was captured in and +with whether it captured DMI producers, so the guard validates only graphs of +the active transport's engine that are bound or carry producers: an unrelated +graph replays untouched, and a later session never adopts a graph captured +for another engine. +""" +from __future__ import annotations + +import contextlib +import itertools +import weakref +from typing import (TYPE_CHECKING, Any, Callable, Iterable, Iterator, List, + Optional, Tuple) + +if TYPE_CHECKING: # avoid a runtime import cycle with the transport + from .hooks.specs import HookSpec + from .transport.ring import RingTransport + + +HookKey = Tuple[int, int] + + +class NodeToggleController: + """Host-side toggle state for one transport. Defaults to inactive.""" + + def __init__(self, engine: Any, transport: "RingTransport") -> None: + self._eng = engine + self._transport = transport + # True once set_active_hooks[_lazy] ran; the transport's + # effective_specs gates on it. + self.gate_active: bool = False + # True when the active gate uses the lazy device-apply path. The + # replay guard reads it to choose ensure_graph_current (lazy) over + # validation only (eager). + self.lazy_active: bool = False + # The enabled-and-captured subset of the transport's active specs, in + # active order (None until the first activation). + self.effective_enabled_specs: "Optional[List[HookSpec]]" = None + # Bumped whenever the enabled set changes. + self.enabled_version: int = 0 + # Reconfiguration caches, both keyed on the engine's registry version + # (bumped by every registry mutation), so neither can serve stale: + # _guard_valid_version: registry version whose validation PASSED + # (failures are never cached); + # _cache: (registry version, enabled frozenset) -> + # (active_specs reference, effective list). The stored reference + # is compared with ``is``; holding it pins the id. + self._guard_valid_version: Optional[int] = None + self._cache: dict = {} + # active_specs reference -> {(hook_type, layer): [active index, ...]}, + # so a reconfiguration queries only the enabled pairs. + self._active_index: Optional[tuple] = None + + # -- reconfiguration --------------------------------------------------- + + def set_active_hooks(self, enabled: Iterable[HookKey]) -> None: + """Enable exactly ``enabled`` (hook_type, layer_no) pairs and disable + the rest, on every bound exec, now. + + Eager apply does not wait for replays in flight: call it only at a + quiescent point. For per-step reconfiguration use + :meth:`set_active_hooks_lazy`. Requires producer nodes recorded at + capture and the graphs bound (see :func:`bind_graph`). + """ + self._validate_registry("set_active_hooks") + pairs = [(int(ht), int(ln)) for (ht, ln) in enabled] + self._eng.set_enabled_hooks(pairs) + err = self._eng.apply_toggle() + if err != 0: + raise RuntimeError(f"apply_toggle failed with CUDA error {err}") + self.gate_active = True + self.lazy_active = False + self._recompute(pairs) + + def set_active_hooks_lazy(self, enabled: Iterable[HookKey]) -> None: + """Update the enabled set and the host gate now, and defer the device + apply to :meth:`ensure_graph_current`, called just before each graph + replays. Only graphs that actually replay are touched.""" + self._validate_registry("set_active_hooks_lazy") + pairs = [(int(ht), int(ln)) for (ht, ln) in enabled] + self._eng.set_enabled_hooks(pairs) # marks every graph stale + self.gate_active = True + self.lazy_active = True + self._recompute(pairs) + + def clear(self) -> None: + """Clear the engine's registry AND deactivate the host gate, so + neither lane is left half-configured. Call before a recapture or + when monitoring stops. Ends the engine's monitoring session, so the + process-wide CUDAGraph patch is restored once no session remains.""" + self._eng.clear_toggle_registry() + _forget_window(self._eng) + end_session(self._eng) + self.gate_active = False + self.lazy_active = False + self.effective_enabled_specs = None + self.enabled_version += 1 + self._guard_valid_version = None + self._cache.clear() + self._active_index = None + + # -- replay-time hooks ------------------------------------------------- + + def is_graph_ready(self, raw_graph: int) -> bool: + """The graph has recorded producer nodes AND a bound exec.""" + return self._eng.is_graph_ready(raw_graph) + + def ensure_graph_current(self, raw_graph: int) -> int: + """Lazy: apply the deferred toggle to the graph about to replay + (no-op when current). Returns a CUDA error code; nonzero is fatal.""" + return self._eng.ensure_graph_current(raw_graph) + + def record_replay_event(self, raw_graph: int) -> int: + """Record the replay event used by the optional pre-apply wait. + Returns a CUDA error code; nonzero is fatal.""" + return self._eng.record_replay_event(raw_graph) + + # -- internals --------------------------------------------------------- + + def _validate_registry(self, who: str) -> None: + """Registry guards for set_active_hooks[_lazy], memoized on the + engine's registry version. The registry changes only at capture, + bind and clear time, so re-running the checks on an unchanged + registry is wasted work. Only a pass is cached.""" + eng = self._eng + version = eng.toggle_registry_version() # read before validating + if version == self._guard_valid_version: + return + # Gate active but the device toggle a no-op: metadata filtered while + # every producer still fires. + if eng.bound_graph_count() == 0 or eng.toggle_node_count() == 0: + raise RuntimeError( + "%s: no producer nodes registered or no graph exec bound " + "(toggle_node_count=%d, bound_graph_count=%d). Capture inside " + "a toggle capture window and bind the captured graphs first." + % (who, eng.toggle_node_count(), eng.bound_graph_count())) + # A hook in one graph but not in the replayed one would push metadata + # with no payload. + if not eng.toggle_registry_uniform(): + raise RuntimeError( + "%s: captured graphs registered different hook sets; the " + "global metadata gate cannot stay aligned. Every toggled " + "graph must register the same hooks (exclude partial-model " + "captures from the capture window)." % who) + # A captured graph without a bound exec replays with every node on + # while the gate filters. + if not eng.toggle_registry_complete(): + raise RuntimeError( + "%s: toggle registry incomplete -- the graphs with recorded " + "producer nodes are not exactly the bound graphs. Every " + "captured graph must be bound, and vice versa." % who) + if eng.capture_anomaly_count() > 0: + raise RuntimeError( + "%s: capture recorded %d producer node(s) the toggle cannot " + "manage (a non-kernel tail dependency, unreadable capture " + "state, or a chunked producer, which the toggle does not " + "manage). Refusing to activate." + % (who, eng.capture_anomaly_count())) + self._guard_valid_version = version + + def _recompute(self, pairs: List[HookKey]) -> None: + """Recompute active AND enabled AND captured with one batched engine + call. Memoized per (registry version, enabled set, active_specs + identity), so repeated presets skip the recompute.""" + active = self._transport._active_specs + key = (self._eng.toggle_registry_version(), frozenset(pairs)) + hit = self._cache.get(key) + if hit is not None and hit[0] is active: + self.effective_enabled_specs = hit[1] + self.enabled_version += 1 + return + # The gate is active here, so the engine's mask is "enabled AND + # registered": a spec outside `pairs` is never effective. Query just + # the active specs whose key is enabled, in active order. + index = self._active_index + if index is None or index[0] is not active: + by_key: dict = {} + for i, s in enumerate(active): + by_key.setdefault((s.hook_type, s.layer_no), []).append(i) + index = self._active_index = (active, by_key) + by_key = index[1] + positions = sorted(i for k in key[1] for i in by_key.get(k, ())) + mask = self._eng.effective_enabled_mask( + [(active[i].hook_type, active[i].layer_no) for i in positions]) + self.effective_enabled_specs = [active[i] for i, m in zip(positions, mask) if m] + self.enabled_version += 1 + if len(self._cache) >= 64: # ad hoc sets: keep the memo bounded + self._cache.clear() + self._cache[key] = (active, self.effective_enabled_specs) + + +def parse_enabled_hooks(value: Any) -> Optional[List[HookKey]]: + """Parse an enabled-hooks setting into (hook_type, layer_no) pairs. + + * ``None`` or ``""`` -> ``None``: gate inactive, every hook fires. + * ``"none"`` -> ``[]``: an explicit empty set, every hook off. + * ``"0:1,0:2"`` -> ``[(0, 1), (0, 2)]``. A token without a layer + (``"14"``) is a global hook (layer ``-1``). + """ + if value is None: + return None + text = str(value).strip() + if text == "": + return None + if text.lower() == "none": + return [] + out: List[HookKey] = [] + for token in text.split(","): + token = token.strip() + if not token: + continue + hook_type, _, layer = token.partition(":") + out.append((int(hook_type), int(layer) if layer != "" else -1)) + return out or None + + +# --------------------------------------------------------------------------- +# CUDA-graph helpers +# --------------------------------------------------------------------------- + +class _CaptureWindow: + engine: Any = None + exclude_capture: Optional[Callable[[], bool]] = None + excluded: int = 0 + + +class _Session: + """One monitoring session of one engine: an owner of the process-wide + CUDAGraph patch while it is open. Holds the engine weakly where it can + (an ended session stamped on a long-lived graph must not pin the ring).""" + + _tokens = itertools.count(1) + + def __init__(self, engine: Any) -> None: + self.token = next(self._tokens) + self.engine_id = id(engine) + try: + self._engine_ref: Optional[Callable[[], Any]] = weakref.ref(engine) + except TypeError: + self._engine_ref = None + # raw cudaGraph_t handles bound to the engine in this session. + self.bound: set = set() + + def engine(self) -> Any: + return self._engine_ref() if self._engine_ref is not None else None + + def owns_engine(self, engine: Any) -> bool: + if engine is None: + return False + if self._engine_ref is not None: + return self._engine_ref() is engine + return self.engine_id == id(engine) + + +_WINDOW = _CaptureWindow() +_REPLAY_GUARD = False +# The patch: the class it replaced (restored when the last session ends) and +# the toggle-aware subclass (built once, so graphs of every session share one +# class). _PATCHED_CLASS is set while the patch is installed. +_ORIG_CUDAGRAPH: Optional[type] = None +_TOGGLE_CLASS: Optional[type] = None +_PATCHED_CLASS: Optional[type] = None +# Open monitoring sessions, keyed by id(engine). +_SESSIONS: dict = {} + + +def set_replay_guard(enabled: bool) -> None: + """Arm or disarm the replay-time guard of graphs created by the patched + ``torch.cuda.CUDAGraph`` (see :func:`install_cudagraph_patch`).""" + global _REPLAY_GUARD + _REPLAY_GUARD = bool(enabled) + + +def excluded_capture_count() -> int: + """Captures excluded from the registry since the window opened.""" + return _WINDOW.excluded + + +def _session_for(engine: Any) -> Optional[_Session]: + session = _SESSIONS.get(id(engine)) if engine is not None else None + if session is not None and not session.owns_engine(engine): + return None # a dead engine's id reused by a new object + return session + + +def _start_session(engine: Any) -> _Session: + """Begin a fresh session for `engine`. A new capture window replaces the + engine's previous session (the window clears its registry).""" + session = _Session(engine) + _SESSIONS[id(engine)] = session + return session + + +def end_session(engine: Any) -> None: + """End `engine`'s monitoring session. When no session remains, restore + the original ``torch.cuda.CUDAGraph`` and disarm the replay guard. + Graphs created while patched keep their class; with the guard disarmed + they replay exactly like unpatched graphs.""" + if engine is not None: + _SESSIONS.pop(id(engine), None) + if not _SESSIONS: + restore_cudagraph_patch() + + +def active_session_count() -> int: + """Open monitoring sessions (owners of the CUDAGraph patch).""" + return len(_SESSIONS) + + +def _probe_engine() -> Any: + from .transport import ring as _rt + engine = getattr(_rt.get_active(), "_ring_engine", None) + if engine is not None: + return engine + for session in list(_SESSIONS.values()): + engine = session.engine() + if engine is not None: + return engine + return None + + +def _launch_count(engine: Any) -> Optional[int]: + getter = getattr(engine, "producer_launch_count", None) + if not callable(getter): + return None + return int(getter()) + + +def _node_count(engine: Any) -> Optional[int]: + getter = getattr(engine, "toggle_node_count", None) + if not callable(getter): + return None + return int(getter()) + + +def _weak(engine: Any) -> Callable[[], Any]: + try: + return weakref.ref(engine) + except TypeError: + return lambda: engine + + +def _engine_of(ref: Any) -> Any: + return ref() if ref is not None else None + + +def _current_exec(graph: Any) -> int: + """The graph's exec, instantiating it first when the template is kept and + not instantiated yet (keep_graph=True defers instantiation).""" + try: + return graph.raw_cuda_graph_exec() + except RuntimeError: + graph.instantiate() + return graph.raw_cuda_graph_exec() + + +def _upload_for_device_toggle(graph: Any, engine: Any) -> None: + """Instantiate if needed and upload the graph's exec before its first + launch (device-updatable producers require it), once per exec. Excluded + captures go through here too: uploading registers nothing. The engine + uploads only execs whose graph captured device-updatable producers.""" + exec_ptr = _current_exec(graph) + if exec_ptr and exec_ptr != getattr(graph, "_dmi_uploaded_exec", None): + err = engine.upload_graph_exec(graph.raw_cuda_graph(), exec_ptr) + if err != 0: + raise RuntimeError( + f"node toggle: cudaGraphUpload failed ({err})") + graph._dmi_uploaded_exec = exec_ptr + + +def install_cudagraph_patch() -> type: + """Replace ``torch.cuda.CUDAGraph`` with a toggle-aware subclass. + + The subclass forces ``keep_graph=True`` -- a framework that creates + ``torch.cuda.CUDAGraph()`` with the default frees the captured template + after instantiation, so the node handles recorded during capture would + dangle. Keeping the template costs host memory only. It also + + * stamps each capture with the monitoring session it was captured in and + whether it captured DMI producers; + * excludes a capture from the registry when the capture window's + ``exclude_capture`` predicate says so (node recording is suspended for + that capture and the graph is tagged ``_dmi_toggle_excluded``); + * in device-toggle mode, instantiates and uploads the exec of a capture + that may hold producers before its first launch (required for + device-updatable nodes), excluded captures included, and uploads a + re-instantiated exec before it replays; + * guards replays while :func:`set_replay_guard` is armed and a gate is + active (see ``replay``). + + Idempotent. Install it before the framework creates its graphs; + :func:`bind_graph` refuses graphs of the unpatched class. The patch is + owned by the monitoring sessions (:func:`begin_capture_window`): when the + last one ends, :func:`restore_cudagraph_patch` puts the original class + back. + """ + global _PATCHED_CLASS, _ORIG_CUDAGRAPH, _TOGGLE_CLASS + import torch + + current = torch.cuda.CUDAGraph + if _PATCHED_CLASS is not None and current is _PATCHED_CLASS: + return current + if getattr(current, "_dmi_keep_graph", False): + # Already a keep-graph class (installed by another copy of DMI). + _PATCHED_CLASS = current + return current + if _TOGGLE_CLASS is not None and current is _ORIG_CUDAGRAPH: + torch.cuda.CUDAGraph = _TOGGLE_CLASS + _PATCHED_CLASS = _TOGGLE_CLASS + return _TOGGLE_CLASS + original = current + + class ToggleCUDAGraph(original): + # keep_graph is forced in both __new__ and __init__: CUDAGraph() with + # no arguments would otherwise reset it to False in __init__. A + # subclass (not a factory) keeps isinstance checks working. + _dmi_keep_graph = True + + def __new__(cls, *args, **kwargs): + kwargs["keep_graph"] = True + return original.__new__(cls, **kwargs) + + def __init__(self, *args, **kwargs): + kwargs["keep_graph"] = True + super().__init__(**kwargs) + + def capture_begin(self, *args, **kwargs): + engine = _WINDOW.engine + predicate = _WINDOW.exclude_capture + # Owner: the session whose window is open now (None outside any + # window). Producer detection reads the process-wide + # producer-launch counter (through the window's engine, else the + # active transport's or any session's) and, inside the window, + # the nodes recorded for this graph. + self._dmi_owner = _session_for(engine) + probe = engine if engine is not None else _probe_engine() + self._dmi_probe = _weak(probe) if probe is not None else None + self._dmi_launches0 = _launch_count(probe) + self._dmi_nodes0 = _node_count(engine) + self._dmi_producers = None + self._dmi_toggle_excluded = bool( + engine is not None and predicate is not None and predicate()) + self._dmi_toggle_suspended = False + if self._dmi_toggle_excluded and engine.toggle_capture_enabled(): + engine.enable_toggle_capture(False) + self._dmi_toggle_suspended = True + _WINDOW.excluded += 1 + try: + super().capture_begin(*args, **kwargs) + except BaseException: + # capture_end will not run: restore node recording now, or the + # rest of the window would record nothing. + if self._dmi_toggle_suspended: + self._dmi_toggle_suspended = False + engine.enable_toggle_capture(True) + raise + + def capture_end(self, *args, **kwargs): + engine = _WINDOW.engine + try: + super().capture_end(*args, **kwargs) + finally: + if getattr(self, "_dmi_toggle_suspended", False): + self._dmi_toggle_suspended = False + if engine is not None: + engine.enable_toggle_capture(True) + self._dmi_producers = _captured_producers(self, engine) + if (engine is not None and engine.toggle_device_mode() + and self._dmi_producers is not False): + # keep_graph=True defers instantiation to the first replay, + # which would launch device-updatable producers without the + # required upload. Instantiate now (what an unpatched graph + # does at capture end) and upload; remember the engine so a + # re-instantiated exec is uploaded before it replays. + self._dmi_upload_engine = _weak(engine) + _upload_for_device_toggle(self, engine) + + def replay(self): + upload_engine = _engine_of(getattr(self, "_dmi_upload_engine", None)) + if upload_engine is not None: + _upload_for_device_toggle(self, upload_engine) + if _REPLAY_GUARD: + from .transport import ring as _rt + transport = _rt.get_active() + if (transport is not None and transport.toggle_gate_active + and _guard_applies(self, transport)): + _guarded_replay(self, transport, super().replay) + return + super().replay() + + ToggleCUDAGraph.__name__ = "CUDAGraph" + ToggleCUDAGraph.__qualname__ = "CUDAGraph" + torch.cuda.CUDAGraph = ToggleCUDAGraph + _ORIG_CUDAGRAPH = original + _TOGGLE_CLASS = ToggleCUDAGraph + _PATCHED_CLASS = ToggleCUDAGraph + return ToggleCUDAGraph + + +def restore_cudagraph_patch() -> None: + """Put the original ``torch.cuda.CUDAGraph`` back and disarm the replay + guard. Called when the last monitoring session ends; graphs created + while patched keep their class, and their guard stays inert.""" + global _PATCHED_CLASS, _REPLAY_GUARD + _REPLAY_GUARD = False + if _PATCHED_CLASS is None: + return + import torch + + if (_ORIG_CUDAGRAPH is not None and _PATCHED_CLASS is _TOGGLE_CLASS + and torch.cuda.CUDAGraph is _PATCHED_CLASS): + torch.cuda.CUDAGraph = _ORIG_CUDAGRAPH + _PATCHED_CLASS = None + + +def _captured_producers(graph: Any, window_engine: Any) -> Optional[bool]: + """Whether the capture that just ended recorded DMI producers. + + True or False when it can be told, None when it cannot (callers treat + None like True). Two independent sources of evidence, and positive + evidence from either one wins: + + * the producer-launch counter, sampled across the capture, answers for + any capture, excluded or outside the window; + * the producer nodes the registry recorded for this graph (inside the + window only; an excluded capture records none). + + An unchanged counter never overrides recorded nodes. Without positive + evidence, the answer is False when either source could answer (a capture + inside the window that recorded no node has no toggle-managed producer; + an unmanaged one is a capture anomaly, which already blocks activation), + and None when neither could. + """ + counted: Optional[bool] = None + probe = _engine_of(getattr(graph, "_dmi_probe", None)) + before = getattr(graph, "_dmi_launches0", None) + if probe is not None and before is not None: + after = _launch_count(probe) + if after is not None: + if after != before: + return True + counted = False + recorded = _recorded_producer_nodes(graph, window_engine) + if recorded: + return True + if counted is False: + return False + return recorded + + +def _recorded_producer_nodes(graph: Any, window_engine: Any) -> Optional[bool]: + """Whether the registry recorded producer nodes for this capture; None + when it cannot tell (no window, or an excluded capture).""" + if window_engine is None or getattr(graph, "_dmi_toggle_excluded", False): + return None + counter = getattr(window_engine, "toggle_node_count_for_graph", None) + if counter is not None: + return int(counter(graph.raw_cuda_graph())) > 0 + nodes0 = getattr(graph, "_dmi_nodes0", None) + nodes1 = _node_count(window_engine) + if nodes0 is None or nodes1 is None: + return None + return nodes1 != nodes0 + + +def _guard_applies(graph: Any, transport: Any) -> bool: + """Whether the active transport's gate governs this graph's replay. + + No for a graph stamped with another engine's session (a later session + never adopts it) and for a graph that captured no DMI producer and is not + bound to the transport's engine -- an unrelated CUDA-graph user in the + process. Yes, so the fail-closed checks run, for every bound graph and + every graph that captured producers or may have. + """ + engine = getattr(transport, "_ring_engine", None) + owner = getattr(graph, "_dmi_owner", None) + if owner is not None and engine is not None and not owner.owns_engine(engine): + return False + session = _session_for(engine) + if session is not None and session.bound and graph.raw_cuda_graph() in session.bound: + return True + return getattr(graph, "_dmi_producers", None) is not False + + +def _guarded_replay(graph: Any, transport: Any, replay: Callable[[], None]) -> None: + """Replay with the toggle's consistency checks. Every failure here is + fatal: this step's metadata has already been published, so replaying + anyway would desynchronize the ring.""" + if getattr(graph, "_dmi_toggle_excluded", False): + # Excluded from the registry by design: its producers run all on, + # which is right exactly on an ungated step, where the full active + # set's metadata was pushed. + if transport.gated_step_requires_registry(): + raise RuntimeError( + "node toggle: a graph excluded from the toggle registry is " + "about to replay on a step classified as gated. Its producers " + "run all on while the metadata gate published only the " + "enabled subset, so the ring would desynchronize. The " + "integration's step classification disagrees with the " + "framework's graph dispatch. The worker must terminate.") + replay() + return + raw = graph.raw_cuda_graph() + rebound = _rebind_if_reinstantiated(graph, transport, raw) + # A graph captured after the capture window closed has no recorded nodes + # and no bound exec: its producers run default-on while the gate filters. + if not transport.is_graph_ready(raw): + raise RuntimeError( + f"node toggle: graph {raw:#x} is about to replay but is not " + "registered and bound in the toggle registry -- it was most " + "likely captured after the capture window closed. Its producers " + "would run all on while the metadata gate filters to the enabled " + "subset. The worker must terminate.") + if rebound and not transport.toggle_lazy_active: + # Eager mode applies only at set_active_hooks time, to the execs bound + # then. A freshly bound exec starts with every producer node in its + # instantiation default (on), so bring it to the enabled set now, + # before its first replay. Nothing of the new exec is in flight yet. + err = transport.ensure_graph_current(raw) + if err != 0: + raise RuntimeError( + f"node toggle: applying the enabled set to re-instantiated " + f"graph {raw:#x} failed with CUDA error {err}; this step's " + "metadata is already published. The worker must terminate.") + if transport.toggle_lazy_active: + err = transport.ensure_graph_current(raw) + if err != 0: + raise RuntimeError( + f"node toggle: lazy apply failed with CUDA error {err} for " + f"graph {raw:#x}; this step's metadata is already published. " + "The worker must terminate.") + replay() + err = transport.record_replay_event(raw) + if err != 0: + raise RuntimeError( + f"node toggle: record_replay_event failed with CUDA error " + f"{err} for graph {raw:#x}. The worker must terminate.") + return + replay() + + +def _rebind_if_reinstantiated(graph: Any, transport: Any, raw: int) -> bool: + """A bound graph re-instantiated since bind_graph has a new exec: bind it + before it replays. The registry then resets that graph's per-exec state, + so the next ensure_graph_current sets every node of the new exec. + + Returns True when a new exec was bound. A graph is bound to exactly one + engine (``_dmi_bound_engine``), so at most one binding changes here.""" + bound = getattr(graph, "_dmi_bound_exec", None) + if bound is None: + return False + engine = getattr(transport, "_ring_engine", None) + if engine is None or _engine_of(getattr(graph, "_dmi_bound_engine", None)) is not engine: + return False + exec_ptr = _current_exec(graph) + if exec_ptr == bound: + return False + engine.bind_graph_exec(raw, exec_ptr) + graph._dmi_bound_exec = exec_ptr + graph._dmi_uploaded_exec = exec_ptr + return True + + +def begin_capture_window( + engine: Any, + *, + exclude_capture: Optional[Callable[[], bool]] = None, + device_toggle: bool = False, +) -> None: + """Open the capture window: clear any stale registry, record producer + nodes from now on, start the engine's monitoring session and install the + CUDA-graph patch it owns. + + ``exclude_capture`` is called at the start of every capture; returning + True keeps that capture out of the registry (for example a framework's + partial-model segment graphs, which would make the registry non-uniform). + ``device_toggle`` asks the engine for device-side toggling; it stays off + when the device or build lacks device graph updates. + """ + engine.clear_toggle_registry() + # Explicit either way: the mode is a per-window choice. + engine.enable_device_toggle(bool(device_toggle)) + engine.enable_toggle_capture(True) + _WINDOW.engine = engine + _WINDOW.exclude_capture = exclude_capture + _WINDOW.excluded = 0 + _start_session(engine) + install_cudagraph_patch() + + +def end_capture_window(engine: Any) -> None: + """Close the capture window (bind graphs before or after, not during a + capture). A window closed with no graph bound ends the engine's session + (restoring the CUDAGraph patch if it was the last); binding a graph + afterwards renews it.""" + engine.enable_toggle_capture(False) + _forget_window(engine) + session = _session_for(engine) + if session is not None and not session.bound: + end_session(engine) + + +@contextlib.contextmanager +def capture_window( + engine: Any, + *, + exclude_capture: Optional[Callable[[], bool]] = None, + device_toggle: bool = False, +) -> Iterator[None]: + """``begin_capture_window`` ... ``end_capture_window``, closed also when + the body raises (so the window never keeps a reference to the engine).""" + begin_capture_window(engine, exclude_capture=exclude_capture, + device_toggle=device_toggle) + try: + yield + finally: + end_capture_window(engine) + + +def _forget_window(engine: Any) -> None: + """Drop the capture window's reference to `engine` (its ring buffers stay + pinned while the window holds it).""" + if _WINDOW.engine is engine: + _WINDOW.engine = None + _WINDOW.exclude_capture = None + + +def bind_graph(engine: Any, graph: Any) -> None: + """Bind one captured ``torch.cuda.CUDAGraph`` to the engine's registry. + + Instantiates the graph if it has no exec yet. Refuses a graph created + from an unpatched ``torch.cuda.CUDAGraph``: its template was freed at + instantiation, so the recorded node handles dangle. + """ + if not getattr(type(graph), "_dmi_keep_graph", False): + raise RuntimeError( + "node toggle: the graph was created from an unpatched " + "torch.cuda.CUDAGraph (keep_graph=False), so its template is " + "already freed and the recorded producer-node handles dangle. " + "Call install_cudagraph_patch() (or begin_capture_window) before " + "the framework creates its graphs.") + if getattr(graph, "_dmi_toggle_excluded", False): + raise RuntimeError( + "node toggle: refusing to bind a capture excluded from the " + "registry") + exec_ptr = _current_exec(graph) + raw = graph.raw_cuda_graph() + engine.bind_graph_exec(raw, exec_ptr) # uploads a device-toggle exec + graph._dmi_bound_exec = exec_ptr + graph._dmi_bound_engine = _weak(engine) + graph._dmi_uploaded_exec = exec_ptr + session = _session_for(engine) + if session is None: + # Bound after a window that closed with nothing bound: renew the + # session the graph was captured in (or a new one). + owner = getattr(graph, "_dmi_owner", None) + if owner is not None and owner.owns_engine(engine): + _SESSIONS[id(engine)] = session = owner + else: + session = _start_session(engine) + install_cudagraph_patch() + session.bound.add(raw) + + +def activate( + transport: Any, + enabled_hooks: Optional[Iterable[HookKey]], + *, + lazy: bool, +) -> None: + """Apply a static enabled set after capture and arm the replay guard. + + ``enabled_hooks=None`` leaves the gate inactive (every hook fires). A + requested set with no bound graph is refused rather than served with the + toggle silently inert. + """ + if enabled_hooks is None: + return + enabled = list(enabled_hooks) + if transport._ring_engine.bound_graph_count() == 0: + raise RuntimeError( + "node toggle: an enabled-hook set was requested but no captured " + "graph is bound, so the toggle would be silently inert.") + set_replay_guard(True) + if lazy: + transport.set_active_hooks_lazy(enabled) + else: + transport.set_active_hooks(enabled) + + +__all__ = [ + "NodeToggleController", + "activate", + "active_session_count", + "begin_capture_window", + "bind_graph", + "capture_window", + "end_capture_window", + "end_session", + "excluded_capture_count", + "install_cudagraph_patch", + "parse_enabled_hooks", + "restore_cudagraph_patch", + "set_replay_guard", +] diff --git a/src/dmi/transport/ring.py b/src/dmi/transport/ring.py index 16eb4305a..4838d68db 100644 --- a/src/dmi/transport/ring.py +++ b/src/dmi/transport/ring.py @@ -7,7 +7,7 @@ """ from __future__ import annotations -from typing import Any, List, Optional, Tuple +from typing import TYPE_CHECKING, Any, Iterable, List, Optional, Tuple import torch import torch.library @@ -25,6 +25,9 @@ compute_hook_shape, ) +if TYPE_CHECKING: + from ..node_toggle import NodeToggleController + # --------------------------------------------------------------------------- # Module-level active transport @@ -249,6 +252,20 @@ def __init__(self, ring_engine: Any) -> None: self._active_specs: List[HookSpec] = [] self._using_forward_hooks: bool = False + # Runtime node-toggle controller (dmi.node_toggle), created on first + # use. None means the toggle is inactive: every active spec fires + # every step. + self._toggle: "Optional[NodeToggleController]" = None + # Per-step gate selector. The toggle (cudaGraphNodeSetEnabled on + # captured producer nodes) governs only producers that run inside a + # replayed graph. On an eager or prefill step producers run + # unconditionally, so metadata publication and capacity reservation + # must use the full active set that step. The integration sets this + # per step (True: a toggled graph replays, the gate applies; False: + # the full active set fires). Default True so a caller that only + # replays toggled graphs keeps gating. + self._gated_step: bool = True + # Hook selection preset name (e.g. "full", "hidden-states", "logits"). # Set by the active adapter before hook installation. self._hook_selection: Optional[str] = None @@ -320,6 +337,81 @@ def set_model_cfg(self, cfg: ModelShapeConfig) -> None: """Set the model shape config for analytical shape computation.""" self._model_cfg = cfg + # ------------------------------------------------------------------ + # Runtime node toggling (see dmi.node_toggle) + + @property + def effective_specs(self) -> "List[HookSpec]": + """The specs that fire on a gated step: the enabled-and-captured + subset while the toggle gate is active, else every active spec.""" + toggle = getattr(self, "_toggle", None) + if (toggle is not None and toggle.gate_active + and toggle.effective_enabled_specs is not None): + return toggle.effective_enabled_specs + return self._active_specs + + def specs_for_step(self) -> "List[HookSpec]": + """The spec set that drives capacity reservation and metadata + publication THIS step: ``effective_specs`` on a gated step, the full + active set on an ungated (eager or prefill) step. Identical when the + toggle is inactive.""" + if getattr(self, "_gated_step", True): + return self.effective_specs + return self._active_specs + + def set_gated_step(self, gated: bool) -> None: + """Declare whether this step replays a toggled graph. + + Integrations call this before planning each step: True when the + step's producers run inside a toggle-managed graph (the enabled + subset fires), False on an eager or partial-graph step (the full + active set fires).""" + self._gated_step = bool(gated) + + def gated_step_requires_registry(self) -> bool: + """True when this step's producers must all be toggle-managed, i.e. + a graph excluded from the registry must not replay.""" + return bool(getattr(self, "_gated_step", True)) + + def _toggle_ctrl(self) -> "NodeToggleController": + if self._toggle is None: + from ..node_toggle import NodeToggleController + self._toggle = NodeToggleController(self._ring_engine, self) + return self._toggle + + @property + def toggle_gate_active(self) -> bool: + toggle = getattr(self, "_toggle", None) + return toggle is not None and toggle.gate_active + + @property + def toggle_lazy_active(self) -> bool: + toggle = getattr(self, "_toggle", None) + return toggle is not None and toggle.lazy_active + + def set_active_hooks(self, enabled: "Iterable[Tuple[int, int]]") -> None: + """Eager node-toggle reconfiguration (quiescent points only). See + ``NodeToggleController.set_active_hooks``.""" + self._toggle_ctrl().set_active_hooks(enabled) + + def set_active_hooks_lazy(self, enabled: "Iterable[Tuple[int, int]]") -> None: + """Lazy node-toggle reconfiguration (device apply deferred per + graph). See ``NodeToggleController.set_active_hooks_lazy``.""" + self._toggle_ctrl().set_active_hooks_lazy(enabled) + + def is_graph_ready(self, raw_graph: int) -> bool: + return self._toggle_ctrl().is_graph_ready(raw_graph) + + def ensure_graph_current(self, raw_graph: int) -> int: + return self._toggle_ctrl().ensure_graph_current(raw_graph) + + def record_replay_event(self, raw_graph: int) -> int: + return self._toggle_ctrl().record_replay_event(raw_graph) + + def clear_toggle(self) -> None: + """Clear the registry and deactivate the host gate together.""" + self._toggle_ctrl().clear() + def pre_push_all_metas(self, batch: int, q_len: int, kv_dim: int, logits_to_keep: int = 0, token_ids_dtype: Optional[torch.dtype] = None, @@ -353,7 +445,9 @@ def pre_push_all_metas(self, batch: int, q_len: int, kv_dim: int, shapes = [] dtypes = [] flags = [] - for spec in self._active_specs: + # The per-step spec set: the metadata published here, the device + # node state and the capacity reservation stay in lockstep. + for spec in self.specs_for_step(): spec_q_len = (actual_q_len if actual_q_len is not None and spec.dim0_is_actual_tokens else q_len) @@ -569,7 +663,16 @@ def activate(transport: RingTransport) -> None: def deactivate() -> None: global _active_transport + transport = _active_transport _active_transport = None + if transport is not None: + # A node-toggle capture window must not keep the engine alive, and + # tearing the transport down ends its engine's monitoring session + # (restoring torch.cuda.CUDAGraph once no session remains). + from .. import node_toggle + engine = getattr(transport, "_ring_engine", None) + node_toggle._forget_window(engine) + node_toggle.end_session(engine) try: from . import native as _ne _ne.ring_clear_active_engine() diff --git a/tests/native/ring/Makefile b/tests/native/ring/Makefile index 5c7018b8b..ee32064eb 100644 --- a/tests/native/ring/Makefile +++ b/tests/native/ring/Makefile @@ -62,14 +62,17 @@ RECORD_CONSUMER_SRC := ../../../native/csrc/ring/record_consumer.cpp CLICKHOUSE_RECORD_SINK_SRC := ../../../native/csrc/clickhouse_record_sink.cpp RING_ENGINE_CU := ../../../native/csrc/ring/ring_engine.cu RING_ENGINE_PY_CU := ../../../native/csrc/ring/ring_engine_py.cu +TOGGLE_REGISTRY_CU := ../../../native/csrc/ring/toggle_registry.cu TARGETS := $(BUILD)/test_rings $(BUILD)/test_producer \ $(BUILD)/test_ring_engine $(BUILD)/test_null_mode \ $(BUILD)/test_record_consumer \ - $(BUILD)/test_clickhouse_record_sink + $(BUILD)/test_clickhouse_record_sink \ + $(BUILD)/test_toggle_registry .PHONY: all test_rings test_producer test_ring_engine test_null_mode \ - test_record_consumer test_clickhouse_record_sink run clean FORCE \ + test_record_consumer test_clickhouse_record_sink test_toggle_registry \ + run clean FORCE \ bench_producer_grid all: $(TARGETS) @@ -90,11 +93,14 @@ $(BUILD)/test_producer: test_producer.cu $(PRODUCER_CU) $(CUDA_CONFIG_STAMP) | $ $(BUILD)/test_ring_engine: test_ring_engine.cu $(PRODUCER_CU) $(DRAIN_SRC) \ $(P2P_SRC) $(RECORD_CONSUMER_SRC) $(RING_ENGINE_CU) \ - $(RING_ENGINE_PY_CU) $(CUDA_CONFIG_STAMP) | $(BUILD) + $(RING_ENGINE_PY_CU) $(TOGGLE_REGISTRY_CU) $(CUDA_CONFIG_STAMP) | $(BUILD) $(CUDA_NVCC) $(NVCC_FLAGS) $(TORCH_INCLUDE_FLAGS) $(TORCH_LIB_FLAGS) \ $(filter-out $(CUDA_CONFIG_STAMP),$^) $(TORCH_RPATH_FLAGS) \ -ltorch -ltorch_cpu -ltorch_cuda -lc10 -lc10_cuda -lpthread -o $@ +$(BUILD)/test_toggle_registry: test_toggle_registry.cu $(TOGGLE_REGISTRY_CU) $(CUDA_CONFIG_STAMP) | $(BUILD) + $(CUDA_NVCC) $(NVCC_FLAGS) $(filter-out $(CUDA_CONFIG_STAMP),$^) -o $@ + $(BUILD)/test_null_mode: test_null_mode.cu $(PRODUCER_CU) $(CUDA_CONFIG_STAMP) | $(BUILD) $(CUDA_NVCC) $(NVCC_FLAGS) $(filter-out $(CUDA_CONFIG_STAMP),$^) -o $@ @@ -126,6 +132,9 @@ test_record_consumer: $(BUILD)/test_record_consumer test_clickhouse_record_sink: $(BUILD)/test_clickhouse_record_sink $(BUILD)/test_clickhouse_record_sink +test_toggle_registry: $(BUILD)/test_toggle_registry + $(BUILD)/test_toggle_registry + # Manual microbenchmark; not part of `all` or `run`. $(BUILD)/bench_producer_grid: bench_producer_grid.cu $(PRODUCER_CU) $(CUDA_CONFIG_STAMP) | $(BUILD) $(CUDA_NVCC) $(NVCC_FLAGS) bench_producer_grid.cu -o $@ @@ -140,6 +149,7 @@ run: all $(BUILD)/test_null_mode $(BUILD)/test_record_consumer $(BUILD)/test_clickhouse_record_sink + $(BUILD)/test_toggle_registry clean: rm -rf $(BUILD) diff --git a/tests/native/ring/cpu_cuda_stub/cuda_runtime.h b/tests/native/ring/cpu_cuda_stub/cuda_runtime.h new file mode 100644 index 000000000..08f6e8e91 --- /dev/null +++ b/tests/native/ring/cpu_cuda_stub/cuda_runtime.h @@ -0,0 +1,71 @@ +// CPU-only stand-in for the CUDA runtime API used by the host half of +// ring/toggle_registry.cu. No CUDA library is linked or initialized: each +// call records what a real exec would hold so tests can compare the +// registry's view with the exec's actual node state. +#pragma once +#include +#include +#include +#include +#include + +using cudaError_t = int; +using cudaGraph_t = void*; +using cudaGraphNode_t = void*; +using cudaGraphExec_t = void*; +using cudaStream_t = void*; +using cudaEvent_t = void*; +using cudaStreamCaptureStatus = int; +constexpr int cudaSuccess = 0; +constexpr unsigned cudaStreamNonBlocking = 1; +constexpr unsigned cudaEventDisableTiming = 2; + +// (exec, node) -> enabled. A node missing from the map is enabled: CUDA +// instantiates every kernel node enabled. +inline std::map, bool> fake_enabled; +inline unsigned fake_set_enabled_calls = 0; +inline int fake_set_enabled_error = 0; +inline std::multiset fake_uploads; +inline unsigned fake_events_live = 0; + +inline bool fake_node_on(uint64_t exec, uint64_t node) { + auto it = fake_enabled.find({exec, node}); + return it == fake_enabled.end() ? true : it->second; +} + +inline int cudaGraphNodeSetEnabled(cudaGraphExec_t ex, cudaGraphNode_t node, + unsigned on) { + ++fake_set_enabled_calls; + if (fake_set_enabled_error) return fake_set_enabled_error; + fake_enabled[{reinterpret_cast(ex), + reinterpret_cast(node)}] = on != 0; + return 0; +} +inline int cudaStreamGetCaptureInfo(cudaStream_t, cudaStreamCaptureStatus*, + unsigned long long*, cudaGraph_t*, + const cudaGraphNode_t**, size_t*) { + return 0; +} +inline int cudaMallocManaged(void** p, size_t n) { + *p = std::malloc(n); + return *p ? 0 : 1; +} +inline int cudaFree(void* p) { std::free(p); return 0; } +inline int cudaStreamCreateWithFlags(cudaStream_t* p, unsigned) { + *p = reinterpret_cast(1); + return 0; +} +inline int cudaGraphUpload(cudaGraphExec_t ex, cudaStream_t) { + fake_uploads.insert(reinterpret_cast(ex)); + return 0; +} +inline int cudaStreamSynchronize(cudaStream_t) { return 0; } +inline int cudaStreamDestroy(cudaStream_t) { return 0; } +inline int cudaEventSynchronize(cudaEvent_t) { return 0; } +inline int cudaEventCreateWithFlags(cudaEvent_t* p, unsigned) { + ++fake_events_live; + *p = reinterpret_cast(2); + return 0; +} +inline int cudaEventRecord(cudaEvent_t, cudaStream_t) { return 0; } +inline int cudaEventDestroy(cudaEvent_t) { --fake_events_live; return 0; } diff --git a/tests/native/ring/test_ring_engine.cu b/tests/native/ring/test_ring_engine.cu index 20c7efd13..fc905b62f 100644 --- a/tests/native/ring/test_ring_engine.cu +++ b/tests/native/ring/test_ring_engine.cu @@ -52,6 +52,10 @@ public: // executable; the extension supplies their real definitions. void ring_diag_reset_host_counters() {} void ring_diag_print_host_counters() {} +void ring_set_toggle_capture(bool) {} +// The producer-launch counter lives with the producer ops in +// ring_torch_op.cpp, which this executable does not link. +uint64_t ring_producer_launch_count() { return 0; } #define CUDA_CHECK(expr) \ do { \ diff --git a/tests/native/ring/test_toggle_registry.cu b/tests/native/ring/test_toggle_registry.cu new file mode 100644 index 000000000..d521a5ea2 --- /dev/null +++ b/tests/native/ring/test_toggle_registry.cu @@ -0,0 +1,472 @@ +// Native tests for ToggleRegistry: capture-time node recording, eager and +// lazy cudaGraphNodeSetEnabled application, registry guards, teardown, and +// the device-side planner (toggle_device.h). +// +// Each "producer" here is a one-block marker kernel that adds its value to +// its own output slot, so a replay's effect shows exactly which nodes ran. + +#include "ring/toggle_registry.h" + +#include + +#include +#include +#include +#include +#include +#include + +static int g_pass = 0; +static int g_fail = 0; + +#define CUDA_CHECK(expr) \ + do { \ + cudaError_t error_ = (expr); \ + if (error_ != cudaSuccess) { \ + std::fprintf(stderr, "CUDA error at %s:%d: %s\n", \ + __FILE__, __LINE__, cudaGetErrorString(error_)); \ + std::exit(1); \ + } \ + } while (0) + +#define EXPECT(condition) \ + do { \ + if (!(condition)) { \ + std::fprintf(stderr, "FAIL %s:%d: %s\n", \ + __FILE__, __LINE__, #condition); \ + ++g_fail; \ + } else { \ + ++g_pass; \ + } \ + } while (0) + +static void banner(const char* name) { std::printf("[ TEST ] %s\n", name); } + +__global__ void marker_kernel(int* out, int slot, int value) { + if (threadIdx.x == 0) out[slot] += value; +} + +// A marker that honours the control-row contract of the producers: a -1 row +// returns without writing. +__global__ void gated_marker_kernel(int* out, int slot, int value, + const int64_t* control_row) { + if (control_row[0] == -1) return; + if (threadIdx.x == 0) out[slot] += value; +} + +static constexpr int kHookType = 3; + +struct Graph { + cudaGraph_t graph{nullptr}; + cudaGraphExec_t exec{nullptr}; + ~Graph() { + if (exec) cudaGraphExecDestroy(exec); + if (graph) cudaGraphDestroy(graph); + } +}; + +// Record the capture's tail dependency exactly as the producer ops do. +static uint64_t tail_node(cudaStream_t stream, cudaGraph_t* graph_out) { + cudaStreamCaptureStatus status; + unsigned long long id; + cudaGraph_t graph; + const cudaGraphNode_t* deps; + size_t n; + CUDA_CHECK(dmx_capture_info(stream, &status, &id, &graph, &deps, &n)); + if (status != cudaStreamCaptureStatusActive || n < 1) return 0; + *graph_out = graph; + return reinterpret_cast(deps[n - 1]); +} + +// Capture one marker node per layer in `layers`; marker for layer L adds +// base + L to out[L]. +static void capture_markers(ring_py::ToggleRegistry& registry, Graph& g, + cudaStream_t stream, int* out, + const std::vector& layers, int base) { + CUDA_CHECK(cudaStreamBeginCapture(stream, cudaStreamCaptureModeThreadLocal)); + for (int layer : layers) { + marker_kernel<<<1, 32, 0, stream>>>(out, layer, base + layer); + cudaGraph_t graph = nullptr; + const uint64_t node = tail_node(stream, &graph); + EXPECT(node != 0); + registry.register_node(reinterpret_cast(graph), kHookType, + layer, node); + } + CUDA_CHECK(cudaStreamEndCapture(stream, &g.graph)); + CUDA_CHECK(cudaGraphInstantiate(&g.exec, g.graph, 0)); +} + +static std::vector read_and_clear(int* out, int n) { + std::vector host(n); + CUDA_CHECK(cudaMemcpy(host.data(), out, n * sizeof(int), cudaMemcpyDeviceToHost)); + CUDA_CHECK(cudaMemset(out, 0, n * sizeof(int))); + return host; +} + +static std::vector> keys(const std::vector& layers) { + std::vector> result; + for (int layer : layers) result.push_back({kHookType, layer}); + return result; +} + +// Expected marker output: enabled layers carry base + layer, others 0. +static std::vector expected(int n, const std::vector& enabled, int base) { + std::vector result(n, 0); + for (int layer : enabled) result[layer] = base + layer; + return result; +} + +static void replay(const Graph& g, cudaStream_t stream) { + CUDA_CHECK(cudaGraphLaunch(g.exec, stream)); + CUDA_CHECK(cudaStreamSynchronize(stream)); +} + +static void test_eager_apply_and_diff() { + banner("eager apply disables exactly the unselected nodes, diff-based"); + constexpr int n = 6; + int* out = nullptr; + CUDA_CHECK(cudaMalloc(&out, n * sizeof(int))); + CUDA_CHECK(cudaMemset(out, 0, n * sizeof(int))); + cudaStream_t stream; + CUDA_CHECK(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking)); + { + ring_py::ToggleRegistry registry; + registry.set_capture(true); + Graph g; + capture_markers(registry, g, stream, out, {0, 1, 2, 3, 4, 5}, 10); + registry.bind_exec(reinterpret_cast(g.graph), + reinterpret_cast(g.exec)); + registry.set_capture(false); + EXPECT(registry.node_count() == 6); + EXPECT(registry.bound_count() == 1); + EXPECT(registry.complete()); + EXPECT(registry.uniform()); + EXPECT(registry.anomaly_count() == 0); + // Inactive gate: every hook counts as enabled. + EXPECT(registry.is_enabled(kHookType, 4)); + + registry.set_enabled(keys({1, 4})); + EXPECT(registry.apply_all() == 0); + EXPECT(registry.last_apply_count() == 4); // four nodes disabled + replay(g, stream); + EXPECT(read_and_clear(out, n) == expected(n, {1, 4}, 10)); + EXPECT(registry.is_enabled(kHookType, 4)); + EXPECT(!registry.is_enabled(kHookType, 0)); + + registry.set_enabled(keys({1, 2, 4})); + EXPECT(registry.apply_all() == 0); + EXPECT(registry.last_apply_count() == 1); // only node 2 flips + replay(g, stream); + EXPECT(read_and_clear(out, n) == expected(n, {1, 2, 4}, 10)); + + // Enabled but never captured -> not effective (no payload behind it). + registry.set_enabled(keys({2}) ); + std::vector> query = {{kHookType, 2}, {kHookType, 9}, {kHookType + 1, 2}}; + EXPECT(registry.effective_mask(query) == std::vector({1, 0, 0})); + } + CUDA_CHECK(cudaStreamDestroy(stream)); + CUDA_CHECK(cudaFree(out)); +} + +static void test_lazy_per_graph_apply() { + banner("lazy apply brings only the graph about to replay up to date"); + constexpr int n = 4; + int* out = nullptr; + CUDA_CHECK(cudaMalloc(&out, n * sizeof(int))); + CUDA_CHECK(cudaMemset(out, 0, n * sizeof(int))); + cudaStream_t stream; + CUDA_CHECK(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking)); + { + ring_py::ToggleRegistry registry; + registry.set_capture(true); + Graph a, b; + capture_markers(registry, a, stream, out, {0, 1, 2, 3}, 0); + capture_markers(registry, b, stream, out, {0, 1, 2, 3}, 100); + const uint64_t ga = reinterpret_cast(a.graph); + const uint64_t gb = reinterpret_cast(b.graph); + registry.bind_exec(ga, reinterpret_cast(a.exec)); + registry.bind_exec(gb, reinterpret_cast(b.exec)); + registry.set_capture(false); + EXPECT(registry.uniform()); + EXPECT(registry.complete()); + + // A sequence of reconfigurations; each graph is only brought current + // when it is about to replay, possibly several versions late. + const std::vector> sets = {{0, 2}, {1, 3}, {0, 1, 2, 3}, {2}, {}}; + const std::vector> order = { + {'B', 'B', 'A'}, {'A', 'B'}, {}, {'A', 'B'}, {'A', 'B'}}; + for (size_t i = 0; i < sets.size(); ++i) { + registry.set_enabled(keys(sets[i])); + EXPECT(!registry.applied_current(ga)); + EXPECT(!registry.applied_current(gb)); + for (char which : order[i]) { + const Graph& g = which == 'A' ? a : b; + const uint64_t raw = which == 'A' ? ga : gb; + EXPECT(registry.ensure_current(raw) == 0); + EXPECT(registry.applied_current(raw)); + EXPECT(registry.record_replay_event(raw, stream) == 0); + replay(g, stream); + EXPECT(read_and_clear(out, n) == + expected(n, sets[i], which == 'A' ? 0 : 100)); + } + } + // Unknown graphs are skipped, not failed. + EXPECT(registry.ensure_current(0x1234) == 0); + } + CUDA_CHECK(cudaStreamDestroy(stream)); + CUDA_CHECK(cudaFree(out)); +} + +static void test_registry_guards() { + banner("uniform/complete/ready guards and anomalies"); + constexpr int n = 4; + int* out = nullptr; + CUDA_CHECK(cudaMalloc(&out, n * sizeof(int))); + cudaStream_t stream; + CUDA_CHECK(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking)); + { + ring_py::ToggleRegistry registry; + registry.set_capture(true); + const uint64_t v0 = registry.version(); + Graph a, b; + capture_markers(registry, a, stream, out, {0, 1, 2}, 0); + capture_markers(registry, b, stream, out, {0, 1}, 0); + EXPECT(registry.version() > v0); + const uint64_t ga = reinterpret_cast(a.graph); + const uint64_t gb = reinterpret_cast(b.graph); + EXPECT(!registry.uniform()); // graph b lacks layer 2 + registry.bind_exec(ga, reinterpret_cast(a.exec)); + EXPECT(!registry.complete()); // b captured, not bound + EXPECT(registry.is_ready(ga)); + EXPECT(!registry.is_ready(gb)); + registry.bind_exec(gb, reinterpret_cast(b.exec)); + EXPECT(registry.complete()); + EXPECT(registry.is_ready(gb)); + const uint64_t v1 = registry.version(); + registry.note_anomaly(); + EXPECT(registry.anomaly_count() == 1); + EXPECT(registry.version() > v1); + } + CUDA_CHECK(cudaStreamDestroy(stream)); + CUDA_CHECK(cudaFree(out)); +} + +static void test_failed_lazy_apply_leaves_graph_stale() { + banner("a failed lazy apply mutates nothing and leaves the graph stale"); + constexpr int n = 3; + int* out = nullptr; + CUDA_CHECK(cudaMalloc(&out, n * sizeof(int))); + CUDA_CHECK(cudaMemset(out, 0, n * sizeof(int))); + cudaStream_t stream; + CUDA_CHECK(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking)); + { + ring_py::ToggleRegistry registry; + registry.set_capture(true); + Graph g; + capture_markers(registry, g, stream, out, {0, 1, 2}, 5); + const uint64_t raw = reinterpret_cast(g.graph); + registry.bind_exec(raw, reinterpret_cast(g.exec)); + registry.set_enabled(keys({1})); + registry.force_apply_error(static_cast(cudaErrorInvalidValue)); + EXPECT(registry.ensure_current(raw) == static_cast(cudaErrorInvalidValue)); + EXPECT(!registry.applied_current(raw)); + replay(g, stream); // nodes untouched: all still enabled + EXPECT(read_and_clear(out, n) == expected(n, {0, 1, 2}, 5)); + registry.force_apply_error(0); + EXPECT(registry.ensure_current(raw) == 0); + EXPECT(registry.applied_current(raw)); + replay(g, stream); + EXPECT(read_and_clear(out, n) == expected(n, {1}, 5)); + } + CUDA_CHECK(cudaStreamDestroy(stream)); + CUDA_CHECK(cudaFree(out)); +} + +static void test_upload_skips_graphs_without_device_nodes() { + banner("graphs without device-updatable producers are never uploaded"); + int* out = nullptr; + CUDA_CHECK(cudaMalloc(&out, sizeof(int))); + cudaStream_t stream; + CUDA_CHECK(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking)); + { + ring_py::ToggleRegistry registry; + Graph g; + capture_markers(registry, g, stream, out, {0}, 0); + // Not a device graph: nothing to upload (the upload would only cost a + // stream create and a synchronize). + EXPECT(registry.upload_device_graph(reinterpret_cast(g.graph), + reinterpret_cast(g.exec)) == 0); + } + CUDA_CHECK(cudaStreamDestroy(stream)); + CUDA_CHECK(cudaFree(out)); +} + +static void test_pre_apply_sync_and_clear() { + banner("pre-apply wait records replay events; clear re-enables nodes"); + constexpr int n = 3; + int* out = nullptr; + CUDA_CHECK(cudaMalloc(&out, n * sizeof(int))); + CUDA_CHECK(cudaMemset(out, 0, n * sizeof(int))); + cudaStream_t stream; + CUDA_CHECK(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking)); + { + ring_py::ToggleRegistry registry; + registry.set_pre_apply_sync(true); + EXPECT(registry.pre_apply_sync()); + registry.set_capture(true); + Graph g; + capture_markers(registry, g, stream, out, {0, 1, 2}, 1); + const uint64_t raw = reinterpret_cast(g.graph); + registry.bind_exec(raw, reinterpret_cast(g.exec)); + registry.set_enabled(keys({0})); + EXPECT(registry.ensure_current(raw) == 0); + CUDA_CHECK(cudaGraphLaunch(g.exec, stream)); + EXPECT(registry.record_replay_event(raw, stream) == 0); + registry.set_enabled(keys({2})); + EXPECT(registry.ensure_current(raw) == 0); // waits on the event + replay(g, stream); + const std::vector two_replays = read_and_clear(out, n); + EXPECT(two_replays == std::vector({1, 0, 3})); + + registry.clear(); + EXPECT(registry.node_count() == 0); + EXPECT(registry.bound_count() == 0); + EXPECT(!registry.capture_enabled()); + EXPECT(registry.is_enabled(kHookType, 1)); // gate inactive again + replay(g, stream); // every node re-enabled + EXPECT(read_and_clear(out, n) == expected(n, {0, 1, 2}, 1)); + } + CUDA_CHECK(cudaStreamDestroy(stream)); + CUDA_CHECK(cudaFree(out)); +} + +#if DMX_HAS_DEVICE_GRAPH_UPDATE +static cudaError_t launch_device_updatable(int* out, int slot, int value, + const int64_t* row, + cudaStream_t stream) { + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3(1); + cfg.blockDim = dim3(32); + cfg.stream = stream; + cudaLaunchAttribute attr[1]; + attr[0].id = cudaLaunchAttributeDeviceUpdatableKernelNode; + attr[0].val.deviceUpdatableKernelNode.deviceUpdatable = 1; + attr[0].val.deviceUpdatableKernelNode.devNode = nullptr; + cfg.attrs = attr; + cfg.numAttrs = 1; + return cudaLaunchKernelEx(&cfg, gated_marker_kernel, out, slot, value, row); +} +#endif + +static void test_device_planner() { + banner("device planner enables each node from its own control row"); +#if DMX_HAS_DEVICE_GRAPH_UPDATE + ring_py::ToggleRegistry registry; + if (!registry.set_device_mode(true)) { + std::printf(" skipped: device graph updates unavailable\n"); + return; + } + constexpr int n = 5; + int* out = nullptr; + int64_t* rows = nullptr; + CUDA_CHECK(cudaMalloc(&out, n * sizeof(int))); + CUDA_CHECK(cudaMemset(out, 0, n * sizeof(int))); + CUDA_CHECK(cudaMalloc(&rows, n * sizeof(int64_t))); + CUDA_CHECK(cudaMemset(rows, 0, n * sizeof(int64_t))); + cudaStream_t stream; + CUDA_CHECK(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking)); + registry.set_capture(true); // allocates the pool + EXPECT(registry.device_armed()); + + Graph g; + CUDA_CHECK(cudaStreamBeginCapture(stream, cudaStreamCaptureModeThreadLocal)); + for (int i = 0; i < n; ++i) { + cudaStreamCaptureStatus st; + unsigned long long id; + cudaGraph_t graph; + const cudaGraphNode_t* deps; + size_t nd; + CUDA_CHECK(dmx_capture_info(stream, &st, &id, &graph, &deps, &nd)); + bool fresh = false; + const int slot = registry.device_slot(reinterpret_cast(graph), + rows + i, &fresh); + EXPECT(slot == 0); + EXPECT(fresh == (i == 0)); + if (fresh) ring::launch_toggle_planner(registry.device_pool(), slot, stream); + CUDA_CHECK(launch_device_updatable(out, i, 20 + i, rows + i, stream)); + CUDA_CHECK(dmx_capture_info(stream, &st, &id, &graph, &deps, &nd)); + cudaGraphNode_t node = deps[nd - 1]; + cudaKernelNodeAttrValue value = {}; + CUDA_CHECK(cudaGraphKernelNodeGetAttribute( + node, cudaKernelNodeAttributeDeviceUpdatableKernelNode, &value)); + EXPECT(value.deviceUpdatableKernelNode.deviceUpdatable); + registry.register_device_entry(slot, reinterpret_cast(graph), + reinterpret_cast(node), rows + i, + value.deviceUpdatableKernelNode.devNode); + } + CUDA_CHECK(cudaStreamEndCapture(stream, &g.graph)); + CUDA_CHECK(cudaGraphInstantiate(&g.exec, g.graph, 0)); + EXPECT(registry.upload_device_graph(reinterpret_cast(g.graph), + reinterpret_cast(g.exec)) == 0); + // At most once: binding the same graph later does not upload it again + // (and an upload failure would count as an anomaly). + registry.bind_exec(reinterpret_cast(g.graph), + reinterpret_cast(g.exec)); + EXPECT(registry.upload_device_graph(reinterpret_cast(g.graph), + reinterpret_cast(g.exec)) == 0); + registry.set_capture(false); + EXPECT(registry.device_node_count() == n); + EXPECT(registry.anomaly_count() == 0); + + const std::vector> row_sets = { + {0, 0, 0, 0, 0}, {-1, 3, -1, 0, -1}, {5, -1, 7, -1, 2}, {-1, -1, -1, -1, -1}, + {0, 0, 0, 0, 0}}; + for (const auto& row_values : row_sets) { + CUDA_CHECK(cudaMemcpyAsync(rows, row_values.data(), n * sizeof(int64_t), + cudaMemcpyHostToDevice, stream)); + replay(g, stream); + std::vector want(n, 0); + for (int i = 0; i < n; ++i) if (row_values[i] != -1) want[i] = 20 + i; + EXPECT(read_and_clear(out, n) == want); + } + EXPECT(registry.device_errors() == 0); + + // A repeated control row starts a new group (the next forward). + cudaGraph_t dummy = nullptr; + bool fresh = false; + EXPECT(registry.device_slot(reinterpret_cast(dummy), rows, &fresh) == 1); + EXPECT(fresh); + + // The mode is per capture window: clear() turns it off (the pool stays + // allocated), and nothing is armed until it is enabled again. + registry.clear(); + EXPECT(!registry.device_mode()); + EXPECT(!registry.device_armed()); + EXPECT(registry.device_slot(0, rows, &fresh) == -1); + EXPECT(registry.set_device_mode(true)); + EXPECT(registry.device_armed()); + EXPECT(registry.set_device_mode(false)); + EXPECT(!registry.device_mode()); + + CUDA_CHECK(cudaStreamDestroy(stream)); + CUDA_CHECK(cudaFree(rows)); + CUDA_CHECK(cudaFree(out)); +#else + std::printf(" skipped: built without device graph updates\n"); +#endif +} + +int main() { + setbuf(stdout, nullptr); + CUDA_CHECK(cudaFree(nullptr)); + std::printf("test_toggle_registry\n"); + test_eager_apply_and_diff(); + test_lazy_per_graph_apply(); + test_registry_guards(); + test_failed_lazy_apply_leaves_graph_stale(); + test_pre_apply_sync_and_clear(); + test_upload_skips_graphs_without_device_nodes(); + test_device_planner(); + std::printf("Results: %d passed, %d failed\n", g_pass, g_fail); + return g_fail == 0 ? 0 : 1; +} diff --git a/tests/native/ring/test_toggle_registry_host.cpp b/tests/native/ring/test_toggle_registry_host.cpp new file mode 100644 index 000000000..a482566e6 --- /dev/null +++ b/tests/native/ring/test_toggle_registry_host.cpp @@ -0,0 +1,184 @@ +// Host-only tests of ToggleRegistry exec rebinding. Appended by +// tests/test_toggle_registry_host.py to the host half of +// ring/toggle_registry.cu (namespace ring_py) and compiled against the CPU +// CUDA stub in cpu_cuda_stub/; the device-kernel half is not compiled. + +#include +#include +#include + +static int failures = 0; +#define EXPECT(cond) \ + do { \ + if (!(cond)) { \ + std::fprintf(stderr, "FAIL %s:%d: %s\n", __FILE__, __LINE__, \ + #cond); \ + ++failures; \ + } \ + } while (0) + +namespace { + +constexpr uint64_t kGraph = 0x10; +constexpr uint64_t kOldExec = 0x20; +constexpr uint64_t kNewExec = 0x21; +const std::vector kNodes = {0x30, 0x31, 0x32}; +// Node i belongs to hook (0, i). + +using Set = std::vector>; + +bool exec_matches(uint64_t exec, const std::set& on_layers) { + for (size_t i = 0; i < kNodes.size(); ++i) { + const bool want = on_layers.count(static_cast(i)) > 0; + if (fake_node_on(exec, kNodes[i]) != want) return false; + } + return true; +} + +ring_py::ToggleRegistry* make_registry() { + auto* r = new ring_py::ToggleRegistry(); + r->set_capture(true); + for (size_t i = 0; i < kNodes.size(); ++i) + r->register_node(kGraph, 0, static_cast(i), kNodes[i]); + r->set_capture(false); + return r; +} + +// Disable every hook on the old exec, then rebind to a fresh exec (all +// nodes enabled, as CUDA instantiates them) and apply `desired` lazily. +void rebind_then_lazy(const Set& desired, const std::set& expect_on, + bool reapply_same_set) { + fake_enabled.clear(); + auto* r = make_registry(); + r->bind_exec(kGraph, kOldExec); + r->set_enabled({}); + EXPECT(r->apply_all() == 0); + EXPECT(exec_matches(kOldExec, {})); + + r->bind_exec(kGraph, kNewExec); + EXPECT(!r->applied_current(kGraph)); // the new exec is stale + if (reapply_same_set) r->set_enabled({}); + r->set_enabled(desired); + const unsigned before = fake_set_enabled_calls; + EXPECT(r->ensure_current(kGraph) == 0); + EXPECT(r->applied_current(kGraph)); + EXPECT(exec_matches(kNewExec, expect_on)); + // Every host-managed node of the new exec was set explicitly. + EXPECT(fake_set_enabled_calls - before == kNodes.size()); + // Current now: no further calls. + const unsigned settled = fake_set_enabled_calls; + EXPECT(r->ensure_current(kGraph) == 0); + EXPECT(fake_set_enabled_calls == settled); + delete r; +} + +void test_rebind_then_lazy_apply() { + rebind_then_lazy({}, {}, false); // empty set + rebind_then_lazy({}, {}, true); // same set again + rebind_then_lazy({{0, 1}}, {1}, false); // subset + rebind_then_lazy({{0, 0}, {0, 1}, {0, 2}}, {0, 1, 2}, false); // all +} + +void test_rebind_without_set_change_is_not_reported_current() { + // The audit's sequence: bind, disable, rebind, ensure with no new set. + fake_enabled.clear(); + auto* r = make_registry(); + r->bind_exec(kGraph, kOldExec); + r->set_enabled({}); + EXPECT(r->apply_all() == 0); + r->bind_exec(kGraph, kNewExec); + EXPECT(!r->applied_current(kGraph)); + EXPECT(r->ensure_current(kGraph) == 0); + EXPECT(exec_matches(kNewExec, {})); + EXPECT(!r->is_enabled(0, 0)); + delete r; +} + +void test_rebind_then_eager_apply() { + fake_enabled.clear(); + auto* r = make_registry(); + r->bind_exec(kGraph, kOldExec); + r->set_enabled({{0, 2}}); + EXPECT(r->apply_all() == 0); + EXPECT(exec_matches(kOldExec, {2})); + r->bind_exec(kGraph, kNewExec); + EXPECT(r->apply_all() == 0); + EXPECT(exec_matches(kNewExec, {2})); + EXPECT(r->applied_current(kGraph)); + delete r; +} + +void test_same_exec_rebind_keeps_state() { + fake_enabled.clear(); + auto* r = make_registry(); + r->bind_exec(kGraph, kOldExec); + r->set_enabled({{0, 0}}); + EXPECT(r->ensure_current(kGraph) == 0); + r->bind_exec(kGraph, kOldExec); // idempotent + EXPECT(r->applied_current(kGraph)); + const unsigned before = fake_set_enabled_calls; + EXPECT(r->ensure_current(kGraph) == 0); + EXPECT(fake_set_enabled_calls == before); + delete r; +} + +void test_failed_apply_after_rebind_retries() { + fake_enabled.clear(); + auto* r = make_registry(); + r->bind_exec(kGraph, kOldExec); + r->set_enabled({}); + EXPECT(r->apply_all() == 0); + r->bind_exec(kGraph, kNewExec); + fake_set_enabled_error = 700; + EXPECT(r->ensure_current(kGraph) == 700); + EXPECT(!r->applied_current(kGraph)); + fake_set_enabled_error = 0; + EXPECT(r->ensure_current(kGraph) == 0); + EXPECT(exec_matches(kNewExec, {})); + delete r; +} + +void test_clear_after_rebind_reenables_the_new_exec() { + fake_enabled.clear(); + auto* r = make_registry(); + r->bind_exec(kGraph, kOldExec); + r->set_enabled({}); + EXPECT(r->apply_all() == 0); + r->bind_exec(kGraph, kNewExec); + // Something disabled a node of the new exec that the registry has no + // record for: clear() must not trust the old exec's cache. + fake_enabled[{kNewExec, kNodes[0]}] = false; + r->clear(); + EXPECT(exec_matches(kNewExec, {0, 1, 2})); + delete r; +} + +void test_rebind_drops_the_old_replay_event() { + fake_enabled.clear(); + auto* r = make_registry(); + r->set_pre_apply_sync(true); + r->bind_exec(kGraph, kOldExec); + EXPECT(r->record_replay_event(kGraph, nullptr) == 0); + EXPECT(fake_events_live == 1); + r->bind_exec(kGraph, kNewExec); + EXPECT(fake_events_live == 0); + delete r; +} + +} // namespace + +int main() { + test_rebind_then_lazy_apply(); + test_rebind_without_set_change_is_not_reported_current(); + test_rebind_then_eager_apply(); + test_same_exec_rebind_keeps_state(); + test_failed_apply_after_rebind_retries(); + test_clear_after_rebind_reenables_the_new_exec(); + test_rebind_drops_the_old_replay_event(); + if (failures != 0) { + std::fprintf(stderr, "%d failure(s)\n", failures); + return 1; + } + std::printf("toggle registry host: ok\n"); + return 0; +} diff --git a/tests/test_adapter_protocol.py b/tests/test_adapter_protocol.py index 104c6c95c..117b9c8be 100644 --- a/tests/test_adapter_protocol.py +++ b/tests/test_adapter_protocol.py @@ -331,6 +331,59 @@ def test_force_eager_cleared_on_normal_step_after_overflow(): assert a.transport.force_eager is False +@pytest.mark.parametrize("needs_eager", [False, True]) +def test_strict_fallback_rejection_precedes_metadata_and_adaptation(needs_eager): + """allow_eager_fallback=False: a step that needs the eager safety net + raises before anything is published. An eager-only plan is refused + before the reservation; an oversized step after prepare_step (which + reserves nothing for it) but before metadata and the fallback hooks.""" + a = _make_adaptor(prepare_result=2, step_plan=(1024, 3, needs_eager)) + a.allow_eager_fallback = False + with pytest.raises(RuntimeError, match="capture eager fallback forbidden"): + a.before_forward(None) + assert a.engine._ring_engine.prepare_step_calls == ( + [] if needs_eager else [(1024, 3)]) + assert a.transport.set_step_context_calls == [] + assert a.transport.pre_push_all_metas_calls == [] + assert a.adapt_for_cpu_direct_calls == [] + assert a.on_capacity_exceeded_calls == [] + # The step was never armed. + assert a.transport.capture_step is False + counts = a.capture_dispatch_state() + assert counts["allow_eager_fallback"] is False + assert counts["rejected_eager_steps"] == 1 + assert counts["forced_eager_steps"] == 0 + assert counts["oversized_steps"] == int(not needs_eager) + assert counts["eager_required_steps"] == int(needs_eager) + + +def test_dispatch_counts_track_fallback_when_allowed(): + a = _make_adaptor(prepare_result=2) + a.before_forward(None) + counts = a.capture_dispatch_state() + assert counts["allow_eager_fallback"] is True + assert counts["prepared_steps"] == 1 + assert counts["oversized_steps"] == 1 + assert counts["forced_eager_steps"] == 1 + assert counts["rejected_eager_steps"] == 0 + + +def test_allow_eager_fallback_must_be_a_bool(): + engine = FakeEngine() + with pytest.raises(ValueError, match="must be a bool"): + StubAdapterStrict(engine, "m", _make_ctx(), allow_eager_fallback=1) + adapter = StubAdapterStrict(engine, "m", _make_ctx(), + allow_eager_fallback=False) + assert adapter.allow_eager_fallback is False + + +class StubAdapterStrict(StubAdapter): + def __init__(self, engine, model_id, ctx, *, allow_eager_fallback): + BackendAdapter.__init__(self, engine, model_id, + allow_eager_fallback=allow_eager_fallback) + self._ctx = ctx + + def test_n_hooks_zero_skips_prepare_step(): """When _compute_step_plan returns (0, 0, False), prepare_step is skipped but set_step_context + pre_push_all_metas still run.""" @@ -424,3 +477,20 @@ def test_resolve_hook_selection_refuses_an_unknown_token_and_an_empty_selection( # The refusal propagates through the attach-path entry point. with pytest.raises(ValueError, match="Unknown hook selection"): selection.select_hook_specs([], "nope") + + +def test_plan_step_needs_no_transport(): + """Integrations build adapters field by field in their own tests; the + plan reads the toggle only through a transport that exists.""" + import torch + from dmi.hooks.specs import HOOK_TYPE_RESID_PRE, HookSpec, ModelShapeConfig + + adaptor = StubAdapter.__new__(StubAdapter) + adaptor.model_cfg = ModelShapeConfig(hidden_dim=8, num_heads=2, num_kv_heads=2, + head_dim=4, dtype=torch.float32) + adaptor.active_specs = [HookSpec(hook_type=HOOK_TYPE_RESID_PRE, module=None, + layer_no=0)] + adaptor._spec_needs_eager = lambda spec: False + ctx = dataclasses.replace(_make_ctx(), q_len=2, batch=0, flattened=True) + plan = BackendAdapter.plan_step(adaptor, ctx) + assert (plan.total_bytes, plan.hook_count, plan.needs_eager) == (64, 1, False) diff --git a/tests/test_node_toggle.py b/tests/test_node_toggle.py new file mode 100644 index 000000000..a067ed0ae --- /dev/null +++ b/tests/test_node_toggle.py @@ -0,0 +1,980 @@ +"""CPU tests for the host side of runtime CUDA-graph node toggling. + +A fake engine stands in for the native ToggleRegistry so the controller's +guards, its effective-set recomputation and caches, the transport's per-step +spec seam, the adapter's reservation set and the HookPoint eager gate can be +checked without a GPU. tests/test_node_toggle_gpu.py drives the same paths +against real captured graphs, and tests/native/ring/test_toggle_registry.cu +covers the native registry. +""" +from __future__ import annotations + +import random +from types import SimpleNamespace + +import pytest + +from dmi.hooks.specs import ( + HOOK_TYPE_FINAL_LN, + HOOK_TYPE_RESID_PRE, + HookSpec, + ModelShapeConfig, +) +from dmi.node_toggle import NodeToggleController, parse_enabled_hooks +from dmi.transport.ring import RingTransport + +pytestmark = pytest.mark.cpu + + +class FakeToggleEngine: + """The engine surface NodeToggleController uses, with a registry state.""" + + def __init__(self, registered=(), *, bound=1, uniform=True, complete=True, + anomalies=0): + self.registered = set(registered) + self.bound = bound + self.uniform = uniform + self.complete = complete + self.anomalies = anomalies + self.version = 7 + self.enabled = set() + self.active = False + self.apply_calls = 0 + self.apply_result = 0 + self.mask_calls = [] + self.cleared = 0 + + def toggle_registry_version(self): + return self.version + + def bound_graph_count(self): + return self.bound + + def toggle_node_count(self): + return len(self.registered) + + def toggle_registry_uniform(self): + return self.uniform + + def toggle_registry_complete(self): + return self.complete + + def capture_anomaly_count(self): + return self.anomalies + + def set_enabled_hooks(self, pairs): + self.enabled = set(map(tuple, pairs)) + self.active = True + + def apply_toggle(self): + self.apply_calls += 1 + return self.apply_result + + def effective_enabled_mask(self, query): + self.mask_calls.append(list(query)) + return [int(not self.active + or (tuple(k) in self.enabled and tuple(k) in self.registered)) + for k in query] + + def is_hook_enabled(self, hook_type, layer_no): + return not self.active or ( + (hook_type, layer_no) in self.enabled + and (hook_type, layer_no) in self.registered) + + def clear_toggle_registry(self): + self.cleared += 1 + self.active = False + self.enabled = set() + + def is_graph_ready(self, raw): + return raw == 1 + + def ensure_graph_current(self, raw): + return 0 + + def record_replay_event(self, raw): + return 0 + + +def _specs(n=6): + keys = [(HOOK_TYPE_RESID_PRE, layer) for layer in range(n)] + keys.append((HOOK_TYPE_FINAL_LN, -1)) + return [HookSpec(hook_type=ht, module=None, layer_no=ln) for ht, ln in keys] + + +def _transport(engine, specs): + transport = RingTransport.__new__(RingTransport) + transport._ring_engine = engine + transport._active_specs = specs + transport._toggle = None + transport._gated_step = True + return transport + + +def _keys(specs): + return [(s.hook_type, s.layer_no) for s in specs] + + +# --------------------------------------------------------------------------- +# Parsing + + +@pytest.mark.parametrize("value, expected", [ + (None, None), + ("", None), + (" ", None), + ("none", []), + ("NONE", []), + ("0:1, 0:2", [(0, 1), (0, 2)]), + ("14", [(14, -1)]), + ("3:4,,5", [(3, 4), (5, -1)]), +]) +def test_parse_enabled_hooks(value, expected): + assert parse_enabled_hooks(value) == expected + + +# --------------------------------------------------------------------------- +# Controller + + +def test_gate_is_inactive_until_activated_and_every_spec_fires(): + specs = _specs() + transport = _transport(FakeToggleEngine(_keys(specs)), specs) + assert transport.toggle_gate_active is False + assert transport.effective_specs is specs + assert transport.specs_for_step() is specs + + +def test_eager_activation_applies_and_selects_enabled_and_captured_in_order(): + specs = _specs() + registered = set(_keys(specs)) - {(HOOK_TYPE_RESID_PRE, 3)} + engine = FakeToggleEngine(registered) + transport = _transport(engine, specs) + transport.set_active_hooks([(HOOK_TYPE_FINAL_LN, -1), + (HOOK_TYPE_RESID_PRE, 3), # not captured + (HOOK_TYPE_RESID_PRE, 1), + (HOOK_TYPE_RESID_PRE, 99)]) # not installed + assert engine.apply_calls == 1 + assert transport.toggle_gate_active and not transport.toggle_lazy_active + assert _keys(transport.effective_specs) == [ + (HOOK_TYPE_RESID_PRE, 1), (HOOK_TYPE_FINAL_LN, -1)] + # Only enabled-and-installed pairs are queried. + assert engine.mask_calls[-1] == [(HOOK_TYPE_RESID_PRE, 1), + (HOOK_TYPE_RESID_PRE, 3), + (HOOK_TYPE_FINAL_LN, -1)] + + +def test_lazy_activation_defers_the_device_apply(): + specs = _specs() + engine = FakeToggleEngine(_keys(specs)) + transport = _transport(engine, specs) + transport.set_active_hooks_lazy([(HOOK_TYPE_RESID_PRE, 0)]) + assert engine.apply_calls == 0 + assert transport.toggle_gate_active and transport.toggle_lazy_active + assert _keys(transport.effective_specs) == [(HOOK_TYPE_RESID_PRE, 0)] + + +def test_failed_eager_apply_raises_and_leaves_the_gate_inactive(): + specs = _specs() + engine = FakeToggleEngine(_keys(specs)) + engine.apply_result = 400 + transport = _transport(engine, specs) + with pytest.raises(RuntimeError, match="CUDA error 400"): + transport.set_active_hooks([(HOOK_TYPE_RESID_PRE, 0)]) + assert transport.toggle_gate_active is False + + +@pytest.mark.parametrize("state, message", [ + (dict(bound=0), "no producer nodes registered or no graph exec bound"), + (dict(uniform=False), "different hook sets"), + (dict(complete=False), "registry incomplete"), + (dict(anomalies=2), "cannot manage"), +]) +def test_activation_fails_closed_on_an_unsafe_registry(state, message): + specs = _specs() + engine = FakeToggleEngine(_keys(specs), **state) + transport = _transport(engine, specs) + for activate in (transport.set_active_hooks, transport.set_active_hooks_lazy): + with pytest.raises(RuntimeError, match=message): + activate([(HOOK_TYPE_RESID_PRE, 0)]) + assert transport.toggle_gate_active is False + + +def test_activation_fails_closed_without_any_registered_node(): + specs = _specs() + transport = _transport(FakeToggleEngine(()), specs) + with pytest.raises(RuntimeError, match="toggle_node_count=0"): + transport.set_active_hooks([(HOOK_TYPE_RESID_PRE, 0)]) + + +def test_guard_verdict_is_cached_per_registry_version(): + specs = _specs() + engine = FakeToggleEngine(_keys(specs)) + transport = _transport(engine, specs) + transport.set_active_hooks_lazy([(HOOK_TYPE_RESID_PRE, 0)]) + # Same registry version: a registry going bad is not re-checked... + engine.uniform = False + transport.set_active_hooks_lazy([(HOOK_TYPE_RESID_PRE, 1)]) + # ...but any mutation (a version bump) re-validates. + engine.version += 1 + with pytest.raises(RuntimeError, match="different hook sets"): + transport.set_active_hooks_lazy([(HOOK_TYPE_RESID_PRE, 2)]) + + +def test_repeated_presets_reuse_the_effective_set(): + specs = _specs() + engine = FakeToggleEngine(_keys(specs)) + transport = _transport(engine, specs) + a = [(HOOK_TYPE_RESID_PRE, 0), (HOOK_TYPE_RESID_PRE, 4)] + b = [(HOOK_TYPE_RESID_PRE, 2)] + transport.set_active_hooks_lazy(a) + first = transport.effective_specs + transport.set_active_hooks_lazy(b) + transport.set_active_hooks_lazy(list(reversed(a))) + assert transport.effective_specs is first + assert len(engine.mask_calls) == 2 + # A different active list (re-attach) is not served from the memo. + transport._active_specs = list(specs) + transport.set_active_hooks_lazy(a) + assert len(engine.mask_calls) == 3 + assert _keys(transport.effective_specs) == _keys(first) + + +def test_randomized_reconfiguration_matches_the_reference_set(): + specs = _specs(32) + keys = _keys(specs) + registered = set(keys) - {keys[7]} + engine = FakeToggleEngine(registered) + transport = _transport(engine, specs) + rng = random.Random(83) + for step in range(200): + enabled = set(rng.sample(keys, rng.randrange(0, 12))) + transport.set_active_hooks_lazy(sorted(enabled)) + expected = [s for s in specs + if (s.hook_type, s.layer_no) in enabled & registered] + assert transport.effective_specs == expected, step + + +def test_ungated_steps_use_the_full_active_set(): + specs = _specs() + transport = _transport(FakeToggleEngine(_keys(specs)), specs) + transport.set_active_hooks_lazy([(HOOK_TYPE_RESID_PRE, 0)]) + transport.set_gated_step(False) + assert transport.specs_for_step() is specs + assert transport.gated_step_requires_registry() is False + transport.set_gated_step(True) + assert _keys(transport.specs_for_step()) == [(HOOK_TYPE_RESID_PRE, 0)] + assert transport.gated_step_requires_registry() is True + + +def test_clear_resets_registry_and_gate_together(): + specs = _specs() + engine = FakeToggleEngine(_keys(specs)) + transport = _transport(engine, specs) + transport.set_active_hooks_lazy([(HOOK_TYPE_RESID_PRE, 0)]) + transport.clear_toggle() + assert engine.cleared == 1 + assert transport.toggle_gate_active is False + assert transport.toggle_lazy_active is False + assert transport.effective_specs is specs + + +def test_minimal_transport_without_a_controller_is_inactive(): + transport = RingTransport.__new__(RingTransport) + transport._active_specs = _specs() + assert transport.toggle_gate_active is False + assert transport.specs_for_step() is transport._active_specs + + +# --------------------------------------------------------------------------- +# Metadata, reservation and the eager gate use the per-step set + + +class _RecordingRingEngine(FakeToggleEngine): + def __init__(self, registered): + super().__init__(registered) + self.pushed = [] + self.prepared = [] + + def push_all_metas(self, hook_types, layer_nos, *args, **kwargs): + self.pushed.append(list(zip(hook_types, layer_nos))) + + def prepare_step(self, total_bytes, n_hooks, reserve=True): + self.prepared.append((total_bytes, n_hooks)) + return 0 + + +def _cfg(): + import torch + return ModelShapeConfig(hidden_dim=8, num_heads=2, num_kv_heads=2, + head_dim=4, dtype=torch.float32) + + +def _real_transport(engine, specs): + transport = RingTransport.__new__(RingTransport) + transport._ring_engine = engine + transport._ring_payload = None + transport._effective_cap = None + transport.null_offload = False + transport.force_eager = False + transport.capture_step = True + transport._model_cfg = _cfg() + transport._active_specs = specs + transport._using_forward_hooks = True + transport._toggle = None + transport._gated_step = True + transport._hook_selection = None + transport._warned_shapes = set() + transport.set_step_context(model_id="m", req_ids=["r"], + token_ranges=[(0, 4)], flattened=True) + return transport + + +def test_metadata_follows_the_per_step_set(): + specs = _specs() + engine = _RecordingRingEngine(_keys(specs)) + transport = _real_transport(engine, specs) + transport.pre_push_all_metas(batch=0, q_len=4, kv_dim=0) + assert engine.pushed[-1] == _keys(specs) + transport.set_active_hooks_lazy([(HOOK_TYPE_RESID_PRE, 2), + (HOOK_TYPE_FINAL_LN, -1)]) + transport.pre_push_all_metas(batch=0, q_len=4, kv_dim=0) + assert engine.pushed[-1] == [(HOOK_TYPE_RESID_PRE, 2), (HOOK_TYPE_FINAL_LN, -1)] + transport._gated_step = False + transport.pre_push_all_metas(batch=0, q_len=4, kv_dim=0) + assert engine.pushed[-1] == _keys(specs) + + +def test_adapter_reservation_counts_the_per_step_set(): + from dmi.adapters.base import BackendAdapter + from dmi.adapters.types import StepContext + + specs = _specs() + engine = _RecordingRingEngine(_keys(specs)) + transport = _real_transport(engine, specs) + + class Adapter(BackendAdapter): + def detect_model_shape(self, model): + raise NotImplementedError + + def detect_parallel_ranks(self): + return (0, 0, 0, 0) + + def is_pp_first(self): + return True + + def is_pp_last(self): + return True + + def build_step_context(self, *raw): + return StepContext(model_id="m", flattened=True, req_ids=["r"], + token_ranges=[(0, 4)], dim0_offsets=[0], + kv_offsets=[0], batch=0, q_len=4, kv_dim=0, + logits_to_keep=0) + + def on_capacity_exceeded(self, ctx): + pass + + host = SimpleNamespace(_ring_transport=transport, _ring_engine=engine) + adapter = Adapter(host, "m") + adapter.model_cfg = _cfg() + adapter.active_specs = specs + ctx = adapter.build_step_context() + full = adapter.plan_step(ctx) + assert full.hook_count == len(specs) + transport.set_active_hooks_lazy([(HOOK_TYPE_RESID_PRE, 1)]) + gated = adapter.plan_step(ctx) + assert gated.hook_count == 1 + assert gated.total_bytes == 4 * 8 * 4 + transport._gated_step = False + assert adapter.plan_step(ctx) == full + # Reservation and metadata agree on every step kind. + for gated_step in (True, False): + transport._gated_step = gated_step + adapter.before_forward() + assert engine.prepared[-1][1] == len(engine.pushed[-1]) + + +def test_hook_point_eager_net_skips_a_disabled_hook_on_a_gated_step(monkeypatch): + import torch + + from dmi.hooks import point as point_module + from dmi.hooks.point import HookPoint + from dmi.transport import ring as ring_module + + specs = _specs() + engine = _RecordingRingEngine(_keys(specs)) + transport = _real_transport(engine, specs) + transport.set_active_hooks_lazy([(HOOK_TYPE_RESID_PRE, 1)]) + transport.force_eager = True + monkeypatch.setattr(ring_module, "_active_transport", transport) + dispatched = [] + monkeypatch.setattr(point_module, "dispatch_producer", + lambda *args: dispatched.append(args[-1])) + + class CudaLike(torch.Tensor): + @property + def is_cuda(self): + return True + + hook = HookPoint() + hook._ring_hook_type = HOOK_TYPE_RESID_PRE + hook._ring_hook_id = 0 # disabled + x = torch.zeros(2, 2).as_subclass(CudaLike) + assert hook(x) is not None + assert dispatched == [] + assert engine.prepared == [] # nothing reserved for the disabled hook + + +# --------------------------------------------------------------------------- +# Capture window lifecycle (fake engine and graph base class; no GPU) + + +class FakeWindowEngine: + def __init__(self): + self.capture = False + self.device_mode = False + self.device_calls = [] + self.cleared = 0 + self.uploads = [] + + def clear_toggle_registry(self): + self.cleared += 1 + self.capture = False + self.device_mode = False + + def enable_device_toggle(self, enabled): + self.device_calls.append(enabled) + self.device_mode = enabled + return enabled + + def enable_toggle_capture(self, enabled): + self.capture = enabled + + def toggle_capture_enabled(self): + return self.capture + + def toggle_device_mode(self): + return self.device_mode + + def upload_graph_exec(self, graph, exec_ptr): + self.uploads.append((graph, exec_ptr)) + return 0 + + +def _reset_patch_state(monkeypatch, node_toggle): + monkeypatch.setattr(node_toggle, "_PATCHED_CLASS", None) + monkeypatch.setattr(node_toggle, "_ORIG_CUDAGRAPH", None) + monkeypatch.setattr(node_toggle, "_TOGGLE_CLASS", None) + monkeypatch.setattr(node_toggle, "_SESSIONS", {}) + monkeypatch.setattr(node_toggle, "_REPLAY_GUARD", False) + + +@pytest.fixture +def window(monkeypatch): + """dmi.node_toggle with a fake torch.cuda.CUDAGraph base class.""" + import torch + + from dmi import node_toggle + + class FakeGraph: + fail_begin = False + + def __init__(self, keep_graph=False): + self.keep_graph = keep_graph + + def __new__(cls, keep_graph=False): + return object.__new__(cls) + + def capture_begin(self, *args, **kwargs): + if FakeGraph.fail_begin: + raise RuntimeError("capture_begin failed") + + def capture_end(self): + pass + + def raw_cuda_graph(self): + return 0x10 + + def raw_cuda_graph_exec(self): + return 0x20 + + monkeypatch.setattr(torch.cuda, "CUDAGraph", FakeGraph) + _reset_patch_state(monkeypatch, node_toggle) + yield node_toggle, FakeGraph + node_toggle._WINDOW.engine = None + node_toggle._WINDOW.exclude_capture = None + + +def test_device_toggle_is_chosen_per_window(window): + node_toggle, _ = window + engine = FakeWindowEngine() + node_toggle.begin_capture_window(engine, device_toggle=True) + assert engine.toggle_device_mode() is True + node_toggle.end_capture_window(engine) + node_toggle.begin_capture_window(engine) + assert engine.device_calls == [True, False] + assert engine.toggle_device_mode() is False + node_toggle.end_capture_window(engine) + + +def test_failed_capture_begin_restores_node_recording(window): + node_toggle, FakeGraph = window + engine = FakeWindowEngine() + node_toggle.begin_capture_window(engine, exclude_capture=lambda: True) + cls = node_toggle.install_cudagraph_patch() + graph = cls() + FakeGraph.fail_begin = True + try: + with pytest.raises(RuntimeError, match="capture_begin failed"): + graph.capture_begin() + finally: + FakeGraph.fail_begin = False + assert engine.toggle_capture_enabled() is True # recording restored + assert graph._dmi_toggle_suspended is False + # A successful excluded capture suspends and restores around the capture. + graph.capture_begin() + assert engine.toggle_capture_enabled() is False + graph.capture_end() + assert engine.toggle_capture_enabled() is True + node_toggle.end_capture_window(engine) + + +def test_capture_end_uploads_through_the_registry_guard(window): + node_toggle, _ = window + engine = FakeWindowEngine() + node_toggle.begin_capture_window(engine, device_toggle=True) + graph = node_toggle.install_cudagraph_patch()() + graph.capture_begin() + graph.capture_end() + assert engine.uploads == [(0x10, 0x20)] # the registry decides per graph + node_toggle.end_capture_window(engine) + node_toggle.begin_capture_window(engine) # host-toggle window: no upload + graph = node_toggle.install_cudagraph_patch()() + graph.capture_begin() + graph.capture_end() + assert engine.uploads == [(0x10, 0x20)] + node_toggle.end_capture_window(engine) + + +def test_capture_window_releases_the_engine_on_exception(window): + node_toggle, _ = window + engine = FakeWindowEngine() + with pytest.raises(KeyError): + with node_toggle.capture_window(engine, exclude_capture=lambda: False): + assert node_toggle._WINDOW.engine is engine + assert engine.toggle_capture_enabled() + raise KeyError("boom") + assert node_toggle._WINDOW.engine is None + assert node_toggle._WINDOW.exclude_capture is None + assert engine.toggle_capture_enabled() is False + + +def test_clear_and_deactivate_release_the_window_engine(window, monkeypatch): + node_toggle, _ = window + from dmi.transport import ring as ring_module + + engine = FakeWindowEngine() + node_toggle.begin_capture_window(engine) + controller = NodeToggleController(engine, SimpleNamespace(_active_specs=[])) + controller.clear() + assert node_toggle._WINDOW.engine is None + + node_toggle.begin_capture_window(engine) + other = FakeWindowEngine() + NodeToggleController(other, SimpleNamespace(_active_specs=[])).clear() + assert node_toggle._WINDOW.engine is engine # a different engine + monkeypatch.setattr(ring_module, "_active_transport", + SimpleNamespace(_ring_engine=engine)) + ring_module.deactivate() + assert node_toggle._WINDOW.engine is None + + +# --------------------------------------------------------------------------- +# Patch ownership, replay-guard scope and exec lifecycle (fake engine and a +# keep_graph-faithful graph base class; no GPU) + + +class LifecycleGraph: + """torch.cuda.CUDAGraph's keep_graph contract: with keep_graph=True the + exec does not exist until instantiate() or the first replay, and + raw_cuda_graph_exec() raises until then.""" + + next_raw = 0x100 + + def __new__(cls, keep_graph=False): + return object.__new__(cls) + + def __init__(self, keep_graph=False): + self.keep_graph = keep_graph + self.exec_ptr = None + self.raw = LifecycleGraph.next_raw + LifecycleGraph.next_raw += 1 + self.instantiations = 0 + self.replays = 0 + self.trace = [] + + def capture_begin(self, *args, **kwargs): + self.trace.append("capture_begin") + + def capture_end(self): + if not self.keep_graph: + self.instantiate() + self.trace.append("capture_end") + + def instantiate(self): + self.instantiations += 1 + self.exec_ptr = self.raw * 0x100 + self.instantiations + self.trace.append("instantiate") + + def raw_cuda_graph(self): + return self.raw + + def raw_cuda_graph_exec(self): + if self.exec_ptr is None: + raise RuntimeError("instantiate must be called when keep_graph=True") + return self.exec_ptr + + def replay(self): + if self.exec_ptr is None: + self.instantiate() + self.replays += 1 + self.trace.append("replay") + + +class LifecycleEngine(FakeWindowEngine): + """Window + registry surface. `produce()` stands in for a producer op: + it bumps the process-wide launch counter and, while node recording is on, + records a node for the graph being captured.""" + + launches = 0 + + def __init__(self, *, launch_counter=True): + super().__init__() + self.nodes = {} + self.bound = {} + self.capturing = None + self.ensured = [] + if not launch_counter: + self.producer_launch_count = None + + def produce(self, layer=0, *, counted=True): + if counted: + LifecycleEngine.launches += 1 + if self.capture and self.capturing is not None: + self.nodes.setdefault(self.capturing.raw, []).append((0, layer)) + + def producer_launch_count(self): + return LifecycleEngine.launches + + def toggle_node_count(self): + return sum(len(v) for v in self.nodes.values()) + + def toggle_node_count_for_graph(self, raw): + return len(self.nodes.get(raw, ())) + + def clear_toggle_registry(self): + super().clear_toggle_registry() + self.nodes.clear() + self.bound.clear() + + def bind_graph_exec(self, raw, exec_ptr): + self.bound[raw] = exec_ptr + + def is_graph_ready(self, raw): + return raw in self.nodes and raw in self.bound + + +class GatedTransport: + def __init__(self, engine, *, lazy=True, requires_registry=True): + self._ring_engine = engine + self.toggle_gate_active = True + self.toggle_lazy_active = lazy + self.requires_registry = requires_registry + + def gated_step_requires_registry(self): + return self.requires_registry + + def is_graph_ready(self, raw): + return self._ring_engine.is_graph_ready(raw) + + def ensure_graph_current(self, raw): + self._ring_engine.ensured.append((raw, self._ring_engine.bound.get(raw))) + return 0 + + def record_replay_event(self, raw): + return 0 + + +@pytest.fixture +def lifecycle(monkeypatch): + import torch + + from dmi import node_toggle + from dmi.transport import ring as ring_module + + monkeypatch.setattr(torch.cuda, "CUDAGraph", LifecycleGraph) + _reset_patch_state(monkeypatch, node_toggle) + monkeypatch.setattr(ring_module, "_active_transport", None) + + def capture(engine, *, producers=0): + graph = torch.cuda.CUDAGraph() + graph.capture_begin() + engine.capturing = graph + for layer in range(producers): + engine.produce(layer) + engine.capturing = None + graph.capture_end() + return graph + + def serve(transport): + monkeypatch.setattr(ring_module, "_active_transport", transport) + + yield node_toggle, torch, capture, serve + node_toggle._WINDOW.engine = None + node_toggle._WINDOW.exclude_capture = None + + +def _monitor(node_toggle, capture, engine, **window): + node_toggle.begin_capture_window(engine, **window) + graph = capture(engine, producers=2) + node_toggle.bind_graph(engine, graph) + node_toggle.end_capture_window(engine) + node_toggle.set_replay_guard(True) + return graph + + +def test_clearing_the_controller_restores_the_cudagraph_class(lifecycle): + node_toggle, torch, capture, serve = lifecycle + engine = LifecycleEngine() + monitored = _monitor(node_toggle, capture, engine) + assert torch.cuda.CUDAGraph is not LifecycleGraph # session open + assert torch.cuda.CUDAGraph(keep_graph=False).keep_graph is True + NodeToggleController(engine, SimpleNamespace(_active_specs=[])).clear() + assert torch.cuda.CUDAGraph is LifecycleGraph + assert node_toggle._REPLAY_GUARD is False + assert node_toggle.active_session_count() == 0 + assert torch.cuda.CUDAGraph(keep_graph=False).keep_graph is False + # The graph created while patched still replays (its guard is inert). + serve(GatedTransport(engine)) + monitored.replay() + assert monitored.replays == 1 + + +def test_an_unactivated_window_and_a_deactivated_transport_end_the_session(lifecycle): + node_toggle, torch, capture, serve = lifecycle + from dmi.transport import ring as ring_module + + engine = LifecycleEngine() + node_toggle.begin_capture_window(engine) + capture(engine) + node_toggle.end_capture_window(engine) # nothing bound + assert torch.cuda.CUDAGraph is LifecycleGraph + with node_toggle.capture_window(engine): + capture(engine) + assert torch.cuda.CUDAGraph is LifecycleGraph + + _monitor(node_toggle, capture, engine) + assert torch.cuda.CUDAGraph is not LifecycleGraph + serve(SimpleNamespace(_ring_engine=engine)) + ring_module.deactivate() + assert torch.cuda.CUDAGraph is LifecycleGraph + assert node_toggle._REPLAY_GUARD is False + + +def test_binding_after_the_window_closed_renews_the_session(lifecycle): + node_toggle, torch, capture, _ = lifecycle + engine = LifecycleEngine() + node_toggle.begin_capture_window(engine) + graph = capture(engine, producers=1) + node_toggle.end_capture_window(engine) + assert torch.cuda.CUDAGraph is LifecycleGraph + node_toggle.bind_graph(engine, graph) + assert torch.cuda.CUDAGraph is not LifecycleGraph + assert node_toggle._session_for(engine) is graph._dmi_owner + + +def test_the_patch_stays_until_the_last_session_ends(lifecycle): + node_toggle, torch, capture, _ = lifecycle + first, second = LifecycleEngine(), LifecycleEngine() + _monitor(node_toggle, capture, first) + _monitor(node_toggle, capture, second) + NodeToggleController(first, SimpleNamespace(_active_specs=[])).clear() + assert torch.cuda.CUDAGraph is not LifecycleGraph + assert node_toggle.active_session_count() == 1 + NodeToggleController(second, SimpleNamespace(_active_specs=[])).clear() + assert torch.cuda.CUDAGraph is LifecycleGraph + + +def test_a_graph_without_producers_replays_unguarded(lifecycle): + node_toggle, _, capture, serve = lifecycle + engine = LifecycleEngine() + node_toggle.begin_capture_window(engine) + monitored = capture(engine, producers=2) + unrelated = capture(engine) # a helper graph, no producer + node_toggle.bind_graph(engine, monitored) + node_toggle.end_capture_window(engine) + after_window = capture(engine) # captured later, no producer + node_toggle.set_replay_guard(True) + serve(GatedTransport(engine)) + monitored.replay() + unrelated.replay() + after_window.replay() + assert (monitored.replays, unrelated.replays, after_window.replays) == (1, 1, 1) + assert [raw for raw, _ in engine.ensured] == [monitored.raw] + + +@pytest.mark.parametrize("launch_counter", [True, False]) +def test_an_unbound_producer_graph_still_fails_closed(lifecycle, launch_counter): + node_toggle, _, capture, serve = lifecycle + engine = LifecycleEngine(launch_counter=launch_counter) + _monitor(node_toggle, capture, engine) + # Captured after the window closed: producers, but no node and no bind. + late = capture(engine, producers=1) + serve(GatedTransport(engine)) + with pytest.raises(RuntimeError, match="not registered and bound"): + late.replay() + assert late.replays == 0 + + +def test_an_unbound_graph_with_recorded_producers_fails_closed(lifecycle): + node_toggle, _, capture, serve = lifecycle + engine = LifecycleEngine(launch_counter=False) + node_toggle.begin_capture_window(engine) + monitored = capture(engine, producers=1) + forgotten = capture(engine, producers=1) # recorded, never bound + node_toggle.bind_graph(engine, monitored) + node_toggle.end_capture_window(engine) + node_toggle.set_replay_guard(True) + serve(GatedTransport(engine)) + with pytest.raises(RuntimeError, match="not registered and bound"): + forgotten.replay() + + +def test_a_later_session_does_not_adopt_another_engines_graph(lifecycle): + node_toggle, _, capture, serve = lifecycle + old_engine, new_engine = LifecycleEngine(), LifecycleEngine() + old_graph = _monitor(node_toggle, capture, old_engine) + NodeToggleController(old_engine, SimpleNamespace(_active_specs=[])).clear() + _monitor(node_toggle, capture, new_engine) + serve(GatedTransport(new_engine)) + old_graph.replay() # its producers write old_engine's ring + assert old_graph.replays == 1 + assert new_engine.ensured == [] + + +def test_a_stale_graph_of_the_same_engine_fails_closed(lifecycle): + node_toggle, _, capture, serve = lifecycle + engine = LifecycleEngine() + stale = _monitor(node_toggle, capture, engine) + _monitor(node_toggle, capture, engine) # recapture: registry cleared + serve(GatedTransport(engine)) + with pytest.raises(RuntimeError, match="not registered and bound"): + stale.replay() + + +@pytest.mark.parametrize("lazy", [True, False], ids=["lazy", "eager"]) +def test_a_reinstantiated_bound_graph_is_rebound_before_it_replays(lifecycle, lazy): + node_toggle, _, capture, serve = lifecycle + engine = LifecycleEngine() + graph = _monitor(node_toggle, capture, engine) + first_exec = engine.bound[graph.raw] + serve(GatedTransport(engine, lazy=lazy)) + graph.replay() + engine.ensured.clear() + graph.instantiate() # a new exec, same template + replays = [] + graph.trace.clear() + real_replay = LifecycleGraph.replay + + def replay_after_ensure(self): + # The new exec must already carry the enabled set when it launches. + replays.append(list(engine.ensured)) + real_replay(self) + + LifecycleGraph.replay = replay_after_ensure + try: + graph.replay() + finally: + LifecycleGraph.replay = real_replay + assert engine.bound[graph.raw] == graph.exec_ptr != first_exec + # Eager mode too: the rebound exec is brought current before its first + # replay, exactly once. + assert replays == [[(graph.raw, graph.exec_ptr)]] + engine.ensured.clear() + graph.replay() # same exec again + assert engine.ensured == ([(graph.raw, graph.exec_ptr)] if lazy else []) + + +def test_eager_replays_of_an_unchanged_exec_do_not_reapply(lifecycle): + node_toggle, _, capture, serve = lifecycle + engine = LifecycleEngine() + graph = _monitor(node_toggle, capture, engine) + serve(GatedTransport(engine, lazy=False)) + graph.replay() + graph.replay() + assert graph.replays == 2 and engine.ensured == [] + + +def test_recorded_producer_nodes_win_over_an_unchanged_launch_counter(lifecycle): + # A producer op that recorded a node but did not bump the launch counter + # must still classify the capture as carrying producers: an unbound + # capture fails closed instead of replaying around the gate. + node_toggle, torch, capture, serve = lifecycle + engine = LifecycleEngine() + node_toggle.begin_capture_window(engine) + monitored = capture(engine, producers=1) + forgotten = torch.cuda.CUDAGraph() + forgotten.capture_begin() + engine.capturing = forgotten + engine.produce(0, counted=False) + engine.capturing = None + forgotten.capture_end() + assert forgotten._dmi_producers is True + node_toggle.bind_graph(engine, monitored) + node_toggle.end_capture_window(engine) + node_toggle.set_replay_guard(True) + serve(GatedTransport(engine)) + with pytest.raises(RuntimeError, match="not registered and bound"): + forgotten.replay() + assert forgotten.replays == 0 + + +@pytest.mark.parametrize("launch_counter", [True, False]) +def test_excluded_device_toggle_capture_uploads_before_first_replay( + lifecycle, launch_counter): + node_toggle, _, capture, serve = lifecycle + engine = LifecycleEngine(launch_counter=launch_counter) + node_toggle.begin_capture_window(engine, exclude_capture=lambda: True, + device_toggle=True) + graph = capture(engine, producers=1) + assert graph.keep_graph is True + # Instantiated and uploaded at capture end, before any launch. + assert graph.trace == ["capture_begin", "capture_end", "instantiate"] + assert engine.uploads == [(graph.raw, graph.exec_ptr)] + assert engine.nodes == {} # never enrolled + with pytest.raises(RuntimeError, match="excluded"): + node_toggle.bind_graph(engine, graph) + node_toggle.end_capture_window(engine) + node_toggle.set_replay_guard(True) + serve(GatedTransport(engine, requires_registry=False)) + graph.replay() + graph.replay() + assert engine.uploads == [(graph.raw, graph.exec_ptr)] # once per exec + graph.instantiate() # new exec + graph.replay() + assert engine.uploads[-1] == (graph.raw, graph.exec_ptr) + assert graph.trace[-2:] == ["instantiate", "replay"] + assert len(engine.uploads) == 2 and graph.replays == 3 + + +def test_device_toggle_capture_without_producers_is_not_instantiated(lifecycle): + node_toggle, _, capture, _ = lifecycle + engine = LifecycleEngine() + node_toggle.begin_capture_window(engine, exclude_capture=lambda: True, + device_toggle=True) + graph = capture(engine) + node_toggle.end_capture_window(engine) + assert graph.exec_ptr is None and engine.uploads == [] + graph.replay() + assert engine.uploads == [] diff --git a/tests/test_node_toggle_gpu.py b/tests/test_node_toggle_gpu.py new file mode 100644 index 000000000..a7f3ad1ca --- /dev/null +++ b/tests/test_node_toggle_gpu.py @@ -0,0 +1,503 @@ +"""Runtime CUDA-graph node toggling on a real legacy ring. + +Each test captures ``torch.ops.ring.producer`` launches into +``torch.cuda.CUDAGraph`` graphs, binds them to the engine's toggle registry, +and checks what an in-process consumer receives after replays. Every source +tensor is filled with a marker (graph offset + layer), so a delivery shows +which graph's node produced it and whether payload and metadata stayed +paired: a stale or misapplied toggle shows up as a missing, extra or +misattributed slice. + +Needs CUDA and the full native backend. Memory use is a few MiB. +""" +from __future__ import annotations + +import time + +import pytest +import torch + +from tests._requirements import require_cuda, require_native_backend + +pytestmark = [ + pytest.mark.gpu, + pytest.mark.native_backend, + require_cuda(), + require_native_backend(), +] + +QLEN, HID = 4, 8 + + +class Ring: + """A legacy ring engine + transport with an in-process consumer.""" + + def __init__(self, n_hooks=4, payload=4 << 20, task_entries=1024): + from dmi.hooks.specs import (HOOK_TYPE_RESID_PRE, HookSpec, + ModelShapeConfig) + from dmi.transport import native as ne + from dmi.transport import ring as rt + + self.ht = HOOK_TYPE_RESID_PRE + self.collected = [] + cfg = ne.RingConfig() + cfg.payload_ring_bytes = payload + cfg.pinned_staging_bytes = payload + cfg.task_ring_entries = task_entries + cfg.drain_poll_timeout_us = 200 + + def consume(model_id, shard_rank, req_id, act_name, layer_no, + start, end, tensor): + self.collected.append( + (int(layer_no), float(tensor.reshape(-1)[0].item()))) + + self.engine = ne.RingEngine(cfg, consume) + self.engine.init(0) + self.engine.start() + self.transport = rt.RingTransport(self.engine) + rt.activate(self.transport) + self.transport.set_model_cfg(ModelShapeConfig( + hidden_dim=HID, num_heads=2, num_kv_heads=2, head_dim=HID // 2, + dtype=torch.float32)) + self.transport._active_specs = [ + HookSpec(hook_type=self.ht, module=None, layer_no=j, + dim0_is_actual_tokens=True) + for j in range(n_hooks)] + self.transport.set_step_context( + model_id="m", req_ids=["r"], token_ranges=[(0, QLEN)], + dim0_offsets=[0], kv_offsets=[0], flattened=True) + self.payload = self.engine.payload_tensor() + self.n = n_hooks + self.keep_alive = [] + + def sources(self, offset): + return [torch.full((QLEN, HID), offset + j, device="cuda", + dtype=torch.float32) for j in range(self.n)] + + def produce(self, sources, layers=None): + for j, src in enumerate(sources): + if layers is None or j in layers: + torch.ops.ring.producer(self.payload, src, self.ht, j) + + def capture(self, sources, graph=None): + """Capture one producer per hook; null mode keeps any eager work out + of the ring.""" + self.engine.set_null_mode(True) + graph = graph if graph is not None else torch.cuda.CUDAGraph(keep_graph=True) + with torch.cuda.graph(graph): + self.produce(sources) + graph.instantiate() + self.keep_alive.append(sources) # the graph reads these addresses + torch.cuda.synchronize() + self.engine.set_null_mode(False) + return graph + + def bind(self, graph): + self.engine.bind_graph_exec(graph.raw_cuda_graph(), + graph.raw_cuda_graph_exec()) + + def drain(self, expected): + torch.cuda.synchronize() + self.engine.flush_and_wait() + deadline = time.monotonic() + 5.0 + while len(self.collected) < expected and time.monotonic() < deadline: + time.sleep(0.005) + time.sleep(0.02) + got = sorted(self.collected) + self.collected.clear() + return got + + def close(self): + from dmi.transport import ring as rt + torch.cuda.synchronize() + self.engine.flush_and_wait() + rt.deactivate() + self.engine.stop() + + +def _keys(ring, layers): + return [(ring.ht, j) for j in layers] + + +def _expect(layers, offset): + return sorted((j, float(offset + j)) for j in layers) + + +@pytest.fixture +def ring(): + r = Ring() + try: + yield r + finally: + r.close() + + +def test_eager_toggle_delivers_exactly_the_enabled_hooks(ring): + ring.engine.enable_toggle_capture(True) + graph = ring.capture(ring.sources(0.0)) + ring.engine.enable_toggle_capture(False) + ring.bind(graph) + assert ring.engine.toggle_node_count() == ring.n + assert ring.engine.toggle_registry_complete() + + # Gate inactive: every hook fires. + ring.transport.pre_push_all_metas(batch=0, q_len=QLEN, kv_dim=0) + graph.replay() + assert ring.drain(ring.n) == _expect(range(ring.n), 0.0) + + for enabled in ([1, 3], [0], [], [0, 1, 2, 3], [2]): + ring.transport.set_active_hooks(_keys(ring, enabled)) + ring.transport.pre_push_all_metas(batch=0, q_len=QLEN, kv_dim=0) + graph.replay() + assert ring.drain(len(enabled)) == _expect(enabled, 0.0), enabled + assert ring.engine.capture_anomaly_count() == 0 + + +def test_lazy_reconfiguration_applies_per_graph_before_replay(ring): + ring.engine.enable_toggle_capture(True) + graphs = {"A": (ring.capture(ring.sources(0.0)), 0.0), + "B": (ring.capture(ring.sources(100.0)), 100.0)} + ring.engine.enable_toggle_capture(False) + for graph, _ in graphs.values(): + ring.bind(graph) + assert ring.engine.toggle_registry_uniform() + + current = {} + + def reconfigure(enabled): + current["enabled"] = sorted(enabled) + ring.transport.set_active_hooks_lazy(_keys(ring, enabled)) + + def replay(name): + graph, offset = graphs[name] + raw = graph.raw_cuda_graph() + ring.transport.pre_push_all_metas(batch=0, q_len=QLEN, kv_dim=0) + assert ring.transport.ensure_graph_current(raw) == 0 + assert ring.engine._test_applied_current(raw) + graph.replay() + assert ring.transport.record_replay_event(raw) == 0 + assert ring.drain(len(current["enabled"])) == _expect( + current["enabled"], offset), (name, current["enabled"]) + + # First-time apply, one version stale, two versions stale, empty. + reconfigure([0, 2]); replay("B"); replay("B"); replay("A") + reconfigure([1, 3]); replay("A"); replay("B") + reconfigure([0, 1, 2, 3]); reconfigure([2]); replay("A"); replay("B") + reconfigure([]); replay("A"); replay("B") + + +def test_pre_apply_sync_mode_waits_for_the_previous_replay(ring): + ring.engine.set_toggle_pre_apply_sync(True) + ring.engine.enable_toggle_capture(True) + graph = ring.capture(ring.sources(0.0)) + ring.engine.enable_toggle_capture(False) + ring.bind(graph) + raw = graph.raw_cuda_graph() + for enabled in ([0, 1], [3], [1, 2, 3]): + ring.transport.set_active_hooks_lazy(_keys(ring, enabled)) + ring.transport.pre_push_all_metas(batch=0, q_len=QLEN, kv_dim=0) + assert ring.transport.ensure_graph_current(raw) == 0 + graph.replay() + assert ring.transport.record_replay_event(raw) == 0 + assert ring.drain(len(enabled)) == _expect(enabled, 0.0) + + +def test_reservation_matches_what_fires_on_gated_and_ungated_steps(): + """Plan each step from specs_for_step(), as BackendAdapter.plan_step + does, over a long sequence of gated replays and ungated eager steps with + a small ring: if reservation and production disagree, the CPU accounting + drifts and the ring never drains back to empty.""" + from dmi.hooks.specs import align_up_py + + r = Ring(payload=64 * 1024, task_entries=16) + try: + r.engine.enable_toggle_capture(True) + sources = r.sources(0.0) + graph = r.capture(sources) + r.engine.enable_toggle_capture(False) + r.bind(graph) + delivered = 0 + expected = 0 + presets = [[0, 1], [3], [0, 1, 2, 3], [], [2, 3]] + for step in range(60): + if step % 3 == 0: + r.transport.set_active_hooks_lazy( + _keys(r, presets[(step // 3) % len(presets)])) + gated = step % 4 != 3 + r.transport._gated_step = gated + specs = r.transport.specs_for_step() + total = sum(align_up_py(QLEN * HID * 4, 16) for _ in specs) + if specs: + result = r.engine.prepare_step(total, len(specs)) + assert result in (0, 1), result + r.transport.pre_push_all_metas(batch=0, q_len=QLEN, kv_dim=0) + if gated: + assert r.transport.ensure_graph_current(graph.raw_cuda_graph()) == 0 + graph.replay() + else: + r.produce(sources) # eager: every hook fires + expected += len(specs) + if step % 7 == 6: + delivered += len(r.drain(0)) + delivered += len(r.drain(expected - delivered)) + assert delivered == expected + assert r.engine.available_capacity() == r.engine.payload_cap() + assert r.engine.available_task_slots() == r.engine.task_cap() + finally: + r.close() + + +def test_capacity_protection_engages_after_unreserved_warmup_producers(): + """Warmup runs of captured producers publish without a reservation, so + the drain's tails end up ahead of the reservation heads. Reservations + must start from the tails: otherwise, for as many tasks as the skew, the + room check reads a full ring as empty and producers overwrite unread + slots.""" + from dmi.hooks.specs import align_up_py + + r = Ring(n_hooks=1, payload=64 * 1024, task_entries=8) + try: + sources = r.sources(0.0) + for _ in range(3): # 21 unreserved tasks, drained + for _ in range(7): + r.transport.pre_push_all_metas(batch=0, q_len=QLEN, kv_dim=0) + r.produce(sources) + assert len(r.drain(7)) == 7 + assert r.engine.available_task_slots() == r.engine.task_cap() + step = align_up_py(QLEN * HID * 4, 16) + results = [r.engine.prepare_step(step, 1) for _ in range(9)] + # Eight reservations fill the 8-entry task ring; the ninth must flush. + assert results[:8] == [0] * 8 + assert results[8] == 1 + assert r.engine.available_task_slots() == 0 + finally: + r.close() + + +def test_activation_refuses_unsafe_registries(ring): + # Nothing captured or bound. + with pytest.raises(RuntimeError, match="no producer nodes registered"): + ring.transport.set_active_hooks(_keys(ring, [0])) + + ring.engine.enable_toggle_capture(True) + graph = ring.capture(ring.sources(0.0)) + partial = torch.cuda.CUDAGraph(keep_graph=True) + partial_sources = ring.sources(50.0) + ring.engine.set_null_mode(True) + with torch.cuda.graph(partial): + ring.produce(partial_sources, layers={0, 1}) + partial.instantiate() + ring.engine.set_null_mode(False) + ring.engine.enable_toggle_capture(False) + ring.bind(graph) + # `partial` holds nodes but is unbound, and registered a different set. + assert not ring.engine.is_graph_ready(partial.raw_cuda_graph()) + assert ring.engine.is_graph_ready(graph.raw_cuda_graph()) + assert not ring.engine.toggle_registry_complete() + with pytest.raises(RuntimeError, match="different hook sets"): + ring.transport.set_active_hooks_lazy(_keys(ring, [0])) + ring.bind(partial) + with pytest.raises(RuntimeError, match="different hook sets"): + ring.transport.set_active_hooks_lazy(_keys(ring, [0])) + assert ring.transport.toggle_gate_active is False + + +def test_chunked_producer_captured_in_the_window_fails_closed(ring): + counts = torch.full((2,), QLEN * HID * 2, dtype=torch.int64, device="cuda") + src = torch.zeros(2 * QLEN, HID, device="cuda") + ring.engine.enable_toggle_capture(True) + graph = ring.capture(ring.sources(0.0)) + ring.engine.set_null_mode(True) + chunked = torch.cuda.CUDAGraph(keep_graph=True) + with torch.cuda.graph(chunked): + torch.ops.ring.producer_chunked(ring.payload, src, counts, ring.ht, 9) + ring.engine.set_null_mode(False) + ring.engine.enable_toggle_capture(False) + ring.bind(graph) + assert ring.engine.capture_anomaly_count() == 1 + with pytest.raises(RuntimeError, match="cannot manage"): + ring.transport.set_active_hooks(_keys(ring, [0])) + del chunked + + +def test_failed_lazy_apply_is_reported_and_retried(ring): + ring.engine.enable_toggle_capture(True) + graph = ring.capture(ring.sources(0.0)) + ring.engine.enable_toggle_capture(False) + ring.bind(graph) + raw = graph.raw_cuda_graph() + ring.transport.set_active_hooks_lazy(_keys(ring, [2])) + ring.engine._test_force_apply_error(999) + assert ring.transport.ensure_graph_current(raw) == 999 + assert not ring.engine._test_applied_current(raw) + ring.engine._test_force_apply_error(0) + assert ring.transport.ensure_graph_current(raw) == 0 + ring.transport.pre_push_all_metas(batch=0, q_len=QLEN, kv_dim=0) + graph.replay() + assert ring.drain(1) == _expect([2], 0.0) + + +def test_clear_re_enables_nodes_and_deactivates_the_gate(ring): + ring.engine.enable_toggle_capture(True) + graph = ring.capture(ring.sources(0.0)) + ring.engine.enable_toggle_capture(False) + ring.bind(graph) + ring.transport.set_active_hooks(_keys(ring, [1])) + ring.transport.clear_toggle() + assert not ring.transport.toggle_gate_active + assert ring.engine.toggle_node_count() == 0 + # The graph keeps replaying after the clear with every producer on. + ring.transport.pre_push_all_metas(batch=0, q_len=QLEN, kv_dim=0) + graph.replay() + assert ring.drain(ring.n) == _expect(range(ring.n), 0.0) + + +def test_eager_safety_net_skips_disabled_hooks_on_gated_steps(ring): + from dmi.hooks.point import HookPoint + + ring.engine.enable_toggle_capture(True) + graph = ring.capture(ring.sources(0.0)) + ring.engine.enable_toggle_capture(False) + ring.bind(graph) + ring.transport.set_active_hooks(_keys(ring, [0, 3])) + hooks = [] + for j in range(ring.n): + hook = HookPoint() + hook._ring_hook_type = ring.ht + hook._ring_hook_id = j + hook._ring_payload = ring.payload + hooks.append(hook) + sources = ring.sources(200.0) + ring.transport.force_eager = True + try: + # Gated step that fell back to eager: only the enabled hooks fire. + ring.transport.pre_push_all_metas(batch=0, q_len=QLEN, kv_dim=0) + for hook, src in zip(hooks, sources): + hook(src) + assert ring.drain(2) == _expect([0, 3], 200.0) + # Ungated step: the full active set fires. + ring.transport._gated_step = False + ring.transport.pre_push_all_metas(batch=0, q_len=QLEN, kv_dim=0) + for hook, src in zip(hooks, sources): + hook(src) + assert ring.drain(ring.n) == _expect(range(ring.n), 200.0) + finally: + ring.transport.force_eager = False + ring.transport._gated_step = True + + +@pytest.fixture +def patched_cudagraph(monkeypatch): + from dmi import node_toggle + + monkeypatch.setattr(torch.cuda, "CUDAGraph", torch.cuda.CUDAGraph) + monkeypatch.setattr(node_toggle, "_PATCHED_CLASS", None) + monkeypatch.setattr(node_toggle, "_ORIG_CUDAGRAPH", None) + monkeypatch.setattr(node_toggle, "_TOGGLE_CLASS", None) + monkeypatch.setattr(node_toggle, "_SESSIONS", {}) + monkeypatch.setattr(node_toggle, "_REPLAY_GUARD", False) + yield node_toggle + if node_toggle._WINDOW.engine is not None: + node_toggle.end_capture_window(node_toggle._WINDOW.engine) + + +def test_capture_window_excludes_partial_captures_and_guards_replays( + ring, patched_cudagraph): + """A capture the integration excludes (e.g. a partial-model segment + graph) records no nodes; the guarded replay allows it only on an + ungated step, refuses it on a gated one, and refuses any graph captured + after the window closed.""" + node_toggle = patched_cudagraph + exclude = {"now": False} + node_toggle.begin_capture_window(ring.engine, + exclude_capture=lambda: exclude["now"]) + cls = torch.cuda.CUDAGraph + assert getattr(cls, "_dmi_keep_graph", False) + + full = ring.capture(ring.sources(0.0), graph=cls()) + exclude["now"] = True + segment = cls() + segment_sources = ring.sources(300.0) + ring.engine.set_null_mode(True) + with torch.cuda.graph(segment): + ring.produce(segment_sources, layers={1, 2}) + segment.instantiate() + ring.engine.set_null_mode(False) + exclude["now"] = False + assert segment._dmi_toggle_excluded and not full._dmi_toggle_excluded + assert node_toggle.excluded_capture_count() == 1 + assert ring.engine.toggle_capture_enabled() # restored after the segment + + node_toggle.bind_graph(ring.engine, full) + with pytest.raises(RuntimeError, match="excluded"): + node_toggle.bind_graph(ring.engine, segment) + node_toggle.end_capture_window(ring.engine) + assert ring.engine.toggle_node_count() == ring.n + assert ring.engine.toggle_registry_complete() + + late = ring.capture(ring.sources(400.0), graph=cls()) # window closed + + node_toggle.activate(ring.transport, _keys(ring, [0, 2]), lazy=True) + # Gated step: the bound full graph replays with the lazy apply. + ring.transport.pre_push_all_metas(batch=0, q_len=QLEN, kv_dim=0) + full.replay() + assert ring.drain(2) == _expect([0, 2], 0.0) + # The excluded segment on a gated step would desynchronize the ring. + with pytest.raises(RuntimeError, match="excluded from the toggle registry"): + segment.replay() + # On an ungated step the full set's metadata is published and the + # excluded segment's all-on producers match it. + ring.transport._gated_step = False + ring.transport._active_specs, saved = ( + [s for s in ring.transport._active_specs if s.layer_no in (1, 2)], + ring.transport._active_specs) + try: + ring.transport.pre_push_all_metas(batch=0, q_len=QLEN, kv_dim=0) + segment.replay() + assert ring.drain(2) == _expect([1, 2], 300.0) + finally: + ring.transport._active_specs = saved + ring.transport._gated_step = True + # A graph captured after the window closed is refused before it runs. + with pytest.raises(RuntimeError, match="not registered and bound"): + late.replay() + node_toggle.set_replay_guard(False) + + +def test_bind_refuses_graphs_of_the_unpatched_class(ring): + from dmi import node_toggle + + class UnpatchedGraph: # what torch.cuda.CUDAGraph() is without the patch + pass + + with pytest.raises(RuntimeError, match="unpatched"): + node_toggle.bind_graph(ring.engine, UnpatchedGraph()) + + +def test_activate_refuses_a_requested_set_without_bound_graphs(ring): + from dmi import node_toggle + + node_toggle.activate(ring.transport, None, lazy=False) # inactive: fine + with pytest.raises(RuntimeError, match="silently inert"): + node_toggle.activate(ring.transport, _keys(ring, [0]), lazy=False) + + +def test_python_callable_consumers_are_validated(): + from dmi.transport import native as ne + + with pytest.raises(TypeError, match="callable consumer"): + ne.RingEngine(ne.RingConfig(), 42) + + +def test_toggle_capture_is_refused_on_a_record_ring(): + from dmi.transport import native as ne + + cfg = ne.RingConfig() + cfg.payload_ring_bytes = 1 << 20 + cfg.pinned_staging_bytes = 1 << 20 + try: + engine = ne.RingEngine.create_record(cfg, None, "raise", 0) + except Exception as exc: # pragma: no cover - record runtime unavailable + pytest.skip(f"record ring unavailable: {exc}") + with pytest.raises(RuntimeError, match="record ring"): + engine.enable_toggle_capture(True) diff --git a/tests/test_toggle_registry_host.py b/tests/test_toggle_registry_host.py new file mode 100644 index 000000000..22b5c300f --- /dev/null +++ b/tests/test_toggle_registry_host.py @@ -0,0 +1,46 @@ +"""Host-only tests of ToggleRegistry exec rebinding (no CUDA toolkit or GPU). + +The host half of native/csrc/ring/toggle_registry.cu (namespace ring_py) is +compiled unchanged against a CPU stand-in for the CUDA runtime API +(tests/native/ring/cpu_cuda_stub/), together with +tests/native/ring/test_toggle_registry_host.cpp. The device-kernel half +(namespace ring) is left out. tests/native/ring/test_toggle_registry.cu +covers the same registry on a GPU. +""" + +from __future__ import annotations + +import shutil +import subprocess +from pathlib import Path + +import pytest + + +@pytest.mark.cpu +def test_toggle_registry_rebind_resets_per_exec_state(tmp_path): + compiler = shutil.which("g++") or shutil.which("c++") + if compiler is None: + pytest.skip("a C++17 compiler is required") + root = Path(__file__).resolve().parents[1] + ring_dir = root / "native" / "csrc" / "ring" + tests_dir = root / "tests" / "native" / "ring" + source = (ring_dir / "toggle_registry.cu").read_text() + host, separator, _ = source.partition("\nnamespace ring {") + assert separator, "toggle_registry.cu no longer splits at namespace ring" + assert host.rstrip().endswith("} // namespace ring_py") + generated = tmp_path / "toggle_registry_host_test.cpp" + generated.write_text( + host + "\n" + (tests_dir / "test_toggle_registry_host.cpp").read_text()) + executable = tmp_path / "toggle_registry_host_test" + build = subprocess.run( + [compiler, "-std=c++17", "-O0", "-DDMX_NO_DEVICE_GRAPH_UPDATE", + f"-I{tests_dir / 'cpu_cuda_stub'}", f"-I{ring_dir}", + str(generated), "-o", str(executable)], + capture_output=True, text=True, check=False, + ) + assert build.returncode == 0, build.stdout + build.stderr + run = subprocess.run([str(executable)], capture_output=True, text=True, + check=False, timeout=60) + assert run.returncode == 0, run.stdout + run.stderr + assert "toggle registry host: ok" in run.stdout