Skip to content

USAGE.md

Latest commit

 

History

History
1574 lines (1231 loc) · 60 KB

USAGE.md

File metadata and controls

1574 lines (1231 loc) · 60 KB
Raw

FlashRT Usage Guide

Complete reference for installation, API parameters, mechanisms, and usage patterns.

Installation

TL;DR — see README.md § Build & install for the full walkthrough (Docker and non-Docker paths). The short version:

git clone https://github.com/LiangSu8899/FlashRT.git
cd FlashRT
git clone --depth 1 --branch v4.4.2 \
    https://github.com/NVIDIA/cutlass.git third_party/cutlass

pip install -e ".[torch]"     # or "[jax]" / "[all]"

# Build — produces flash_rt_kernels.so AND (on RTX) flash_rt_fa2.so.
# Thor builds produce only the former (FA2 skipped, Thor uses fvk
# cuBLAS-decomposed attention).
mkdir build && cd build
cmake ..                       # auto-detects GPU arch from nvidia-smi
make -j$(nproc)
cp flash_rt_kernels*.so flash_rt_fa2*.so ../flash_rt/ 2>/dev/null || \
   cp flash_rt_kernels*.so ../flash_rt/     # Thor path
cd ..

Crucially — no pip install flash-attn required. FlashRT vendors FA2 v2.7.4.post1 (fp16 + bf16) at source level and builds it into flash_rt/flash_rt_fa2.so. Zero pip flash-attn wheel dependency at runtime.

After installation, import flash_rt works from any directory.

Pi0 / Pi0.5 prerequisite — PaliGemma tokenizer

TL;DR — required one-time download for Pi0 / Pi0.5 only. Qwen3-8B and other text-only models do NOT need this.

Pi0 and Pi0.5 prefix prompts via the PaliGemma SentencePiece tokenizer (paligemma_tokenizer.model, ~4.1 MiB). This file is not bundled inside the openpi pi0 / pi05 checkpoints — it lives in Google's public big_vision storage bucket, and you have to fetch it once before running any Pi0 / Pi0.5 inference.

One-shot download (recommended)

bash scripts/download_paligemma_tokenizer.sh
# → ~/.cache/flash_rt/paligemma_tokenizer.model

