Model Demo: PI0.5 End-to-End VLA Inference

NOTE: The Jupyter Notebook below is included in the Chimera SDK and can be run interactively by running the following CLI command:

$ quadric sdk notebook

From the Jupyter Notebook window in your browser, select the notebook named /quadric/sdk-cli/examples/models/pi0_5/pi05_e2e.ipynb.

PI0.5 VLA — fused vision→language→action on Chimera GPNPU

PI0.5 is a vision-language-action (VLA) model for general-purpose robot control. It conditions on camera images plus a natural-language instruction and produces a chunk of denoised joint commands via Euler refinement of noisy actions.

This notebook runs the entire PI0.5 INT8 pipeline as a single CCL kernel, with the vision tower running live on the GPNPU. It runs the SigLIP vision tower through chipy.infer_onnx, bridges its image embeddings with the baked language tail, and flows straight into LM prefill + Action Expert — vision → language → action, end to end, in one binary.

Note: we do not run the full PI0.5 VLA model here. We run 1) a smoke test involving 1 SigLIP transformer/FFN block + 1 prefill block + 1 action expert block and 2) a more representative test involving a larger subset of the graph (i.e. 5 SigLIP blocks, 2 prefill blocks, 10 action expert blocks).

Model stages

• SigLIP vision tower (live) — pixel_values[3,3,224,224] (one LIBERO frame per camera) compiled by CGC via chipy.infer_onnx(siglip_vision_b3_int8.onnx) → image_embeddings[3,256,2048].
• Vision→prefill bridge — the vision_prefill_bridge CCL custom op concatenates the 768 image tokens with the 200-row baked language tail into the [1, 968, 2048] prefill sequence.
• LM Decoder (prefill) — 18-layer Gemma2 with causal attention; writes K/V to a shared flat DDR cache via the sdpaPersistOffset CCL custom op.
• Action Expert (denoising) — 18-layer expert with bidirectional cross-attention reading the prefix K/V via the sdpaCrossFromCacheOffset CCL custom op; AdaRMSNorm carries the time embedding.
• Single CCL kernel — vision, prefill, and denoise all compile into one binary; no host-CPU trampoline anywhere in the chain. The only runtime input is the camera frame.

Model specs

1. Setup

Make the pi0_5 package importable when running this notebook from its own directory, then silence verbose CGC compilation logs.

## pi0.5 needs transformers==4.55 (tvm-side modeling_gemma uses imports added in 4.55).
## Install into a notebook-local directory so the kernel picks it up without touching
## the global transformers pin.
import os
import sys

!{sys.executable} -m pip install --no-deps --target=pi05_t455 --quiet transformers==4.55.0
sys.path.insert(0, os.path.abspath("pi05_t455"))

WARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv


import logging
import pathlib
import sys

_pkg_parent = str(pathlib.Path.cwd().parent.resolve())
if _pkg_parent not in sys.path:
    sys.path.insert(0, _pkg_parent)

from pi0_5.common import MODEL_PATH, TEST_VECTORS_PATH

logging.getLogger("epu").setLevel(logging.ERROR)
logging.getLogger("epu").handlers = [logging.NullHandler()]

print(f"Model:        {MODEL_PATH}")
print(f"Test vectors: {TEST_VECTORS_PATH}")

Model:        /quadric/sdk-cli/examples/models/pi0_5/checkpoint/full_model.pth
Test vectors: /quadric/sdk-cli/examples/models/pi0_5/checkpoint/test_vectors.safetensors

2. Download Model

The quantized LM weights, test vectors, and per-layer quant params live on S3 (too large for git) — run the next cell once to populate the local checkpoint/ directory. The SigLIP vision tower (siglip_vision_b3_int8.onnx + its .tranges, plus the 1-block _smoke pair) also lives on S3 (under the pi_0_5/siglip/ prefix) and is fetched automatically by vision_prefill_pipeline.py the first time a kernel is built — no manual step needed.

from pathlib import Path
from urllib.request import urlretrieve

