SolitaryThinker
4a33540c08
The torch adapter boundary marshalled the latent host<->device on EVERY denoise step: _t uploaded the latent (and re-uploaded the text embeds) and _n downloaded the velocity with a forced CUDA sync — 2*N PCIe copies + N syncs per generation, buying nothing, since the latent could stay resident on the GPU the whole loop. Root cause was a numpy loop surface. But the loop MATH is already array-agnostic (CFG combine + flow-match Euler are pure arithmetic; the solver kernel already passes torch through). So introduce a per-platform array namespace (v2/platform/ array_ns): numpy on CPU (torch-free — the parity mini is unchanged), torch-on- device on cuda. The latent is seeded with numpy and uploaded ONCE; it then stays resident through forward -> CFG combine -> solver -> next step; a single host marshal happens at the request/output boundary (engine._to_artifact). Opt-in per recipe via ModelCard.device_io (set on the Wan cards). When set on a GPU box, build_component flips the components' TorchComponent.device_io so _out keeps tensors on-device (in fp32, matching the old _to_numpy cast so the combine dtype is unchanged). Un-migrated families and the CPU toy keep numpy in/out. Also: PrecisionPolicy.cast is array-preserving; _t accepts resident tensors. Verified BIT-IDENTICAL on real Wan2.1-1.3B / H100: the on-device latent equals the pre-change numpy-path latent exactly (max_abs_diff 0.0, np.array_equal True). CPU mini holds 237 passed / 5 pre-existing; pre-commit clean. Other WanDenoiseLoop families can flip device_io next (per-family GPU re-verify); non-Wan loops migrate to the xp namespace later.
205 lines
9.0 KiB
Python
205 lines
9.0 KiB
Python
"""Platform — the detected ``(device, arch)`` that resolves both registries.
|
||
|
||
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()`` returns a CPU/numpy platform unless torch+CUDA are actually present. The
|
||
(device, arch) machinery is exercised on a CPU box by the pure-python ``accel`` stand-in backend,
|
||
proving the dispatch is generic without a real GPU.
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
from typing import Any
|
||
from collections.abc import Iterator
|
||
|
||
from v2.platform.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 v2.platform.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 or rank >= 0 and ar < rank: # portable terminal (ptx/generic): always valid
|
||
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
|
||
self._xp: Any = 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
|
||
|
||
# --- the array namespace (numpy on CPU, torch-on-device on a GPU box) ----- #
|
||
@property
|
||
def xp(self) -> Any:
|
||
"""Birth/cast/marshal helpers so the denoise loop is array-agnostic: numpy on CPU (torch-free),
|
||
torch-on-device on cuda (the latent stays resident — no per-step host<->device round-trip)."""
|
||
if self._xp is None:
|
||
from v2.platform.array_ns import get_array_ns
|
||
self._xp = get_array_ns(self.device)
|
||
return self._xp
|
||
|
||
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()
|