File size: 3,878 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 | #!/usr/bin/env python3
"""Persistent resident QNN worker for Fast-WAM (one NPU session per process).
Loads the `action_step` HTP context ONCE, stays alive, serves inference over a TCP
socket (no per-inference reload). Fast-WAM's on-device unit here is the MoT **action
expert** denoise step — it cross-attends to the video KV cache produced by the video
world-model DiT (`video_prefill`), which runs host-side for now (see README).
role:
action_step : one flow-matching denoise step of the action expert.
inputs : latents_action [1,32,7], timestep [1], video_keys [30,1,98,3072],
video_values [30,1,98,3072], context [1,129,4096], context_mask [1,129]
output : pred_noise [1,32,7] (host applies the Euler/flow-match update)
Usage: resident_worker.py <role> <port>
"""
import sys, os, socket, struct, time
import numpy as np
WS = os.environ.get("FASTWAM_WS", "/root/fastwam_workspace")
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)
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 feed(ctx, m):
"""Order inputs to match the context; QNN consumes float32 raws (bool mask too)."""
out = []
for nm, shp in zip(ctx.getInputName(), ctx.getInputShapes()):
a = np.asarray(m[nm], np.float32)
if tuple(a.shape) != tuple(shp) and a.size == int(np.prod(shp)):
a = a.reshape(shp)
out.append(np.ascontiguousarray(a, np.float32))
return out
def named(ctx, o):
return {nm: np.asarray(o[i], np.float32) for i, nm in enumerate(ctx.getOutputName())}
def main():
role, port = sys.argv[1], int(sys.argv[2])
from qai_appbuilder import QNNConfig, QNNContext, PerfProfile, Runtime, LogLevel
QNNConfig.Config(qnn_lib_path=f"{WS}/qnn_libs", runtime=Runtime.HTP, log_level=LogLevel.ERROR)
suffix = os.environ.get("ACTION_CTX_SUFFIX", "")
ctx = QNNContext(role, f"{WS}/ctx/{role}{suffix}_socid77_archv73.bin")
PerfProfile.SetPerfProfileGlobal(PerfProfile.BURST)
srv = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
srv.bind(("127.0.0.1", port)); srv.listen(1)
open(f"{WS}/tmp/w_{role}.ready", "w").write("1")
conn, _ = srv.accept()
conn.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
while True:
try:
req = recv_tensors(conn)
except (ConnectionError, struct.error):
break
_t = time.time()
o = ctx.Inference(feed(ctx, req))
print(f"[infer-ms] {role} {(time.time()-_t)*1000:.1f}", flush=True)
send_tensors(conn, named(ctx, o))
os._exit(0) # dodge an appbuilder<->libs teardown double-free (see README device gotchas)
if __name__ == "__main__":
main()
|