Single-GPU inference for PI0

How PhyAI runs pi0 inference on a single GPU

Overview

pi0 turns camera views, task text, and robot state into a chunk of future actions. Its ws1 path keeps the entire inference path on one GPU. The vision and language stacks build a reusable prefix; conditioned on that prefix and the robot state, the action expert turns noise into an action chunk through flow matching.

PI0WS1Scheduler implements only this single-GPU path. It has no tensor parallelism, continuous batching, or preemption.

Robot state enters pi0 as a numeric token on the expert side. In pi0.5, discretized state bins become part of the language prompt instead.

Architecture

The request enters PhyAI through its engine + plugin contract. PI0Entry builds the model and scheduler, while three runners own the vision, prefix, and expert work. Their implementation lives in these files:

phyai/src/phyai/models/pi0

main_pi0.py

scheduler_ws1_pi0.py

model_runner_pi0.py

modeling_pi0.py

configuration_pi0.py

Component	Responsibility
PI0Entry	Registers the "pi0" engine plugin, builds PI0Model, loads weights, and creates the scheduler
PI0Config	Stores vision, text, expert, action chunk, tokenizer, and camera-count geometry
PI0Model	Owns the SigLIP/PaliGemma vision tower, PaliGemma text stack, Gemma expert stack, RoPE, and action/time heads
PI0VisionRunner	Runs the vision tower, with optional CUDA graph capture
PI0LLMRunner	Runs the PaliGemma prefix pass and writes prefix K/V into the shared cache
PI0ExpertRunner	Runs the expert state/action passes for each flow-matching step
PI0WS1Scheduler	Orchestrates one complete inference request on a single GPU

Model layout

The code spans several modules, but the data path stays linear. SigLIP turns camera images into tokens, PaliGemma combines them with the task text, and the smaller Gemma expert produces the action trajectory.

Stack	Default shape	Notes
Vision	SigLIP, 27 layers, 224×224 images, 14×14 patches	Produces image tokens projected into the PaliGemma text width
Text	PaliGemma/Gemma, 18 layers, hidden size 2048	Processes image + language prefix and writes prefix K/V
Expert	Gemma action expert, 18 layers, hidden size 1024	Processes one state token plus the full action chunk

These stacks share a small set of top-level geometry settings:

Field	Default	Meaning
chunk_size	50	Number of action tokens returned per engine step
max_state_dim	32	Padded robot-state width
max_action_dim	32	Padded action width
num_inference_steps	10	Flow-matching Euler steps
tokenizer_max_length	48	Right-padded PaliGemma task prompt length
empty_cameras	0	num_images = 3 - empty_cameras; pi0 supports 2 or 3 cameras

params_dtype controls the language and expert stacks. The vision tower keeps its own vision_params_dtype and defaults to fp32 to match the reference implementation. Pass PI0Args(vision_params_dtype=torch.bfloat16) only when bf16 vision is part of the experiment.

Request contract

At runtime, PI0Request carries the tensors that connect preprocessing to the scheduler:

Field	Shape	Notes
pixel_values	(B, num_images, 3, image_size, image_size)	Already resized and normalized camera tensors
input_ids	(B, tokenizer_max_length) int64	Right-padded PaliGemma token ids
lang_lens	(B,) int64	Real task-prompt length for each sample
state	(B, max_state_dim)	Numeric robot state, padded before the expert
noise	(B, chunk_size, max_action_dim) or None	Optional initial action noise; when None, the scheduler samples Gaussian noise

B may range from 1 to max_batch_size. Smaller batches are padded to the captured shape inside the scheduler, then sliced back to actual_B before the result is returned.

Scheduler phases

Once the request enters engine.step, the scheduler moves through these phases in order:

