| |
| """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) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
