From cd0fbbd3a011018d8a9dee0819ec052816734551 Mon Sep 17 00:00:00 2001 From: SolitaryThinker Date: Wed, 17 Jun 2026 13:28:56 -0700 Subject: [PATCH] [feat] static-buffer capture form for the cudagraph step body (Path A) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Close the loudest deferred gap from the cudagraph audit: the capturable step now binds its I/O to address-stable static buffers (modeling real CUDA static I/O buffers), instead of allocating fresh arrays per call. - StaticWorkspace: address-stable buffers allocated once per capture key; bind() copies the current step's inputs in place via np.copyto, which RAISES on a shape/dtype mismatch — turning the weak peak_activation_bytes proxy into a real key-soundness backstop (a step whose shape doesn't fit can't replay an incompatible graph). Output written into a static buffer too. - WorkPlan.graph_fn / graph_inputs: a capturable step exposes its deterministic op-structure as graph_fn(model, workspace) reading EVERY per-step input (latent, sigmas, conditioning, scale) from the workspace — never from closure over per-step data — plus the dict of current values. Loops without both stay on the eager path. wan21 factors a shared _velocity() so graph_fn and the eager run stay bit-identical. - Capturer dispatch captures/replays via graph_fn against the keyed workspace; the workspace is shared per key on the instance. Correct under the engine's synchronous step execution (bind+graph_fn atomic per dispatch, output returned as a copy) — proven by the batch-of-N interleave gate running two same-key requests through the shared workspace bit-identically. A concurrent/multi-stream executor would need a per-stream pool (documented). - 2 new tests (no-static-form eager-break, static-buffer shape-mismatch raises); workspace collision test split into bytes-proxy vs shape-backstop. 197 tests pass. --- v2/README.md | 15 +++--- v2/loop/contracts.py | 9 ++++ v2/models/wan21/loop.py | 36 +++++++++---- v2/runtime/cudagraph.py | 104 +++++++++++++++++++++++++++---------- v2/tests/test_cudagraph.py | 30 ++++++++--- 5 files changed, 146 insertions(+), 48 deletions(-) diff --git a/v2/README.md b/v2/README.md index ad6b35783..14c87bb26 100644 --- a/v2/README.md +++ b/v2/README.md @@ -24,12 +24,15 @@ pinned to the rollout's kernel for C2 correctness. The `dit` and `vae` component overrides (`text_encoder` demonstrates the device→cpu fallback). wan21's denoise also runs under **piecewise CUDA-graph capture/replay** at the step boundary (Path A): the capture *lifecycle* — key (device/arch/shape/resident-weights/branch-set+expert), eager-break for data-dependent steps -(SDE / interceptor override), and version-eviction on weight sync — is wired and tested. Not yet -done (deferred, and bites only a real GPU): the static-buffer refactor of the step body, admission -budgeting of capture cost (`GRAPH_CAPTURE`), the AR/vocoder op families (omni/qwen-omni) and the -full torch/CUDA path. Replay here re-runs the step thunk — it models the lifecycle, not the speedup. +(SDE / interceptor override), version-eviction on weight sync, and the **static-buffer** discipline +(the capturable step reads every per-step input from an address-stable `StaticWorkspace`; rebinding +in place raises on a shape mismatch — a real key-soundness check) — is wired and tested. Not yet +done (deferred, and bites only a real GPU): admission budgeting of capture cost (`GRAPH_CAPTURE`), +a per-stream workspace pool for a concurrent executor, the AR/vocoder op families (omni/qwen-omni), +and the full torch/CUDA path. Replay here re-runs the step's `graph_fn` against the captured static +buffers — it models the capture lifecycle + buffer discipline, not the GPU speedup. -**195 tests, 31 files**, run two ways (`pytest` and a zero-dependency runner). Python 3.10+ and numpy only. +**197 tests, 31 files**, run two ways (`pytest` and a zero-dependency runner). Python 3.10+ and numpy only. ## Scope @@ -105,7 +108,7 @@ Stress tests (the design held — see **designv4 §9**): ```bash cd /Users/willlin/src/FastVideo-mini -python3 -m pytest v2/tests/ -q # 195 tests +python3 -m pytest v2/tests/ -q # 197 tests python3 v2/run_tests.py # same suite, ZERO deps (no pytest needed) python3 -m v2.examples # the worked examples ``` diff --git a/v2/loop/contracts.py b/v2/loop/contracts.py index 8e57c6bd6..40b8544c6 100644 --- a/v2/loop/contracts.py +++ b/v2/loop/contracts.py @@ -121,6 +121,15 @@ class WorkPlan: # the captured graph's *shape of computation*. Part of the capture key so a step with a different # branch set / expert never replays an incompatible graph. Empty ⇒ structure fixed by shape alone. graph_key: tuple = () + # The static-buffer capture form (design_v3 §6.2). A capturable step exposes its deterministic op + # structure as ``graph_fn(model, workspace) -> StepResult`` reading EVERY per-step-varying input + # (latent, sigmas, conditioning) from ``workspace`` — never from closure — plus ``graph_inputs``, + # the dict of those current values. The runtime allocates address-stable buffers once per key and + # rebinds ``graph_inputs`` into them in place each step (modeling CUDA static I/O buffers). Loops + # that don't provide both stay on the eager path (the runtime eager-breaks them). ``run`` remains + # the eager thunk (override / stochastic paths, and the no-capture baseline). + graph_fn: Any = None + graph_inputs: dict[str, Any] | None = None @dataclass diff --git a/v2/models/wan21/loop.py b/v2/models/wan21/loop.py index 260107d3e..3bbc569d3 100644 --- a/v2/models/wan21/loop.py +++ b/v2/models/wan21/loop.py @@ -80,18 +80,21 @@ class WanDenoiseLoop: cfg, precision = self.cfg, self.precision sde, noise_scale, rng = st.scratch.get("sde", False), st.scratch.get("sde_noise_scale", 0.7), st.rng + def _velocity(model, x_, sigma_t_, pe_, ne_, scale_): + # The conditioned forward + CFG combine. Solver/forward dispatch through the platform's + # kernel table (numpy on CPU, the device kernel on a GPU/accel backend — design_v3 §17). + dit = model.component(expert_id) + preds = {b: dit(x_, pe_ if b == "cond" else ne_, sigma_t_) for b in branches} + return precision.cast(cfg.combine(preds, scale_, sctx, cfg_state)) + def run(model, override=None): - # Solver primitives are dispatched through the platform's kernel table: on CPU these - # resolve to the numpy reference; on a GPU/accel backend they resolve to that device's - # kernel — same loop, same StepResult shape (design_v3 §17). + # The EAGER thunk: handles the override (cached-prediction) path and the stochastic SDE + # rollout path — both eager-break under capture. Allocates freely. kernels = model.platform.kernels if override is not None and "noise_pred" in override: - velocity = np.asarray(override["noise_pred"], dtype="float32") + velocity = precision.cast(np.asarray(override["noise_pred"], dtype="float32")) else: - dit = model.component(expert_id) - preds = {b: dit(x, pe if b == "cond" else ne, sigma_t) for b in branches} - velocity = cfg.combine(preds, scale, sctx, cfg_state) - velocity = precision.cast(velocity) + velocity = _velocity(model, x, sigma_t, pe, ne, scale) if sde: # FlowGRPO rollout: stochastic + log-prob noise = rng.standard_normal(x.shape) x_next, logp, _m, _s = kernels.get(FLOW_SDE_STEP)( @@ -103,6 +106,17 @@ class WanDenoiseLoop: return StepResult(output={"noise_pred": np.asarray(velocity, dtype="float32"), "latents": x_next.astype("float32")}) + def graph_fn(model, ws): + # The CAPTURABLE op-structure: reads EVERY per-step input from the static workspace (never + # from closure over per-step data) and writes into the static output buffer — so the same + # captured graph replays correctly with rebound buffer contents (design_v3 §6.2). + st_t, st_n = float(ws["sigma_t"]), float(ws["sigma_next"]) + velocity = _velocity(model, ws["x"], st_t, ws["pe"], ws["ne"], float(ws["scale"])) + x_next = model.platform.kernels.get(FLOW_MATCH_STEP)(precision.cast(ws["x"]), velocity, st_t, st_n) + np.copyto(ws["out"], x_next.astype("float32")) # write the static output buffer + return StepResult(output={"noise_pred": np.asarray(velocity, dtype="float32"), + "latents": np.array(ws["out"], copy=True)}) + cond_bytes = sum(int(np.asarray(e).nbytes) for e in (pe, ne) if e is not None) res = ResourceRequest( compute_seconds=self.cost.predict(int(np.prod(x.shape)), float(len(branches))), @@ -127,7 +141,11 @@ class WanDenoiseLoop: # is sound for branch-set/expert-determined policies (ClassicCFG); a policy whose op # structure forks on other state (PerModalityCFG modality) would need a graph_key hook. capturable=not sde, - graph_key=(tuple(sorted(branches)), expert_id, precision.scheduler_step_in_fp32)) + graph_key=(tuple(sorted(branches)), expert_id, precision.scheduler_step_in_fp32), + # the static-buffer capture form: graph_fn reads all per-step inputs from the workspace + graph_fn=graph_fn, + graph_inputs={"x": x, "sigma_t": sigma_t, "sigma_next": sigma_next, + "pe": pe, "ne": ne, "scale": scale}) def advance(self, st: LoopState, result: StepResult) -> LoopState: st.latents["video"] = result.output["latents"] diff --git a/v2/runtime/cudagraph.py b/v2/runtime/cudagraph.py index e68980e4a..cdc13e828 100644 --- a/v2/runtime/cudagraph.py +++ b/v2/runtime/cudagraph.py @@ -24,17 +24,27 @@ correctness-critical and bug-prone regardless of device: AUTOMATIC: the resident-weight version is *in* the key, so a synced component's old graphs simply become unreachable and the next step recaptures. -The capturer NEVER executes a *stored* thunk — it always runs the *current* step's thunk (so it -cannot smear per-request state across interleaved requests). The ``CapturedGraph`` is the keyed -record used to account capture-vs-replay and to detect a same-key/different-workspace clash. +**Static buffers (the capture-safety discipline).** A captured graph references FIXED memory +addresses — replay copies the current step's inputs *into* those buffers in place, then relaunches; +only the contents change, never the op-structure or the addresses. So a capturable step must expose +its computation as ``graph_fn(model, workspace)`` that reads EVERY per-step-varying input (latent, +sigmas, conditioning) from the workspace — never from closure — and must not allocate fresh I/O each +call. ``StaticWorkspace`` models those buffers: ``alloc`` reserves them once per key; ``bind`` copies +the current step's values in via ``np.copyto``, which RAISES on a shape/dtype mismatch — so an +ill-fitting input cannot silently replay an incompatible graph (a real key-soundness backstop, not +the weak ``peak_activation_bytes`` proxy). Because all per-step/per-request inputs flow through the +workspace, the SAME captured buffers are correctly reused across interleaved requests; the capturer +re-invokes the *current* step's ``graph_fn`` (CPU can't bake op-structure into a replayable object), +which is sound because the key guarantees every step under it is structurally identical. -Honest scope — what this mini does NOT yet model (these bite a real GPU, not the CPU tests): - * **Static-buffer refactor (not done).** The wan21 step body still allocates per call - (``np.asarray``/``.astype``/host RNG). On real CUDA a capturable region must bind its I/O and - scratch to reused static buffers first; ``workspace_bytes`` here is a *declared number*, not a - buffer binding. So the workspace check below is a shape-change proxy (it compares the step's - ``peak_activation_bytes``), NOT the kernel-level capture-safety backstop a real integration needs - — the registered per-kernel ``workspace_bytes`` is declared in the matrix but not yet consulted. +The per-key workspace is SHARED across requests (it lives on the instance). That is correct under +this engine's synchronous, one-step-at-a-time execution: ``dispatch`` rebinds then runs ``graph_fn`` +atomically within a single step, and returns a copy of the output buffer, so an interleaved request's +next step overwrites the buffers only after the prior step fully consumed them. (The batch-of-N +interleave-parity gate exercises two same-key requests sharing this workspace and stays bit-identical.) +A truly concurrent / multi-stream executor would instead need a per-stream workspace pool. + +Honest scope — what this mini still does NOT model (bites a real GPU, not the CPU tests): * **Capture cost is not budgeted.** ``WorkUnitKind.GRAPH_CAPTURE`` and ``ResourceRequest.graph_capture_size`` are declared placeholders; admission does not yet account the one-time capture (warmup) cost. @@ -46,11 +56,47 @@ from __future__ import annotations from dataclasses import dataclass from typing import Any, Callable +import numpy as np + + +class StaticWorkspace: + """Address-stable buffers for a captured region (models CUDA static I/O buffers). Allocated once + per capture key; ``bind`` updates contents IN PLACE (``np.copyto``), raising on a shape mismatch + so an ill-fitting input can't replay an incompatible graph.""" + + def __init__(self, buffers: dict[str, Any]): + self._b = buffers + + @staticmethod + def _mk(v: Any): + return None if v is None else np.array(v) # scalars become 0-d arrays; arrays are copied + + @classmethod + def alloc(cls, inputs: dict[str, Any]) -> "StaticWorkspace": + b = {k: cls._mk(v) for k, v in inputs.items()} + x = inputs.get("x") + if x is not None: + b["out"] = np.array(x) # static output buffer, latent-shaped + return cls(b) + + def bind(self, inputs: dict[str, Any]) -> None: + for k, v in inputs.items(): + cur = self._b.get(k) + if cur is None: + if v is not None: + raise ValueError(f"workspace[{k!r}] was None at capture but has a value on replay") + continue + np.copyto(cur, v) # in place; raises if shapes/dtypes don't fit + + def __getitem__(self, k: str): + return self._b[k] + @dataclass class CapturedGraph: """A keyed record standing in for a captured ``torch.cuda.CUDAGraph`` + its static buffers.""" key: tuple + workspace: Any # StaticWorkspace (the reused fixed buffers), or None workspace_bytes: int replays: int = 0 @@ -78,37 +124,41 @@ class GraphCapturer: def dispatch(self, plan: Any, instance: Any, override: Any, eager: Callable[[Any], Any]) -> Any: - """Run one step under the capture lifecycle. ``eager(override)`` executes the step thunk and - returns its StepResult|dict. Returns that result unchanged on every path (capture, replay, - eager-break) — replay re-runs the current thunk, so the result is always correct for *this* - step; the cache only governs accounting + structural-compatibility checks.""" + """Run one step under the capture lifecycle. ``eager(override)`` runs the eager step thunk + (override / stochastic / no-capture paths). A capturable step runs its ``graph_fn`` against + the keyed ``StaticWorkspace``; replay rebinds the current inputs into the SAME buffers and + re-invokes the current ``graph_fn`` (sound because the key fixes the op-structure).""" if not self.enabled: return eager(override) - # Eager-break: a model-declared non-capturable step (host RNG / data-dependent) or a step an - # interceptor overrode this iteration never enters the graph cache. - if not getattr(plan, "capturable", True) or override is not None: + # Eager-break: a model-declared non-capturable step (host RNG / data-dependent), an + # interceptor override this iteration, or a step with no static graph form never enters the + # graph cache. + has_static = plan.graph_fn is not None and plan.graph_inputs is not None + if not getattr(plan, "capturable", True) or override is not None or not has_static: self.stats["eager_breaks"] += 1 return eager(override) key = self.key(plan, instance) - ws = int(plan.resources.peak_activation_bytes) + ws_bytes = int(plan.resources.peak_activation_bytes) g = self._graphs.get(key) if g is None: - # MISS → capture (cheap warmup): record the keyed graph, then run the step once. - self._graphs[key] = CapturedGraph(key=key, workspace_bytes=ws) + # MISS → capture (cheap warmup): allocate the static buffers, bind inputs, run once. + ws = StaticWorkspace.alloc(plan.graph_inputs) + self._graphs[key] = CapturedGraph(key=key, workspace=ws, workspace_bytes=ws_bytes) self.stats["captures"] += 1 - return eager(None) + return plan.graph_fn(instance, ws) - # HIT → replay. Safety net: everything under one key must need the same static workspace; a - # mismatch means the key is insufficient (it would replay an incompatible graph on a GPU). - if g.workspace_bytes != ws: + # HIT → replay. peak-workspace proxy guard (declared static workspace must match)... + if g.workspace_bytes != ws_bytes: raise RuntimeError( f"cudagraph key collision: key {key} captured a {g.workspace_bytes}-byte workspace " - f"but this step needs {ws} — the capture key is insufficient and would corrupt a " - f"real graph. This is a keying bug, not a runtime condition.") + f"but this step needs {ws_bytes} — the capture key is insufficient and would corrupt " + f"a real graph. This is a keying bug, not a runtime condition.") + # ...and the real backstop: rebinding into the fixed buffers raises if shapes don't fit. + g.workspace.bind(plan.graph_inputs) g.replays += 1 self.stats["replays"] += 1 - return eager(None) + return plan.graph_fn(instance, g.workspace) def invalidate(self, component_ids: Any = None) -> int: """Evict captured graphs whose resident weights just changed (RL weight sync), mirroring diff --git a/v2/tests/test_cudagraph.py b/v2/tests/test_cudagraph.py index d800e3bf7..9e25046b0 100644 --- a/v2/tests/test_cudagraph.py +++ b/v2/tests/test_cudagraph.py @@ -126,13 +126,15 @@ class _Res: class _Plan: - def __init__(self, ws=64, gk=(), cap=True, bk=("step",)): + def __init__(self, ws=64, gk=(), cap=True, bk=("step",), gin=None, gfn=None): self.loop_id = "L" self.shape_sig = _Sig(bk) self.resources = _Res(ws) self.graph_key = gk self.capturable = cap self.run = None + self.graph_inputs = gin if gin is not None else {"x": np.zeros((2, 2), dtype="float32")} + self.graph_fn = gfn if gfn is not None else (lambda model, w: {"ok": True}) class _Plat: @@ -164,23 +166,39 @@ def test_override_and_noncapturable_eager_break(): assert cap.stats["eager_breaks"] == 2 and cap.stats["captures"] == 0 +def test_no_static_graph_form_eager_breaks(): + cap, inst = GraphCapturer(), _Inst() + p = _Plan(); p.graph_fn = None # a plan with no static capture form + cap.dispatch(p, inst, None, lambda ov: {"r": 1}) + assert cap.stats["eager_breaks"] == 1 and cap.stats["captures"] == 0 + + def test_invalidate_drops_only_matching_component_graphs(): from v2.runtime.cudagraph import CapturedGraph cap = GraphCapturer() # two graphs whose keys carry different resident components (key[4] = resident ((cid, ver), ...)) k_dit = ("accel", "sm90", "L", ("a",), (("transformer", "v0"),), ()) k_txt = ("accel", "sm90", "L", ("a",), (("text_encoder", "v0"),), ()) - cap._graphs[k_dit] = CapturedGraph(k_dit, 1) - cap._graphs[k_txt] = CapturedGraph(k_txt, 1) + cap._graphs[k_dit] = CapturedGraph(k_dit, None, 1) + cap._graphs[k_txt] = CapturedGraph(k_txt, None, 1) dropped = cap.invalidate({"transformer"}) assert dropped == 1 and cap.captured == 1 and k_txt in cap._graphs # only the dit graph evicted -def test_workspace_collision_raises(): - """Same key, different static workspace ⇒ the key is insufficient and would corrupt a real - graph. The capturer refuses (the safety net), rather than silently replaying.""" +def test_workspace_bytes_collision_raises(): + """Same key, different declared static-workspace size ⇒ the key is insufficient. Refuse.""" cap, inst = GraphCapturer(), _Inst() eager = lambda ov: {"ok": True} cap.dispatch(_Plan(ws=100), inst, None, eager) # capture: key→100B workspace with pytest.raises(RuntimeError, match="key collision"): cap.dispatch(_Plan(ws=200), inst, None, eager) # same key, 200B ⇒ refuse + + +def test_static_buffer_shape_mismatch_raises(): + """The real backstop: rebinding an input whose shape doesn't fit the captured static buffer + raises (np.copyto) — catching a key that let incompatible shapes share a graph.""" + cap, inst = GraphCapturer(), _Inst() + eager = lambda ov: {} + cap.dispatch(_Plan(gin={"x": np.zeros((2, 2), dtype="float32")}), inst, None, eager) # 2x2 buffer + with pytest.raises(ValueError): # 4x4 won't fit the captured 2x2 buffer + cap.dispatch(_Plan(gin={"x": np.zeros((4, 4), dtype="float32")}), inst, None, eager)