xpuenabler/molmoact2-drift-keystone-libero-qnn

MolmoAct2-LIBERO Drift-Keystone — QNN / HTP context binaries

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

The source policy is allenai/MolmoAct2-LIBERO retrained with:

technique	change	conversion effect
Drift	one-step drifting loss	action expert Euler decode 10-step → 1-step (dt=1.0)
KeyStone	inference-time geometric best-of-K selection	host-side option — the per-step graph is unchanged (run action_step K times, select on host)

Architecture / checkpoint format are identical to the original MolmoAct2 (no grid sampler, 25-layer ViT, prompt seq ~488, +0 params). This bundle contains the 7 prebuilt .bin contexts, the ONNX graphs they were built from, and the golden reference tensors. Conversion code: nota-github/xpu-molmoact2-qnn-htp (scripts_drift_keystone/, branch support-drift-keystone, PR #5).

---

What's in here

ctx/      7 HTP context binaries (soc_id 77 / Hexagon v73), fp16 weights
  vision_socid77_archv73.bin           926 MB   (ViT-25L + pooling)
  llm_split0_socid77_archv73.bin       1825 MB
  llm_split1_socid77_archv73.bin       1825 MB
  llm_split2_socid77_archv73.bin       1825 MB
  llm_split3_socid77_archv73.bin       1825 MB
  action_context_socid77_archv73.bin   4 MB
  action_step_socid77_archv73.bin      1164 MB  (1-step flow-matching, Drift)
golden/   reference I/O (PyTorch fp32), the parity gate for every context
  vision_io.npz  llm_split{0..3}_io.npz  action_context_io.npz  action_step_io.npz  trace.json
onnx/     per-component ONNX graphs (opset 20) that feed the DLC/context build; large weights
          stored as external data next to each .onnx
  vision/  llm_split0..3/  action_context/  action_step/

---

Architecture — 3 components

            image ─┐  ViT-25L + pooling      ┌──────────────────┐  36-layer KV  ┌──────────────────┐
 text ─────────────┼──────────────────────▶  │   LLM (Qwen-ish) │ ────────────▶ │  Action Expert   │─▶ actions
            state ─┘                          │  single prefill  │  (as context) │ flow-matching    │  [1,10,7]
                                              └──────────────────┘               └──────────────────┘
                                            1 pass, KV = OUTPUT only        context ×1 + Euler step ×1 (Drift)

• Drift: the flow-matching action expert runs 1 host-looped Euler step (action_step), not
  10. (num_inference_steps=1.)
• KeyStone (optional): run action_step with K different noise seeds and select the medoid on the host (geometric best-of-K). The context binary is identical; only the host loop count changes.
• The 3.5B LLM is split into 4 layer-range contexts (9 layers each) for the HTP per-session limit.

Per-component I/O contract

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

Host pipeline: vision (NPU) → scatter image embeds into the prompt scaffold → llm_split0..3 (chain, collect 36-layer KV) → action_context (NPU) → 1-step Euler with action_step → de-normalize. (KeyStone: repeat the action step K× and select.)

---

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 KV / vision embeds span ~1e2–1e4) and is informational.

context	output	cosine
vision	image_embeds	0.99999361
llm_split0	hidden_out / keys / values	1.0 / 0.99999999 / 0.99999998
llm_split1	hidden_out / keys / values	1.0 / 1.0 / 1.0
llm_split2	hidden_out / keys / values	1.0 / 1.0 / 1.0
llm_split3	hidden_out / keys / values	1.0 / 1.0 / 1.0
action_context	k_ctx / v_ctx	1.0 / 0.99999998
action_step	next_trajectory	1.0

ONNX-vs-PyTorch parity (ORT CPU) was cosine 1.0 for every component. On-device (HTP fp16) end-to-end action cosine and latency are the runtime deliverable.

---

Build provenance

• Source policy: xpuenabler/MolmoAct2-LIBERO-Drift-Keystone (LeRobot checkpoint; base allenai/MolmoAct2-LIBERO).
• Loaded via the nota-github/xpu-lerobot @ feat/molmoact2-drift-keystone fork (parses drift/keystone config); one real LIBERO inference captured golden I/O at the four boundaries.
• Drift → action_step host-looped 1 step; KeyStone is host-side best-of-K (graph unchanged).
• Toolchain: QAIRT 2.47.0, opset 20, float DLC (no quant), offline HTP context-binary 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). See docs/drift_keystone_qnn.md in the conversion repo.