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.
70 lines
2.5 KiB
Python
70 lines
2.5 KiB
Python
"""Array namespace — numpy on a CPU box, torch-on-device on a GPU box.
|
||
|
||
The denoise loop's math is already array-agnostic: CFG combine (``uncond + s·(cond−uncond)``) and the
|
||
flow-match Euler step (``x + (σ_next−σ_t)·v``) are pure arithmetic that runs identically on numpy and
|
||
torch. This namespace supplies the few NON-arithmetic helpers the loop still needs — birthing the
|
||
latent, dtype casts, and the single host round-trip at the request boundary — so on a GPU box the
|
||
latent stays resident on-device for the whole loop (no per-step host<->device copy) while the CPU/toy
|
||
path stays pure numpy (torch-free; the parity mini is unchanged).
|
||
|
||
The latent is still *seeded* with numpy (``rng.standard_normal``) and uploaded once via ``from_host``,
|
||
so a GPU run is bit-identical to the pre-on-device path (the noise values are the same; upload is
|
||
lossless). ``Platform.xp`` picks the namespace from the platform's device.
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
from typing import Any
|
||
|
||
import numpy as np
|
||
|
||
|
||
class NumpyNS:
|
||
"""The CPU/toy namespace: every op is a numpy identity (no torch dependency)."""
|
||
|
||
device = "cpu"
|
||
|
||
def from_host(self, a: Any) -> np.ndarray:
|
||
return np.asarray(a, dtype=np.float32)
|
||
|
||
def to_host(self, a: Any) -> np.ndarray:
|
||
return np.asarray(a)
|
||
|
||
def to_f32(self, a: Any) -> np.ndarray:
|
||
return np.asarray(a, dtype=np.float32)
|
||
|
||
def is_native(self, a: Any) -> bool:
|
||
return isinstance(a, np.ndarray)
|
||
|
||
|
||
class TorchNS:
|
||
"""The GPU namespace: keeps arrays as device tensors; marshals host<->device only on demand."""
|
||
|
||
def __init__(self, device: str = "cuda") -> None:
|
||
import torch
|
||
self._torch = torch
|
||
self.device = device
|
||
|
||
def from_host(self, a: Any) -> Any:
|
||
t = self._torch
|
||
if t.is_tensor(a):
|
||
return a.to(self.device)
|
||
return t.as_tensor(np.asarray(a, dtype=np.float32), device=self.device)
|
||
|
||
def to_host(self, a: Any) -> np.ndarray:
|
||
t = self._torch
|
||
if t.is_tensor(a):
|
||
return a.detach().to("cpu", t.float32).numpy()
|
||
return np.asarray(a)
|
||
|
||
def to_f32(self, a: Any) -> Any:
|
||
t = self._torch
|
||
return a.float() if t.is_tensor(a) else np.asarray(a, dtype=np.float32)
|
||
|
||
def is_native(self, a: Any) -> bool:
|
||
return bool(self._torch.is_tensor(a))
|
||
|
||
|
||
def get_array_ns(device: str) -> Any:
|
||
"""The array namespace for a platform device: torch-on-device for cuda, numpy otherwise."""
|
||
return TorchNS(device) if device == "cuda" else NumpyNS()
|