xpuenabler/molmoact2-libero_grid_sampler_qnn

MolmoAct2-LIBERO + Grid Sampler — QNN / HTP context binaries

Offline-compiled Qualcomm QNN (HTP) context binaries for xpuenabler/molmoact2-libero_grid_sampler_fine_tuned, ready to run resident on-device on a Qualcomm Dragonwing IQ‑9075 (QCS9075, Hexagon v73, soc_id 77).

This is the Step 3 (runtime) deployment bundle: 7 prebuilt .bin contexts, the golden reference tensors used as inputs / parity references, and the resident multi-process runtime. The conversion pipeline (PyTorch → ONNX → QNN DLC → HTP context binary) lives in the xpu-molmoact2-qnn-htp repo.

Grid Sampler variant. Each camera contributes 16 visual tokens instead of 196 (pruned by the ActiveTokenSampler / F.grid_sample in the vision backbone), so the LIBERO 2‑camera prompt is 128 tokens (the original allenai/MolmoAct2-LIBERO checkpoint was 488). Everything downstream — the action cross-attention context, the runtime reshapes — uses 128. See Grid Sampler & seq=128 before running.

---

What's in here

ctx/      7 HTP context binaries (soc_id 77 / Hexagon v73), fp16 weights
  vision_socid77_archv73.bin           925 MB
  llm_split0_socid77_archv73.bin       1.8 GB
  llm_split1_socid77_archv73.bin       1.8 GB
  llm_split2_socid77_archv73.bin       1.8 GB
  llm_split3_socid77_archv73.bin       1.8 GB
  action_context_socid77_archv73.bin   3.6 MB
  action_step_socid77_archv73.bin      1.2 GB
golden/   reference I/O (PyTorch fp32), consumed by the runtime as inputs + the parity gate
  vision_io.npz  llm_split{0..3}_io.npz  action_context_io.npz  action_step_io.npz
  step0_boundary.npz   trace.json   norm_stats.json
runtime/  resident multi-process runtime (Step 3)
  resident_run.py      orchestrator (host glue: scatter, split chaining, denorm, latency, parity)
  resident_worker.py   one NPU session per process (Grid-Sampler PROMPT_SEQ=128)
  profile_device.sh    one-command host→device profiler (SSH; reads device creds from IQ9_info.txt)

You also need, from the QAIRT 2.47.x SDK (not redistributed here — see Device setup): qnn_libs/ (aarch64 QNN runtime .sos) and dsp_libs/ (Hexagon v73 skel libs).

---

Architecture — 3 components, 7 co-resident NPU sessions

            image ─┐
                   ▼
        ┌──────────────────┐  16 tok/cam  ┌──────────────────┐  36-layer KV  ┌──────────────────┐
 text ─▶│  Vision backbone │ ───────────▶ │   LLM (Qwen-ish) │ ────────────▶ │  Action Expert   │─▶ actions
        │ ViT+GridSampler  │  host scatter│  single prefill  │  (as context) │ flow-matching    │  [1,10,7]
        └──────────────────┘              └──────────────────┘               └──────────────────┘
            once / image                1 pass, KV = OUTPUT only        context ×1 + Euler step ×10

• The LLM runs once (single prefix-encode, no autoregression). Its 36-layer KV cache is an output, consumed by the Action Expert as cross-attention context.
• The 3.5B LLM exceeds the HTP per-session weight limit, so it is split into 4 layer-range contexts (9 layers each).
• The Action Expert is flow-matching diffusion: action_context (project KV → cross-attn context, once) + action_step (one Euler denoise step, looped 10× on host).
• Resident set = vision(1) + llm_split0..3(4) + action_context(1) + action_step(1) = 7 sessions, exactly the device's ~7 concurrent-session cap. The action worker process holds two contexts (action_context + action_step) and projects KV→context on the NPU in-process.

Per-component I/O contract (Grid-Sampler, seq = 128)

context	inputs (name: shape, f32)	outputs (name: shape, f32)
vision	pixel_values: [1,2,729,588]	image_embeds: [32,2560] (16 tok × 2 cams)
llm_split{i}	hidden_states: [1,128,2560], attention_bias: [1,1,128,128]	hidden_out: [1,128,2560], keys: [9,1,128,1024], values: [9,1,128,1024]
action_context	keys: [36,1,128,1024], values: [36,1,128,1024]	k_ctx: [36,1,128,8,96], v_ctx: [36,1,128,8,96]
action_step	trajectory: [1,10,32], timestep: [1], k_ctx: [36,1,128,8,96], v_ctx: [36,1,128,8,96], encoder_attention_mask: [1,128], action_dim_is_pad: [1,32]	next_trajectory: [1,10,32]

The host pipeline: vision (NPU) → scatter the 32 image embeds into the prompt embedding scaffold → llm_split0..3 (chain, collect 36-layer KV) → action worker (NPU action_context → 10-step Euler with action_step) → de-normalize [...,:7] with norm_stats.json quantiles.

---

Inference guide

1. Device workspace

On the IQ‑9075, create /root/molmoact2_workspace/ with:

/root/molmoact2_workspace/
  ctx/        ← the 7 .bin files from ctx/ in this repo
  golden/     ← the golden/ folder from this repo (incl. norm_stats.json)
  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

Get qnn_libs/ and dsp_libs/ from the QAIRT 2.47.x SDK matching the build SoC/arch (soc_id 77, v73). 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.

Device gotchas (from prior bring-up): reboot before a clean resident run (crashed cycles leak DSP sessions); the workers end with os._exit() to dodge an appbuilder↔libs teardown double-free.

