xpuenabler/Fast-WAM-libero-qnn

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, 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 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.

---

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 .sos) 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); 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

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)

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

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.
• 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).