From 2f044c02dde46503bd31cd0ab43901efd18f0415 Mon Sep 17 00:00:00 2001 From: SolitaryThinker Date: Wed, 17 Jun 2026 12:26:11 -0700 Subject: [PATCH] [feat] multi-backend dispatch substrate (device/arch/kernel registries) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add v2/platform/: the (recipe, runtime) backend membrane that lets CPU, GPU, and other devices coexist behind one dispatch substrate. - Two tuple-keyed registries: COMPONENTS(kind, device, variant) for weight-bearing components and KERNELS(op, device, arch, variant) for stateless primitives, each with an availability predicate and an enumerable manifest (declared-but-unavailable cells listed without importing torch). - Platform: detected (device, arch) owning the device + arch fallback chains and a per-platform cached KernelTable; detect() is honest (CPU/numpy unless torch+CUDA are actually present). Arch fallback is monotonic — only degrades to older/portable archs, never a newer binary-incompatible one. - Seams wired: ModelInstance.component() -> platform.build_component(spec,self) with spec.factory as the numpy terminal rung (existing cards untouched); the wan21 denoise thunk dispatches solver ops through model.platform.kernels. - Three backends: cpu (numpy terminal + parity oracle), accel (pure-python stand-in proving cross-device resolution, arch fallback, and the oracle), torch_cuda (declared-but-unavailable; no faked GPU). - test_platform.py: 16 tests — detection, terminal rung, arch-fallback walk + monotonicity, device precedence, variant fallback, component override + per-kind device fallback, the parity oracle (accel == cpu, bit_identical via the C1 ladder), and matrix enumeration without importing torch. Scope is honestly bounded in the README/docstrings: only wan21-denoise routes through the kernel table and only the dit kind is overridden today (each a one-line adoption); the torch/CUDA path and a cudagraph workspace-safety contract are declared/deferred, not implemented. 183 tests pass. --- v2/README.md | 33 +++-- v2/card/instance.py | 26 ++-- v2/models/wan21/loop.py | 10 +- v2/platform/__init__.py | 42 +++++++ v2/platform/backends/__init__.py | 8 ++ v2/platform/backends/accel.py | 81 ++++++++++++ v2/platform/backends/cpu.py | 22 ++++ v2/platform/backends/torch_cuda.py | 53 ++++++++ v2/platform/platform.py | 196 +++++++++++++++++++++++++++++ v2/platform/registry.py | 128 +++++++++++++++++++ v2/tests/test_platform.py | 196 +++++++++++++++++++++++++++++ 11 files changed, 777 insertions(+), 18 deletions(-) create mode 100644 v2/platform/__init__.py create mode 100644 v2/platform/backends/__init__.py create mode 100644 v2/platform/backends/accel.py create mode 100644 v2/platform/backends/cpu.py create mode 100644 v2/platform/backends/torch_cuda.py create mode 100644 v2/platform/platform.py create mode 100644 v2/platform/registry.py create mode 100644 v2/tests/test_platform.py diff --git a/v2/README.md b/v2/README.md index 6431a7cab..8c7330c57 100644 --- a/v2/README.md +++ b/v2/README.md @@ -7,12 +7,24 @@ elaborates). It is "mini" in two senses: (1) it supports a deliberate subset of models/methods, and (2) the heavy neural forwards are small numpy stand-ins so the whole thing runs and is tested on a laptop with no -GPU and no torch. **The architecture is real; the kernels are toys.** On a GPU box you swap -`ComponentSpec.factory` for the torch adapters (which wrap the real `fastvideo.models.*` modules) and -nothing else in the runtime, scheduler, caches, parity, training, or workflow code changes — that is -the whole point of the (recipe, runtime) separation. +GPU and no torch. **The architecture is real; the kernels are toys.** Backends are selected through +the `platform/` dispatch substrate: two tuple-keyed registries — `COMPONENTS(kind, device, variant)` +and `KERNELS(op, device, arch, variant)` — that a detected `Platform` resolves, with the numpy +reference as the terminal fallback rung and parity oracle (`§17`). On a GPU box you register the +torch component/kernel backends (the `cuda` cells are declared in `platform/backends/torch_cuda.py`, +`available=False` until torch+CUDA are present) and `Platform.detect()` resolves them — the loops, +scheduler, caches, parity, training, and workflow code are unchanged. -**127 tests, 20 files**, run two ways (`pytest` and a zero-dependency runner). Python 3.10+ and numpy only. +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. + +**183 tests, 30 files**, run two ways (`pytest` and a zero-dependency runner). Python 3.10+ and numpy only. ## Scope @@ -98,9 +110,14 @@ python3 -m v2.examples # the worked examples Each `ComponentSpec` records the real `load_id` (e.g. `fastvideo.models.dits.wanvideo:WanTransformer3DModel`). To run real Wan/LTX/etc.: 1. install torch + the parent `fastvideo` package + weights; -2. replace each `ComponentSpec.factory` with a lazy adapter that loads the real module and exposes the - same call surface the loops use (a velocity forward; `vae.decode`; `text_encoder.encode`); -3. swap the numpy flow-match/SDE samplers for the real schedulers inside the loop's step body. +2. implement the `cuda` (or other-device) cells in `platform/backends/` — a component builder that + loads the real module and exposes the same call surface the loops use (a velocity forward; + `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). 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/card/instance.py b/v2/card/instance.py index a603e6ddc..2773ecd01 100644 --- a/v2/card/instance.py +++ b/v2/card/instance.py @@ -13,6 +13,7 @@ from __future__ import annotations from typing import Any +from ..platform import Platform from .specs import ModelCard @@ -21,11 +22,16 @@ class ModelInstance: passed to ``Loop.init`` (the loop reads components through it).""" def __init__(self, card: ModelCard, parallel_plan: Any = None, - cache_manager: Any = None, weights_version: str = "v0"): + cache_manager: Any = None, weights_version: str = "v0", + platform: Any = None): self.card = card self.parallel_plan = parallel_plan self.caches = cache_manager self.weights_version = weights_version + # The detected (device, arch). Resolves component/kernel implementations through the two + # backend registries; defaults to CPU/numpy. Swapping this to a GPU platform is the whole + # "change only the backend, not the loops/policies/training" story (design_v3 §17). + self.platform = platform if platform is not None else Platform.cpu() self.adapter_versions: dict[str, str] = {} # per-component weight versions (§7.1): a component's version changes only when IT is synced, # so a transformer-only RL sync never invalidates the frozen text-encoder's feature cache. @@ -42,9 +48,9 @@ class ModelInstance: spec = self.card.components.get(component_id) if spec is None: raise KeyError(f"component {component_id!r} not declared on card {self.card.model_id!r}") - if spec.factory is None: - raise RuntimeError(f"component {component_id!r} has no factory (cannot instantiate)") - self._components[component_id] = spec.factory(self) + # The single materialization seam: the platform resolves (kind, device, variant) through + # the COMPONENTS registry, falling back to spec.factory as the cpu/numpy terminal rung. + self._components[component_id] = self.platform.build_component(spec, self) return self._components[component_id] def has_component(self, component_id: str) -> bool: @@ -93,8 +99,14 @@ class ModelInstance: def load_card(card: ModelCard, parallel_plan: Any = None, cache_manager: Any = None, - *, validate: bool = True) -> ModelInstance: - """The card-as-factory entrypoint (design_v3 §4.1: card is a runtime factory).""" + *, validate: bool = True, platform: Any = None) -> ModelInstance: + """The card-as-factory entrypoint (design_v3 §4.1: card is a runtime factory). + + ``platform`` selects the backend (device, arch); when omitted it is detected (CPU/numpy here, + CUDA on a torch+GPU box). The same card loads on any backend — only the resolved component/kernel + implementations differ. + """ if validate: card.validate() - return ModelInstance(card, parallel_plan=parallel_plan, cache_manager=cache_manager) + return ModelInstance(card, parallel_plan=parallel_plan, cache_manager=cache_manager, + platform=platform if platform is not None else Platform.detect()) diff --git a/v2/models/wan21/loop.py b/v2/models/wan21/loop.py index b27cb4d85..4ff8579b3 100644 --- a/v2/models/wan21/loop.py +++ b/v2/models/wan21/loop.py @@ -21,7 +21,7 @@ from ...loop.contracts import ( StepResult, WorkPlan, ) -from ...loop.sampler import flow_match_euler_step, flow_sde_step_with_logprob +from ...platform import FLOW_MATCH_STEP, FLOW_SDE_STEP from ...request.streams import StreamChunk from ..backend import LATENT_CHANNELS @@ -81,6 +81,10 @@ class WanDenoiseLoop: sde, noise_scale, rng = st.scratch.get("sde", False), st.scratch.get("sde_noise_scale", 0.7), st.rng 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). + kernels = model.platform.kernels if override is not None and "noise_pred" in override: velocity = np.asarray(override["noise_pred"], dtype="float32") else: @@ -90,12 +94,12 @@ class WanDenoiseLoop: velocity = precision.cast(velocity) if sde: # FlowGRPO rollout: stochastic + log-prob noise = rng.standard_normal(x.shape) - x_next, logp, _m, _s = flow_sde_step_with_logprob( + x_next, logp, _m, _s = kernels.get(FLOW_SDE_STEP)( precision.cast(x), velocity, sigma_t, sigma_next, noise=noise, noise_scale=noise_scale) return StepResult(output={"noise_pred": np.asarray(velocity, dtype="float32"), "latents": x_next.astype("float32"), "sde_logprob": logp, "prev": np.asarray(x, dtype="float32")}) - x_next = flow_match_euler_step(precision.cast(x), velocity, sigma_t, sigma_next) + x_next = kernels.get(FLOW_MATCH_STEP)(precision.cast(x), 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/__init__.py b/v2/platform/__init__.py new file mode 100644 index 000000000..2dc780bb8 --- /dev/null +++ b/v2/platform/__init__.py @@ -0,0 +1,42 @@ +"""The multi-backend dispatch substrate (design_v3 §17). + +Two tuple-keyed registries (``COMPONENTS``, ``KERNELS``) + a ``Platform`` that detects ``(device, +arch)`` and resolves both — with the numpy reference as the terminal fallback rung and parity +oracle. This package is the membrane that lets CPU, GPU, and other backends coexist: a model's +loops, policies, scheduler, caches, and training code never name a device; they call through the +component/kernel seams and the resolved ``Platform`` decides the implementation. + +Importing this package is cheap and cycle-free (registry + platform classes only). The concrete +backend registrations in ``backends/`` are imported lazily on first ``Platform`` use. +""" +from __future__ import annotations + +from .platform import ( + KernelTable, + Platform, + component_matrix, + ensure_backends_loaded, + kernel_matrix, +) +from .registry import ( + COMPONENTS, + FLOW_MATCH_STEP, + FLOW_SDE_STEP, + KERNELS, + register_component, + register_kernel, +) + +__all__ = [ + "Platform", + "KernelTable", + "COMPONENTS", + "KERNELS", + "register_component", + "register_kernel", + "ensure_backends_loaded", + "kernel_matrix", + "component_matrix", + "FLOW_MATCH_STEP", + "FLOW_SDE_STEP", +] diff --git a/v2/platform/backends/__init__.py b/v2/platform/backends/__init__.py new file mode 100644 index 000000000..fd362b2ba --- /dev/null +++ b/v2/platform/backends/__init__.py @@ -0,0 +1,8 @@ +"""Backend registration modules. Imported for side effects by ``Platform.ensure_backends_loaded``. + +Each module calls ``register_component`` / ``register_kernel`` at import time. They live behind the +lazy loader (not imported by ``platform/__init__`` or ``card/``) so the registry stays a pure leaf +and there are no import cycles. On disk these mirror the kernel-colocation answer: ``cpu`` is the +unified numpy reference; a real ``torch_cuda`` backend would register both unified primitives and +model-co-located fusion *variants* through this same API. +""" diff --git a/v2/platform/backends/accel.py b/v2/platform/backends/accel.py new file mode 100644 index 000000000..75832b08e --- /dev/null +++ b/v2/platform/backends/accel.py @@ -0,0 +1,81 @@ +"""``accel`` — a pure-python stand-in accelerator backend (design_v3 §17 honesty note). + +There is no GPU in this environment, so to *prove* the dispatch substrate is genuinely +device-generic — not a single-backend abstraction with a CPU special case — this backend registers +a second device, ``accel``, entirely in CPU-resident Python. It exercises everything a real GPU +backend would: + + * a **component** override — ``AccelDiT`` is built via ``COMPONENTS[("dit", "accel", …)]`` instead + of the card's numpy factory, while components it does NOT register (text_encoder, vae) fall back + over the device chain ``accel → cpu`` to the numpy factory; + * **kernels at different arch rungs** — ``flow_sde_step`` is registered at ``sm90`` directly, while + ``flow_match_step`` is registered only at ``generic``, so resolving it on an ``sm90`` accel + platform must walk the arch fallback ``sm90 → sm80 → generic`` to find it; + * the **parity oracle** — the ODE solver op (``_accel_flow_match_step``) is an *independently + written* kernel that nonetheless matches the numpy reference bit-for-bit, so a denoise run on the + accel platform equals the cpu run. That independent-impl-vs-reference equality is the oracle a + real backend must pass; the stand-in shows the check has teeth. (The SDE op below deliberately + *reuses* the reference impl — a backend may legitimately not specialize every op — so the SDE + test exercises the dispatch path, not an independent implementation.) + +On a real box this file is where a torch/cuda backend would register its adapters and fused kernels +(unified primitives + model-co-located fusion variants) through the same two functions. +""" +from __future__ import annotations + +from typing import Any + +from ...models.backend import ToyDiT +from ..registry import FLOW_MATCH_STEP, FLOW_SDE_STEP, register_component, register_kernel + + +# --------------------------------------------------------------------------- # +# Component: an accel-resident DiT (stand-in for a torch module placed on the # +# device). Wraps the same toy compute so the parity oracle is bit-identical. # +# --------------------------------------------------------------------------- # +class AccelDiT: + """A device-tagged DiT wrapper. On a real box this would hold a torch module on the accelerator; + here it delegates to the same numpy toy so accel ≡ cpu numerically (the oracle).""" + + def __init__(self, inner: Any): + self._inner = inner + self.device = "accel" + + def __call__(self, *args, **kwargs): + return self._inner(*args, **kwargs) + + def __getattr__(self, name: str): + # delegate the trainable surface (clone/copy_from/blend_from/mse_grad_step) to the inner toy + return getattr(self._inner, name) + + +def _build_accel_dit(spec: Any, instance: Any, platform: Any) -> AccelDiT: + # Build the same component the cpu rung would (via the card's factory), then place it "on device". + inner = spec.factory(instance) if getattr(spec, "factory", None) is not None else ToyDiT() + return AccelDiT(inner) + + +register_component("dit", _build_accel_dit, device="accel", source="accel(stand-in):AccelDiT") + + +# --------------------------------------------------------------------------- # +# Kernels: independent accel implementations that match the numpy reference. # +# --------------------------------------------------------------------------- # +def _accel_flow_match_step(x_t, velocity, sigma_t: float, sigma_next: float): + """Stand-in 'fused' accel solver step. Identical math to the numpy reference (the oracle).""" + return x_t + (sigma_next - sigma_t) * velocity + + +# flow_match_step lives only at the `generic` arch rung → resolving on sm90 must fall back through +# the arch chain (sm90 → sm80 → generic). This proves the arch-fallback walk. +register_kernel(FLOW_MATCH_STEP, _accel_flow_match_step, + device="accel", arch="generic", source="accel(stand-in):fused_flow_match") + +# flow_sde_step is registered at sm90 directly. This backend does NOT specialize the SDE op — it +# reuses the numpy reference, so the accel SDE path is the same function as cpu (the SDE parity test +# therefore checks dispatch-path equivalence, not an independent impl; the ODE op above is the real +# independent oracle). A device that DID specialize SDE would be checked the same way the ODE op is. +from ...loop.sampler import flow_sde_step_with_logprob # noqa: E402 + +register_kernel(FLOW_SDE_STEP, flow_sde_step_with_logprob, + device="accel", arch="sm90", source="accel(stand-in):sde_logprob") diff --git a/v2/platform/backends/cpu.py b/v2/platform/backends/cpu.py new file mode 100644 index 000000000..42f3d974c --- /dev/null +++ b/v2/platform/backends/cpu.py @@ -0,0 +1,22 @@ +"""CPU / numpy backend — the terminal rung and parity oracle (design_v3 §17). + +This backend registers the numpy reference kernels. It is the bottom of every device fallback +chain: any op a richer backend hasn't implemented resolves here, and every other backend's output +is checked against this one (the consistency ladder's C0/C1 oracle). + +Components are NOT registered here: the card's ``ComponentSpec.factory`` *is* the cpu/numpy +component rung (``Platform.build_component`` falls back to it when no ``(kind, device, …)`` cell +matches). That keeps every existing card working with zero registration. +""" +from __future__ import annotations + +from ...loop.sampler import flow_match_euler_step, flow_sde_step_with_logprob +from ..registry import FLOW_MATCH_STEP, FLOW_SDE_STEP, register_kernel + +# The numpy solver primitives, registered as the terminal (cpu, numpy) kernels. These are the exact +# same functions the loops used to call directly — so dispatching through the registry on a CPU box +# is bit-for-bit identical to the pre-registry behavior. +register_kernel(FLOW_MATCH_STEP, flow_match_euler_step, + device="cpu", arch="numpy", source="loop.sampler:flow_match_euler_step") +register_kernel(FLOW_SDE_STEP, flow_sde_step_with_logprob, + device="cpu", arch="numpy", source="loop.sampler:flow_sde_step_with_logprob") diff --git a/v2/platform/backends/torch_cuda.py b/v2/platform/backends/torch_cuda.py new file mode 100644 index 000000000..c91ef26fd --- /dev/null +++ b/v2/platform/backends/torch_cuda.py @@ -0,0 +1,53 @@ +"""``cuda`` — the real torch/GPU backend, *declared but unavailable* on this box (design_v3 §17). + +This environment has no torch and no GPU, so these cells cannot run here. But they are registered +anyway, gated by an ``available`` predicate (a ``find_spec`` check that never imports torch), so the +backend matrix is **enumerable without a GPU and without importing torch**: ``kernel_matrix()`` / +``component_matrix()`` list these rows with ``available=False``. (Registration is via the fixed +backend import list, not setuptools entry points — see the scope caveat in ``registry.py``; full +entry-point discovery is the real-package mechanism, not wired in this mini.) + +On a real box (torch + CUDA present) ``available`` flips to ``True`` and ``Platform.detect()`` +returns a ``cuda`` platform that resolves these instead of the numpy/accel rungs. The bodies below +are deliberately not implemented — they exist to anchor the matrix and to fail loudly if ever +resolved on a box that lacks the toolchain. +""" +from __future__ import annotations + +import importlib.util +from typing import Any + +from ..registry import FLOW_MATCH_STEP, FLOW_SDE_STEP, register_component, register_kernel + + +def _torch_present() -> bool: + return importlib.util.find_spec("torch") is not None + + +def _cuda_available() -> bool: + if not _torch_present(): + return False + try: + import torch # type: ignore + return bool(torch.cuda.is_available()) + except Exception: + return False + + +def _unavailable(what: str): + def _fn(*_a, **_k): + raise RuntimeError( + f"{what} requires a GPU box with torch + the fastvideo-kernel toolchain; " + f"this cell is declared for matrix enumeration only") + return _fn + + +# Component: the real DiT torch adapter (loads `spec.load_id`, places weights on the device). +register_component("dit", _unavailable("torch DiT adapter"), device="cuda", + available=_cuda_available, source="fastvideo.models.dits (torch adapter)") + +# Kernels: the real fused solver primitives from fastvideo-kernel, at the Hopper rung. +register_kernel(FLOW_MATCH_STEP, _unavailable("cuda flow_match kernel"), device="cuda", arch="sm90", + available=_cuda_available, source="fastvideo-kernel:flow_match (cuda)") +register_kernel(FLOW_SDE_STEP, _unavailable("cuda flow_sde kernel"), device="cuda", arch="sm90", + available=_cuda_available, source="fastvideo-kernel:flow_sde (cuda)") diff --git a/v2/platform/platform.py b/v2/platform/platform.py new file mode 100644 index 000000000..c39d74bf2 --- /dev/null +++ b/v2/platform/platform.py @@ -0,0 +1,196 @@ +"""Platform — the detected ``(device, arch)`` that resolves both registries (design_v3 §17). + +One ``Platform`` per pool. It owns: + * the **device fallback chain** (e.g. ``accel → cpu``) — so a component/kernel a backend hasn't + implemented falls back to the numpy reference (the terminal rung + parity oracle); + * the **arch fallback chain** per device (e.g. cuda ``sm100→sm90→sm80→ptx→generic``) — so a kernel + built for an older arch still resolves on a newer card; + * ``build_component(spec, instance)`` — the single materialization seam (replaces a bare + ``spec.factory`` call); and + * ``kernels`` — a per-platform ``KernelTable`` that resolves an op once and caches the callable. + +``Platform.detect()`` is honest: it returns a CPU/numpy platform unless torch+CUDA are actually +present (they are not in this environment). The (device, arch) machinery is fully exercised on the +CPU box by the pure-python ``accel`` stand-in backend — proving the dispatch is generic without +faking a GPU. +""" +from __future__ import annotations + +from typing import Any, Iterator + +from .registry import COMPONENTS, KERNELS + +_BACKENDS_LOADED = False + + +def ensure_backends_loaded() -> None: + """Import the backend modules once so they self-register. Lazy → import-cycle-free.""" + global _BACKENDS_LOADED + if _BACKENDS_LOADED: + return + _BACKENDS_LOADED = True + # imported for side effects (each module calls register_component/register_kernel at import) + from .backends import accel, cpu, torch_cuda # noqa: F401 + + +# Arch fallback chains (newest → oldest → portable terminal). Resolution walks these in order. +_CUDA_ARCH_FALLBACK = ("sm100", "sm90", "sm80", "ptx", "generic") +_ACCEL_ARCH_FALLBACK = ("sm90", "sm80", "generic") +_PORTABLE_ARCHS = ("ptx", "generic") # vendor-neutral terminals — valid on any device + + +def _arch_rank(arch: str) -> int: + """``smXY`` → integer ``XY`` for comparison; portable terminals (ptx/generic/unknown) → -1.""" + if arch.startswith("sm") and arch[2:].isdigit(): + return int(arch[2:]) + return -1 + + +def _suffix_from(chain: tuple[str, ...], arch: str) -> tuple[str, ...]: + """The arch fallback tail for ``arch``: itself, then ONLY older real archs from ``chain``, then + portable terminals. Never falls back to a *newer* arch — a kernel built for sm90 is binary- + incompatible on an sm70 device, so resolution must only ever degrade to older/more-portable code. + (A newer device, e.g. sm120, may still run older-arch kernels — those are kept.) + """ + rank = _arch_rank(arch) + tail: list[str] = [arch] + for a in chain: + if a in tail: + continue + ar = _arch_rank(a) + if ar < 0: # portable terminal (ptx/generic): always valid + tail.append(a) + elif rank >= 0 and ar < rank: # strictly-older real arch: a safe fallback + tail.append(a) + # else: newer-or-equal real arch → skip (would be binary-incompatible on this device) + return tuple(tail) + + +class KernelTable: + """A platform's resolved op→callable map. Resolves each op once (over the fallback chain), + then caches — so the hot loop pays the registry walk at most once per (op, variant).""" + + def __init__(self, platform: "Platform"): + self.platform = platform + self._cache: dict[tuple[str, str], Any] = {} + + def get(self, op: str, variant: str = "default") -> Any: + ck = (op, variant) + if ck not in self._cache: + reg = self.platform.resolve_kernel(op, variant) + if reg is None: + raise KeyError( + f"no kernel for op={op!r} variant={variant!r} on device chain " + f"{self.platform.device_chain} (and no numpy terminal registered)") + self._cache[ck] = reg.fn + return self._cache[ck] + + def resolved_source(self, op: str, variant: str = "default") -> str: + """Which registered cell an op resolves to — used by parity/diagnostics tests.""" + reg = self.platform.resolve_kernel(op, variant) + return "" if reg is None else f"{reg.key}@{reg.source}" + + +class Platform: + def __init__(self, device: str, arch: str, device_chain: tuple[str, ...], + arch_chains: dict[str, tuple[str, ...]]): + self.device = device + self.arch = arch + self.device_chain = tuple(device_chain) + self.arch_chains = dict(arch_chains) + self._kernels: KernelTable | None = None + + # --- the kernel table (lazy, cached) ------------------------------------- # + @property + def kernels(self) -> KernelTable: + if self._kernels is None: + ensure_backends_loaded() + self._kernels = KernelTable(self) + return self._kernels + + def arch_chain(self, device: str) -> tuple[str, ...]: + return self.arch_chains.get(device, ("generic",)) + + # --- candidate-key generators (the fallback chains, in priority order) --- # + def _component_keys(self, kind: str, variant: str) -> Iterator[tuple]: + for dev in self.device_chain: + yield (kind, dev, variant) + if variant != "default": + yield (kind, dev, "default") + + def _kernel_keys(self, op: str, variant: str) -> Iterator[tuple]: + for dev in self.device_chain: + for arch in self.arch_chain(dev): + yield (op, dev, arch, variant) + if variant != "default": + yield (op, dev, arch, "default") + + # --- the two resolution seams -------------------------------------------- # + def build_component(self, spec: Any, instance: Any, variant: str = "default") -> Any: + """Materialize a component. Registry first (device chain), then ``spec.factory`` as the + terminal numpy rung — so existing cards (factory-only) are unchanged, while a GPU/accel + backend overrides by registering ``(kind, device, …)``.""" + ensure_backends_loaded() + reg = COMPONENTS.resolve_first(self._component_keys(spec.kind, variant)) + if reg is not None: + return reg.fn(spec, instance, self) + if getattr(spec, "factory", None) is not None: + return spec.factory(instance) # the (kind, "cpu", "numpy") terminal rung + raise RuntimeError( + f"component kind={spec.kind!r} ({getattr(spec, 'component_id', '?')!r}) has no impl on " + f"device chain {self.device_chain} and no factory terminal") + + def resolve_kernel(self, op: str, variant: str = "default"): + ensure_backends_loaded() + return KERNELS.resolve_first(self._kernel_keys(op, variant)) + + # --- constructors -------------------------------------------------------- # + @classmethod + def cpu(cls) -> "Platform": + return cls("cpu", "numpy", ("cpu",), {"cpu": ("numpy",)}) + + @classmethod + def accel(cls, arch: str = "sm90") -> "Platform": + """A pure-python stand-in accelerator (CPU-resident) used to prove cross-device dispatch, + arch fallback, and the parity oracle without a real GPU.""" + return cls("accel", arch, ("accel", "cpu"), + {"accel": _suffix_from(_ACCEL_ARCH_FALLBACK, arch), "cpu": ("numpy",)}) + + @classmethod + def cuda(cls, arch: str = "sm90") -> "Platform": + """A real CUDA platform (resolves torch adapters + fastvideo-kernel cells). Only usable on a + box where those cells are ``available``; declared here so detection has a target.""" + return cls("cuda", arch, ("cuda", "cpu"), + {"cuda": _suffix_from(_CUDA_ARCH_FALLBACK, arch), "cpu": ("numpy",)}) + + @classmethod + def detect(cls) -> "Platform": + """Honest detection: CUDA iff torch+CUDA are actually importable, else CPU/numpy.""" + ensure_backends_loaded() + import importlib.util + if importlib.util.find_spec("torch") is not None: + try: + import torch # type: ignore + if torch.cuda.is_available(): + major, minor = torch.cuda.get_device_capability() + return cls.cuda(f"sm{major}{minor}") + except Exception: + pass + return cls.cpu() + + def __repr__(self) -> str: + return f"Platform(device={self.device!r}, arch={self.arch!r}, chain={self.device_chain})" + + +# --------------------------------------------------------------------------- # +# Matrix dumps — enumerate the full (op/kind × device × arch × variant) grid, # +# including unavailable backends, without importing torch or touching a GPU. # +# --------------------------------------------------------------------------- # +def kernel_matrix() -> list[dict[str, Any]]: + ensure_backends_loaded() + return KERNELS.manifest() + + +def component_matrix() -> list[dict[str, Any]]: + ensure_backends_loaded() + return COMPONENTS.manifest() diff --git a/v2/platform/registry.py b/v2/platform/registry.py new file mode 100644 index 000000000..1eadc0d1a --- /dev/null +++ b/v2/platform/registry.py @@ -0,0 +1,128 @@ +"""The two tuple-keyed backend registries — the dispatch membrane (design_v3 §17; multi-backend RFC). + +Everything device/arch-specific resolves through one of two registries, keyed by a tuple: + + * ``COMPONENTS`` — keyed ``(kind, device, variant)``. Resolves which *weight-bearing* component + implementation to materialize (the toy numpy module, a torch adapter on a GPU box, …). This is + the generalization of ``ComponentSpec.factory``: the card's ``factory`` is simply the terminal + ``(kind, "cpu", "numpy")`` rung that every other backend falls back to. + * ``KERNELS`` — keyed ``(op, device, arch, variant)``. Resolves which *stateless primitive* kernel + implements an op (the flow-match solver step, an SDE step, a gemm/norm on a real box). Arch is a + fallback axis (sm100→sm90→…→generic), with the numpy reference as the terminal rung. + +Two properties this module guarantees (the kernel-colocation answer made concrete): + + 1. **Location is decoupled from dispatch.** A registration carries a key, not a path. Where the .cu + (or .py) lives on disk — unified ``fastvideo-kernel/`` or co-located with a model — is invisible + here; both call the same ``register_*``. + 2. **The matrix is enumerable without a GPU and without importing torch.** ``manifest()`` lists + every declared cell — including ones whose backend is *unavailable* here (e.g. a torch/cuda cell + on a CPU box, gated by a ``find_spec``-based ``available`` predicate that never imports torch) — + so the parity/fallback matrix can be dumped on a laptop. ``resolve`` skips unavailable cells; + ``manifest`` lists them. + + Scope caveat (honest): in this in-tree mini the cells are populated by a fixed backend import + list (``Platform.ensure_backends_loaded`` imports ``backends/{cpu,accel,torch_cuda}``), not by + setuptools entry points. So the matrix is complete for the backends in that list, but an + *out-of-tree* or *model-co-located* backend would not self-enumerate until added to it — full + entry-point discovery is the real-package mechanism and is deliberately not wired here. + +This is a pure leaf module: stdlib only. The actual backend registrations live in ``backends/`` and +are imported lazily (``Platform.ensure_backends_loaded``) to keep this import-cycle-free. +""" +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Callable, Iterable + +# --------------------------------------------------------------------------- # +# Op names — the stable identifiers a loop dispatches on (not free functions). # +# --------------------------------------------------------------------------- # +FLOW_MATCH_STEP = "flow_match_step" # deterministic flow-match Euler solver step (ODE serve) +FLOW_SDE_STEP = "flow_sde_step" # FlowGRPO stochastic step + log-prob (RL rollout) + + +def _always_available() -> bool: + return True + + +@dataclass +class Registration: + """One cell in a registry: a key, the implementation, and an availability predicate. + + ``available`` lets a cell be *declared* (so it shows up in ``manifest()``) while being + *unresolvable* in this environment — e.g. a torch/cuda kernel on a CPU box declares + ``available=lambda: torch.cuda.is_available()`` and is listed-but-skipped here. + """ + key: tuple + fn: Callable[..., Any] + available: Callable[[], bool] = _always_available + source: str = "" + meta: dict[str, Any] = field(default_factory=dict) + + def is_available(self) -> bool: + try: + return bool(self.available()) + except Exception: + return False + + +class TupleRegistry: + """A registry keyed by an ordered field tuple, with availability-aware resolution.""" + + def __init__(self, name: str, key_fields: tuple[str, ...]): + self.name = name + self.key_fields = key_fields + self._reg: dict[tuple, Registration] = {} + + def put(self, key: tuple, fn: Callable[..., Any], *, + available: Callable[[], bool] | None = None, source: str = "", + meta: dict[str, Any] | None = None) -> None: + if len(key) != len(self.key_fields): + raise ValueError(f"{self.name}: key {key} does not match fields {self.key_fields}") + self._reg[key] = Registration(key=key, fn=fn, + available=available or _always_available, + source=source, meta=dict(meta or {})) + + def lookup(self, key: tuple) -> Registration | None: + """Exact lookup, ignoring availability (used by ``manifest``/diagnostics).""" + return self._reg.get(key) + + def resolve_first(self, candidate_keys: Iterable[tuple]) -> Registration | None: + """First candidate key that is both registered AND available wins (the fallback chain).""" + for key in candidate_keys: + reg = self._reg.get(key) + if reg is not None and reg.is_available(): + return reg + return None + + def manifest(self) -> list[dict[str, Any]]: + """Every declared cell as a flat dict — enumerable without importing/calling the kernels. + + This is what makes the parity matrix dumpable on a CPU box: it lists torch/cuda cells too, + each with ``available=False`` here, so coverage is never silently truncated to 'what imported'. + """ + rows: list[dict[str, Any]] = [] + for key, reg in sorted(self._reg.items(), key=lambda kv: tuple(map(str, kv[0]))): + row = dict(zip(self.key_fields, key)) + row["available"] = reg.is_available() + row["source"] = reg.source + rows.append(row) + return rows + + +# The two registries (design_v3 §17). Populated by ``backends/`` modules (lazily imported). +COMPONENTS = TupleRegistry("components", ("kind", "device", "variant")) +KERNELS = TupleRegistry("kernels", ("op", "device", "arch", "variant")) + + +def register_component(kind: str, fn: Callable[..., Any], *, device: str, variant: str = "default", + available: Callable[[], bool] | None = None, source: str = "") -> None: + """Register a component builder. ``fn(spec, instance, platform) -> live component``.""" + COMPONENTS.put((kind, device, variant), fn, available=available, source=source) + + +def register_kernel(op: str, fn: Callable[..., Any], *, device: str, arch: str, variant: str = "default", + available: Callable[[], bool] | None = None, source: str = "") -> None: + """Register a stateless kernel. ``fn(*args, **kwargs)`` — same signature as its numpy reference.""" + KERNELS.put((op, device, arch, variant), fn, available=available, source=source) diff --git a/v2/tests/test_platform.py b/v2/tests/test_platform.py new file mode 100644 index 000000000..54f11cfbd --- /dev/null +++ b/v2/tests/test_platform.py @@ -0,0 +1,196 @@ +"""Multi-backend dispatch substrate: device/arch resolution, fallback, and the parity oracle (§17). + +These tests prove the two tuple-keyed registries (COMPONENTS, KERNELS) + the Platform resolve +implementations generically across devices, with the numpy reference as the terminal rung — and +that a second backend (the pure-python ``accel`` stand-in) is bit-for-bit identical to it. The +``cuda`` cells are declared-but-unavailable here, so the matrix is enumerable without a GPU and +without importing torch. +""" +from __future__ import annotations + +import sys + +import numpy as np + +from v2._enums import ExecutionProfile +from v2.cache import CacheManager +from v2.card import load_card +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.parity import bit_identical # the typed C1 instrument (ladder.py) +from v2.platform import ( + FLOW_MATCH_STEP, + FLOW_SDE_STEP, + Platform, + component_matrix, + kernel_matrix, +) +from v2.platform.backends.accel import AccelDiT +from v2.request import DiffusionParams, TaskType, make_request +from v2.request.streams import Stream +from v2.runtime import Engine +from v2.runtime.context import RuntimeLoopContext + +_HUB = Engine() # borrow its observer/interceptor hubs for the loop context + + +def _instance(platform): + card = build_wan21_card() + 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.""" + req = make_request(TaskType.T2V, inst.card.model_id, prompt, + diffusion=DiffusionParams(num_steps=steps, seed=seed, sde_rollout=sde)) + slots: dict = {} + text_encode_node_fn(inst, slots, req, None) + profile = ExecutionProfile.ROLLOUT if sde else ExecutionProfile.SERVE + 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) + while not runner.done: + runner.step() + return np.asarray(runner.result.outputs["latents"]) + + +# --------------------------------------------------------------------------- # +# Detection + the default (unchanged) CPU path # +# --------------------------------------------------------------------------- # +def test_detect_is_cpu_without_torch(): + p = Platform.detect() + assert p.device == "cpu" and p.arch == "numpy" + assert p.device_chain == ("cpu",) + + +def test_default_instance_is_cpu_and_factory_terminal_unchanged(): + inst = load_card(build_wan21_card(), cache_manager=None) # no platform → detect() → cpu + assert inst.platform.device == "cpu" + # the card's factory is the cpu/numpy terminal rung: component build is the toy, untouched + assert isinstance(inst.component("transformer"), ToyDiT) + assert isinstance(inst.component("text_encoder"), ToyTextEncoder) + assert isinstance(inst.component("vae"), ToyVAE) + + +# --------------------------------------------------------------------------- # +# Kernel resolution: terminal rung, arch fallback, variant fallback # +# --------------------------------------------------------------------------- # +def test_cpu_kernels_resolve_to_numpy_reference(): + p = Platform.cpu() + src_fm = p.kernels.resolved_source(FLOW_MATCH_STEP) + src_sde = p.kernels.resolved_source(FLOW_SDE_STEP) + assert "('flow_match_step', 'cpu', 'numpy', 'default')" in src_fm + assert "loop.sampler:flow_match_euler_step" in src_fm + assert "('flow_sde_step', 'cpu', 'numpy', 'default')" in src_sde + + +def test_accel_arch_fallback_walk(): + """flow_match is only at the `generic` arch rung → resolving on sm90 must walk sm90→sm80→generic; + flow_sde is at sm90 directly.""" + p = Platform.accel("sm90") + assert p.arch_chain("accel") == ("sm90", "sm80", "generic") + assert "('flow_match_step', 'accel', 'generic', 'default')" in p.kernels.resolved_source(FLOW_MATCH_STEP) + assert "('flow_sde_step', 'accel', 'sm90', 'default')" in p.kernels.resolved_source(FLOW_SDE_STEP) + + +def test_arch_fallback_is_monotonic_never_newer(): + """An OLDER device must never fall back to a newer-arch (binary-incompatible) kernel; a NEWER + device may still run older-arch kernels (forward-compatible).""" + old = Platform.accel("sm70") # older than the accel chain's sm90/sm80 + chain = old.arch_chain("accel") + assert chain[0] == "sm70" + assert "sm90" not in chain and "sm80" not in chain # newer archs excluded + assert "generic" in chain # portable terminal kept + new = Platform.cuda("sm120") # newer than the cuda chain + assert {"sm100", "sm90", "sm80"} <= set(new.arch_chain("cuda")) # older archs still usable + + +def test_device_precedence_accel_before_cpu(): + """flow_sde exists on BOTH accel(sm90) and cpu(numpy) — the accel device wins (chain order).""" + p = Platform.accel("sm90") + assert "accel" in p.kernels.resolved_source(FLOW_SDE_STEP) + + +def test_unknown_variant_falls_back_to_default(): + p = Platform.cpu() + # asking for a variant that was never registered still resolves to the "default" cell + assert p.kernels.get(FLOW_MATCH_STEP, variant="does_not_exist") is p.kernels.get(FLOW_MATCH_STEP) + + +def test_cuda_platform_falls_back_to_cpu_when_unavailable(): + """A cuda platform on a torchless box: the cuda cells are unavailable, so resolution skips them + and walks the device chain cuda→cpu to the numpy terminal.""" + p = Platform.cuda("sm90") + assert "('flow_match_step', 'cpu', 'numpy', 'default')" in p.kernels.resolved_source(FLOW_MATCH_STEP) + + +# --------------------------------------------------------------------------- # +# Component resolution: device override + per-kind fallback to the factory # +# --------------------------------------------------------------------------- # +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 + assert isinstance(inst.component("text_encoder"), ToyTextEncoder) + assert isinstance(inst.component("vae"), ToyVAE) + + +def test_cuda_component_falls_back_to_factory_when_unavailable(): + inst = _instance(Platform.cuda("sm90")) # cuda dit cell unavailable on this box + assert isinstance(inst.component("transformer"), ToyDiT) # → spec.factory terminal + + +# --------------------------------------------------------------------------- # +# THE PARITY ORACLE: a new backend must match the numpy reference bit-for-bit # +# (checked with the typed C1 instrument `parity.bit_identical`, not a bare ==) # +# --------------------------------------------------------------------------- # +def test_parity_oracle_accel_equals_cpu_ode(): + """accel's ODE solver kernel is an INDEPENDENT impl (`_accel_flow_match_step`) — this is the real + oracle: a separately-written backend kernel must match the numpy reference, C1 bit-identical.""" + cpu_latent = _drive_denoise(_instance(Platform.cpu()), seed=7, steps=8) + accel_latent = _drive_denoise(_instance(Platform.accel("sm90")), seed=7, steps=8) + assert bit_identical(cpu_latent, accel_latent) + + +def test_accel_sde_dispatch_path_equivalent(): + """The SDE op is NOT specialized by accel (it reuses the numpy reference), so this checks the + stochastic path still resolves + runs through the kernel table and matches — a dispatch-path + test, not an independent-impl oracle (see accel.py).""" + cpu_latent = _drive_denoise(_instance(Platform.cpu()), seed=11, steps=6, sde=True) + accel_latent = _drive_denoise(_instance(Platform.accel("sm90")), seed=11, steps=6, sde=True) + assert bit_identical(cpu_latent, accel_latent) + + +def test_accel_denoise_still_produces_valid_output(): + out = _drive_denoise(_instance(Platform.accel("sm90")), seed=3, steps=4) + assert out.shape[0] == 4 and np.all(np.isfinite(out)) + + +# --------------------------------------------------------------------------- # +# Enumerable matrix — including unavailable backends, without importing torch # +# --------------------------------------------------------------------------- # +def test_kernel_matrix_enumerates_all_backends_including_unavailable_cuda(): + rows = kernel_matrix() + devices = {r["device"] for r in rows} + assert {"cpu", "accel", "cuda"} <= devices + cuda_rows = [r for r in rows if r["device"] == "cuda"] + assert cuda_rows and all(r["available"] is False for r in cuda_rows) # declared, not runnable here + assert all(r["available"] is True for r in rows if r["device"] == "cpu") + + +def test_component_matrix_includes_declared_unavailable_torch_adapter(): + rows = component_matrix() + cuda = [r for r in rows if r["device"] == "cuda" and r["kind"] == "dit"] + assert cuda and cuda[0]["available"] is False + accel = [r for r in rows if r["device"] == "accel" and r["kind"] == "dit"] + assert accel and accel[0]["available"] is True + + +def test_matrix_enumeration_does_not_import_torch(): + """The manifest is built from declarations + a find_spec predicate — never an import of torch.""" + kernel_matrix(); component_matrix() + assert "torch" not in sys.modules