CKPT_DIR = Path("checkpoint")
CKPT_DIR.mkdir(exist_ok=True)
S3_BASE = "https://sdk-cli-models.s3.us-east-2.amazonaws.com/pi_0_5"

## Checkpoint assets (LM weights, test vectors, quant params) -> checkpoint/.
## The SigLIP vision tower (siglip_vision_b3_int8.onnx + .tranges) is also on S3
## (pi_0_5/siglip/) and is fetched automatically by vision_prefill_pipeline.py on first use.
for name in ["full_model.pth", "test_vectors.safetensors", "quant_params.pth"]:
    local = CKPT_DIR / name
    if local.exists():
        print(f"  {name} ({local.stat().st_size / 1e6:.0f} MB) — cached")
    else:
        print(f"  Downloading {name}...")
        urlretrieve(f"{S3_BASE}/{name}", str(local))
        print(f"  {name} ({local.stat().st_size / 1e6:.0f} MB) — done")

  Downloading full_model.pth...
  full_model.pth (5463 MB) — done
  Downloading test_vectors.safetensors...
  test_vectors.safetensors (6004 MB) — done
  quant_params.pth (3 MB) — cached

3. Compile & Run

Build a single CCL kernel combining the live SigLIP vision tower, LM prefill, and Action Expert, then compile it with CGC and run it on the ISS.

validate_fused_kernel_ccl handles the full flow in one call: it loads one LIBERO frame per real camera as pixel_values, builds the fused @chipy.func() (infer_onnx(SigLIP) → vision_prefill_bridge → calibrated prefill + expert chain), runs one CPU forward pass through the ChiPy chain to harvest frac_bits for every FixedPoint32 intermediate, compiles the kernel via CGC, runs it on the ISS, and reports cosine + match% for both the prefill (informational) and the terminal expert (the gated result) against the ChiPy CPU reference, plus a cycle profile.

