[feat] static-buffer capture form for the cudagraph step body (Path A)

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

@@ -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
 ```

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

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"]

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

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)