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