Phase	Work
pi0.vision_loop	Move camera tensors to the vision dtype and run PI0VisionRunner once per real batch item
pi0.lang_pack	Embed language ids, then pack image tokens and language tokens into the per-sample prefix buffer
pi0.llm_prefix_plan	Reset static caches and prepare ragged prefix attention metadata
pi0.llm_prefix_fwd	Run the PaliGemma text stack and write prefix K/V into KVCachePool
pi0.expert_plan	Prepare state and action expert attention metadata over prefix + suffix slots
pi0.expert_state_prefill	Embed the robot state and cache its per-layer K/V once
pi0.expert_loop	Initialize or copy action noise and run flow-matching integration
pi0.expert_step	One expert velocity prediction and Euler update inside pi0.expert_loop

The prefix remains unchanged during denoising, so the scheduler caches it once. State and action queries then see different portions of the cache:

state query  -> prefix + state
action query -> prefix + state + action chunk

The suffix therefore contains 1 + chunk_size tokens: one state token followed by the action chunk.

CUDA graphs

The fixed request geometry makes CUDA graph capture possible. With RuntimeConfig(use_cuda_graph=True), each runner captures its graph during scheduler.setup():

Runner	Captured shape
PI0VisionRunner	(num_images, 3, image_size, image_size)
PI0LLMRunner	(max_batch_size * n_per_sample, text_hidden_size)
PI0ExpertRunner	state, x_t, and time buffers at fixed max_batch_size

For each request, the runners refresh their static input buffers and replay the captured graphs. Attention metadata stays outside the captured region and is written through the backend’s capture-aware buffers.

Disable CUDA graphs when you need an expanded Nsight Systems trace:

uv run python benchmark/bench_n_batch_ws1_pi0.py \
    --batch-sizes 4 \
    --no-cuda-graph

Running pi0

1. Prepare weights

   Use a HF-style pi0 PyTorch checkpoint directory containing config.json and model.safetensors. Omit --checkpoint only for a random-weight smoke test.

2. Construct the engine

   Select the "pi0" plugin. Engine construction loads the optional checkpoint, prepares the runners, and captures CUDA graphs when enabled.

   import torch
   from pathlib import Path
   
   from phyai.engine import Engine, EngineArgs
   from phyai.engine_config import DeviceConfig, EngineConfig, RuntimeConfig
   from phyai.models.pi0.main_pi0 import PI0Args
   
   engine = Engine(
       EngineArgs(
           plugin="pi0",
           plugin_args=PI0Args(
               checkpoint_dir=Path("/path/to/pi0_pytorch"),
               max_batch_size=4,
               vision_params_dtype=torch.float32,
           ),
           config=EngineConfig(
               device=DeviceConfig(target="cuda", params_dtype=torch.bfloat16),
               runtime=RuntimeConfig(use_cuda_graph=True),
           ),
       )
   )
   

   max_batch_size becomes part of the captured shape. Changing it requires a new engine.

3. Build a request

   PI0Processor, provided by phyai-utils-tools, converts raw robot observations into the tensors expected by PI0Request.

   from phyai.models.pi0.scheduler_ws1_pi0 import PI0Request
   from phyai_utils_tools.models.pi0 import PI0Processor
   
   processor = PI0Processor(
       image_size=224,
       num_channels=3,
       num_images=3,
       tokenizer_max_length=48,
       max_state_dim=32,
       action_dim=7,
       device="cuda",
       params_dtype=torch.bfloat16,
   )
   
   processed = processor.preprocess(
       {
           "images": [cam0, cam1, cam2],
           "task": ["pick up the object"],
           "state": state,
       }
   )
   
   request = PI0Request(
       pixel_values=processed.pixel_values,
       input_ids=processed.input_ids,
       lang_lens=processed.lang_lens,
       state=processed.state,
   )
   
4. Run one step

   actions = engine.step(request)  # (B, chunk_size, max_action_dim)
   

   When the processor has an action_dim, processor.postprocess(actions) removes the padded action width and applies dataset unnormalization when statistics are available.

5. Close the engine

   engine.close()
   

