Fast-WAM-libero-qnn / README.md
xpuenabler's picture
Fast-WAM action-expert QNN HTP context binary + resident runtime + README
b324599 verified
|
Raw
History Blame Contribute Delete
8.3 kB
---
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).