File size: 7,515 Bytes
b324599 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 | #!/usr/bin/env python3
"""Launch the resident Fast-WAM runtime and run the action-expert denoise loop on the NPU
WITHOUT reload, verify parity vs the PyTorch golden, and measure latency.
Fast-WAM's on-device unit here is the MoT **action expert** (`action_step`), which cross-
attends to the video KV cache produced by the video world-model DiT (`video_prefill`).
The video prefill runs host-side for now (its context binary needs a MolmoAct2-style
layer-split, see README), so the host:
1. loads the video KV cache + context (here: from golden/, i.e. the video_prefill output),
2. runs the flow-matching Euler loop: for each of num_inference_steps, feed the current
latents_action + timestep -> pred_noise (NPU), then update latents with the Wan
continuous flow-match step latents += pred * delta,
3. verifies the single-step output vs golden (parity gate), reports latency.
Pipeline (on-device portion): action_step ×N (resident). No per-inference reload.
"""
import sys, os, socket, struct, subprocess, time
import numpy as np
WS = os.environ.get("FASTWAM_WS", "/root/fastwam_workspace")
RT = f"{WS}/runtime"
ROLES = ["action_step"]
PORT0 = 5801
STEPS = int(os.environ.get("FASTWAM_STEPS", "10")) # num_inference_steps
SHIFT = float(os.environ.get("FASTWAM_SHIFT", "5.0")) # action_scheduler infer_shift
NUM_TRAIN = int(os.environ.get("FASTWAM_NUM_TRAIN_TIMESTEPS", "1000"))
def recvall(s, n):
b = bytearray()
while len(b) < n:
c = s.recv(min(n - len(b), 8 << 20))
if not c:
raise ConnectionError()
b += c
return bytes(b)
def recv_tensors(s):
n = struct.unpack("<I", recvall(s, 4))[0]
out = {}
for _ in range(n):
nl = struct.unpack("<I", recvall(s, 4))[0]; name = recvall(s, nl).decode()
dl = struct.unpack("<I", recvall(s, 4))[0]; dt = recvall(s, dl).decode()
nd = struct.unpack("<I", recvall(s, 4))[0]; shape = struct.unpack(f"<{nd}I", recvall(s, nd * 4))
nb = struct.unpack("<Q", recvall(s, 8))[0]; data = recvall(s, nb)
out[name] = np.frombuffer(data, dtype=dt).reshape(shape).copy()
return out
def send_tensors(s, d):
parts = [struct.pack("<I", len(d))]
for name, a in d.items():
a = np.ascontiguousarray(a, np.float32)
nb = name.encode(); dt = str(a.dtype).encode()
parts.append(struct.pack("<I", len(nb)) + nb)
parts.append(struct.pack("<I", len(dt)) + dt)
parts.append(struct.pack("<I", a.ndim) + struct.pack(f"<{a.ndim}I", *a.shape))
parts.append(struct.pack("<Q", a.nbytes) + a.tobytes())
s.sendall(b"".join(parts))
def call(sock, d):
send_tensors(sock, d)
return recv_tensors(sock)
def cos(r, g):
r = np.asarray(r, np.float64).ravel(); g = np.asarray(g, np.float64).ravel()
return float(r @ g / (np.linalg.norm(r) * np.linalg.norm(g) + 1e-12))
def build_schedule(steps, shift, num_train):
"""Wan continuous flow-match inference schedule (vendored from video_dit.py):
sigma = linspace(1,0,steps+1)[:steps]; sigma = shift*sigma/(1+(shift-1)*sigma).
timesteps = sigma*num_train; deltas = sigma_next - sigma (sigma_next[-1]=0)."""
sigma = np.linspace(1, 0, steps + 1)[:steps]
sigma = shift * sigma / (1 + (shift - 1) * sigma)
timesteps = (sigma * num_train).astype(np.float32)
sigma_next = np.concatenate([sigma[1:], np.zeros(1)])
deltas = (sigma_next - sigma).astype(np.float32)
return timesteps, deltas
def main():
os.makedirs(f"{WS}/tmp", exist_ok=True)
for r in ROLES:
try: os.remove(f"{WS}/tmp/w_{r}.ready")
except OSError: pass
env = os.environ.copy()
env["LD_LIBRARY_PATH"] = f"{WS}/qnn_libs:" + env.get("LD_LIBRARY_PATH", "")
env["ADSP_LIBRARY_PATH"] = f"{WS}/dsp_libs"
procs, ports = {}, {}
t0 = time.time()
for i, r in enumerate(ROLES):
ports[r] = PORT0 + i
logf = open(f"{WS}/tmp/w_{r}.log", "w")
procs[r] = subprocess.Popen(["python3", "-u", f"{RT}/resident_worker.py", r, str(ports[r])],
env=env, stdout=logf, stderr=subprocess.STDOUT)
print(f"[launch] {len(ROLES)} worker spawning ...", flush=True)
for r in ROLES:
for _ in range(240):
if os.path.exists(f"{WS}/tmp/w_{r}.ready"):
break
if procs[r].poll() is not None:
print(f"[ERR] worker {r} died during load (rc={procs[r].returncode}); see tmp/w_{r}.log")
return 3
time.sleep(0.5)
else:
print(f"[ERR] worker {r} not ready"); return 3
print(f"[launch] worker RESIDENT in {time.time()-t0:.1f}s (1 NPU session)", flush=True)
socks = {}
for r in ROLES:
s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
s.connect(("127.0.0.1", ports[r])); s.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
socks[r] = s
g = np.load(f"{WS}/golden/action_step_io.npz")
vk = g["video_keys"].astype(np.float32) # [30,1,98,3072] (video world-model KV cache)
vv = g["video_values"].astype(np.float32)
ctx = g["context"].astype(np.float32) # [1,129,4096]
cmask = g["context_mask"].astype(np.float32) # [1,129]
la0 = g["latents_action"].astype(np.float32) # [1,32,7] initial noise
ts0 = g["timestep"].astype(np.float32) # [1] step-0 timestep
pred_gold = g["pred"].astype(np.float32) # [1,32,7] one-step reference
const = {"video_keys": vk, "video_values": vv, "context": ctx, "context_mask": cmask}
# ---- parity: single step with golden inputs vs golden pred ----
out1 = call(socks["action_step"], {**const, "latents_action": la0, "timestep": ts0})
key = "pred_noise" if "pred_noise" in out1 else list(out1)[0]
c = cos(pred_gold, out1[key])
print(f"[verify] action_step single-step pred cos={c:.8f} (PASS={c>=0.999})", flush=True)
# ---- full N-step flow-match loop (latency) ----
timesteps, deltas = build_schedule(STEPS, SHIFT, NUM_TRAIN)
def run_loop(timing=None):
la = la0.copy()
for i in range(STEPS):
ts = np.array([timesteps[i]], np.float32)
s = time.time()
o = call(socks["action_step"], {**const, "latents_action": la, "timestep": ts})
if timing is not None:
timing["action_step"] = timing.get("action_step", 0) + time.time() - s
pred = np.ascontiguousarray(o[key], np.float32).reshape(la.shape)
la = la + pred * deltas[i] # Wan continuous flow-match step: x += pred*delta
return la
run_loop() # warmup
N = 5
timing = {}
t0 = time.time()
for _ in range(N):
run_loop(timing)
total = (time.time() - t0) / N
print(f"\n=== RESIDENT LATENCY (avg of {N}, 1 session, NO reload, {STEPS}-step flow-match) ===", flush=True)
print(f" action_step (x{STEPS}) {timing['action_step']/N*1000:7.1f} ms (incl. video-KV TCP each step)", flush=True)
print(f" {'per-step (wall)':18} {timing['action_step']/N/STEPS*1000:7.1f} ms", flush=True)
print(f" {'TOTAL':18} {total*1000:7.1f} ms", flush=True)
print(" (pure NPU ctx.Inference per step -> tmp/w_action_step.log [infer-ms])", flush=True)
for s in socks.values():
try: s.close()
except OSError: pass
time.sleep(1)
for p in procs.values():
try: p.terminate()
except OSError: pass
return 0
if __name__ == "__main__":
sys.exit(main())
|