[feat] multi-backend dispatch substrate (device/arch/kernel registries)
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.
This commit is contained in:
SolitaryThinker
committed by
Peiyuan Zhang
Peiyuan Zhang
parent
431f4daddb
commit
2f044c02dd
+25
-8
@@ -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)
|
||||
|
||||
+19
-7
@@ -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())
|
||||
|
||||
@@ -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")})
|
||||
|
||||
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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.
|
||||
"""
|
||||
@@ -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")
|
||||
@@ -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")
|
||||
@@ -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)")
|
||||
@@ -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()
|
||||
@@ -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)
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user