Skip to content

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

Merged
zhaoyboo merged 1 commit into
mainfrom
feat/cuda-graph-node-toggle
Nov 5, 2026
Merged

zhaoyboo merged 1 commit into
mainfrom
feat/cuda-graph-node-toggle

Conversation

@zhaoyboo

@zhaoyboo zhaoyboo commented Nov 4, 2026

Copy link
Copy Markdown
Collaborator

Depends on #170. Supersedes #69.

Summary

Lets an integration enable or disable individual HookPoint producers inside
already-captured CUDA graphs, per step, without recapture: producer kernel
nodes are recorded during capture and toggled with cudaGraphNodeSetEnabled
(eagerly, or lazily at each graph's next replay), or by an in-graph planner on
CUDA 12.4+ (device-side toggling).

Implementation details

  • Native: ToggleRegistry (node recording during a capture window, graph and
    exec binding, eager/lazy apply with replay events, pre-apply sync,
    completeness checks), device-side mode (toggle_device.h: device-updatable
    producer nodes with a planner kernel; each exec uploaded at most once),
    RingEnginePy delegates and ring_torch_op node recording for static and
    prefix producers.
  • DrainThread::reserve snaps lagging heads up to the committed tails, so
    capacity protection engages immediately after unreserved warmup producers.
  • Python: dmi.node_toggle (NodeToggleController, install_cudagraph_patch
    with an exclusion predicate and replay guard, begin_capture_window /
    end_capture_window / capture_window, bind_graph, activate,
    parse_enabled_hooks); RingTransport.effective_specs, specs_for_step,
    set_gated_step and toggle delegates; BackendAdapter.plan_step reserves
    for the per-step set; strict mode allow_eager_fallback=False; the eager
    safety net skips disabled hooks on gated steps.
  • docs/node-toggle.md.

Compatibility and scope

Inactive unless an integration opens a capture window. The record ring is
refused. A RingEngine also accepts a Python callable consumer.

Validation

  • Native: test_toggle_registry 151 assertions (device planner, upload-once,
    device-mode reset on clear); all ring suites pass.
  • Python: test_node_toggle.py, test_node_toggle_gpu.py,
    test_adapter_protocol.py, with a regression test for each of: reservation
    debt after unreserved warmup, sticky device mode, capture_begin exception
    safety, capture-window lifecycle, upload-once, and planning without a
    transport.
  • CPU tests: 2280 passed, 458 skipped. test_node_toggle_gpu.py: 15 passed.

@Samfisheryu

Copy link
Copy Markdown
Collaborator

Independent audit: three required changes (CPU evidence)

Reviewed commit: 2323c9945544212fe67e256ad6ba6d3113ceb19d.

The findings below exercise the actual Python module and the actual C++ host implementation. Python graph/engine objects and CUDA APIs are replaced with CPU doubles; no real GPU execution, CUDA initialization, inference, or performance benchmark was performed. The scripts and observed outputs are embedded below so the evidence is reproducible without access to the auditor's machine. GPU validation remains pending, as requested by the repository owner.

T171-1 · P2 · Restore the process-wide CUDAGraph patch when monitoring is cleared

Code: forced keep_graph=True, global class replacement, and NodeToggleController.clear().

Trigger: open a capture window, bind a monitored graph, activate toggling, then clear the controller and remove the active transport. An unrelated component subsequently creates torch.cuda.CUDAGraph(keep_graph=False).

Expected: after the monitoring session ends, newly created unrelated graphs use the original constructor semantics. A controlled global patch needs an explicit ownership and restoration lifecycle.

Observed: the class remains ToggleCUDAGraph, _REPLAY_GUARD remains true, and the new graph's explicitly requested keep_graph=False is overwritten with true. Clearing the native registry and host gate does not undo the global patch. The Python reproduction reports:

{"class_restored": false, "replay_guard": true, "new_graph_keep_graph": true}

Impact: enabling and then disabling this feature does not return the process to its original behavior. For unrelated graphs created afterward, keep_graph=True retains the graph template and defers instantiation until an explicit instantiate or the first replay. The constructor change is reproduced on CPU; its GPU memory and first-replay cost have not been measured.

Requested change: give the patch a monitoring-session owner, preserve the original class, and restore the class and guard when the last owner ends. Existing graph instances also need owner-aware replay behavior so a later session cannot attach them to a different active transport. Add an enable → clear/close → unrelated graph regression.

T171-2 · P2 · Limit the replay guard to graphs owned by this monitoring session

Code: global active-transport dispatch and unregistered-graph rejection.

Trigger: toggling is active for a correctly bound monitored graph. A separate helper/component in the same process replays a graph containing no DMI producer, so that graph legitimately has no DMI registry entry.

Expected: the unrelated graph replays normally. Fail-closed validation still applies to graphs that contain this session's producers but were captured or bound incorrectly.

Observed: every instance of the patched class is routed through the globally active transport. is_graph_ready() is false for the unrelated graph, and replay raises the “worker must terminate” error before the underlying replay executes. The same reproduction first successfully replays the bound monitored graph, then reports unrelated_graph.replays == 0 and:

node toggle: graph 0x101 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.

Impact: enabling firewall monitoring can stop unrelated CUDA graph users in the same process; this does not require multiple DMI engines. The fake graph confirms the Python rejection and absence of a forwarded replay, not an executed GPU failure.

Requested change: track graph ownership and whether a capture contains managed producers, then pass explicitly unrelated graphs through unchanged. Keep the consistency check for owned producer graphs; removing the guard globally would hide real metadata/payload mismatches.

CPU reproduction for T171-1 and T171-2 (complete Python script)

Save the following as pr171_python.py. Run against an unchanged checkout of the reviewed commit:

CUDA_VISIBLE_DEVICES='' OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 PYTHONNOUSERSITE=1 python pr171_python.py /path/to/DMI

This imports the real src/dmi/node_toggle.py directly, without importing the dmi package or PyTorch. Graph models the keep_graph/instantiate lifecycle; Engine models a valid registry for the monitored graph and no registry entry for the unrelated graph. Assertions below confirm the current buggy behavior and should be replaced with the intended behavior in regression tests.

"""CPU-only execution of the real node_toggle.py against graph/engine doubles.

No torch or CUDA library is imported. Graph keep_graph/instantiate behavior
matches the locally installed torch/cuda/graphs.py documented contract.
"""
import importlib.util
import json
import pathlib
import sys
import types


class Graph:
    next_raw = 0x100
    fail_end = False

    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.replays = 0
        self.raw = Graph.next_raw
        Graph.next_raw += 1

    def capture_begin(self):
        pass

    def capture_end(self):
        if Graph.fail_end:
            raise RuntimeError("simulated capture_end failure")
        if not self.keep_graph:
            self.instantiate()

    def instantiate(self):
        self.exec_ptr = self.raw + 0x1000

    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


class Engine:
    def __init__(self):
        self.capture = False
        self.device = False
        self.uploads = []
        self.bound = {}
        self.version = 0
        self.enabled = set()

    def clear_toggle_registry(self):
        self.capture = self.device = False
        self.bound.clear()
        self.version += 1

    def enable_device_toggle(self, enabled):
        self.device = enabled

    def toggle_device_mode(self):
        return self.device

    def enable_toggle_capture(self, enabled):
        self.capture = enabled

    def toggle_capture_enabled(self):
        return self.capture

    def upload_graph_exec(self, raw, exec_ptr):
        self.uploads.append((raw, exec_ptr))
        return 0

    def bind_graph_exec(self, raw, exec_ptr):
        self.bound[raw] = exec_ptr

    def toggle_registry_version(self):
        return self.version

    def bound_graph_count(self):
        return len(self.bound)

    def toggle_node_count(self):
        return 2

    def toggle_registry_uniform(self):
        return True

    def toggle_registry_complete(self):
        return True

    def capture_anomaly_count(self):
        return 0

    def set_enabled_hooks(self, pairs):
        self.enabled = set(pairs)

    def effective_enabled_mask(self, query):
        return [int(pair in self.enabled) for pair in query]

    def ensure_graph_current(self, raw):
        return 0

    def record_replay_event(self, raw):
        return 0

    def is_graph_ready(self, raw):
        return raw in self.bound


repo = pathlib.Path(sys.argv[1])
fake_torch = types.ModuleType("torch")
fake_torch.cuda = types.SimpleNamespace(CUDAGraph=Graph)
sys.modules["torch"] = fake_torch
pkg = types.ModuleType("audit_dmi")
pkg.__path__ = []
sys.modules[pkg.__name__] = pkg
transports = types.ModuleType("audit_dmi.transport")
ring = types.SimpleNamespace(active=None)
ring.get_active = lambda: ring.active
transports.ring = ring
sys.modules[transports.__name__] = transports
spec = importlib.util.spec_from_file_location(
    "audit_dmi.node_toggle", repo / "src/dmi/node_toggle.py")
nt = importlib.util.module_from_spec(spec)
sys.modules[spec.name] = nt
spec.loader.exec_module(nt)


class Transport:
    def __init__(self, engine):
        self._ring_engine = engine
        self._active_specs = [types.SimpleNamespace(hook_type=0, layer_no=i)
                              for i in (0, 1)]
        self.ctrl = nt.NodeToggleController(engine, self)
        self.requires_registry = True

    @property
    def toggle_gate_active(self):
        return self.ctrl.gate_active

    @property
    def toggle_lazy_active(self):
        return self.ctrl.lazy_active

    def set_active_hooks_lazy(self, pairs):
        self.ctrl.set_active_hooks_lazy(pairs)

    def is_graph_ready(self, raw):
        return self.ctrl.is_graph_ready(raw)

    def gated_step_requires_registry(self):
        return self.requires_registry

    def ensure_graph_current(self, raw):
        return self.ctrl.ensure_graph_current(raw)

    def record_replay_event(self, raw):
        return self.ctrl.record_replay_event(raw)


def capture():
    g = fake_torch.cuda.CUDAGraph()
    g.capture_begin()
    g.capture_end()
    return g


result = {"repo": str(repo), "cuda_imported": False}
engine = Engine()
transport = Transport(engine)
ring.active = transport
nt.begin_capture_window(engine)
monitored = capture()
nt.bind_graph(engine, monitored)
unrelated = capture()  # no DMI producer, so no registry entry
nt.end_capture_window(engine)
nt.activate(transport, [(0, 0)], lazy=True)
monitored.replay()
try:
    unrelated.replay()
except RuntimeError as exc:
    result["unrelated_graph"] = {"error": str(exc), "replays": unrelated.replays}
else:
    raise AssertionError("expected current implementation to reject unrelated graph")

transport.ctrl.clear()
ring.active = None
result["after_clear"] = {
    "class_restored": fake_torch.cuda.CUDAGraph is Graph,
    "replay_guard": nt._REPLAY_GUARD,
    "new_graph_keep_graph": fake_torch.cuda.CUDAGraph(keep_graph=False).keep_graph,
}
assert result["after_clear"] == {
    "class_restored": False, "replay_guard": True, "new_graph_keep_graph": True}
print(json.dumps(result, indent=2))

T171-3 · P2 · Reset per-executable state when rebinding a graph to a new exec

Code: bind_exec(), applied-version fast path, per-node state cache, and default instantiated state.

Trigger: bind graph G to executable E1, disable its producer and apply that state, then instantiate a replacement executable E2 for the same retained graph and call the supported bind API again. E2 begins with its producer enabled. This is a valid new exec being bound, not a replay through a destroyed old handle.

Expected: changing the executable invalidates the applied version and the per-node actual-state cache. Before E2 replays, its node state must be synchronized with the current enabled-hook set.

Observed: bind_exec() changes reg_exec_[g] and bumps registry version_, but preserves both applied_version_[g] and RegEntry.last_enabled. The first ensure_current() returns at the version fast path. Even explicitly setting the same empty enabled set again does not repair E2: the per-node loop skips the producer because the old exec's cached last_enabled is already false.

The reproduction compiles the unaltered host portion of toggle_registry.cu, using a CPU API double for CUDA. It produces:

{
  "reported_current_after_rebind": 1,
  "host_hook_enabled": 0,
  "actual_new_exec_after_ensure": 1,
  "actual_new_exec_after_reapply": 1,
  "set_enabled_calls_after_rebind": 0
}

Here actual_new_exec_* denotes the executable state held by the CUDA API double, not a GPU measurement. The directly demonstrated defect is that the real host implementation makes zero cudaGraphNodeSetEnabled calls for the replacement exec while reporting its state current and the hook disabled.

Impact: the host metadata/reservation gate can omit a producer that the replacement exec still runs, violating the module's three-lane consistency invariant. The resulting payload/ring behavior needs the deferred GPU integration test.

Requested change: when exec identity changes, invalidate its applied version and reset or invalidate every host-controlled node's actual-state cache before synchronizing the desired set. Review exec-specific replay-event/upload bookkeeping at the same lifecycle boundary. Cover empty, subset, and all-enabled sets after rebind; resetting only applied_version_ is insufficient because the node-level cache also suppresses the update.

CPU reproduction for T171-3 (complete portable runner and CUDA stub)

Save as pr171_native.py and run against the same unchanged commit. Requirements are Python 3 and a C++17-capable g++; no CUDA toolkit or library is required.

CUDA_VISIBLE_DEVICES='' OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 PYTHONNOUSERSITE=1 python pr171_native.py /path/to/DMI

The runner copies the original host namespace without changing its methods, omits the untested device-kernel namespace, and compiles against the repository's original headers with a temporary CUDA API stub. The modeled default-on state is explicit in the test. Generated files are isolated in a temporary directory and removed afterward.

"""Run the actual ToggleRegistry host methods against CPU CUDA API doubles.
Only the untested device-kernel namespace is omitted; no CUDA is loaded.
Usage: python pr171_native.py /path/to/DMI
Requires Python 3 and g++ with C++17 support.
"""
import json
import pathlib
import subprocess
import sys
import tempfile

repo = pathlib.Path(sys.argv[1]).resolve()
source = repo / "native/csrc/ring/toggle_registry.cu"
host, separator, _ = source.read_text().partition("\nnamespace ring {")
assert separator and host.rstrip().endswith("}  // namespace ring_py")
stub = r'''
#pragma once
// CPU-only CUDA API double. No CUDA library is linked or initialized.
#include <cstdlib>
#include <cstdint>
#include <map>
#include <utility>
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;
inline std::map<std::pair<uint64_t, uint64_t>, bool> fake_enabled;
inline int fake_set_enabled_error = 0;
inline unsigned fake_set_enabled_calls = 0;
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<uint64_t>(ex), reinterpret_cast<uint64_t>(node)}] = on;
    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<void*>(1); return 0; }
inline int cudaGraphUpload(cudaGraphExec_t, cudaStream_t) { 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) { *p = reinterpret_cast<void*>(2); return 0; }
inline int cudaEventRecord(cudaEvent_t, cudaStream_t) { return 0; }
inline int cudaEventDestroy(cudaEvent_t) { return 0; }
'''
test = r'''

#include <cassert>
#include <iostream>
int main() {
    ring_py::ToggleRegistry r;
    const uint64_t graph = 0x10, old_exec = 0x20, new_exec = 0x21, node = 0x30;
    fake_enabled[{old_exec, node}] = true;
    fake_enabled[{new_exec, node}] = true; // freshly instantiated default
    r.register_node(graph, 0, 0, node);
    r.bind_exec(graph, old_exec);
    r.set_enabled({});
    assert(r.apply_all() == 0);
    assert((fake_enabled[{old_exec, node}] == false));
    const auto before_rebind_calls = fake_set_enabled_calls;
    r.bind_exec(graph, new_exec);
    const bool reported_current_after_rebind = r.applied_current(graph);
    assert(r.ensure_current(graph) == 0);
    const bool actual_after_ensure = fake_enabled[{new_exec, node}];
    r.set_enabled({}); // even explicitly reapplying the desired empty set
    assert(r.ensure_current(graph) == 0);
    const bool actual_after_reapply = fake_enabled[{new_exec, node}];
    std::cout << "{\"reported_current_after_rebind\":" << reported_current_after_rebind
        << ",\"host_hook_enabled\":" << r.is_enabled(0, 0)
        << ",\"actual_new_exec_after_ensure\":" << actual_after_ensure
        << ",\"actual_new_exec_after_reapply\":" << actual_after_reapply
        << ",\"set_enabled_calls_after_rebind\":"
        << fake_set_enabled_calls - before_rebind_calls;
    assert(reported_current_after_rebind && actual_after_ensure && actual_after_reapply);
    assert(!r.is_enabled(0, 0));
    std::cout << "}\n";
}
'''
with tempfile.TemporaryDirectory(prefix="dmi-pr171-cpu-") as name:
    work = pathlib.Path(name)
    (work / "cuda_runtime.h").write_text(stub)
    generated = work / "repro.cpp"
    binary = work / "repro"
    generated.write_text(host + "\n" + test)
    subprocess.run([
        "g++", "-std=c++17", "-O0", "-DDMX_NO_DEVICE_GRAPH_UPDATE",
        "-I", str(work), "-I", str(source.parent),
        str(generated), "-o", str(binary),
    ], check=True)
    output = subprocess.check_output([str(binary)], text=True)
print(json.dumps(json.loads(output), indent=2))

The separate excluded device-toggle graph upload finding is tracked under #172, where the gather path makes it reachable. This comment intentionally contains only the three #171 findings above. No GPU correctness or performance sign-off is implied by these CPU reproductions.

@zhaoyboo
zhaoyboo force-pushed the feat/cuda-graph-node-toggle branch from 2323c99 to 4e234b4 Compare November 4, 2026 21:41
@zhaoyboo
zhaoyboo force-pushed the perf/producer-grid-geometry branch from a20b6da to 38da7fb Compare November 4, 2026 21:41
@zhaoyboo

zhaoyboo commented Nov 4, 2026

Copy link
Copy Markdown
Collaborator Author

@Samfisheryu thanks for the audit and the runnable reproductions. All three findings are fixed in 4e234b4, which is now stacked on the revised #170. The commit also fixes T172-4, because the upload helper it touches lives in node_toggle.py here.

T171-1 — the CUDAGraph patch is owned and restored

  • The torch.cuda.CUDAGraph replacement now belongs to monitoring sessions, one per engine, started by begin_capture_window.
  • A session ends when its window closes with nothing bound, on NodeToggleController.clear(), or when the transport deactivates.
  • When the last session ends, the original class is restored and _REPLAY_GUARD is disarmed, so a later CUDAGraph(keep_graph=False) keeps its argument.
  • Graph instances created during a session are stamped with that session and engine. A later session cannot attach them to a different transport; a stale graph of the same engine fails closed.
  • Your pr171_python.py, run unchanged against this commit, now fails its "current behavior" assertion. The corrected expectation reports:
    {"class_restored": true, "replay_guard": false, "new_graph_keep_graph": false}.

T171-2 — the replay guard covers only this session's producer graphs

  • Each capture records whether it contains managed producers. That comes from a process-wide producer-launch counter sampled across capture_begin/capture_end, or from nodes recorded for that graph.
  • The guard checks only bound graphs and producer-carrying graphs of the active transport's engine.
  • A graph without producers, or one captured for another engine, replays untouched. In your reproduction, unrelated_graph.replays == 1 after the monitored graph replays.
  • Unbound graphs that do contain producers still fail closed, so real metadata/payload mismatches are not hidden.

T171-3 — rebinding resets per-exec state

  • When a different exec is bound, ToggleRegistry::bind_exec resets that graph's applied version, every host-managed node's last_enabled cache, and the exec-specific replay event and upload record. The next ensure_current() then sets every node of the new exec.
  • The replay guard also rebinds a bound graph that was re-instantiated.
  • Your pr171_native.py, run unchanged, now aborts on its "bug present" assertion. The corrected expectation reports:
    reported_current_after_rebind=0, actual_new_exec_after_ensure=0, actual_new_exec_after_reapply=0, set_enabled_calls_after_rebind=1.
  • tests/test_toggle_registry_host.py compiles the unmodified host half of toggle_registry.cu against a CPU CUDA stub (tests/native/ring/cpu_cuda_stub/). It covers rebind with empty, subset and all-enabled sets.

Regression tests (tests/test_node_toggle.py, CPU fakes)

test_clearing_the_controller_restores_the_cudagraph_class, test_an_unactivated_window_and_a_deactivated_transport_end_the_session, test_the_patch_stays_until_the_last_session_ends, test_a_graph_without_producers_replays_unguarded, test_an_unbound_producer_graph_still_fails_closed, test_an_unbound_graph_with_recorded_producers_fails_closed, test_a_later_session_does_not_adopt_another_engines_graph, test_a_stale_graph_of_the_same_engine_fails_closed, test_a_reinstantiated_bound_graph_is_rebound_before_it_replays, plus the T172-4 tests listed on #172.

Verification for this revision, CPU only: the native extension builds, and the -m cpu suite passes (2294 passed, 459 skipped). The one failure, test_cpu_native_build.py::test_every_torch_including_compile_requests_cxx20[host], also fails on main. No GPU tests, CUDA initialization or benchmarks were run. GPU validation of toggling and device-mode upload waits for the owner's go-ahead.

@Samfisheryu

Copy link
Copy Markdown
Collaborator

Independent follow-up review of 4e234b4. I reran the original cases with corrected expectations, compared the previous revision, and checked the new lifecycle paths. This review remains CPU-only; GPU correctness and performance validation is still deferred until the owner's go-ahead.

Previous findings:

  • T171-1 resolved: clearing the last session restores the original CUDAGraph, disarms the guard, and preserves a subsequent explicit keep_graph=False. Multiple-session ownership and binding after an unbound window closes also work in the independent CPU lifecycle checks.
  • T171-2 resolved for the reported producer-free graph case: unrelated graphs replay; a different engine's old graph is not adopted; stale same-engine producer graphs and late static-producer graphs still fail closed. The process-wide launch counter's behavior under overlapping captures/other-stream eager work remains an isolation boundary to verify, not an additional confirmed blocker here.
  • T171-3 partially resolved: the native bind_exec -> ensure_current cache bug is fixed. The new automatic rebind path still has an eager-mode gap described below.

Targeted repository tests passed: tests/test_node_toggle.py tests/test_toggle_registry_host.py — 46 passed. Passing these tests does not cover the following two issues.

1. P2 — Apply the target state before replaying an automatically rebound exec in eager mode

Code: automatic rebind followed by the lazy-only apply, rebind helper, and native cache invalidation.

Trigger: capture/bind a graph, disable a producer using the supported eager set_active_hooks([]) path at a quiescent point, then call graph.instantiate() again and replay. _rebind_if_reinstantiated() binds the new executable, but _guarded_replay() calls ensure_graph_current() only when toggle_lazy_active is true. The eager branch forwards replay immediately.

Expected: any newly bound executable must receive the desired node state before its first replay, or replay must fail closed until the caller reapplies it.

Observed: the host metadata/reservation gate stays disabled while the newly instantiated executable retains its default-on producer. The CPU model reports:

lazy:  host_enabled=false, new_exec_producer_on=false, ensure_calls=2
eager: host_enabled=false, new_exec_producer_on=true,  ensure_calls=1

new_exec_producer_on is the CPU API double's state, not a GPU measurement. Separately compiling the unchanged native host implementation against CUDA API doubles confirms the missing step: after the new bind, actual_new_exec_after_rebind=1 and set_enabled_calls_only_rebind=0; only a subsequent ensure_current() changes that node to off. The original native defect is fixed: reported_current_after_rebind=0, and the explicit ensure now makes one SetEnabled call and yields actual_new_exec_after_ensure=0 (the old revision incorrectly yielded 1 with zero calls).

The existing automatic reinstantiate regression uses the transport double's default lazy=True. The host-side eager test explicitly invokes apply_all(), unlike the Python replay path.

Requested change: synchronize the new exec before replay whenever automatic rebinding occurs, including eager mode, or explicitly reject the eager replay until reapplication. Add the eager counterpart to the reinstantiate regression.

Runnable CPU call-sequence reproduction using this revision's test doubles

Save as pr171_rebind_cpu.py and run from the checkout root:

CUDA_VISIBLE_DEVICES='' OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 PYTHONNOUSERSITE=1 PYTHONPATH=src nice -n 15 python pr171_rebind_cpu.py

This imports PyTorch but replaces its graph class before constructing any graph; it does not initialize CUDA or execute GPU work. The test transport represents an already active gate. The probe demonstrates the real helper forwarding the eager replay without an apply call; it does not claim to execute a GPU producer.

import json
import runpy
import pytest
import torch
from dmi import node_toggle as nt
from dmi.transport import ring

fixtures = runpy.run_path('tests/test_node_toggle.py')
results = {}
for lazy in (True, False):
    with pytest.MonkeyPatch.context() as patch:
        patch.setattr(torch.cuda, 'CUDAGraph', fixtures['LifecycleGraph'])
        fixtures['_reset_patch_state'](patch, nt)
        patch.setattr(ring, '_active_transport', None)
        engine = fixtures['LifecycleEngine']()
        nt.begin_capture_window(engine)
        graph = torch.cuda.CUDAGraph()
        graph.capture_begin()
        engine.capturing = graph
        engine.produce()
        engine.capturing = None
        graph.capture_end()
        nt.bind_graph(engine, graph)
        nt.end_capture_window(engine)
        # The author's transport double represents an already active gate.
        transport = fixtures['GatedTransport'](engine, lazy=lazy)
        patch.setattr(ring, '_active_transport', transport)
        nt.set_replay_guard(True)
        graph.replay()
        old_exec = engine.bound[graph.raw]
        engine.ensured.clear()
        graph.instantiate()
        graph.replay()
        results['lazy' if lazy else 'eager'] = {
            'new_exec_was_bound': engine.bound[graph.raw] == graph.exec_ptr != old_exec,
            'ensure_calls_before_new_exec_replay': len(engine.ensured),
            'replays_forwarded': graph.replays,
        }
print(json.dumps(results, indent=2))

Output:

{
  "lazy": {
    "new_exec_was_bound": true,
    "ensure_calls_before_new_exec_replay": 1,
    "replays_forwarded": 2
  },
  "eager": {
    "new_exec_was_bound": true,
    "ensure_calls_before_new_exec_replay": 0,
    "replays_forwarded": 2
  }
}

2. P2 — Supply the new producer-counter symbol to the standalone native ring test

The new RingEnginePy::producer_launch_count() calls ring_producer_launch_count(), whose definition lives in ring_torch_op.cpp.

The test_ring_engine link target includes ring_engine_py.cu but not ring_torch_op.cpp. Its existing standalone stubs cover the older diagnostic/capture hooks, but omit the new counter getter. Consequently this target no longer links.

This is observed in the latest stacked revision's native CI job, not inferred from an unavailable local GPU test:

undefined reference to `ring_producer_launch_count()'

The failing job is on #174; the new unresolved dependency is introduced here in #171 and inherited downstream. This is a native test build failure, distinct from a production-extension link failure. Building the extension and passing the CPU pytest suite do not validate this standalone target.

Requested change: provide the appropriate standalone counter stub, following the existing diagnostic stubs, or factor/link the implementation so this target resolves the symbol. Verify the native target links in CI before considering the native regression gate green.

The excluded device-toggle upload issue T172-4 is tracked in the #172 follow-up: the helper correction here is bypassed by #172's missing gather/head/fused producer-counter updates. It is not counted again as a separate #171 issue.

@zhaoyboo

zhaoyboo commented Nov 5, 2026

Copy link
Copy Markdown
Collaborator Author

@Samfisheryu thanks for the follow-up. Both findings are fixed in 541b10d.

1. Eager replay of an automatically rebound exec

  • _rebind_if_reinstantiated() now returns True when it binds a new exec. _guarded_replay() calls ensure_graph_current() for that graph before the replay in eager mode too, not only when lazy mode is on. A graph is bound to exactly one engine (_dmi_bound_engine), so a re-instantiation rebinds one binding; that is the one brought current. Nothing of the new exec has launched yet, so applying to it is safe under the eager contract.
  • Ordinary eager replays of an unchanged exec still make no ensure call.
  • In the same commit, _captured_producers() no longer lets an unchanged launch counter override recorded producer nodes. A capture counts as carrying producers if the counter moved or the registry recorded nodes for it. This is the helper half of T172-4; the op half is on [Core] Add per-hook request-row gather, PackedDemandPlanner, and fused inline projection heads #172.
  • Tests: the reinstantiate regression is parametrized over lazy and eager. It records the ensure calls as the new exec's first replay begins, so it checks order as well as count. New tests: test_eager_replays_of_an_unchanged_exec_do_not_reapply and test_recorded_producer_nodes_win_over_an_unchanged_launch_counter. Against the old helper, both the eager case and the node-evidence test fail.
  • Your pr171_rebind_cpu.py, unchanged:
    {"lazy":  {"new_exec_was_bound": true, "ensure_calls_before_new_exec_replay": 1, "replays_forwarded": 2},
     "eager": {"new_exec_was_bound": true, "ensure_calls_before_new_exec_replay": 1, "replays_forwarded": 2}}

2. ring_producer_launch_count() in the standalone ring test

  • tests/native/ring/test_ring_engine.cu now has a uint64_t ring_producer_launch_count() stub, next to the existing diagnostic and capture stubs. It has the same global-namespace signature as ring_torch_op.h.
  • make -C tests/native/ring clean all CXX_STD=c++20 links every target on this branch and on the full stack ([Core] Expose ring preparation, graph toggle, gather upload, and staging wait telemetry #174). test_record_consumer reports 151 passed and test_clickhouse_record_sink 21 passed, with CUDA_VISIBLE_DEVICES empty.

tests/test_node_toggle.py, tests/test_toggle_registry_host.py and tests/test_producer_grid_geometry.py pass (50 tests). As before, this is CPU-only: no GPU tests, CUDA initialization or benchmarks were run. GPU validation stays deferred until the owner gives the go-ahead.

@Samfisheryu

Copy link
Copy Markdown
Collaborator

Independent third review of 541b10d815fbb0c5cd5adc45a841a441d012e90f: both round-two items are resolved at the reviewed CPU/host evidence level.

  • R2-171-eager / T171-3: the original short reproduction now records one ensure call before the replacement exec replays in both lazy and eager modes. The independent lifecycle model also confirms the replacement node receives the intended off state. An unchanged eager exec does not acquire an extra apply call. An apply error stops replay.
  • R2-171-link: the standalone ring test now provides the correct global counter stub. The latest CI successfully links the native tests; its CPU and live jobs pass as well.

Corrected output from the unchanged short reproduction:

lazy:  new_exec_was_bound=true, ensure_calls_before_new_exec_replay=1, replays_forwarded=2
eager: new_exec_was_bound=true, ensure_calls_before_new_exec_replay=1, replays_forwarded=2

The earlier class-restoration, session ownership and unrelated-graph cases remain correct in the independent checks. Recorded producer nodes now also take precedence over an unchanged launch counter.

The combined toggle/registry/real-op host test selection passed 85 CPU tests on the dependent #172 revision containing this exact helper. It uses CPU recording boundaries; it does not execute device kernels. No additional confirmed issue was found in the revised helper. Actual graph replay, device-side behavior and performance remain for the separately authorized GPU phase.

@Samfisheryu Samfisheryu left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Approved at 541b10d815fbb0c5cd5adc45a841a441d012e90f following the completed independent code and CPU/host review. The issues raised for this PR have been addressed and independently rechecked. The reviewed commit's CI is successful.

Review record: #171 (comment)

GPU execution and end-to-end performance benchmarks remain a separate follow-up, pending the owner's explicit instruction as agreed. This approval records the completed code/CPU review and does not claim those measurements.

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.
@zhaoyboo
zhaoyboo changed the base branch from perf/producer-grid-geometry to main November 5, 2026 23:12
@zhaoyboo
zhaoyboo force-pushed the feat/cuda-graph-node-toggle branch from 541b10d to ffbedf5 Compare November 5, 2026 23:12
@zhaoyboo
zhaoyboo merged commit 202ec5a into main Nov 5, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants