From 634f0828a2bc4bb8034c49e1dea2ae309afc4e44 Mon Sep 17 00:00:00 2001 From: SolitaryThinker Date: Wed, 17 Jun 2026 12:55:52 -0700 Subject: [PATCH] [feat] route all diffusion loops + RL recompute through the kernel table MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Finish the kernel seam across the board so the platform's KernelTable is the universal solver-dispatch path, not just wan21. - Loops: ltx2, wan_causal, adapters, adaptive now resolve the flow-match solver via model.platform.kernels.get(FLOW_MATCH_STEP) instead of importing the numpy sampler directly (wan21 already did). On CPU this is bit-identical (the cpu kernel IS the old function); a GPU/accel backend now overrides every loop. - RL: the FlowGRPO log-prob recompute in unified_rl / joint_multi_rl / workflow_rl dispatches FLOW_SDE_STEP through the platform, pinned to the SAME kernel the rollout used (C2 kernel-pinning — otherwise the PPO ratio biases on a real GPU where rollout and recompute kernels could differ). - accel backend: add a kind-generic AccelComponent wrapper and override the vae component too (text_encoder left unregistered to keep the device->cpu fallback demonstrated), closing the "only dit is overridden" gap. - tests: vae override assertion; a second-loop (wan-causal chunk rollout) parity oracle proving accel == cpu bit-identical beyond wan21. README scope updated. 184 tests pass. --- v2/README.md | 23 ++++++++++++----------- v2/models/adapters/loop.py | 5 +++-- v2/models/adaptive/loop.py | 4 ++-- v2/models/ltx2/loop.py | 7 ++++--- v2/models/wan_causal/loop.py | 5 +++-- v2/platform/backends/accel.py | 25 +++++++++++++++++++++++++ v2/tests/test_platform.py | 25 +++++++++++++++++++------ v2/training/methods/joint_multi_rl.py | 8 +++++--- v2/training/methods/unified_rl.py | 9 ++++++--- v2/training/methods/workflow_rl.py | 18 ++++++++++++------ 10 files changed, 91 insertions(+), 38 deletions(-) diff --git a/v2/README.md b/v2/README.md index 8c7330c57..0ad2c6156 100644 --- a/v2/README.md +++ b/v2/README.md @@ -17,14 +17,15 @@ scheduler, caches, parity, training, and workflow code are unchanged. Honest scope of what is *wired* today (the substrate supports more than it demonstrates): a pure- python `accel` stand-in backend proves the dispatch is genuinely device-generic (cross-device -resolution + arch fallback + the parity oracle, bit-identical to numpy) without a GPU. The wan21 -denoise loop routes its solver ops (`flow_match_step`/`flow_sde_step`) through the kernel table, and -the `dit` component has a backend override; the other model loops still call the numpy samplers -directly and only the `dit` kind is overridden — each is a one-line change to adopt. The full -torch/CUDA path and a cudagraph capture-safety (`workspace_bytes`) contract are declared/​deferred, -not implemented. +resolution + arch fallback + the parity oracle, bit-identical to numpy) without a GPU. All diffusion +loops (wan21, ltx2, wan-causal, adapters, adaptive) and the FlowGRPO RL log-prob recompute route +their solver ops (`flow_match_step`/`flow_sde_step`) through the kernel table — the recompute is +pinned to the rollout's kernel for C2 correctness. The `dit` and `vae` components have `accel` +overrides (`text_encoder` demonstrates the device→cpu fallback). Not yet routed: the AR/vocoder op +families (omni/qwen-omni) have no KERNELS entries, the full torch/CUDA path is declared-unavailable, +and a cudagraph capture-safety (`workspace_bytes`) contract is deferred. -**183 tests, 30 files**, run two ways (`pytest` and a zero-dependency runner). Python 3.10+ and numpy only. +**184 tests, 30 files**, run two ways (`pytest` and a zero-dependency runner). Python 3.10+ and numpy only. ## Scope @@ -100,7 +101,7 @@ Stress tests (the design held — see **designv4 §9**): ```bash cd /Users/willlin/src/FastVideo-mini -python3 -m pytest v2/tests/ -q # 127 tests +python3 -m pytest v2/tests/ -q # 184 tests python3 v2/run_tests.py # same suite, ZERO deps (no pytest needed) python3 -m v2.examples # the worked examples ``` @@ -115,9 +116,9 @@ Each `ComponentSpec` records the real `load_id` (e.g. `vae.decode`; `text_encoder.encode`), and the device kernels for the solver/primitive ops. They flip to `available=True` when torch+CUDA are present, and `Platform.detect()` resolves them; the numpy `ComponentSpec.factory` stays as the CPU terminal fallback rung; -3. for the loops that don't yet route through the kernel table, swap the direct numpy - flow-match/SDE sampler calls for `model.platform.kernels.get(...)` (the wan21 loop shows the - pattern — a one-line change per call site). +3. the diffusion loops + RL recompute already dispatch their solver ops via + `model.platform.kernels.get(...)`, so the registered `cuda` kernels are picked up automatically; + only the not-yet-routed op families (AR sampling, vocoder synth) need the same one-line swap. The loops, policies, scheduler, caches, parity gates, training methods, and workflows are unchanged. ## Serving & fleet (our own version — Dynamo is an option, not a dependency) diff --git a/v2/models/adapters/loop.py b/v2/models/adapters/loop.py index f160efea7..6f2c8e22c 100644 --- a/v2/models/adapters/loop.py +++ b/v2/models/adapters/loop.py @@ -11,7 +11,7 @@ from __future__ import annotations import numpy as np from ...loop.contracts import Done -from ...loop.sampler import flow_match_euler_step +from ...platform import FLOW_MATCH_STEP from ..wan21.loop import WanDenoiseLoop @@ -39,7 +39,8 @@ class AdapterDenoiseLoop(WanDenoiseLoop): for aid in adapters: # apply each active adapter's delta v = v + model.component(aid).delta(x, control) res.output["noise_pred"] = v - res.output["latents"] = flow_match_euler_step(prec.cast(x), v, sigma_t, sigma_next).astype("float32") + fm = model.platform.kernels.get(FLOW_MATCH_STEP) # solver dispatched per (device, arch) + res.output["latents"] = fm(prec.cast(x), v, sigma_t, sigma_next).astype("float32") return res plan.run = run diff --git a/v2/models/adaptive/loop.py b/v2/models/adaptive/loop.py index d7ee9e2c0..9086dd9a0 100644 --- a/v2/models/adaptive/loop.py +++ b/v2/models/adaptive/loop.py @@ -17,7 +17,7 @@ from __future__ import annotations import numpy as np from ...loop.contracts import Done, StepResult -from ...loop.sampler import flow_match_euler_step +from ...platform import FLOW_MATCH_STEP from ..wan21.loop import WanDenoiseLoop @@ -46,7 +46,7 @@ class CacheDiTDenoiseLoop(WanDenoiseLoop): v = np.asarray(last_v, dtype="float32") def run(model, override=None): - x_next = flow_match_euler_step(prec.cast(x), v, sigma_t, sigma_next) + x_next = model.platform.kernels.get(FLOW_MATCH_STEP)(prec.cast(x), v, sigma_t, sigma_next) return StepResult(output={"noise_pred": v, "latents": x_next.astype("float32"), "skipped": True}) plan.run = run diff --git a/v2/models/ltx2/loop.py b/v2/models/ltx2/loop.py index 170a81950..beb4a3ae5 100644 --- a/v2/models/ltx2/loop.py +++ b/v2/models/ltx2/loop.py @@ -28,7 +28,7 @@ from ...loop.contracts import ( StepResult, WorkPlan, ) -from ...loop.sampler import flow_match_euler_step +from ...platform import FLOW_MATCH_STEP from ..backend import LATENT_CHANNELS BASE_SIGMAS = [1.0, 0.99375, 0.9875, 0.98125, 0.975, 0.909375, 0.725, 0.421875, 0.0] @@ -99,6 +99,7 @@ class LTX2DenoiseLoop: au = st.latents.get("audio") def run(model, override=None): + fm = model.platform.kernels.get(FLOW_MATCH_STEP) # solver dispatched per (device, arch) if override is not None and "noise_pred" in override: velocity = np.asarray(override["noise_pred"], dtype="float32") else: @@ -110,14 +111,14 @@ class LTX2DenoiseLoop: velocity = (pos + (v_guidance - 1.0) * (pos - neg) + stg_scale * (pos - ptb)) - x_next = flow_match_euler_step(x, np.asarray(velocity, dtype="float32"), sigma_t, sigma_next) + x_next = fm(x, np.asarray(velocity, dtype="float32"), sigma_t, sigma_next) out = {"noise_pred": np.asarray(velocity, dtype="float32"), "latents": x_next.astype("float32")} if want_audio and au is not None: # joint audio denoise, conditioned on video dit = model.component("transformer") a_pos = dit(au, pe, sigma_t, context=x) # context=x ⇒ audio synced to the video a_neg = dit(au, ne, sigma_t, context=x) a_vel = a_pos + (a_guidance - 1.0) * (a_pos - a_neg) # audio's OWN guidance scale - out["audio_latents"] = flow_match_euler_step(au, a_vel, sigma_t, sigma_next).astype("float32") + out["audio_latents"] = fm(au, a_vel, sigma_t, sigma_next).astype("float32") return StepResult(output=out) res = ResourceRequest( diff --git a/v2/models/wan_causal/loop.py b/v2/models/wan_causal/loop.py index c324fecd2..b9b923c89 100644 --- a/v2/models/wan_causal/loop.py +++ b/v2/models/wan_causal/loop.py @@ -24,7 +24,7 @@ from ...loop.contracts import ( StepResult, WorkPlan, ) -from ...loop.sampler import flow_match_euler_step +from ...platform import FLOW_MATCH_STEP from ...request.streams import StreamChunk from ..backend import LATENT_CHANNELS @@ -80,13 +80,14 @@ class ChunkRolloutLoop: cfg, precision = self.cfg, self.precision def run(model, override=None): + fm = model.platform.kernels.get(FLOW_MATCH_STEP) # solver dispatched per (device, arch) if override is not None and "noise_pred" in override: velocity = np.asarray(override["noise_pred"], dtype="float32") else: dit = model.component("transformer") preds = {b: dit(x, pe if b == "cond" else ne, sigma_t, context=context) for b in branches} velocity = cfg.combine(preds, scale, sctx, cfg_state) - x_next = flow_match_euler_step(precision.cast(x), precision.cast(velocity), sigma_t, sigma_next) + x_next = fm(precision.cast(x), precision.cast(velocity), sigma_t, sigma_next) return StepResult(output={"noise_pred": np.asarray(velocity, dtype="float32"), "latents": x_next.astype("float32")}) diff --git a/v2/platform/backends/accel.py b/v2/platform/backends/accel.py index 75832b08e..03e53a859 100644 --- a/v2/platform/backends/accel.py +++ b/v2/platform/backends/accel.py @@ -58,6 +58,31 @@ def _build_accel_dit(spec: Any, instance: Any, platform: Any) -> AccelDiT: register_component("dit", _build_accel_dit, device="accel", source="accel(stand-in):AccelDiT") +class AccelComponent: + """Generic device-tagged wrapper for non-callable component kinds (vae, audio_vae, …). Delegates + the full call surface (``decode``/``encode``/…) to the inner toy, so accel ≡ cpu numerically.""" + + def __init__(self, inner: Any): + self._inner = inner + self.device = "accel" + + def __getattr__(self, name: str): + return getattr(self._inner, name) + + +def _build_accel_component(spec: Any, instance: Any, platform: Any) -> AccelComponent: + if getattr(spec, "factory", None) is None: + raise RuntimeError(f"accel: no factory to wrap for kind={spec.kind!r}") + return AccelComponent(spec.factory(instance)) + + +# Override more than one component kind on accel (proving the seam is kind-generic, not dit-special). +# text_encoder is deliberately left UNregistered so the device→cpu fallback path stays demonstrated. +for _kind in ("vae", "audio_vae"): + register_component(_kind, _build_accel_component, device="accel", + source="accel(stand-in):AccelComponent") + + # --------------------------------------------------------------------------- # # Kernels: independent accel implementations that match the numpy reference. # # --------------------------------------------------------------------------- # diff --git a/v2/tests/test_platform.py b/v2/tests/test_platform.py index 54f11cfbd..133f630bc 100644 --- a/v2/tests/test_platform.py +++ b/v2/tests/test_platform.py @@ -19,6 +19,7 @@ from v2.loop.driver import LoopRunner from v2.models.backend import ToyDiT, ToyTextEncoder, ToyVAE from v2.models.common import text_encode_node_fn from v2.models.wan21 import build_wan21_card +from v2.models.wan_causal import build_wan_causal_card from v2.parity import bit_identical # the typed C1 instrument (ladder.py) from v2.platform import ( FLOW_MATCH_STEP, @@ -27,7 +28,7 @@ from v2.platform import ( component_matrix, kernel_matrix, ) -from v2.platform.backends.accel import AccelDiT +from v2.platform.backends.accel import AccelComponent, AccelDiT from v2.request import DiffusionParams, TaskType, make_request from v2.request.streams import Stream from v2.runtime import Engine @@ -41,8 +42,8 @@ def _instance(platform): return load_card(card, cache_manager=CacheManager.from_card(card), platform=platform) -def _drive_denoise(inst, prompt="a fox", seed=1, *, steps=8, sde=False): - """Drive the wan21 denoise loop step-by-step on ``inst``'s platform; return the final latent.""" +def _drive_denoise(inst, prompt="a fox", seed=1, *, steps=8, sde=False, loop_id="diffusion_denoise"): + """Drive a loop step-by-step on ``inst``'s platform; return the final latent.""" req = make_request(TaskType.T2V, inst.card.model_id, prompt, diffusion=DiffusionParams(num_steps=steps, seed=seed, sde_rollout=sde)) slots: dict = {} @@ -51,7 +52,7 @@ def _drive_denoise(inst, prompt="a fox", seed=1, *, steps=8, sde=False): ctx = RuntimeLoopContext(inst, observers=_HUB.observers, interceptors=_HUB.interceptors, slots=slots, stream=Stream(req.request_id), cancel_scope=None, profile=profile, metrics={}, request_id=req.request_id) - runner = LoopRunner(inst.loop("diffusion_denoise"), ctx, req, inst) + runner = LoopRunner(inst.loop(loop_id), ctx, req, inst) while not runner.done: runner.step() return np.asarray(runner.result.outputs["latents"]) @@ -134,9 +135,10 @@ def test_accel_component_override_and_device_fallback(): inst = _instance(Platform.accel("sm90")) dit = inst.component("transformer") # registered for accel → AccelDiT wrapper assert isinstance(dit, AccelDiT) and dit.device == "accel" - # text_encoder / vae are NOT registered for accel → fall back over accel→cpu to the factory toy + vae = inst.component("vae") # vae is ALSO overridden on accel (kind-generic) + assert isinstance(vae, AccelComponent) and vae.device == "accel" + # text_encoder is NOT registered for accel → falls back over accel→cpu to the numpy factory toy assert isinstance(inst.component("text_encoder"), ToyTextEncoder) - assert isinstance(inst.component("vae"), ToyVAE) def test_cuda_component_falls_back_to_factory_when_unavailable(): @@ -170,6 +172,17 @@ def test_accel_denoise_still_produces_valid_output(): assert out.shape[0] == 4 and np.all(np.isfinite(out)) +def test_parity_oracle_accel_equals_cpu_second_loop(): + """A SECOND loop family (wan-causal chunk rollout) also routes its solver through the kernel + table — so the seam is finished across loops, not just wan21 denoise. accel ≡ cpu, bit-identical.""" + def drive(platform): + card = build_wan_causal_card() + inst = load_card(card, cache_manager=CacheManager.from_card(card), platform=platform) + loop_id = next(iter(inst.card.loops)) # the chunk-rollout loop + return _drive_denoise(inst, seed=5, loop_id=loop_id) + assert bit_identical(drive(Platform.cpu()), drive(Platform.accel("sm90"))) + + # --------------------------------------------------------------------------- # # Enumerable matrix — including unavailable backends, without importing torch # # --------------------------------------------------------------------------- # diff --git a/v2/training/methods/joint_multi_rl.py b/v2/training/methods/joint_multi_rl.py index 304791dfd..0140f17ff 100644 --- a/v2/training/methods/joint_multi_rl.py +++ b/v2/training/methods/joint_multi_rl.py @@ -28,8 +28,8 @@ import numpy as np from ..._enums import ConsistencyLevel, ExecutionProfile from ...loop.sampler import flow_sde_ml_velocity as _ml_velocity -from ...loop.sampler import flow_sde_step_with_logprob from ...models.common import cached_text_encode +from ...platform import FLOW_SDE_STEP from ...request import DiffusionParams, TaskType, make_request from ..rollout import rollout_loop from ..weight_sync import WeightRole, WeightSyncPlan @@ -116,6 +116,8 @@ class JointMultiExpertRL(TrainingMethod): def _dit_ppo(self, behavior, emb, advantage): ref_dit = self.reference.component(self.generator_id) + # Pin the log-prob recompute to the SAME SDE kernel the rollout used (C2 kernel-pinning). + sde_step = self.student.platform.kernels.get(FLOW_SDE_STEP) ppo_losses, kls, ratios, gnorms = [], [], [], [] for rec in behavior or []: if "sde_logprob" not in rec: @@ -123,10 +125,10 @@ class JointMultiExpertRL(TrainingMethod): prev, sample = rec["prev"], rec["sample"] st, sn = float(rec["sigma_t"]), float(rec["sigma_next"]) v_cur = self.dit(prev, emb, st) - _, logp_cur, mean_cur, eff_std = flow_sde_step_with_logprob( + _, logp_cur, mean_cur, eff_std = sde_step( prev, v_cur, st, sn, prev_sample=sample, noise_scale=self.sde_noise_scale) v_ref = ref_dit(prev, emb, st) - _, _, mean_ref, _ = flow_sde_step_with_logprob( + _, _, mean_ref, _ = sde_step( prev, v_ref, st, sn, prev_sample=sample, noise_scale=self.sde_noise_scale) ratio = float(np.exp(np.clip(logp_cur - float(rec["sde_logprob"]), -10.0, 10.0))) clipped = float(np.clip(ratio, 1.0 - self.ppo_clip, 1.0 + self.ppo_clip)) diff --git a/v2/training/methods/unified_rl.py b/v2/training/methods/unified_rl.py index 4dd3e6855..64fa62f22 100644 --- a/v2/training/methods/unified_rl.py +++ b/v2/training/methods/unified_rl.py @@ -32,8 +32,8 @@ import numpy as np from ..._enums import ConsistencyLevel, ExecutionProfile from ...loop.sampler import flow_sde_ml_velocity as _ml_velocity -from ...loop.sampler import flow_sde_step_with_logprob from ...models.common import cached_text_encode +from ...platform import FLOW_SDE_STEP from ...request import DiffusionParams, TaskType, make_request from ..rollout import rollout_loop from ..weight_sync import WeightRole, WeightSyncPlan @@ -112,6 +112,9 @@ class UnifiedRLMethod(TrainingMethod): """FlowGRPO PPO over one rollout's SDE trajectory: per-step ratio + KL, then a toy DiT step. Returns (mean_ppo_loss, mean_kl, mean_ratio, grad_norm).""" ref_dit = self.reference.component("transformer") + # Pin the log-prob recompute to the SAME SDE kernel the rollout used (C2 kernel-pinning): + # on a GPU box the rollout's logp and this ratio must come from one kernel or the PPO ratio biases. + sde_step = self.student.platform.kernels.get(FLOW_SDE_STEP) ppo_losses, kls, ratios, gnorms = [], [], [], [] for rec in behavior or []: if "sde_logprob" not in rec: @@ -119,10 +122,10 @@ class UnifiedRLMethod(TrainingMethod): prev, sample = rec["prev"], rec["sample"] st, sn = float(rec["sigma_t"]), float(rec["sigma_next"]) v_cur = self.dit(prev, emb, st) - _, logp_cur, mean_cur, eff_std = flow_sde_step_with_logprob( + _, logp_cur, mean_cur, eff_std = sde_step( prev, v_cur, st, sn, prev_sample=sample, noise_scale=self.sde_noise_scale) v_ref = ref_dit(prev, emb, st) - _, _, mean_ref, _ = flow_sde_step_with_logprob( + _, _, mean_ref, _ = sde_step( prev, v_ref, st, sn, prev_sample=sample, noise_scale=self.sde_noise_scale) ratio = float(np.exp(np.clip(logp_cur - float(rec["sde_logprob"]), -10.0, 10.0))) clipped = float(np.clip(ratio, 1.0 - self.ppo_clip, 1.0 + self.ppo_clip)) diff --git a/v2/training/methods/workflow_rl.py b/v2/training/methods/workflow_rl.py index 151c71556..f03dd7158 100644 --- a/v2/training/methods/workflow_rl.py +++ b/v2/training/methods/workflow_rl.py @@ -21,16 +21,20 @@ import numpy as np from ..._enums import ConsistencyLevel, ExecutionProfile from ...loop.sampler import flow_sde_ml_velocity, flow_sde_step_with_logprob from ...models.common import cached_text_encode +from ...platform import FLOW_SDE_STEP from ...request import DiffusionParams, TaskType, make_request from ..rollout import rollout_loop from ..weight_sync import WeightRole, WeightSyncPlan from .base import TrainingMethod, new_instance -def _flowgrpo_ppo_step(dit, ref_dit, behavior, emb, advantage, *, lr, ppo_clip, beta, noise_scale): +def _flowgrpo_ppo_step(dit, ref_dit, behavior, emb, advantage, *, lr, ppo_clip, beta, noise_scale, + sde_step=None): """One FlowGRPO PPO pass over an SDE trajectory: per-step ratio + KL, then move the policy toward the max-likelihood velocity of each realized sample (advantage-weighted). Returns - (ppo_loss, kl, ratio, grad_norm). Shared by both stages.""" + (ppo_loss, kl, ratio, grad_norm). Shared by both stages. ``sde_step`` is the platform-resolved SDE + kernel (pins the recompute to the rollout's kernel, C2); defaults to the numpy reference.""" + sde_step = sde_step or flow_sde_step_with_logprob ppo_losses, kls, ratios, gnorms = [], [], [], [] for rec in behavior or []: if "sde_logprob" not in rec: @@ -38,10 +42,10 @@ def _flowgrpo_ppo_step(dit, ref_dit, behavior, emb, advantage, *, lr, ppo_clip, prev, sample = rec["prev"], rec["sample"] st, sn = float(rec["sigma_t"]), float(rec["sigma_next"]) v_cur = dit(prev, emb, st) - _, logp_cur, mean_cur, eff_std = flow_sde_step_with_logprob( + _, logp_cur, mean_cur, eff_std = sde_step( prev, v_cur, st, sn, prev_sample=sample, noise_scale=noise_scale) v_ref = ref_dit(prev, emb, st) - _, _, mean_ref, _ = flow_sde_step_with_logprob( + _, _, mean_ref, _ = sde_step( prev, v_ref, st, sn, prev_sample=sample, noise_scale=noise_scale) ratio = float(np.exp(np.clip(logp_cur - float(rec["sde_logprob"]), -10.0, 10.0))) clipped = float(np.clip(ratio, 1.0 - ppo_clip, 1.0 + ppo_clip)) @@ -141,11 +145,13 @@ class WorkflowRLMethod(TrainingMethod): pt, _kt, rt, gt = _flowgrpo_ppo_step( self.t2i.component("transformer"), self.t2i_ref.component("transformer"), res_t.behavior, emb_t, a, lr=self.t2i_lr, ppo_clip=self.ppo_clip, - beta=self.beta, noise_scale=self.sde_noise_scale) + beta=self.beta, noise_scale=self.sde_noise_scale, + sde_step=self.t2i.platform.kernels.get(FLOW_SDE_STEP)) pi_, _ki, ri, gi = _flowgrpo_ppo_step( self.i2v.component("transformer"), self.i2v_ref.component("transformer"), res_i.behavior, emb_i, a, lr=self.i2v_lr, ppo_clip=self.ppo_clip, - beta=self.beta, noise_scale=self.sde_noise_scale) + beta=self.beta, noise_scale=self.sde_noise_scale, + sde_step=self.i2v.platform.kernels.get(FLOW_SDE_STEP)) t2i_loss.append(pt); i2v_loss.append(pi_) t2i_gn.append(gt); i2v_gn.append(gi) ratios.extend([rt, ri])