Fast-WAM action-expert QNN HTP context binary + resident runtime + README

README.md

@@ -0,0 +1,176 @@
+---
+license: apache-2.0
+base_model: ZibinDong/fastwam_libero_uncond_2cam224
+tags:
+- fast-wam
+- world-model
+- mot
+- vla
+- robotics
+- libero
+- qnn
+- qualcomm
+- htp
+- hexagon
+- npu
+- edge
+pipeline_tag: robotics
+---
+
+# Fast-WAM (LIBERO) — QNN / HTP context binary (action expert)
+
+Offline-compiled **Qualcomm QNN (HTP) context binary** for the MoT **action expert** of
+[`ZibinDong/fastwam_libero_uncond_2cam224`](https://huggingface.co/ZibinDong/fastwam_libero_uncond_2cam224),
+ready to run **resident on-device** on a **Qualcomm Dragonwing IQ‑9075 (QCS9075, Hexagon v73, soc_id 77)**.
+
+This is the **runtime deployment bundle**: the prebuilt `action_step` `.bin` context, the golden
+reference tensors (inputs + parity reference, including the video KV cache), and the resident
+runtime. The conversion pipeline (`PyTorch → ONNX → QNN DLC → HTP context binary`) lives in the
+[`nota-github/xpu-molmoact2-qnn-htp`](https://github.com/nota-github/xpu-molmoact2-qnn-htp) repo
+(`fastwam_qnn/`, branch `support_Fast-WAM`, PR #2).
+
+> **What's on the NPU.** Fast-WAM is a Mixture-of-Transformers (MoT) world model: a frozen Wan2.2
+> video VAE + UMT5 text encoder feed a **video DiT** (world model) and an **action DiT** (action
+> expert). This bundle ships the **action expert** denoise step — after a **real-valued-RoPE +
+> constant-mask surgery** it converts to HTP with **cosine 1.000000** CPU-backend parity. The
+> **video world-model prefill** (which produces the video KV cache the action expert cross-attends
+> to) converts to a QNN DLC with cosine 1.0 too, but its fp16 context is ~9.7 GB (> the device's
+> ~2 GB/session limit) so it needs a MolmoAct2-style layer-split before it can go on-NPU; for now
+> it runs **host-side** and the runtime feeds its output (the golden video KV cache). The frozen
+> Wan2.2 VAE stays host-side (Conv3d 5D exceeds QNN's rank-5 limit). See [Remaining](#remaining).
+
+---
+
+## What's in here
+
+```
+ctx/      1 HTP context binary (soc_id 77 / Hexagon v73), fp16 weights
+  action_step_socid77_archv73.bin       2.06 GB
+golden/   reference I/O (PyTorch fp32): action-step inputs (incl. the video KV cache) + parity ref
+  action_step_io.npz     video_prefill_io.npz     trace.json     convert_report.json
+runtime/  resident runtime
+  resident_run.py      orchestrator (host glue: flow-match Euler loop, parity, latency)
+  resident_worker.py   one NPU session per process (loads the action_step context once)
+  profile_device.sh    one-command host→device profiler (SSH; reads creds from IQ9_info.txt)
+```
+
+You also need, **from the QAIRT 2.47.x SDK** (not redistributed here):
+`qnn_libs/` (aarch64 QNN runtime `.so`s) and `dsp_libs/` (Hexagon **v73** skel libs).
+
+---
+
+## Architecture — 1 on-NPU component
+
+```
+  image ─┐  Wan2.2 VAE      ┌───────────────┐  video KV cache   ┌────────────────────┐
+         ├───(host)────────▶│  video DiT    │ ────────────────▶ │  action DiT (NPU)  │─▶ actions
+ text ───┘  UMT5 (host)     │ world model   │  [30,1,98,3072]   │  flow-matching     │   [1,32,7]
+                            │ prefill (host)│  (×2 k/v)         │  step ×N (resident)│
+                            └───────────────┘                   └────────────────────┘
+```
+
+- The **action expert** (`action_step`) is a single HTP context, **looped N× on the host** for
+  flow-matching denoising. Only `latents_action` and `timestep` change across the loop; the video
+  KV cache + context are constant per chunk.
+- The **video world-model prefill** and the **VAE** run host-side for now (see [Remaining](#remaining));
+  the runtime feeds their outputs from `golden/`.
+
+### I/O contract
+
+| context | inputs (name: shape, f32) | output (name: shape, f32) |
+|---|---|---|
+| `action_step` | `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]` | `pred_noise: [1,32,7]` |
+
+Host loop (Wan continuous flow-match): schedule `sigma = linspace(1,0,N+1)[:N]; sigma = shift·sigma/(1+(shift−1)·sigma)`,
+`timesteps = sigma·1000`, `deltas = sigma_next − sigma`; for step `i`, feed `latents_action` + `timesteps[i]`,
+take `pred_noise`, update `latents_action += pred_noise · deltas[i]`. (shift = 5.0.)
+
+---
+
+## Inference guide
+
+### 1. Device workspace
+
+On the IQ‑9075, create `/root/fastwam_workspace/` with:
+
+```
+/root/fastwam_workspace/
+  ctx/        ← action_step_socid77_archv73.bin from this repo
+  golden/     ← the golden/ folder from this repo
+  runtime/    ← resident_run.py, resident_worker.py  (profile_device.sh pushes these for you)
+  qnn_libs/   ← aarch64 QNN runtime .so from QAIRT 2.47.x: lib/aarch64-oe-linux-gcc11.2/*.so
+  dsp_libs/   ← Hexagon v73 skel from QAIRT 2.47.x:        lib/hexagon-v73/unsigned/*.so
+  tmp/        ← created automatically
+```
+
+The runtime loads them via `LD_LIBRARY_PATH=$WS/qnn_libs` and `ADSP_LIBRARY_PATH=$WS/dsp_libs`.
+Device Python needs **`qai_appbuilder`** (QAI AppBuilder for QNN) + `numpy`.
+
+> **action_step is 2.06 GB** — right at the device's ~2 GB single-session fastrpc weights-buffer
+> limit. **Reboot before a clean resident run** (crashed cycles leak DSP sessions); the worker ends
+> with `os._exit()` to dodge an appbuilder↔libs teardown double-free.
+
+### 2a. Run directly on the device
+
+```bash
+cd /root/fastwam_workspace
+python3 runtime/resident_run.py          # FASTWAM_STEPS=10 (default), FASTWAM_SHIFT=5.0
+```
+
+It spawns the worker, keeps it **resident**, runs the N-step flow-match loop, prints the
+**single-step pred cosine** vs the golden reference, then a latency summary (avg of 5) plus the
+pure per-step `[infer-ms]` (excl. TCP) in `tmp/w_action_step.log`.
+
+### 2b. Or profile from the host (one command)
+
+```bash
+runtime/profile_device.sh
+```
+
+Copies the runtime scripts to an ephemeral `/tmp` workspace, symlinks the heavy assets, runs, prints
+the latency table + per-step `[infer-ms]`, then deletes the temp dir. Device creds are read from
+`IQ9_info.txt` (`IP:` / `passwd:`; override with `IQ9_INFO=...`).
+
+### fp16 / quantized A/B switch
+
+```bash
+ACTION_CTX_SUFFIX=_w4a16 python3 runtime/resident_run.py   # if you add a quantized bin
+```
+
+---
+
+## Validated parity (host, QNN CPU backend vs PyTorch fp32)
+
+Verified device-free on the QNN **CPU** backend against the PyTorch golden. **Cosine is the gate**
+(≥ 0.9999).
+
+| stage | output | cosine |
+|---|---|---|
+| ONNX (ORT CPU) — reference | — | action expert exports clean after real-RoPE surgery |
+| QNN DLC (CPU backend) vs PyTorch | pred_noise | **1.000000** |
+| video_prefill DLC (CPU backend) vs PyTorch | video_keys / video_values | 1.000000 / 1.000000 |
+
+The on-device (HTP fp16) single-step pred cosine + latency are what `resident_run.py` reports.
+
+---
+
+## Remaining
+
+- **video world-model prefill — layer-split.** The 30-layer video DiT converts to a QNN DLC with
+  cosine 1.0, but its fp16 context (~9.7 GB) exceeds the device's ~2 GB/session limit. Split it into
+  layer-range contexts (MolmoAct2-LLM pattern) to run on-NPU; until then the runtime uses the
+  host-computed video KV cache from `golden/`.
+- **Wan2.2 VAE — host-side.** The frozen `AutoencoderKLWan` is Conv3d (5D); an internal transpose
+  exceeds QNN HTP's rank-5 limit. Keep it host-side (feed latents), or refactor the T=1 path to rank-4.
+
+---
+
+## Build provenance
+
+- Source policy: [`ZibinDong/fastwam_libero_uncond_2cam224`](https://huggingface.co/ZibinDong/fastwam_libero_uncond_2cam224).
+- Action expert: MoT action DiT (30 layers, hidden 1024, 24 heads × 128), flow-matching, action chunk `[1,32,7]`.
+- Surgery for ONNX/HTP: complex `view_as_complex`/`view_as_real` RoPE → real rotate-half (`.freqs`
+  stored as real cos/sin); MoT boolean attention masks precomputed as constant buffers.
+- Toolchain: QAIRT **2.47.0**, opset 20, float DLC (no quant), offline HTP context-binary for
+  soc_id 77 / dsp_arch v73, `O3`.
+- Conversion code: `nota-github/xpu-molmoact2-qnn-htp` → `fastwam_qnn/`, `scripts_fastwam/`, `fastwam_qnn/surgery.py` (branch `support_Fast-WAM`, PR #2).

ctx/action_step_socid77_archv73.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:97ca5d72e1c5b79af2bfef17c006b091ecbb9385a47c9e3f55f8e6fd82e1ee51
+size 2064179200

golden/action_step_io.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:894e8dd3bf10ac8cafaf74deda1acdc4b3cbba916c960398f4ef58c7e3c0ac83
+size 74370681

golden/convert_report.json

@@ -0,0 +1,66 @@
+{
+  "vae_encoder": {
+    "status": "BLOCKED",
+    "in": [
+      "pixels"
+    ],
+    "out": [
+      "latent"
+    ],
+    "error": "CalledProcessError: Command '['/ssd6/geonmin.kim/qairt/2.47.0.260601/bin/x86_64-linux-clang/qairt-converter', '--input_network', '/ssd6/geonmin.kim/xpu-molmoact2-qnn-htp/outputs_fastwam/onnx/vae_encoder/vae_encoder.onnx', '-o', '/ssd6/geonmin.kim/xpu-molmoact2-qnn-htp/outputs_fastwam/qnn/vae_encoder/vae_encoder.dlc', '--enable_tensor_deduplication']' returned non-zero exit status 1.",
+    "trace": "Traceback (most recent call last):\n  File \"/ssd6/geonmin.kim/xpu-molmoact2-qnn-htp/scripts_fastwam/02_convert_qnn.py\", line 54, in convert_and_verify\n    qnn.to_dlc(s, onnx_path, dlc)\n  File \"/ssd6/geonmin.kim/xpu-molmoact2-qnn-htp/fastwam_qnn/qnn.py\", line 45, in to_dlc\n    _run(cmd)\n  File \"/ssd6/geonmin.kim/xpu-molmoact2-qnn-htp/fastwam_qnn/qnn.py\", line 28, in _run\n    subprocess.run(cmd, check=True)\n  File \"/usr/lib/python3.12/subprocess.py\", line 571, in run\n    raise CalledProcessError(retcode, process.args,\nsubprocess.CalledProcessError: Command '['/ssd6/geonmin.kim/qairt/2.47.0.260601/bin/x86_64-linux-clang/qairt-converter', '--input_network', '/ssd6/geonmin.kim/xpu-molmoact2-qnn-htp/outputs_fastwam/onnx/vae_encoder/vae_encoder.onnx', '-o', '/ssd6/geonmin.kim/xpu-molmoact2-qnn-htp/outputs_fastwam/qnn/vae_encoder/vae_encoder.dlc', '--enable_tensor_deduplication']' returned non-zero exit status 1.\n"
+  },
+  "video_prefill": {
+    "status": "OK",
+    "in": [
+      "first_frame_latent",
+      "context",
+      "context_mask"
+    ],
+    "out": [
+      "video_keys",
+      "video_values"
+    ],
+    "dlc": "/ssd6/geonmin.kim/xpu-molmoact2-qnn-htp/outputs_fastwam/qnn/video_prefill/video_prefill.dlc",
+    "parity": {
+      "video_keys": {
+        "cosine": 0.9999999982130974,
+        "max_abs": 0.0009984970092773438,
+        "max_rel": 407.65616383779053,
+        "within_tol": 1.0
+      },
+      "video_values": {
+        "cosine": 0.9999999892041477,
+        "max_abs": 0.001104697585105896,
+        "max_rel": 66029.9847368151,
+        "within_tol": 0.9999998892786282
+      }
+    },
+    "cpu_verify": "OK"
+  },
+  "action_step": {
+    "status": "OK",
+    "in": [
+      "latents_action",
+      "timestep",
+      "video_keys",
+      "video_values",
+      "context",
+      "context_mask"
+    ],
+    "out": [
+      "pred_noise"
+    ],
+    "dlc": "/ssd6/geonmin.kim/xpu-molmoact2-qnn-htp/outputs_fastwam/qnn/action_step/action_step.dlc",
+    "parity": {
+      "pred_noise": {
+        "cosine": 0.9999999855798641,
+        "max_abs": 0.0006463229656219482,
+        "max_rel": 0.007894328519392351,
+        "within_tol": 1.0
+      }
+    },
+    "cpu_verify": "OK",
+    "context_binary_mb": 2064.2
+  }
+}

golden/trace.json

@@ -0,0 +1,30 @@
+{
+  "vae": {
+    "out": [
+      [
+        1,
+        48,
+        1,
+        14,
+        28
+      ],
+      "torch.float32"
+    ]
+  },
+  "video_pre": {
+    "out": "<class 'dict'>"
+  },
+  "video_prefill": {
+    "out": "<class 'list'>"
+  },
+  "action_step": {
+    "out": [
+      [
+        1,
+        32,
+        7
+      ],
+      "torch.float32"
+    ]
+  }
+}

golden/video_prefill_io.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a5e8267ac75d67687de913105df5e3c259c262b8811499950c3f6ae340aa0f97
+size 74443669

runtime/profile_device.sh

@@ -0,0 +1,40 @@
+#!/usr/bin/env bash
+# Profile the resident Fast-WAM runtime on the IQ-9075 — from the host, one command.
+#
+# Copies the LATEST runtime/*.py from this repo to an EPHEMERAL workspace under the
+# device's /tmp, symlinks the heavy assets (ctx/ golden/ qnn_libs/ dsp_libs/) from the
+# persistent /root/fastwam_workspace (no multi-GB copy), runs resident_run.py, prints the
+# latency table + per-step pure [infer-ms], then deletes the temp dir. Nothing persistent
+# is modified.
+#
+# Usage (fp16 default; A/B a quantized action bin via env):
+#   runtime/profile_device.sh
+#   runtime/profile_device.sh ACTION_CTX_SUFFIX=_w4a16
+#
+# Device creds are read from IQ9_info.txt (override path with IQ9_INFO=...).
+set -euo pipefail
+
+HERE="$(cd "$(dirname "$0")" && pwd)"
+INFO="${IQ9_INFO:-$HERE/../../IQ9_info.txt}"
+[ -f "$INFO" ] || { echo "device info not found: $INFO (set IQ9_INFO=...)" >&2; exit 1; }
+IP=$(awk -F'[: ]+' '/^IP/{print $2}' "$INFO")
+PW=$(awk -F'[: ]+' '/passwd/{print $2}' "$INFO")
+PERSIST=/root/fastwam_workspace
+ENVS="$*"
+
+SSH="sshpass -p $PW ssh -o StrictHostKeyChecking=no -o ConnectTimeout=15 root@$IP"
+SCP="sshpass -p $PW scp -o StrictHostKeyChecking=no"
+
+echo "[profile] device $IP  envs: ${ENVS:-<fp16>}"
+
+TMP=$($SSH "mktemp -d /tmp/fastwam_run.XXXXXX")
+cleanup() { $SSH "rm -rf $TMP" >/dev/null 2>&1 || true; echo "[profile] cleaned $TMP"; }
+trap cleanup EXIT
+$SSH "mkdir -p $TMP/runtime $TMP/tmp && for d in ctx golden qnn_libs dsp_libs; do ln -sfn $PERSIST/\$d $TMP/\$d; done"
+
+$SCP "$HERE"/resident_run.py "$HERE"/resident_worker.py root@"$IP":"$TMP"/runtime/ >/dev/null
+
+$SSH "cd $TMP && FASTWAM_WS=$TMP $ENVS python3 -u runtime/resident_run.py"
+
+echo "[profile] pure ctx.Inference times ([infer-ms], excl. TCP):"
+$SSH "grep -h infer-ms $TMP/tmp/w_*.log 2>/dev/null | awk '{a[\$2]+=\$3;n[\$2]++} END{for(k in a) printf \"  %-12s %.1f ms\n\", k, a[k]/n[k]}' | sort" || true

runtime/resident_run.py

@@ -0,0 +1,176 @@
+#!/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())

runtime/resident_worker.py

@@ -0,0 +1,102 @@
+#!/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()