Note on accuracy. FixedPoint32 intermediates compound error across layers. The pass/fail gate is on the terminal (Action Expert) output only — cosine ≥ 0.98 OR match% ≥ 98% vs the ChiPy CPU reference. The prefill cosine/match are printed for information (a non-terminal intermediate, not gated — its cosine is expected to run lower than the terminal's).

Cost. The smoke test (§3.1) is fast (1-block tower + 1P/1E); the representative run (§3.2) is the heavier one (5-block tower + 2P/10E — compile + ISS, several GB of host RAM). The patch_large_consts.sh post-pass auto-triggers for the >2 GB const blob at seq=968.

3.1 Smoke test (1-block SigLIP + 1P + 1E)

A fast end-to-end sanity check: a 1-block cut of the SigLIP tower feeds 1 prefill + 1 expert LM layer. Small enough to compile + run on the ISS quickly while still exercising the entire vision→bridge→language→action path.

• SigLIP (1-block cut) on pixel_values[3,3,224,224] (LIBERO frame on cameras 0/1, camera 2 zero-padded) → image_embeddings[3,256,2048] → in-ONNX bridge → hidden_prefill[1,968,2048].
• 1 LM prefill layer (KV write) + 1 expert layer (KV read) at seq=968.

The cosine is self-consistent: EPU-ISS vs the ChiPy CPU reference of the same reduced graph.

from pi0_5.siglip.vision_prefill_pipeline import validate_fused_kernel_ccl, SIGLIP_B3_SMOKE_ONNX

passed = validate_fused_kernel_ccl(
    num_prefill_layers=1,
    num_expert_layers=1,
    prefill_seq_len=968,
    siglip_onnx=SIGLIP_B3_SMOKE_ONNX,
)
assert passed

============================================================
ISS Validation: Fused VLA Kernel (1P + 1E, seq=968)
============================================================
Building kernel: 1P + 1E (prefill_seq=968)...


WARNING:root:Detected 1 tied-parameter group(s); quantization will break the tie (each site quantized independently):
  - paligemma_with_expert.paligemma.model.language_model.embed_tokens.weight is paligemma_with_expert.paligemma.lm_head.weight


Calibration: armed 30 ops, stats observed during chipy CPU forward
  Prefill [informational, non-terminal]: cosine=0.953321, max_diff=152.4329, mean_diff=3.5932, match=19.60%
  Expert [TERMINAL, gated]: cosine=0.990761, max_diff=4.1287, mean_diff=0.1013, match=98.01%  [PASS]

PASS

ISS profile: 46,698,696 (46.70M) cycles
  top stalls: STORE=9,237,273, MEU=8,194,231, LOAD=7,916,303

3.2 Representative run (5-block SigLIP + 2P + 10E)

The representative validation we run here is a subset of the full model to avoid the multi-hour compilation/ISS run: a 5-block cut of the SigLIP tower feeding 2 prefill + 10 expert LM layers at production sequence length (seq=968).

• SigLIP (5-block cut) → image_embeddings[3,256,2048] → in-ONNX bridge → hidden_prefill[1,968,2048].
• 2 prefill layers (KV write to the shared flat DDR buffer) + 10 expert layers (cross-attention reading that prefix KV cache).

The terminal Expert: cosine=… is the gated result; Prefill: cosine=… is printed for information (non-terminal, not gated — expected to run lower than the terminal).

from pi0_5.siglip.vision_prefill_pipeline import SIGLIP_B3_ONNX

passed = validate_fused_kernel_ccl(
    num_prefill_layers=2, num_expert_layers=10, prefill_seq_len=968, siglip_onnx=SIGLIP_B3_ONNX
)
assert passed

============================================================
ISS Validation: Fused VLA Kernel (2P + 10E, seq=968)
============================================================
Building kernel: 2P + 10E (prefill_seq=968)...
Calibration: armed 180 ops, stats observed during chipy CPU forward
  Prefill [informational, non-terminal]: cosine=0.861109, max_diff=562.0262, mean_diff=7.5279, match=10.44%
  Expert [TERMINAL, gated]: cosine=0.988454, max_diff=16.6106, mean_diff=0.8360, match=73.51%  [PASS]

PASS

ISS profile: 112,410,536 (112.41M) cycles
  top stalls: LOAD=24,546,069, STORE=19,238,634, MEU=17,982,735

Summary

Key takeaways

1. One fused binary, vision included. The SigLIP vision tower, LM prefill, and Action Expert compile into a single CCL binary. The vision tower runs live via chipy.infer_onnx — there is no pre-baked hidden_prefill and no host-CPU trampoline anywhere in the chain.
2. Vision is the only runtime input. The language block is a build-time baked constant (the language tail of the merged test-vector sequence), baked exactly like the LM weights. The only tensor that varies at runtime is the camera frame (pixel_values).
3. Bridged in-graph. The vision_prefill_bridge CCL custom op concatenates the 768 SigLIP image tokens with the 200-row baked language tail into the [1, 968, 2048] prefill sequence — the same shape the existing prefill path consumes — so vision flows into language flows into action with no glue code on the host.

Citation

@article{intelligence2025pi05,
  title   = {{$\pi_{0.5}$}: a Vision-Language-Action Model with Open-World Generalization},
  author  = {{Physical Intelligence} and Black, Kevin and Brown, Noah and Driess, Danny and Esmail, Adnan and Equi, Michael and Finn, Chelsea and Fusai, Niccolo and Galliker, Manuel Y. and Ghosh, Dibya and Groom, Lachy and Hausman, Karol and Ichter, Brian and Jakubczak, Szymon and Jones, Tim and Ke, Liyiming and LeBlanc, Devin and Levine, Sergey and Li-Bell, Adrian and Mothukuri, Mohith and Nair, Suraj and Pertsch, Karl and Ren, Allen Z. and Shi, Lucy Xiaoyang and Smith, Laura and Springenberg, Jost Tobias and Stachowicz, Kyle and Tanner, James and Vuong, Quan and Walke, Homer and Walling, Anna and Wang, Haohuan and Yu, Lili and Zhilinsky, Ury},
  journal = {arXiv preprint arXiv:2504.16054},
  year    = {2025}
}