2a. Run directly on the device

cd /root/molmoact2_workspace
python3 runtime/resident_run.py

It spawns the 7 workers, keeps them resident (no per-inference reload), runs the full pipeline, prints the final-action cosine vs the golden reference, then a per-component latency table (avg of 5).

2b. Or profile from the host (one command)

runtime/profile_device.sh copies the runtime scripts to an ephemeral /tmp workspace on the device, symlinks the heavy assets (ctx/ golden/ qnn_libs/ dsp_libs/) from the persistent workspace, runs, prints the latency table + pure [infer-ms], then deletes the temp dir (nothing persistent is modified). Device creds are read from IQ9_info.txt (IP: ... / passwd: ...; override path with IQ9_INFO=...).

runtime/profile_device.sh

fp16 / quantized A/B switch (per component)

The runtime selects a context per component by env var (default "" = fp16). Only fp16 bins are shipped here; if you add quantized bins named *_w4a16_socid77_archv73.bin, the same runtime loads them:

LLM_CTX_SUFFIX=_w4a16    python3 runtime/resident_run.py   # quantize LLM only
VISION_CTX_SUFFIX=_w4a16 python3 runtime/resident_run.py   # vision only
ACTION_CTX_SUFFIX=_w4a16 python3 runtime/resident_run.py   # action_step only

---

Grid Sampler & why seq=128

The vision backbone prunes each camera's pooled feature grid to 16 tokens via a trained ActiveTokenSampler (predict 16 sample coordinates with an MLP → F.grid_sample bilinear → add a coordinate embedding). It is fully static (fixed 16 tokens/image, no data-dependent shapes), exports to ONNX GridSample (opset 20), and compiles cleanly to HTP v73 — no graph surgery needed.

Consequence: the prompt is 128 tokens (16×2 cameras + text + state), not 488. The runtime already reflects this:

• golden/*.npz were captured at seq 128.
• resident_worker.py reshapes the action context with PROMPT_SEQ (default 128); override with MOLMOACT2_PROMPT_SEQ=... if you ever rebuild at a different token budget.

If you start from the original (non-Grid-Sampler) runtime, the one required change is the action_context reshape [36,1,488,8,96] → [36,1,128,8,96] — already applied here.

---

Validated parity (host, QNN CPU backend vs PyTorch fp32)

Each context was verified device-free on the QNN CPU backend against the PyTorch golden. Cosine is the correctness gate (≥ 0.9999); within_tol is a known massive-activation fp artifact (LLM hidden_states/KV span ~1e4) and is informational.

context	output	cosine
vision	image_embeds	0.99999735
llm_split0	hidden_out / keys / values	1.0 / 0.99999999 / 0.99999983
llm_split1	hidden_out / keys / values	0.99999998 / 0.99999998 / 0.99999993
llm_split2	hidden_out / keys / values	1.0 / 0.99999998 / 0.99999992
llm_split3	hidden_out / keys / values	0.99999998 / 0.99999998 / 0.99999992
action_context	k_ctx / v_ctx	1.0 / 0.99999996
action_step	next_trajectory	1.0

ONNX-vs-PyTorch parity (Step 1, ORT CPU) was equally tight (vision incl. GridSample: cosine 0.99999999997). On-device (HTP fp16) end-to-end action cosine and latency are measured by resident_run.py and are the Step-3 deliverable — not included here.

---

Build provenance

• Source policy: xpuenabler/molmoact2-libero_grid_sampler_fine_tuned (LeRobot checkpoint), base allenai/MolmoAct2-LIBERO.
• Toolchain: QAIRT 2.47.0, opset 20, float DLC (no quant), offline HTP context-binary generation for soc_id 77 / dsp_arch v73, O3, weight-sharing per multi-graph context.
• Path: PyTorch → ONNX (per component) → qairt-converter DLC → CPU-backend parity → qnn-context-binary-generator (.bin).

---

onnx/ — per-component ONNX graphs (pre-QNN)

The ONNX exports that feed the QNN/DLC context build above (PyTorch → ONNX → DLC → .bin), one folder per component, opset 20. Large weights are stored as external data next to each .onnx.

onnx/
  vision/           vision backbone (ViT + GridSampler)          vision.onnx (self-contained)
  action_context/   KV → cross-attention context projection
  action_step/      one flow-matching Euler step (+ external weights)
  llm/              full 36-layer LLM (reference; not split)
  llm_split0..3/    LLM layer-range contexts, 9 layers each — the ones actually compiled

Operator-fusion variants — parallel projections weight-concatenated into a single MatMul + Split (numerically identical to the originals; verified max_abs_diff = 0 on the QNN/ORT CPU backend):

  vision_fused/                QKV (25 ViT blocks) + projector gate_up (SwiGLU) + pooling KV fused
  action_step_fused/           MLP gate_proj + up_proj → gate_up fused (36 blocks)
  fusion_viz/                  netron before/after screenshots of the fused subgraphs
  operator_fusion_report.md    fusion-opportunity survey across all 9 modules
  operator_fusion_applied.md   what was fused + numerical-equivalence verification

model_libero/ — MolmoAct2-LIBERO model code

The transformers custom-code model definition / processors for allenai/MolmoAct2-LIBERO, used for ONNX export and reference inference. Source .py only (weights come from the base checkpoint): configuration_molmoact2.py, modeling_molmoact2.py, processing_molmoact2.py, image_processing_molmoact2.py, video_processing_molmoact2.py, inference.py.