The script downloads from the public GCS HTTP endpoint (https://storage.googleapis.com/big_vision/paligemma_tokenizer.model), verifies the MD5, and is idempotent (re-runs are no-ops if the file is already valid).

Custom location

# Set this env var to point at any local copy:
export FLASH_RT_PALIGEMMA_TOKENIZER=/data/checkpoints/paligemma_tokenizer.model

The frontend's resolution order is:

  1. $FLASH_RT_PALIGEMMA_TOKENIZER (explicit override)
  2. ~/.cache/flash_rt/paligemma_tokenizer.model
  3. ~/.cache/openpi/big_vision/paligemma_tokenizer.model (compatible with openpi.shared.download.maybe_download)
  4. /workspace/paligemma_tokenizer.model (legacy container path)
  5. openpi auto-download via gs://big_vision/... (only if gcsfs is installed)

If none resolves, the loader raises a FileNotFoundError whose message contains the exact curl command above — no silent failure modes, no segfaults.

Q&A

Q: I downloaded the openpi Pi0 / Pi0.5 checkpoint. Don't I already have this? A: No. openpi ships only model.safetensors + assets/ + configs; the tokenizer is a separate fetch by design (one tokenizer file, many checkpoints).

Q: Can I use the HuggingFace google/paligemma-3b-pt-224 mirror instead? A: That repo contains the same tokenizer.model file, but PaliGemma is a gated HF model (you must accept the license + supply an HF token). The GCS public bucket is auth-free and a one-line curl — strictly easier for fresh setups. The frontend accepts a tokenizer loaded from either source as long as you point $FLASH_RT_PALIGEMMA_TOKENIZER at it.

Pi0.5 State Prompts

Pi0.5 follows openpi's discrete-state prompt format. Passing state through the stable API discretizes normalized state values with OpenPI's np.digitize(state, np.linspace(-1, 1, 257)[:-1]) - 1 bins and embeds: values in the normalized range usually become 0..255, while values below -1 become -1.

actions = model.predict(
    images=[base_img, wrist_img],
    prompt="pick up the red block",
    state=robot_state,
)

State changes update the prompt embeddings. Because the discretized state is rendered into the prompt, its token length drifts with the state values. The Pi0.5 RTX/Thor frontends offer two strategies via state_prompt_mode:

"exact" (default) — per-length capture + manual warmup. A separate pipeline is captured per exact prompt length and cached, so a recurring length is not rebuilt. Front-load the lengths you expect with warm_state_prompt_buckets() so the control loop does not pay a mid-loop capture the first time a new length appears:

model = flash_rt.load_model(
    checkpoint="/path/to/pi05", framework="torch", config="pi05")  # exact
# Pre-warm representative state lengths once at startup:
warmed = model.warm_state_prompt_buckets(
    images=[base_img, wrist_img],
    prompt="pick up the red block",
    states=[s_min, s_typical, s_max, ...])   # a few representative states
# control loop: predict() reuses the warmed graphs, no new capture

"fixed" (opt-in) — one graph, no warmup. ONE pipeline and ONE captured CUDA Graph at the max prompt length serve every length: padded prefix keys are masked with a device-side valid length and the decoder's action K/V are appended right after the valid prefix, so a changing state-token length never re-captures a graph or reruns autotune — no warmup needed. The trade-off is latency: every inference runs at the padded max length. On Thor, keeping state_prompt_fixed_max_len close to the actual maximum (for example 120 for a 117-token prompt) normally costs about 1 ms versus a warmed exact graph; larger caps pay for the extra padded tokens. That makes "fixed" the low-friction choice when the state-token length drifts: zero recapture, no length enumeration. Prefer "exact" + warm_state_prompt_buckets() only when you need absolute peak latency at a known, stable set of lengths. Enable "fixed" like so:

model = flash_rt.load_model(
    checkpoint="/path/to/pi05", framework="torch", config="pi05",
    state_prompt_mode="fixed",
    state_prompt_fixed_max_len=120)  # choose a cap >= your observed max
# or, without code changes: FLASHRT_PI05_STATE_PROMPT_MODE=fixed
# set FLASHRT_PI05_STATE_PROMPT_FIXED_MAX_LEN=120 to tune the cap in serving
# predict() handles any state length with a single captured graph

Thor updates same-length prompt embeddings in place after graph capture.

Q: I'm getting FileNotFoundError: paligemma_tokenizer.model not found. A: Run the download script above (or curl directly into one of the search paths). The error message is itself the fix — it spells out the exact command.

Q: Old setups had /workspace/paligemma_tokenizer.model — do I need to migrate? A: No. The legacy path is still in the search list.

Quick Start

import flash_rt

model = flash_rt.load_model(
    checkpoint="/path/to/checkpoint",
    framework="torch",
)

actions = model.predict(
    images=[base_img, wrist_img],
    prompt="pick up the red block",
)
# actions: numpy (10, 7) — 10 future steps, 7 DOF

API Reference

flash_rt.load_model()

model = flash_rt.load_model(
    checkpoint,                # str: path to checkpoint directory
    framework="torch",         # "torch" or "jax"
    num_views=2,               # number of camera views (2 or 3)
    autotune=3,                # CUDA Graph autotune trials
    recalibrate=False,         # force fresh FP8 calibration
    config="pi05",             # model config
)
Parameter Type Default Description
checkpoint str required Path to checkpoint directory. Safetensors for torch, Orbax for jax.
framework str "torch" "torch" (safetensors) or "jax" (Orbax). Both use the same C++/CUDA kernels.
num_views int 2 Number of camera views. LIBERO uses 2 (base + wrist).
autotune int|bool 3 CUDA Graph autotune intensity. See Autotune.
recalibrate bool False Force fresh FP8 calibration (and weight cache for JAX), ignoring cache. See Calibration.
weight_cache bool True Cache FP8-quantized weights to disk. JAX only — reduces cold start from ~42s to ~6s. Torch loads in ~3s and ignores this. See Weight Cache.
config str "pi05" Model architecture config: "pi05", "pi0", "groot", "groot_n17", "pi0fast", "motus", "wan22_ti2v_5b", "cosmos3_video". "cosmos3_video" is a non-VLA text2video denoise model — drive it with set_prompt(ref=...) + infer(...), not predict().
decode_cuda_graph bool False Pi0-FAST only. Capture action-phase decode as CUDA Graph. Trades startup time for per-token speed. See Pi0-FAST.
decode_graph_steps int 80 Pi0-FAST only. Number of action tokens to capture in the decode graph. Should cover your longest expected action sequence.
use_fp4 bool False Pi0.5 torch + jax on Thor. Enable NVFP4 quantization on the encoder FFN stack. When True, resolves to the production preset (fp4_layers=tuple(range(18)) + use_awq=True + use_p1_split_gu=True). Requires SM100+ GPU. Other configs emit a warning and fall back to FP8. See NVFP4.
fp4_layers tuple[int] | None None Encoder layer indices to FP4-quantize. None → resolved by the use_fp4 preset. Passing an explicit tuple overrides the preset. Only (7,8,9) and range(18)+AWQ are task-level validated.
use_awq bool | None None Activation-aware weight quant. Required for 18-layer FP4 scope (without it, cos collapses to ~0.33). None → resolved by the use_fp4 preset.
awq_alpha float 0.5 AWQ scaling exponent s[k] = (a[k]/a.mean())^alpha.
use_p1_split_gu bool | None None P1 split-GU 2-GEMM path (parity on Pi0.5, kernel reusable for Pi0.6). None → resolved by the use_fp4 preset.
use_fp8 bool True Enable the selected frontend's FP8 path when available. Set False for frontends with a documented BF16 fallback or for explicit full-FP16 reference paths. groot_n17 on RTX/Thor is stricter: use_fp8=False alone raises, and the non-quantized reference path is use_fp16=True, use_fp8=False.
use_fp16 bool False Opt-in non-quantized full-FP16 path (A/B reference against the FP8/bf16 default). Supported: Pi0.5 torch on RTX SM120/SM89, GROOT N1.6 torch on Thor / RTX SM120, and GROOT N1.7 torch on Thor / RTX SM120/SM89. Requires use_fp8=False; other hardware/configs raise a clear error. See the RTX and Thor full-FP16 sections below.
num_steps int|None None Pi0/Pi0.5 torch frontends. Flow-matching denoise steps; None uses the frontend default.
vision_pool_factor int|None None Pi0.5 torch RTX/Orin. Spatial pooling factor for vision tokens. The FP16 RTX path currently supports only 1.
vision_num_layers int|None None Pi0.5 torch RTX/Orin. Number of SigLIP vision layers to run.
cache_frames int|None None Pi0.5 torch RTX/Orin. Temporal encoder K/V cache period; 1 means no temporal reuse.
state_prompt_mode str "exact" Pi0.5 RTX/Thor. State-in-prompt graph strategy: "exact" (exact prompt length; RTX can pre-warm recurring buckets) or "fixed" (one graph at a configured max length, no recapture). Env override FLASHRT_PI05_STATE_PROMPT_MODE. See Pi0.5 State Prompts.
state_prompt_fixed_max_len int|None None Pi0.5 Thor fixed mode. None keeps the default 200-token cap; set a lower cap when serving can bound the state-prompt length. The cap must be >= the actual token length. Env override FLASHRT_PI05_STATE_PROMPT_FIXED_MAX_LEN.

Pi0.5 state prompt bucket warmup

Pi0.5 follows the OpenPI contract where robot state is discretized into the language prefix with np.digitize(state, np.linspace(-1, 1, 257)[:-1]) - 1. Normalized in-range values usually become 0..255; values below -1 become -1.

Task: <prompt>, State: <openpi state bins>;
Action:

Because those state bins are tokenized as text, different state values can produce different prompt lengths. FlashRT keeps the original OpenPI text format for accuracy, and caches RTX Pi0.5 runtime pipelines by prompt length. If your control loop passes a changing state every step, warm the representative prompt-length buckets before entering the realtime loop:

states = [
    np.zeros(8, dtype=np.float32),
    np.linspace(-1.0, 1.0, 8, dtype=np.float32),
    rollout_state_sample,
]

model = flash_rt.load_model(
    "/path/to/pi05_libero_pytorch",
    framework="torch",
    config="pi05",
    hardware="auto",
    num_views=2,
)

warmed_lengths = model.warm_state_prompt_buckets(
    images=[base_img, wrist_img],
    prompt="pick up the red block",
    states=states,
)
print("warmed prompt lengths:", warmed_lengths)

actions = model.predict(
    images=[base_img, wrist_img],
    prompt="pick up the red block",
    state=current_state,
)

This front-loads graph capture/autotune for the token lengths reached by the supplied states. It does not zero-pad or reformat state tokens, so the model still sees the same prompt distribution as OpenPI. For best results, use several states sampled from the deployment domain or from the first few rollout frames.

Serving note: FlashRT reuses the cached pipeline for recurring prompt token lengths. A previously unseen state-token length still needs one first-time bucket build and graph capture. In production, keep state serialization stable (same state dimension, same numeric precision/rounding policy, no ad-hoc extra fields) and prewarm representative states from the deployment range. This avoids turning continuous state changes into an unbounded stream of one-off prompt lengths.

Example bucket prewarm for states that may tokenize to different lengths:

# Use realistic states from reset / mid-rollout / near-goal observations.
# These examples deliberately include values with different digit counts after
# discretization, which can land in different prompt-token buckets.
state_bucket_samples = [
    np.zeros(8, dtype=np.float32),
    np.array([0.0, 0.1, -0.1, 0.25, -0.25, 0.5, -0.5, 1.0], dtype=np.float32),
    np.array([-1.0, -0.75, -0.5, -0.25, 0.25, 0.5, 0.75, 1.0], dtype=np.float32),
    np.asarray(first_rollout_obs["state"], dtype=np.float32),
]

warmed_lengths = model.warm_state_prompt_buckets(
    images=[base_img, wrist_img],
    prompt="pick up the red block",
    states=state_bucket_samples,
)
print("prewarmed state prompt token lengths:", warmed_lengths)

RTX full-FP16 path (Pi0.5 / GROOT)

Pi0.5, GROOT N1.6, and GROOT N1.7 all default to a quantized path on RTX (FP8 backbone; bf16 DiT for GROOT N1.7). Full-FP16 is an opt-in non-quantized baseline — useful to A/B against the quantized path. Enable it by explicitly disabling FP8 and enabling FP16. The FP16 routes are additive and leave the default quantized paths untouched.

Pi0.5

The default Pi0.5 RTX path remains FP8/BF16. To run full-FP16:

import flash_rt

model = flash_rt.load_model(
    "/path/to/pi05_libero_pytorch",
    framework="torch",
    config="pi05",
    hardware="auto",  # or "rtx_sm120" / "rtx_sm89"
    num_views=3,
    num_steps=10,
    cache_frames=1,
    use_fp8=False,
    use_fp16=True,
)

actions = model.predict(
    images=[base_img, wrist_img, wrist_right_img],
    prompt="pick up the red block and place it in the tray",
)

Benchmark helper:

python examples/blackwell/bench_pi05_fp16.py \
  --checkpoint /path/to/pi05_libero_pytorch \
  --num-views 3 \
  --steps 10 \
  --warmup 10 \
  --iters 100

This path is intentionally isolated from Thor, Orin, and the default RTX FP8 path. Passing use_fp16=True without use_fp8=False, or on unsupported hardware/configs, raises a ValueError.

Officially enabled FP16 RTX routes are:

Hardware string Example GPUs Status
rtx_sm120 RTX 5090 validated reference path
rtx_sm89 RTX 4090, L40 enabled through the same RTX FP16 path; validate benchmark/accuracy on the target card

GROOT N1.6

GROOT N1.6 defaults to FP8 on RTX. The full-FP16 baseline runs every GEMM in FP16 with no activation calibration; it is reachable from load_model (or examples/quickstart.py --use_fp16) and exposes the same set_prompt / predict contract as the FP8 path:

model = flash_rt.load_model(
    "/path/to/GR00T-N1.6-3B",
    framework="torch",
    config="groot",
    hardware="rtx_sm120",
    num_views=2,
    embodiment_tag="gr1",
    use_fp8=False,
    use_fp16=True,
)

calibrate() is a no-op on this path and precision_spec is None (no FP8 scales). Reference: FP16-vs-FP8 cosine ≈ 0.999 on RTX 5090; FP16 is slower than FP8 (precision↔speed trade-off), so use it as an A/B reference, not for production latency.

GROOT N1.7

GROOT N1.7 defaults to the FP8 path on RTX (see GROOT N1.7 RTX). The RTX reference variant is the opt-in non-default A/B path:

  • on rtx_sm120, it is the existing GrootN17TorchFrontendRtxFP16
  • on rtx_sm89, it is the dedicated GrootN17TorchFrontendRtxSm89FP16

Both paths replace the default FP8 route selected by load_model(), but the exact kernel/dtype mix remains hardware-specific.

N1.7 uses the aux-driven set_prompt(aux=...) / infer(state, initial_noise=...) contract (not the predict(images=...) quickstart flow), so construct the frontend directly when you want the explicit class. Example for rtx_sm120:

from flash_rt.frontends.torch.groot_n17_rtx_fp16 import (
    GrootN17TorchFrontendRtxFP16,
)

fe = GrootN17TorchFrontendRtxFP16(
    "/path/to/GR00T-N1.7-3B",
    num_views=2,
    embodiment_tag="oxe_droid_relative_eef_relative_joint",
)
fe.set_prompt(aux=aux)                       # aux from the N1.7 preprocessing path
actions = fe.infer(state_normalized, initial_noise=noise)

On rtx_sm89, import and construct flash_rt.frontends.torch.groot_n17_rtx_sm89_fp16.GrootN17TorchFrontendRtxSm89FP16 instead.

load_model(config="groot_n17", hardware="rtx_sm120", use_fp16=True, use_fp8=False) returns this frontend; on hardware="rtx_sm89" the same API returns the dedicated GrootN17TorchFrontendRtxSm89FP16 reference frontend. Reference on RTX 5090: FP16-vs-bf16 action cosine ≈ 0.99999, combined E2E-vs-reference cosine ≈ 0.9999; infer latency ≈ 10.7 ms (≈ the bf16 path — the DiT GEMMs are small, so FP16 and bf16 throughput match).

The underlying FP16 kernels are SM80-family friendly, but the public support matrix is narrower: FlashRT exposes the FP16 RTX route for Pi0.5 on rtx_sm120 / rtx_sm89, for GROOT N1.6 on rtx_sm120, and for GROOT N1.7 on rtx_sm120 / rtx_sm89.

Thor non-FP8 reference paths (GROOT N1.6 / N1.7)

Thor (SM110) defaults to FP8 for both GROOT N1.6 and N1.7. The non-FP8 route is the opt-in, fully non-quantized A/B reference for those FP8 paths: it runs the identical fully-kernelized pipeline with no FP8 anywhere — the backbone GEMMs go through the graph-safe fp16_nn path instead of per-tensor FP8, and the DiT action head stays bf16 (its native ckpt precision) — with no activation calibration. No PyTorch matmul touches the feature path. Enable it with use_fp8=False, use_fp16=True.

GROOT N1.6 (Thor)

Same set_prompt / predict contract as the FP8 path:

model = flash_rt.load_model(
    "/path/to/GR00T-N1.6-3B",
    framework="torch",
    config="groot",
    hardware="thor",
    num_views=2,
    embodiment_tag="gr1",
    use_fp8=False,
    use_fp16=True,
)
actions = model.predict(images=[base_img, wrist_img], prompt="pick up the block")

Reference (Jetson AGX Thor, CUDA-event E2E P50): FP8 default ≈ 44 ms, FP16 ≈ 65 ms; FP16-vs-FP8 action cosine ≈ 0.99996 (fixed denoise noise). FP16 is slower (precision↔speed trade-off) — use it as an accuracy reference, not for production latency.

GROOT N1.7 (Thor)

N1.7 uses the aux-driven set_prompt(aux=...) / infer(state, initial_noise=...) contract (not predict(images=...)). The default on Thor is FP8 (use_fp8=True): the vision backbone (ViT / DeepStack / LLM / VL-self-attn) and the DiT action head both run per-tensor FP8 through the host-alpha GEMM epilogues. The whole per-frame path is two static CUDA graphs — a vision backbone graph and an action graph (cross-KV refresh + state encode + all diffusion steps) — with no Python/torch compute on the hot path.

import flash_rt

# FP8 default (no flags needed)
fe = flash_rt.load_model(
    "/path/to/GR00T-N1.7-3B",
    framework="torch", config="groot_n17", hardware="thor",
    num_views=2, embodiment_tag="oxe_droid_relative_eef_relative_joint",
)
fe.set_prompt(aux=aux)                          # aux from the N1.7 preprocessing path
actions = fe.infer(state_normalized, initial_noise=noise)

Reference (Jetson AGX Thor, wall-clock per-frame image→action e2e): ≈ 49 ms — vision backbone graph ≈ 24 ms + action graph ≈ 26 ms. Action cosine vs the FP16-backbone reference ≈ 0.9999; replays are bit-identical.

Full-FP16 reference (opt-in A/B). use_fp8=False, use_fp16=True returns a fully non-quantized reference: the backbone runs every GEMM in fp16 and the DiT action head stays bf16 (no FP8 anywhere). It is the accuracy baseline for the whole model (e2e ≈ 74 ms on Thor; action cosine vs the FP8 default ≈ 0.9999), not a production-latency path.

fe = flash_rt.load_model(
    "/path/to/GR00T-N1.7-3B",
    framework="torch", config="groot_n17", hardware="thor",
    use_fp8=False, use_fp16=True,
    num_views=2, embodiment_tag="oxe_droid_relative_eef_relative_joint",
)
# equivalently: GrootN17TorchFrontendThorFP16(checkpoint, num_views=2, embodiment_tag=...)

GROOT N1.7 RTX

GROOT N1.7 is registered for the RTX SM120 / SM89 torch path:

import flash_rt

model = flash_rt.load_model(
    "/path/to/GR00T-N1.7-3B",
    framework="torch",
    config="groot_n17",
    hardware="rtx_sm120",
    num_views=2,
    embodiment_tag="oxe_droid_relative_eef_relative_joint",
)

model.set_prompt(aux=aux, prompt="put the blue block in the green bowl")
actions_normalized = model.infer(
    state_normalized,
    initial_noise=initial_noise,
    use_dit_graph=True,
)

aux is the precomputed Qwen3-VL setup bundle for N1.7: LLM input embeddings, visual masks, M-RoPE tables, pixel features, and grid_thw. infer() expects normalized state and returns normalized actions; denormalization remains the caller's responsibility for this N1.7 path.

On RTX the default path is FP8, but the hardware split is explicit:

  • rtx_sm120 uses the shared RTX frontend / pipeline with the established fp8_descale_fp16 backbone path
  • rtx_sm89 uses a dedicated RTX SM89 frontend / pipeline with runtime weights stored as [N, K] and GEMMs dispatched through fp8_nt_dev

In both cases the VLM backbone feature path runs through FlashRT kernels: backbone attention uses the vendored FA2 / FMHA kernels in set_prompt, and DiT self/cross attention uses FlashRT's vendored FA2 slots during infer. The action-head mix is not identical across the two RTX families:

  • rtx_sm120 keeps the established RTX semantics documented by its frontend
  • rtx_sm89 keeps FlashRT kernel execution for the SM89 FP8 backbone and quantized DiT hot-path GEMMs, while some action-head-adjacent helpers stay in the inherited bf16/fp16 route

Activation scales follow the FlashRT calibration convention (Calibration): weight scales are baked at load, activation scales are calibrated once and cached to ~/.flash_rt/calibration/<ckpt_hash>_n17_Se<N>.json. A warm set_prompt loads the cache and runs FP8 kernels only; the torch reference shadow runs solely on a cold cache miss (or recalibrate=True) to extract activation amax — never as the inference backbone. calibrate(aux_list) refines scales across N samples.

Supported hardware is rtx_sm120 and rtx_sm89. rtx_sm120 remains the validated Blackwell reference path in the public docs; rtx_sm89 is validated locally on RTX 4090 using the same aux/reference-fixture methodology and should still be benchmarked on the target card before making hardware-specific claims.

Build FA2 with at least the N1.7 head dimensions and dtypes:

cmake -S . -B build \
  -DFA2_HDIMS="96;128" \
  -DFA2_DTYPES="fp16;bf16"
cmake --build build --target flash_rt_fa2 -j

RTX 5090 validation against the N1.7 reference fixture (denormalized action cosine vs HuggingFace):

Metric FP8 (default) full-FP16 (opt-in)
Combined 0.999849 0.999852
EEF 9D 0.999901 0.999895
Gripper 0.928382 0.913727
Joints 0.999832 0.999838

The gripper DOF is a near-constant scalar whose cosine is intrinsically noisy (the bf16 reference itself scores ≈ 0.94); the combined action cosine clears the ≥ 0.999 bar on both paths.

RTX 5090 latency (warm, P50):

Stage FP8 (default) full-FP16 (opt-in)
Backbone (ViT+LLM+VL self-attn, per observation) 12.8 ms 15.3 ms
Infer (DiT 4-step loop, CUDA graph) 10.7 ms 10.9 ms
E2E (backbone + infer) 24.0 ms 26.6 ms

FP8 saves ≈ 2.7 ms E2E (1.11×), concentrated in the backbone GEMMs (1.19×); the DiT is never quantized, so infer is effectively the same on both paths. When one observation drives several action chunks, the backbone is computed once and reused, so per-chunk cost is just the infer row.

This path does not change CMake targets, C++ bindings, or existing Pi0/Pi0.5/GROOT N1.6 runtime dispatch.

Wan2.2 TI2V-5B

Wan2.2 TI2V-5B is exposed as an RTX SM120 official-pipeline baseline:

import flash_rt

model = flash_rt.load_model(
    "/path/to/Wan2.2-TI2V-5B",
    framework="torch",
    config="wan22_ti2v_5b",
    hardware="rtx_sm120",
)

model.set_prompt("A blue sphere rolls across a wooden table")
video = model.infer(
    mode="t2v",
    width=832,
    height=480,
    frames=81,
    steps=20,
    shift=5.0,
    guide_scale=5.0,
    seed=1234,
)

This route uses the official Wan Python pipeline and original ModelScope checkpoint layout. It is separate from ComfyUI; ComfyUI integration should be provided by an external custom-node package. See docs/wan22_usage.md.

Cosmos3-Nano text2video

Cosmos3-Nano text2video is a self-contained, kernelized FP8 denoise model (non-VLA) on RTX SM120. Conditioning enters via set_prompt(ref=...) and the denoise runs through infer(...); predict() is not part of this API.

import flash_rt

model = flash_rt.load_model(
    "/path/to/cosmos3_video_weights.safetensors",
    framework="torch",
    config="cosmos3_video",
    hardware="rtx_sm120",
    use_fp8=True,                       # default; use_fp8=False -> bf16 reference
)

model.set_prompt(ref="/path/to/tensors.safetensors")   # official reference dump
out = model.infer(
    teacache_skip="3,5,7",             # training-free step caching ("" = off)
    shift=10.0,
    compare_ref=True,
    return_metadata=True,
)
# out["latent"] -> [1,48,T,H,W] denoised vision latent (VAE decode is downstream)

The model-local kernel is built once on the target GPU (cd flash_rt/models/cosmos3_video/kernels && python3 setup.py build_ext --inplace). See docs/cosmos3_video_usage.md.

model.predict()

actions = model.predict(
    images=[base_img, wrist_img],   # list of (224,224,3) uint8/float16 numpy
    prompt="pick up the red block", # text prompt (required on first call)
)
Parameter Type Default Description
images list or dict required Camera images. List: [base, wrist] or [base, wrist, wrist_right]. Dict: {"image": ..., "wrist_image": ...}.
prompt str|None None Task prompt. Required on first call. Omit to reuse previous prompt (no recalibration, no graph recapture).

Returns: np.ndarray shape (10, 7) — 10 future action steps, 7 DOF.

Prompt reuse: When prompt is the same as the previous call, predict() skips set_prompt() entirely — only graph replay happens (~44ms). Changing the prompt triggers recalibration + graph recapture (~4s).

model.recalibrate()

model.recalibrate()
# Next predict() will force fresh calibration
actions = model.predict(images=..., prompt="...")

Clears the calibration cache for this checkpoint and forces recalibration on the next predict() call. Use after:

  • Fine-tuning the model (activation distributions changed)
  • Switching deployment domains (different image statistics)
  • Debugging precision issues

Pi0-FAST

Pi0-FAST is an autoregressive model — actions are generated as discrete FAST tokens (Gemma 2B, 18 layers), not via diffusion. Total latency = prefill + N × per_token_decode where N is typically 30–80 tokens.

Performance Modes

Mode set_prompt (cold) set_prompt (cached) Per-token 50-tok E2E
Default ~2.5 s ~0.1 s 9.5 ms ~480 ms
Max-perf (decode_cuda_graph=True) ~4.0 s ~1.5 s 8.3 ms ~447 ms 
# Default: good for interactive / multi-prompt scenarios
model = flash_rt.load_model(
    checkpoint="/path/to/pi0_fast_base",
    config="pi0fast",
)
actions = model.predict(
    images=[base_img, wrist_img],
    prompt="pick up the red block",
    state=robot_state,
)

# Max-performance: best for fixed-prompt 24h continuous control
model = flash_rt.load_model(
    checkpoint="/path/to/pi0_fast_base",
    config="pi0fast",
    decode_cuda_graph=True,       # capture decode loop as CUDA Graph
    decode_graph_steps=80,        # covers sequences up to 80 action tokens
)

Default mode: Each decode token runs through a Python loop with per-step kernel launches. Lowest startup cost — set_prompt loads calibration from cache in ~0.1s on subsequent runs.

Max-performance mode: The action-phase decode loop is captured as a single CUDA Graph (same technique as Pi0's diffusion loop). All autoregressive token dependencies (argmax → embedding lookup → next step) run entirely on GPU with zero host-device synchronization per step.

  • Adds ~1.5s to set_prompt for decode graph capture (on top of the cached calibration load)
  • Saves ~1.2 ms/token (–13% per-token latency)
  • Break-even at ~40 inferences per prompt
  • For long sequences (80 tokens), saves ~96 ms per inference

set_prompt Timing Breakdown

Component Cold (first ever) Cached (same ckpt + Se)
SigLIP graph capture ~0.2 s ~0.2 s
FP8 calibration (16 samples) ~2.4 s 0 s (loaded from cache)
Prefill graph capture ~0.1 s ~0.1 s
Decode graph capture (decode_cuda_graph=True) ~1.5 s ~1.5 s
Total (default) ~2.5 s ~0.1 s
Total (max-perf) ~4.0 s ~1.5 s

The 2.4s calibration is the largest cold-start cost. It runs 16 forward passes with noise perturbations to find robust FP8 activation scales, then saves to ~/.flash_rt/calibration/. On subsequent runs with the same checkpoint and sequence length, calibration loads from cache instantly.

Calibration Cache

Pi0-FAST follows the same calibration cache mechanism as Pi0/Pi0.5:

  1. set_prompt(): Check cache → hit: load scales (0s) / miss: run 16-sample calibration (~2.4s) and save to cache
  2. First predict(): Always recalibrates with real image data (regardless of cache) for optimal FP8 scales matching actual deployment images. The result is saved back to cache.
  3. Subsequent predict(): Pure inference, no recalibration

Same-scene optimization: When running the same scene repeatedly (same prompt, same checkpoint), the cached calibration from step 2 above is reused. No need to force recalibration — the system handles it automatically.

Force recalibration:

rm -rf ~/.flash_rt/calibration/   # clear all cached scales

Checkpoint Conversion

The Torch frontend requires safetensors format. Convert from Orbax:

python examples/convert_pi0fast_orbax_to_safetensors.py \
    --orbax_dir /path/to/pi0_fast_base \
    --output_dir /path/to/pi0_fast_base_converted

The JAX frontend reads Orbax directly.

Autotune

CUDA Graph instantiation on Thor is non-deterministic — the same kernels can produce different execution schedules with ~2ms latency variance. Autotune recaptures the graph multiple times and keeps the fastest schedule.

How it works

  1. Capture CUDA Graph (SigLIP + Encoder + Decoder)
  2. Benchmark the graph with CUDA events
  3. If latency is within the fast regime (< 38.5ms for Enc+AE), accept
  4. Otherwise, recapture and retry
  5. After autotune trials, use whatever was captured last

Parameter values

autotune= Behavior Extra startup time When to use
0 or False Single capture, no retry 0 Debugging, fastest startup
3 (default) Up to 3 retries ~1s Production (Torch almost always finds fast graph on trial 0)
5 Up to 5 retries ~2.5s JAX or when you need guaranteed best latency
True Same as 3 ~1s Shorthand

Framework differences

  • Torch: Almost always gets the fast graph on trial 0 (43-44ms total)
  • JAX: XLA GPU memory state sometimes prevents fast graph. jax.clear_caches() is called before capture to help, but 5+ trials may be needed

Autotune runs once per prompt

Autotune is part of set_prompt(). If you call predict() with the same prompt, the cached graph is replayed — no autotune overhead.

RTX 4090 / SM89 FP8 NT safety policy

Pi0.5 on RTX 4090 / Ada uses FP8 weights in nk layout and dispatches cuBLASLt FP8 NT GEMMs. On some SM89 driver / cuBLASLt runtime combinations, the default heuristic top-1 algorithm is stable but the top-N autotune benchmark can expose a non-default candidate that raises CUDA illegal memory access. Because that failure poisons the CUDA context, FlashRT avoids the risky benchmark path before it runs.

By default (FLASHRT_FP8_NT_AUTOTUNE=auto), Pi0.5 skips FP8 NT top-N autotune on SM89 and keeps the cuBLASLt heuristic top-1 algorithm. FP8 inference remains enabled; only the FP8 NT autotune benchmark is skipped.

Override for debugging / local benchmarking:

FLASHRT_FP8_NT_AUTOTUNE=force  # force FP8 NT top-N autotune
FLASHRT_FP8_NT_AUTOTUNE=safe   # always skip FP8 NT top-N autotune

Calibration

FP8 inference requires calibrated activation scales — per-layer maximum values that determine the FP8 quantization range. Incorrect scales cause precision loss.

How it works

Phase 1: Initial calibration (during set_prompt()):

  1. Check disk cache: ~/.flash_rt/calibration/{ckpt_hash}_Se{N}.json
  2. Cache hit: Load scales from JSON (instant)
  3. Cache miss: Run encoder_forward_calibrate() + decoder_forward_calibrate() with warmup data (~3-4s), save to cache

Phase 2: Real-data recalibration (on first predict() call):

  1. After SigLIP processes the first real image, the encoder input (enc_x) contains realistic activation distributions
  2. Rerun calibration with this real data for more accurate scales
  3. Recapture CUDA Graph with updated scales
  4. This happens only once — subsequent predict() calls skip this step

Cache details

Property Value
Location ~/.flash_rt/calibration/
Key SHA256(checkpoint_first_64KB + file_size)[:16] + _Se{sequence_length}
Format JSON with enc_scales, ae_scales, enc_alpha, enc_w_scales
Invalidation Automatic per-checkpoint hash. Different checkpoints or finetunes get different caches.

Cache isolation

  • Multi-model safe: Each checkpoint produces a unique hash → separate cache files
  • Multi-framework safe: Torch (safetensors) and JAX (Orbax) hash different files → separate caches
  • Multi-prompt safe: Different prompt lengths produce different Se → separate caches
  • Fine-tune safe: Modified checkpoint → different first 64KB → different hash → fresh calibration

Forcing recalibration

Three ways to force fresh calibration:

# Method 1: At load time
model = flash_rt.load_model(checkpoint, recalibrate=True)

# Method 2: At runtime
model.recalibrate()
model.predict(images=..., prompt="...")

# Method 3: CLI
python examples/quickstart.py --checkpoint /path/to/ckpt --recalibrate

Or manually delete the cache:

rm -rf ~/.flash_rt/calibration/

Explicit multi-sample calibration — model.calibrate()

model.calibrate() lets you seed the FP8 activation scales with a list of real observations instead of the default single-frame / zero-input path. Use it when runtime images cover a wider distribution (different lighting, object poses, camera angles) than any one calibration frame could represent.

The API is the same everywhere — one call, any number of observations:

model.calibrate(observations)

Everything else has a sensible default. You do not need to change any other setting: the cache, graph capture, prompt, and predict() all continue to behave as before.

I have a dataset — how do I call calibrate?

The common case: you have a robot-rollout dataset, and you want to calibrate on N representative frames before deployment. Three lines:

from flash_rt.datasets.libero import load_calibration_obs

# Pick 8 stratified frames (episode × frame-position) from a LIBERO-
# format dataset. Returns list[dict] — each dict has 'image',
# 'wrist_image', 'state' ready to hand straight to calibrate().
obs_list = load_calibration_obs("/path/to/libero_dataset", n=8)

model.calibrate(obs_list)

load_calibration_obs expects the LeRobot-v2 LIBERO layout (meta/info.json + data/chunk-NNN/episode_NNNNNN.parquet). If your dataset has a different layout, either:

  1. Load your observations yourself and pass the list:
obs_list = [
    {"image": ..., "wrist_image": ..., "state": ...},
    {"image": ..., "wrist_image": ..., "state": ...},
    # ... any N >= 1
]
model.calibrate(obs_list)
  1. Or build a pandas DataFrame with (task_index, episode_index, frame_index, index) columns and use the lower-level helper:
from flash_rt.core.calibration import stratified_sample

obs_list = stratified_sample(my_dataframe, my_load_fn, n=8)
model.calibrate(obs_list)

N = 1 — the default, keeps using the disk cache

If you pass a single observation (or skip calibrate() entirely — the first predict() call triggers the same path), FlashRT uses the legacy calibration pipeline:

model.calibrate([obs])        # explicit N=1
# or simply:
model.predict(obs)            # first call auto-calibrates

The first run computes scales from one forward pass and writes them to ~/.flash_rt/calibration/{ckpt_hash}_Se{N}.json. Every subsequent process with the same checkpoint reads the cache — no forward pass required.

N >= 2 — multi-sample dataset calibration

Same API, just pass more observations:

model.calibrate(obs_list)                     # default percentile
model.calibrate(obs_list, percentile=99.9)    # explicit (same default)

Under the hood: each observation runs one calibration forward pass, per-layer activation maxima are reduced by taking the 99.9 th percentile across samples (so one outlier frame cannot inflate every scale), and the reduced scales are written to the device before the graph is recaptured.

Cache is still used: N >= 2 writes the same JSON format as N = 1, so a second process run with the same checkpoint picks up the reduced scales for free.

Parameters

Parameter Type Default Meaning
observations list[dict] required Obs dicts matching the predict() contract (image, wrist_image, optional state). N = 1 → single-frame path; N >= 2 → multi-sample percentile path.
percentile float 99.9 Percentile applied to per-sample per-tensor amax. 100.0 == traditional max. Lower (e.g. 95) clips more aggressively — useful when your dataset has known outliers.
max_samples int | None None Upper bound, lets you pass an iterator without materializing the whole list.
verbose bool False Log a one-line dispersion summary (median amax, outlier cutback, etc.).

Which frontends support it

Frontend N = 1 N >= 2 (dataset)
pi05_thor (torch) ✅ ✅
pi0_thor (torch) ✅ ✅
pi0fast (torch) ✅ ✅
pi05_thor_fp4 (torch, FP4 encoder active) ✅ ✅ (two-phase: FP8 + AWQ refit)
pi05_thor_fp4 (jax, FP4 encoder active) ✅ ✅ (two-phase: FP8 + AWQ refit)
pi0_rtx, pi05_rtx (torch + jax) ✅ ✅
groot_rtx (torch) ✅ ✅
groot_n17_thor (torch, aux samples) ✅ ✅
groot_thor (torch) ✅ Implemented; validate with N1.6 fixtures before deployment
pi05_thor / pi0_thor / pi0fast (jax, non-FP4) ✅ Implemented; validate with target fixtures before deployment

GROOT N1.7 Thor is aux-driven rather than predict()-driven: its calibration samples use the same aux schema consumed by set_prompt(aux=...). The obs-list examples above apply to predict()-style frontends.

pi05_thor_fp4 uses a two-phase multi-sample flow: Phase 1 reduces FP8 activation scales across N samples (same loop the base class uses), then Phase 2 collects AWQ per-channel activation amax with the new FP8 scales active, percentile-reduces across samples, and refits the NVFP4 packed weights + AWQ inv_s buffers in place — no second graph capture. Roughly 2× the wall-clock of non-FP4 N samples because both phases walk the dataset.

Choosing N and percentile

Scenario Recommended
Fixed environment, fine-tuned model N = 8, percentile = 99.9
Fixed environment, base / non-fine-tuned model N = 1 at a representative frame
Runtime drifts (lighting / outdoor / scene switches) N = 64–256, percentile = 99.9
Input contains sensor outliers N >= 256, percentile <= 99.0

Measured cosine and calibration-time numbers, plus the precision-vs-coverage trade-off, live in docs/calibration.md §10. The short version: a LIBERO-fine-tuned Pi0.5 gains cos ≈ +0.0003 and halves max per-channel deviation going from N = 1 to N = 8, at the cost of a few seconds of one-off calibration.

Low-level helper

If your dataset cannot go through load_calibration_obs, the lower-level helper takes a pandas DataFrame + a per-index loader:

flash_rt.core.calibration.stratified_sample(
    metadata,        # pandas.DataFrame with {index, task_index,
                     #   episode_index, frame_index} columns
    load_fn,         # callable: index -> obs dict
    n=8,
    *,
    task_filter=None,        # narrow to one task_index
    exclude=None,            # skip specific global indices
) -> list[dict]

It picks N global frame indices stratified across episode × frame position (not uniform-random, which would over-sample steady-state frames) and applies load_fn to each.

Diagnostic: outlier-scale warning

After calibration, FlashRT scans the produced per-layer FP8 scales and logs a WARNING via flash_rt.core.calibration.check_scale_ceiling if any scale is more than 20 × the median of the same calibration. This catches the case where a single outlier frame in the calibration set stretched the FP8 scale on one layer far beyond its peers — the typical sign that the dataset contains a glitched / overexposed / occluded sample.

Full Parameter Examples

import flash_rt

# === Production deployment (recommended) ===
model = flash_rt.load_model(
    checkpoint="/path/to/pi05_libero_pytorch",
    framework="torch",
    autotune=3,          # stable 44ms
)
actions = model.predict(images=[img, wrist], prompt="pick up the red block")


# === JAX with Orbax checkpoint ===
model = flash_rt.load_model(
    checkpoint="/path/to/orbax_checkpoint",
    framework="jax",
    autotune=5,          # JAX may need more trials
)


# === After fine-tuning: force recalibration ===
model = flash_rt.load_model(
    checkpoint="/path/to/finetuned_checkpoint",
    framework="torch",
    recalibrate=True,    # ignore old cache
)


# === Fast iteration during development ===
model = flash_rt.load_model(
    checkpoint="/path/to/checkpoint",
    framework="torch",
    autotune=0,          # skip autotune for fastest startup
)


# === 3-camera setup ===
model = flash_rt.load_model(
    checkpoint="/path/to/checkpoint",
    framework="torch",
    num_views=3,
)
actions = model.predict(
    images=[base_img, wrist_left, wrist_right],
    prompt="pick up the cup",
)


# === Runtime recalibration (domain shift) ===
model = flash_rt.load_model(checkpoint="/path/to/checkpoint")
actions = model.predict(images=[img1, img2], prompt="task A")

# ... deployment domain changed ...
model.recalibrate()
actions = model.predict(images=[img3, img4], prompt="task B")  # fresh calibration

Weight Cache (JAX only)

JAX (Orbax) checkpoint loading takes ~42s due to OCDBT deserialization + weight transform + FP8 quantization. The weight cache saves the final FP8-quantized engine weights to disk after first load, so subsequent loads skip all three steps.

Why JAX only?

Framework Cold start Bottleneck
Torch (safetensors) ~3s mmap load — already fast
JAX (Orbax) ~42s → ~6s with cache OCDBT deserialize + transform + FP8 quant

Torch uses safetensors which is essentially a flat binary mmap — there's nothing to cache. JAX's Orbax format requires complex deserialization that the weight cache eliminates.

How it works

  1. First load: Orbax → transform → FP8 quantize → upload to GPU → save binary cache (~/.flash_rt/weights/{hash}_nv{N}.bin, ~4 GB)
  2. Subsequent loads: Read binary cache → upload to GPU directly (~6s)

Parameters

# Default: cache enabled (recommended)
model = flash_rt.load_model(checkpoint, framework="jax")

# Disable cache (always re-quantize from Orbax)
model = flash_rt.load_model(checkpoint, framework="jax", weight_cache=False)

# Force re-quantize (clears both weight cache and calibration cache)
model = flash_rt.load_model(checkpoint, framework="jax", recalibrate=True)

When to disable or clear weight cache

Situation Action
First deploy Automatic — cache miss triggers full load + save
Normal restart Automatic — cache hit, ~6s load
After fine-tuning recalibrate=True or weight_cache=False
Checkpoint updated Automatic — different hash → new cache
Debugging precision weight_cache=False to rule out cache issues

Cache isolation

Each cache file is keyed by SHA256(checkpoint_manifest)[:16] + num_views. Different checkpoints, different fine-tunes, different num_views all produce separate cache files. No cross-contamination.

HTTP Server

FlashRT includes a FastAPI server for production deployment. The model loads once at startup; all subsequent requests are pure graph replay (~44ms).

Quick start

pip install fastapi uvicorn

# Torch
python examples/server.py --checkpoint /path/to/ckpt

# JAX
python examples/server.py --checkpoint /path/to/ckpt --framework jax

# Custom port + thorough autotune
python examples/server.py --checkpoint /path/to/ckpt --port 9000 --autotune 5

Endpoints

GET /health — Health check

{"status": "ok", "framework": "torch", "version": "2.2.0", "prompt": "pick up the red block"}

POST /predict — Run inference

{
    "prompt": "pick up the red block",
    "images": ["<base64_raw_uint8>", "<base64_raw_uint8>"],
    "image_shape": [224, 224, 3]
}

Response:

{
    "actions": [[0.1, -0.2, ...], ...],
    "latency_ms": 44.3,
    "shape": [10, 7]
}

If images is omitted, dummy random images are used (for testing).

Test with curl

# Simple test (dummy images)
curl -X POST http://localhost:8000/predict \
    -H "Content-Type: application/json" \
    -d '{"prompt": "pick up the red block"}'

# With real images (Python)
python -c "
import requests, base64, numpy as np
img = np.random.randint(0, 255, (224, 224, 3), dtype=np.uint8)
b64 = base64.b64encode(img.tobytes()).decode()
resp = requests.post('http://localhost:8000/predict', json={
    'prompt': 'pick up the red block',
    'images': [b64, b64],
})
print(resp.json()['shape'], resp.json()['latency_ms'], 'ms')
"

Thread safety

The server uses an asyncio lock to ensure only one inference runs at a time (single GPU). Concurrent requests are queued automatically.

Qwen3.6 OpenAI-compatible server

Qwen3.6-27B NVFP4 is served by the production agent server serving/qwen36_agent/ — session prefix reuse, committed-stream SSE, tool calls, and execution-state capsules. See its README.md for the full run guide, the parameter table, the measured serving numbers, and the differentiation vs vLLM/SGLang. Install the server extra, point FLASHRT_QWEN36_MTP_CKPT_DIR at a paired Qwen3.6 MTP checkpoint, and choose a --max-seq that covers the largest prompt plus completion you intend to serve.

pip install -e ".[torch,server]"   # or: pip install fastapi uvicorn

export FLASHRT_QWEN36_MTP_CKPT_DIR=/path/to/qwen36_mtp_ckpt
export FLASHRT_QWEN36_LONG_KV_CACHE=fp8

python -m serving.qwen36_agent.server \
  --checkpoint /path/to/qwen36_nvfp4 \
  --max-seq 262208 \
  --port 8000
# --route-min-seq 0, --warmup-preset agent, and --graph-cache-max (auto-scaled
# by --max-seq) are the defaults; /health reports "speculative": true/false.

The server exposes /v1/models, /health, /v1/chat/completions, and /v1/sessions. Each completion response carries a flashrt telemetry block (prefill_ms, decode_ms, decode_tok_per_s, cached_tokens, prefix_action, …) and the server logs one metric line per request, so decode tok/s is directly comparable to TPOT-style LLM serving numbers and excludes TTFT. For the parameter reference and the 128-token to 256K context sweep, see docs/qwen36_usage.md and docs/qwen36_nvfp4.md. The earlier single-file example is retired to examples/archive/.

CLI Reference

# Basic inference
python examples/quickstart.py --checkpoint /path/to/ckpt

# JAX framework
python examples/quickstart.py --checkpoint /path/to/ckpt --framework jax

# Benchmark 20 iterations
python examples/quickstart.py --checkpoint /path/to/ckpt --benchmark 20

# Thorough autotune
python examples/quickstart.py --checkpoint /path/to/ckpt --autotune 5

# Force recalibration
python examples/quickstart.py --checkpoint /path/to/ckpt --recalibrate

# LIBERO evaluation
python examples/thor/eval_libero.py \
    --checkpoint /path/to/ckpt \
    --task_suite libero_spatial \
    --framework torch

# Quick LIBERO test (3 tasks × 3 episodes)
python examples/thor/eval_libero.py \
    --checkpoint /path/to/ckpt \
    --task_suite libero_spatial --quick

# HTTP server
python examples/server.py --checkpoint /path/to/ckpt --port 8000

Startup Timeline

Typical load_model() + first predict() timing on Jetson AGX Thor:

Pi0 / Pi0.5

Phase Torch JAX (no cache) JAX (cached)
Load checkpoint + FP8 quantize ~3s ~42s ~6s
set_prompt(): tokenize + RoPE + time conditioning ~0.1s ~0.1s ~0.1s
set_prompt(): SigLIP graph capture ~0.5s ~0.5s ~0.5s
set_prompt(): calibration (cache miss) ~3s ~3s ~3s
set_prompt(): calibration (cache hit) 0s 0s 0s
set_prompt(): autotune=3 ~1s ~1.5s ~1.5s
First predict(): real-data recalibration ~1.5s ~1.5s ~1.5s
Subsequent predict(): graph replay ~44ms ~44ms ~44ms

After the first predict(), all subsequent calls are pure CUDA Graph replay at ~44ms. With weight cache + calibration cache, JAX warm start is ~6s (vs ~42s cold start).

Pi0-FAST

Phase Default Max-perf (decode_cuda_graph=True)
Load checkpoint + FP8 quantize ~5s ~5s
set_prompt(): calibration (cache miss) ~2.4s ~2.4s
set_prompt(): calibration (cache hit) 0s 0s
set_prompt(): SigLIP + prefill graph ~0.3s ~0.3s
set_prompt(): decode graph capture — ~1.5s
First predict(): real-data recalibration ~2.8s ~2.8s
Subsequent predict() (50 tokens) ~480ms ~447ms

With calibration cache, default-mode set_prompt() drops from ~2.5s to ~0.1s. Max-perf mode is ~1.5s (decode graph capture dominates after cache hit).

NVFP4 (Pi0.5 only)

Optional NVFP4 (Blackwell block-scaled FP4) quantization on the Pi0.5 encoder FFN stack, enabled via a single flag use_fp4=True. It is supported for Pi0.5 torch and JAX on Thor. Other configs (pi0 / groot / pi0fast) log a warning and use the FP8 route.

# Production-recommended — single flag, best-known config:
model = flash_rt.load_model(
    checkpoint,
    config="pi05",
    use_fp4=True,
)
# Equivalent to passing:
#   fp4_layers=tuple(range(18))   # full encoder FFN (18 layers)
#   use_awq=True                   # required for 18-layer scope
#   use_p1_split_gu=True           # production P1 path

# Advanced: override sub-flags for A/B or debug:
model = flash_rt.load_model(
    checkpoint, config="pi05",
    use_fp4=True,
    fp4_layers=(7, 8, 9),       # conservative subset
    use_awq=False,
    use_p1_split_gu=False,
)

Preset resolution

When use_fp4=True and a sub-flag (fp4_layers, use_awq, use_p1_split_gu) is left as None, it resolves to the production preset:

Sub-flag Preset value Reason
fp4_layers tuple(range(18)) Full encoder FFN coverage
use_awq True Required for 18-layer scope (without AWQ, full-scope FP4 cos collapses to ~0.33)
use_p1_split_gu True Split Gate+Up → 2× fp4out GEMM + combiner (Pi0.5 parity, Pi0.6 reusable)

What it does

  • GEMMs: all Gate+Up / Down proj GEMMs across the 18 FFN layers run in NVFP4 (block-size 16, UE4M3 scales, Sm1xxBlockScaledConfig tile layout) instead of FP8. Attention (QKV, O) stays FP8 fp16-output.
  • P1 split-GU: gate_proj and up_proj run as two separate cutlass_fp4_gemm_fp4out GEMMs (FP4-packed output + SFA via the LinCombBlockScaleFactor epilogue — proven CUTLASS pattern). A fused geglu_two_mul_fp4_to_fp4 combiner reads both FP4 inputs, applies GELU + Down-AWQ inv_s + per-block FP4 quant, and writes the packed FP4 + SFA directly for the Down GEMM. Eliminates ~31 MB/layer of fp16 DRAM round-trip vs the merged-GU path.
  • AWQ (activation-aware weight quant): per-input-channel pre-scale s[k] = (a[k]/mean(a))^awq_alpha on each NVFP4 weight, with the matching inverse scale fused into the pre-GEMM kernels (residual_add_rms_norm_mul_fp4_sfa for Gate+Up input, geglu_two_mul_fp4_to_fp4 for Down input). Calibrated on first predict() with real images, requantized in-place so the captured CUDA Graph remains valid.
  • Residual stream: stays fp16 through the FP4 region (NVIDIA enable_llm_nvfp4 design — output_quantizer disabled).
  • Non-FP4 paths: attention, decoder, SigLIP are unchanged (bit-identical to the FP8 baseline).

Weight loading

When use_fp4=True, the FP4 layer weights are loaded directly as fp16 from the selected framework checkpoint and NVFP4-quantized offline (no FP8 intermediate). Torch uses safetensors checkpoints. JAX uses Orbax checkpoints and depends on the Orbax loader accepting current StepMetadata checkpoint metadata. This matches the NVIDIA modelopt design and avoids a double-lossy FP8 → fp16 → FP4 round-trip.

Requirements

  • SM100+ GPU with Blackwell Tensor Cores (validated on Thor SM110). Hardware without NVFP4 support logs a warning and uses the FP8 route.
  • flash_rt_fp4.so extension built alongside flash_rt_kernels.so (automatic in standard install).

Validation

Pi0.5 on Jetson AGX Thor, LIBERO Spatial 10 tasks × 50 trials = 500 episodes:

Config Success E2E P50 (normal regime)
FP8 baseline 491 / 500 (98.2%) ~43.5 ms
NVFP4 full-18 + AWQ + P1 (--use_fp4) 491 / 500 (98.2%) ~43.5 ms

Task-level parity with the FP8 baseline. In-graph kernel-sum profile shows -2.1 ms/infer theoretical saving, but CUDA Graph per-node scheduling overhead on the full ~60-kernel pipeline absorbs most of it — net wall-clock P50 is at parity in the production regime. The kernels are designed to scale linearly on Pi0.6 (~2× compute) where per-node overhead stays constant while per-kernel savings grow.

Multi-model precision regression (tests/test_all_models_precision.py):

Model Config cos vs reference P50
Pi0.5 FP8 baseline vs_prod=0.9984, vs_old_torch=0.9999 44.5 ms
Pi0.5 use_fp4=True preset vs_pytorch_ref=0.9989, vs_prod=0.9974 43.3 ms
Pi0 (unchanged) vs_pytorch_ref=0.9972 46.7 ms
Pi0 JAX (unchanged) vs_pytorch_ref=0.9983 45.1 ms
GROOT N1.6 (unchanged) vs_pytorch_ref=0.9986 46.2 ms

Thor replay-latency / precision check for the JAX Orbax path:

Frontend Views P50 P95 cos vs same-origin PyTorch ref max diff
jax_fp8 3 55.03 ms 55.23 ms 0.999358 0.0638
jax_fp4 3 50.30 ms 56.58 ms 0.998351 0.0994

The JAX reference was generated from the same pi05_libero_jax Orbax checkpoint via the upstream OpenPI conversion script and PI0Pytorch.sample_actions; do not compare this validation against a separate HF / LeRobot safetensors checkpoint. These numbers are correctness / availability evidence for the Thor JAX FP4 path, not a broad performance claim.

Layer selection

fp4_layers accepts any subset of encoder layer indices 0-17. Two configurations are task-level LIBERO-validated:

  • tuple(range(18)) + AWQ (production preset — --use_fp4 default)
  • (7, 8, 9) without AWQ (the conservative subset, from the first FP4 drop)

Other subsets are simulation-only (see the internal precision report).

Known limits / roadmap

  • Decoder FP4 — precision simulation (S2: decoder all-proj) has cos 0.9985 and passed LIBERO quick 9/9. Full kernel integration planned (est. -6 ms E2E).
  • SigLIP FFN FP4 — precision simulation favorable, integration requires fp16-native SigLIP weight loader.
  • FP4 on Pi0 / GROOT / Pi0-FAST — architecture supports it; frontend subclasses not yet written. Kernels are reusable without change.
  • Full fp16 fallback (use_fp8=False) — requires a complete fp16 GEMM / RMSNorm pipeline, not yet implemented.

RL inference (classifier-free guidance)

Opt-in CFG inference path for advantage-conditioned VLA policies trained with the RECAP recipe (π*0.6, arXiv:2511.14759). The default infer() path is unchanged; RL mode is activated explicitly.

rt = Pi05TorchFrontendRtx("/path/to/pi05_libero_pytorch", num_views=2)

# Opt-in to advantage-conditioned CFG (rebuilds pipeline as Pi05CFGPipeline)
rt.set_rl_mode(cfg_enable=True, cfg_beta=1.5)
rt.set_prompt("fold the t-shirt")   # cond/uncond prompt pair embedded internally
rt.calibrate([obs])
actions = rt.infer(obs)["actions"]  # runs encoder x2 + decoder x2 per step

# Revert to the standard path any time
rt.set_rl_mode(cfg_enable=False)

Support matrix (v0.1.0):

Frontend CFG inference
pi05_rtx (torch) ✅
pi0_rtx, pi0fast, groot_rtx, all Thor / JAX frontends ❌ (pi0-RTX pattern ports planned)

Measured latency (RTX 5090, pi05_libero_pytorch, FP8):

path median (ms) note
Baseline (no CFG) 19.0 single forward, graph replay
CFG β=1.5 serial (Phase 2) 37.1 2× sequential forward, graph replay
CFG β=1.5 fused batched (Phase 3b) 25.9 cond+uncond in one B=2 forward

For sustained 50 Hz real-robot control, combine RL mode with batched mode to get the fused CFG path:

rt.set_batched_mode(enable=True)
rt.set_rl_mode(cfg_enable=True, cfg_beta=1.5)
rt.set_prompt("fold the t-shirt")
rt.calibrate([obs])
actions = rt.infer(obs)["actions"]   # 25.9 ms / call

Full latency table, numerical contract, and generic batched (RL rollout) roadmap are in docs/rl_inference.md.