End-to-end example

examples/pi0/run_pi0.py can drive either a processed request or the raw smoke path. A checkpoint-backed run looks like this:

uv run python examples/pi0/run_pi0.py \
    --checkpoint /path/to/pi0_pytorch \
    --batch-size 1

To check only the execution path, omit the checkpoint and use random weights:

uv run python examples/pi0/run_pi0.py --raw --batch-size 1

A checkpoint with one empty camera expects two image streams:

uv run python examples/pi0/run_pi0.py \
    --checkpoint /path/to/pi0_pytorch \
    --num-images 2

Benchmarking and profiling

Use benchmark/bench_n_batch_ws1_pi0.py to compare batch sizes under the same scheduler configuration:

uv run python benchmark/bench_n_batch_ws1_pi0.py \
    --batch-sizes 1 2 4 \
    --n-warmup 5 \
    --n-timed 30 \
    --result-file ./pi0_ws1_results.jsonl

The same script can bracket a short Nsight Systems capture:

nsys profile \
    --capture-range=cudaProfilerApi \
    --capture-range-end=stop \
    -o ./prof/pi0_ws1 \
    uv run python benchmark/bench_n_batch_ws1_pi0.py \
        --batch-sizes 4 \
        --profile-backend nsys \
        --profile-start-step 5 \
        --profile-num-steps 3

The default leaves the vision tower in fp32. Add --vision-dtype bfloat16 only when the experiment calls for bf16 vision timing.

Current limitations

Graph reuse depends on fixed geometry. PI0WS1Scheduler fixes max_batch_size when the single-GPU engine is built, while CUDA graph capture also fixes the camera count, image size, and tokenizer length. Changing any of them requires a new engine. The vision tower replays once for each real batch item rather than batching those items together.

The scheduler begins after preprocessing. Image resize, tokenization, state padding, and action unnormalization belong to PI0Processor; requests passed directly to the engine must already satisfy that contract.

Full example

from pathlib import Path

import torch

from phyai.engine import Engine, EngineArgs
from phyai.engine_config import DeviceConfig, EngineConfig, RuntimeConfig
from phyai.models.pi0.configuration_pi0 import PI0Config
from phyai.models.pi0.main_pi0 import PI0Args
from phyai.models.pi0.scheduler_ws1_pi0 import PI0Request
from phyai.utils import load_config

CHECKPOINT_DIR = Path("/path/to/pi0_pytorch")
BATCH_SIZE = 1

cfg = load_config(CHECKPOINT_DIR, PI0Config)
device = torch.device("cuda")
dtype = torch.bfloat16

engine = Engine(
    EngineArgs(
        plugin="pi0",
        plugin_args=PI0Args(
            checkpoint_dir=CHECKPOINT_DIR,
            max_batch_size=BATCH_SIZE,
            vision_params_dtype=torch.float32,
        ),
        config=EngineConfig(
            device=DeviceConfig(target="cuda", params_dtype=dtype),
            runtime=RuntimeConfig(use_cuda_graph=True),
        ),
    )
)

try:
    input_ids = torch.zeros(
        BATCH_SIZE, cfg.tokenizer_max_length, dtype=torch.int64, device=device
    )
    input_ids[:, 0] = 2

    request = PI0Request(
        pixel_values=torch.rand(
            BATCH_SIZE,
            cfg.num_images,
            cfg.vision.num_channels,
            cfg.vision.image_size,
            cfg.vision.image_size,
            dtype=torch.float32,
            device=device,
        ),
        input_ids=input_ids,
        lang_lens=torch.ones(BATCH_SIZE, dtype=torch.int64, device=device),
        state=torch.rand(BATCH_SIZE, cfg.max_state_dim, dtype=dtype, device=device),
    )

    actions = engine.step(request)
    print(f"action chunk shape={tuple(actions.shape)}, dtype={actions.dtype}")
finally:
    engine.close()