The torch adapter boundary marshalled the latent host<->device on EVERY denoise
step: _t uploaded the latent (and re-uploaded the text embeds) and _n downloaded
the velocity with a forced CUDA sync — 2*N PCIe copies + N syncs per generation,
buying nothing, since the latent could stay resident on the GPU the whole loop.
Root cause was a numpy loop surface. But the loop MATH is already array-agnostic
(CFG combine + flow-match Euler are pure arithmetic; the solver kernel already
passes torch through). So introduce a per-platform array namespace (v2/platform/
array_ns): numpy on CPU (torch-free — the parity mini is unchanged), torch-on-
device on cuda. The latent is seeded with numpy and uploaded ONCE; it then stays
resident through forward -> CFG combine -> solver -> next step; a single host
marshal happens at the request/output boundary (engine._to_artifact).
Opt-in per recipe via ModelCard.device_io (set on the Wan cards). When set on a
GPU box, build_component flips the components' TorchComponent.device_io so _out
keeps tensors on-device (in fp32, matching the old _to_numpy cast so the combine
dtype is unchanged). Un-migrated families and the CPU toy keep numpy in/out.
Also: PrecisionPolicy.cast is array-preserving; _t accepts resident tensors.
Verified BIT-IDENTICAL on real Wan2.1-1.3B / H100: the on-device latent equals
the pre-change numpy-path latent exactly (max_abs_diff 0.0, np.array_equal True).
CPU mini holds 237 passed / 5 pre-existing; pre-commit clean.
Other WanDenoiseLoop families can flip device_io next (per-family GPU re-verify);
non-Wan loops migrate to the xp namespace later.
Sweep all of v2/ through pre-commit (was previously only run on changed
files). Fixes surfaced across untouched modules:
- ruff: isinstance-tuple -> X | Y (UP038), try/except/pass ->
contextlib.suppress (SIM105), negated-return (SIM103), %-format ->
f-string (UP031).
- mypy: add annotations for no-untyped-call + var-annotated across recipes,
training methods, torch/toy backends, and serving.
- yapf reflow of the SF-Wan KV-cache call sites (semantics unchanged).
yapf/ruff/codespell/mypy all pass; v2 tests 237 passed / 5 pre-existing
(bucket_c torch-absent on CPU venv).
Sweep all 296 v2 modules: simplify verbose docstrings/comments and remove 376 dangling
"(design_vN §X)" citations to the now-deleted design docs (v2/README.md is the source of
truth). Comment/docstring-only — AST-verified code-identical; the CPU suite holds at 237
passed / 5 pre-existing. Also applies yapf + ruff --fix auto-fixes (import ordering,
forward-ref annotation de-quoting under `from __future__ import annotations`; behavior-
neutral, suite-confirmed) and adds the legitimate domain terms mot/clen/te to the codespell
ignore-list. Remaining ruff (24) + mypy (68 no-untyped-call) findings are pre-existing v2
debt, untouched here.
Frees v2/models/ to become the vendored-architecture namespace that mirrors fastvideo/models
(part of making v2 self-contained / able to replace fastvideo). The v2 recipe layer (per-family
card.py/loop.py/program.py + common.py + the build_*_engine re-exports) is the recipe, not the
architectures, so it moves to v2/recipes/. The pure-numpy toy/parity implementations (ToyDiT etc.)
move from v2/models/backend.py to v2/platform/backends/toy.py (alongside cpu.py/accel.py/torch_*).
Mechanical: all imports are absolute, so v2.models.<x> -> v2.recipes.<x> and v2.models.backend ->
v2.platform.backends.toy across v2/ + examples/ (89 files, 178 refs). No behavior change.
CPU mini green (202 passed, 2 skipped); all v2 files compile.
Mechanical conversion of every relative import under v2/ to an absolute v2.* path
(from .x / ..x / ...x -> from v2.<pkg>.x) so imports are unambiguous, grep-able, and
stable when code is copied/moved between entrypoints (VideoGenerator, CLI, server).
Surgical prefix-only rewrite: only the 'from <dots><module>' prefix changed — import
names, parentheses, multi-line formatting, comments, and ordering are byte-for-byte
preserved (no collapsing, no reorder, no unrelated reformatting).
- 471 imports across 125 files; v2/tests/ was already absolute (untouched).
- Validated: all 125 files compile, every 'from v2.* import' target resolves to a real
module/package, zero relative imports remain (full sweep), CPU mini suite green
(202 passed, 2 skipped).
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.
Wire the capture/replay lifecycle into the driven-loop step boundary — the
other half of Path A (hand-fused kernels behind the registry + piecewise
cudagraphs, no compiler). Models and tests the correctness-critical control
logic; replay re-runs the current step thunk (CPU models the lifecycle, not the
GPU speedup).
- GraphCapturer on the instance (cross-request cache), wired into
RuntimeLoopContext.execute and gated by LoopSpec.graph_capture ==
"breakable_cudagraph". Capture key = (device, arch, loop, shape_sig,
resident-weight-versions, graph_key). Eager-break for non-capturable (SDE) and
interceptor-overridden steps. Never executes a stored thunk, so interleaved
requests can't smear state.
- WorkPlan.capturable / graph_key. wan21 sets capturable=not sde and folds
compute dtype (shape_sig.dtype) + CFG branch set + expert + scheduler-precision
into the key — closing a cross-precision key-collision corruption path a
non-fp32 build would otherwise hit on a real GPU (audit finding).
- Version-in-key auto-invalidation + real eviction: set_weights_version evicts
the synced component's graphs (duck-typed, so card/ imports no runtime),
preventing a GPU graph leak across FlowGRPO's per-iteration syncs.
- register_kernel gains the workspace_bytes capture-safety contract (declared in
the matrix; cuda cells declare real scratch, numpy reference is 0).
- 11 tests: capture-once/replay-many, eager-break (SDE + override), capture ≡
pure-eager bit-identical, recapture on shape + weight-version change with
eviction, eager-loop gating, accel-backend capture, and capturer unit tests
(key discrimination, eager-break, workspace-collision safety net, eviction).
Honestly deferred (bite a real GPU, not the CPU tests): the static-buffer
refactor of the step body, admission budgeting of capture cost (GRAPH_CAPTURE),
and per-card opt-in beyond wan21 — all documented in cudagraph.py + README.
195 tests pass.
The remaining design_v3 probes (all except 8b speculative decoding). All fit with no new runtime
primitive (148 tests pass).
#6 LTX-2 joint audio+video (§9.11) — LTX-2 declared an audio_vae required_for t2vs but never used
it; now a single 2-stage denoise carries a synchronized audio latent (conditioned on video),
applies per-modality CFG (guidance_per_modality), and decodes via video VAE + audio VAE → video +
audio. Gated on requesting audio, so the T2V path is byte-identical (existing tests untouched).
ToyAudioVAE; build_ltx2_av_program. 5 tests.
#4 Live weight-sync under in-flight serving (§9.14) — WeightSyncController makes the freeze → drain →
transfer → bump version + invalidate → resume lifecycle explicit. Tests: a mid-flight swap corrupts
(the hazard); draining first leaves the in-flight request bit-identical to baseline while a
post-sync request reflects new weights; transformer-only sync, so the frozen text-encoder cache
survives. The RL flywheel's hardest correctness. 3 tests.
#5 Reward-model-as-a-served-card (§9.15) — a reward model is a card (scorer + a score loop emitting
REWARD_BATCH units); ServedRewardScorer drop-in-replaces the numpy scorer so any RL method becomes
RLHF/RLAIF with no method change. ToyRewardModel; models/reward/. 4 tests.
#7 Content-adaptive control flow (§9.12) — CacheDiTDenoiseLoop (isolated WanDenoiseLoop subclass)
reuses the cached velocity when predictions barely change (cache-dit skip) and early-exits on
convergence — variable step count; interleave parity holds across ragged loops. models/adaptive/. 4 tests.
#8a Nested workflows (§9.13) — a workflow stage can invoke another workflow (engine.run routes ids);
requires/validate recurse; cycles caught at registration + a run-time guard (engine._wf_running).
build_t2i_i2v_extend_workflow. 5 tests.
designv4 §9.11–§9.15 + falsifier/layout/closing updates (148 tests, 25 files). Also removed a
pre-existing unused import in ltx2/loop.py.
Targets the three design_v3 claims that were most load-bearing AND least
exercised (sessions/realtime, training-plane boundary, the WorkUnit-generality
falsifier). All fit with no new runtime primitive (127 tests pass).
1. Interactive world-model session (runtime/session.py, §9.8) — the Session
plane had ZERO coverage. WorldModelSession drives the causal chunk_rollout
loop as a long-lived session: persistent cross-request world state on the
Session.kv_handle, frame streaming, transactional step-boundary cancellation
(a cancelled act leaves the world resumable), no cross-session smearing.
Added only a continuation seam to the chunk loop (init seeds context from a
world_context slot; default empty = unchanged one-shot path). 5 tests.
2. End-to-end RL over a cross-model workflow (training/methods/workflow_rl.py,
§9.9) — trains BOTH flux-t2i and wan-i2v from ONE final-video reward. Rolls
out the whole workflow with SDE capture in both instances; the same final
advantage drives FlowGRPO PPO on each stage's transformer; two WeightSyncPlans
on two instances. The earlier model (T2I) is trained by a reward on the final
video — end-to-end credit across a model boundary — proven causal by a control
(constant reward => zero advantage => nothing moves). 4 tests.
3. Heterogeneous WorkUnit co-scheduling (models/tiled/, §9.10) — the §17
falsifier. VAETileLoop makes VAE decode a loop of VAE_TILE units; tiling is
exact (== one-shot, C0), and VAE_TILE + DIFFUSION_STEP pipelines interleave
bit-identically and co-run in one batch. Validates the mechanism; the
economic half (does it pay) stays a port-time measurement. 4 tests.
designv4 §9.8–§9.10 + falsifier/layout/closing updates (127 tests, 20 files).
Answers "what's the right way to name/register custom pipelines like T2I→I2V":
treat a Workflow like a card — a stable namespaced id in the same servable
namespace, declared dependencies, and a two-level registry. No new concepts;
mirrors how cards are registered (and vllm-omni's pipeline_registry).
- program/workflow.py: Workflow gains `requires` (the cards it composes, derived
from stages) and `validate(engine)` (fail-fast if a required card is absent,
P7). New WorkflowRegistry: declarative workflow_id -> builder catalog for
out-of-tree/ad hoc use.
- runtime/engine.py: `_workflows` registry + register_workflow (validates deps,
rejects id collision with a model_id) + serves(); engine.run routes a request
whose model_id is a workflow to workflow.run — addressable exactly like a model.
Single-model hot path untouched.
- runtime/async_engine.py + serving/server.py: serves() and /models include
workflows (discoverable as servables).
- models/__init__.py: declarative `_WORKFLOWS` catalog (the cross-model analog of
_BUILDERS) + register_workflows() helper; build_image_video_engine now registers
the workflow too. Adding a custom pipeline = one catalog line.
- Naming convention: dotted/namespaced workflow_id (`image_video.t2i_i2v`),
distinct from kebab model ids, collision-checked. Renamed from `t2i_then_i2v`.
- tests (+5, 12 total in the file): addressable by id, requires/validate, id
collision, registry catalog, register_workflows helper. Full suite 114 passed.
- designv4 §9.6: the naming & registration convention documented.
Directory rename (git mv, history preserved) plus rewrite of all references: absolute imports in
tests, the zero-dep runner, docstrings, comments, and the README. No behavior change.
Run: python3 -m pytest v2/tests/ -q ; python3 v2/run_tests.py ; python3 -m v2.examples