File size: 8,301 Bytes

b324599

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