feat(phyai): Add PhyAI native LIBERO pi0.5 benchmark support

benchmark/pi05/README_external_pi05_benchmarks.md

@@ -0,0 +1,168 @@
+# External PI0.5 benchmark wrappers
+
+This directory contains wrappers for benchmarking external PI0.5 inference runtimes with the same high-level settings and common `bench_n_batch.py` runner used by the PhyAI PI0.5 benchmarks.
+
+These wrappers do **not** call PhyAI's engine. Each wrapper calls the target runtime directly, and all machine-local paths must be passed by command-line flags or environment variables.
+
+| File | Runtime measured | Main timing scope |
+| --- | --- | --- |
+| `bench_flashrt_pi05.py` | FlashRT | direct `Pi05TorchFrontendRtx.infer(...)` hot path after `set_prompt`, calibration, and the first graph-building infer |
+| `bench_realtime_vla_pi05.py` | realtime-vla | `Pi05Inference.forward()` hot path |
+| `bench_vlacpp_pi05_client.py` | vla.cpp | ZMQ client request wall time; server phase timings are stored in JSONL `extras` |
+
+## Common settings
+
+Use the same settings across machines when possible:
+
+| Setting | Recommended value |
+| --- | --- |
+| batch size | 1 via `--batch-sizes 1` |
+| views | 2 synthetic views / camera streams for latency-only runs |
+| chunk size | 50 for strict comparison rows |
+| warmup | 100 via `--n-warmup` |
+| timed iterations | 100 via `--n-timed` |
+| result format | JSONL via `--result-file` |
+| prompt | fixed prompt or fixed prompt length, recorded in JSONL `extras` |
+| precision | BF16 fair row; optimized rows must be labeled separately |
+
+Before running, verify the GPU is idle:
+
+```bash
+nvidia-smi
+uptime
+```
+
+The examples below use placeholders rather than host-specific paths:
+
+| Placeholder | Meaning |
+| --- | --- |
+| `<PHYAI_REPO>` | This PhyAI checkout containing `benchmark/pi05/` |
+| `<FLASHRT_REPO>` | FlashRT repository clone |
+| `<REALTIME_VLA_REPO>` | realtime-vla repository clone |
+| `<VLACPP_REPO>` | vla.cpp repository clone |
+| `<PI05_CHECKPOINT>` | PI0.5 checkpoint directory or file for the selected runtime |
+| `<TOKENIZER_DIR_OR_ID>` | Local tokenizer directory or HF id used by vla.cpp client |
+| `<LIBERO_STATS_JSON>` | Local LIBERO `meta/stats.json` for PI0.5 state tokenization |
+
+## Environment setup
+
+These scripts assume the target runtime can already be imported or executed. Keep each official repo at a known commit and record it with your results.
+
+FlashRT:
+
+```bash
+git clone https://github.com/flashrt-project/FlashRT <FLASHRT_REPO>
+cd <FLASHRT_REPO>
+# Install/build FlashRT following its official README for your GPU/CUDA stack.
+export FLASHRT_ROOT=<FLASHRT_REPO>
+```
+
+Notes: on RTX 5090/SM120, make sure the CUDA toolkit used to build extensions supports the GPU. For the BF16 fair row, the wrapper sets `FVK_PI05_RTX_FORCE_BF16=1`. Do not use `load_model(..., num_steps=50)` for chunk size; FlashRT `num_steps` means denoise steps.
+
+realtime-vla:
+
+```bash
+git clone https://github.com/Dexmal/realtime-vla <REALTIME_VLA_REPO>
+cd <REALTIME_VLA_REPO>
+# Install realtime-vla following its official README.
+export REALTIME_VLA_ROOT=<REALTIME_VLA_REPO>
+```
+
+If you pass a PI0.5 `model.safetensors`, also set `FLASHRT_ROOT` because the wrapper reuses FlashRT's checkpoint conversion helper. Loading `.pkl/.pickle` checkpoints requires `--trust-pickle-checkpoint` and should only be used for trusted files.
+
+vla.cpp:
+
+```bash
+git clone https://github.com/VinRobotics/vla.cpp <VLACPP_REPO>
+cd <VLACPP_REPO>
+# Build vla-server following its official README, for example into build_sm120/.
+export VLACPP_ROOT=<VLACPP_REPO>
+```
+
+Prepare the PI0.5 GGUF model, `mmproj` GGUF, tokenizer, and LIBERO `stats.json` before running. Prefer local tokenizer and stats paths to avoid network/auth issues.
+
+Quick smoke test after setup: use the runtime-specific command below, replace the real output path with a scratch file, and set `--n-warmup 1 --n-timed 1`.
+
+## FlashRT
+
+`--flashrt-root` can be omitted if `FLASHRT_ROOT=<FLASHRT_REPO>` is set.
+
+```bash
+cd <PHYAI_REPO>
+python benchmark/pi05/bench_flashrt_pi05.py \
+  --flashrt-root <FLASHRT_REPO> \
+  --checkpoint <PI05_CHECKPOINT> \
+  --precision bf16 \
+  --num-views 2 \
+  --chunk-size 50 \
+  --batch-sizes 1 \
+  --n-warmup 100 \
+  --n-timed 100 \
+  --result-file results/flashrt_pi05_external.jsonl
+```
+
+For FlashRT optimized precision, use `--precision fp8_bf16` and keep it in a separate table row.
+
+## realtime-vla
+
+`--realtime-vla-root` can be omitted if `REALTIME_VLA_ROOT=<REALTIME_VLA_REPO>` is set. If `--checkpoint` points to a PI0.5 `model.safetensors` file or a directory containing `model.safetensors`, pass `--flashrt-root <FLASHRT_REPO>` as well because the wrapper reuses FlashRT's PI0.5 safetensors conversion helper.
+
+```bash
+cd <PHYAI_REPO>
+python benchmark/pi05/bench_realtime_vla_pi05.py \
+  --realtime-vla-root <REALTIME_VLA_REPO> \
+  --flashrt-root <FLASHRT_REPO> \
+  --checkpoint <PI05_CHECKPOINT> \
+  --num-views 2 \
+  --chunk-size 50 \
+  --prompt-len 16 \
+  --batch-sizes 1 \
+  --n-warmup 100 \
+  --n-timed 100 \
+  --result-file results/realtime_vla_pi05_external.jsonl
+```
+
+If you already have a realtime-vla `.pkl`, `.pt`, or `.pth` checkpoint, `--flashrt-root` is not needed.
+
+## vla.cpp
+
+Start `vla-server` first. The PI0.5 GGUF server reports server phase timing when started with `--timing-detail phase`.
+
+Example server:
+
+```bash
+<VLACPP_REPO>/build_sm120/vla-server \
+  --bind tcp://127.0.0.1:5555 \
+  --timing-detail phase \
+  <MM_PROJ_GGUF> \
+  <PI05_GGUF>
+```
+
+Then run the client benchmark. `--vlacpp-root` can be omitted if `VLACPP_ROOT=<VLACPP_REPO>` is set.
+
+```bash
+cd <PHYAI_REPO>
+python benchmark/pi05/bench_vlacpp_pi05_client.py \
+  --vlacpp-root <VLACPP_REPO> \
+  --addr tcp://127.0.0.1:5555 \
+  --arch pi05 \
+  --tokenizer <TOKENIZER_DIR_OR_ID> \
+  --stats-json <LIBERO_STATS_JSON> \
+  --num-views 2 \
+  --chunk-size 50 \
+  --batch-sizes 1 \
+  --n-warmup 100 \
+  --n-timed 100 \
+  --result-file results/vlacpp_pi05_external.jsonl
+```
+
+For reproducibility, prefer a local tokenizer path for `--tokenizer`. The default PI0.5 tokenizer id, `google/paligemma-3b-pt-224`, may require HuggingFace access and authentication. Passing `--stats-json` avoids implicit network fetches for LIBERO state quantile statistics.
+
+## Notes
+
+- `bench_flashrt_pi05.py` and `bench_realtime_vla_pi05.py` reuse the PhyAI `NBatchBenchRunner` JSONL schema.
+- `bench_flashrt_pi05.py` records wall latency with the runner's perf-counter path because FlashRT may run work on an internal CUDA stream; the step synchronizes CUDA before returning.
+- `bench_realtime_vla_pi05.py` uses the runner's CUDA event timing around `Pi05Inference.forward()`.
+- `bench_vlacpp_pi05_client.py` requires a running `vla-server`; the runner records client wall latency, and server phase latency is recorded in JSONL `extras.server_phase_latency_ms`.
+- vla.cpp phase timing may expose `vision` and combined `inference`, not necessarily separate prefix/expert timing.
+- These wrappers are intended for latency reproduction and support only `--batch-sizes 1` for now. Component MFU tables still require the shared PI0.5 FLOP model and component timings, as described in `thor_pi05_benchmark_plan.md`.

benchmark/pi05/bench_flashrt_pi05.py

@@ -0,0 +1,229 @@
+#!/usr/bin/env python3
+"""FlashRT PI0.5 latency benchmark using the PhyAI bench runner.
+
+This is a thin adapter around FlashRT's direct ``Pi05TorchFrontendRtx`` path.
+The direct frontend is used because action chunk size is a frontend constructor
+argument; ``flash_rt.load_model(..., num_steps=...)`` controls denoise steps,
+not action chunk size. It mirrors ``benchmark/bench_n_batch_ws1_pi05.py``:
+framework-specific setup lives here, while warmup, timed iterations, JSONL
+output, and optional profiling are handled by ``benchmark/bench_n_batch.py``.
+
+Run::
+
+    python benchmark/pi05/bench_flashrt_pi05.py \
+        --flashrt-root <FLASHRT_REPO> \
+        --checkpoint <PI05_CHECKPOINT> \
+        --batch-sizes 1 --n-warmup 100 --n-timed 100 \
+        --result-file results/flashrt_pi05.jsonl
+
+Only batch size 1 is supported because the FlashRT PI0.5 direct frontend path
+used here takes one robot request at a time.
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+from pathlib import Path
+import statistics
+import sys
+from typing import Any
+
+import numpy as np
+import torch
+
+_BENCHMARK_DIR = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(_BENCHMARK_DIR))
+import bench_n_batch as bnb  # noqa: E402
+from phyai.utils.profile import (  # noqa: E402
+    add_profile_cli_args,
+    install_profiler,
+    profile_config_from_args,
+)
+
+
+def env_path(name: str) -> Path | None:
+    value = os.environ.get(name)
+    return Path(value) if value else None
+
+
+class LazySummary(dict):
+    def __init__(self, values_fn_or_items):
+        if callable(values_fn_or_items):
+            super().__init__()
+            self._values_fn = values_fn_or_items
+        else:
+            super().__init__(values_fn_or_items)
+            self._values_fn = None
+
+    def items(self):
+        if self._values_fn is None:
+            return super().items()
+        summary = summarize(self._values_fn())
+        return (summary or {}).items()
+
+
+def summarize(values: list[float]) -> dict[str, float] | None:
+    if not values:
+        return None
+    xs = sorted(float(x) for x in values)
+    return {
+        "count": len(xs),
+        "mean_ms": float(statistics.fmean(xs)),
+        "median_ms": float(statistics.median(xs)),
+        "p50_ms": float(np.percentile(xs, 50)),
+        "p90_ms": float(np.percentile(xs, 90)),
+        "p99_ms": float(np.percentile(xs, 99)),
+        "min_ms": xs[0],
+        "max_ms": xs[-1],
+    }
+
+
+def make_observation(num_views: int, seed: int, prompt: str) -> dict[str, Any]:
+    """Deterministic synthetic observation for latency-only runs."""
+    rng = np.random.default_rng(seed)
+    obs: dict[str, Any] = {
+        "image": rng.integers(0, 256, size=(224, 224, 3), dtype=np.uint8),
+        "state": rng.standard_normal(8).astype(np.float32),
+        "task": prompt,
+        "prompt": prompt,
+    }
+    if num_views >= 2:
+        obs["wrist_image"] = rng.integers(0, 256, size=(224, 224, 3), dtype=np.uint8)
+    if num_views >= 3:
+        obs["wrist_image_right"] = rng.integers(
+            0, 256, size=(224, 224, 3), dtype=np.uint8
+        )
+    return obs
+
+
+def import_flashrt_frontend(repo: Path):
+    # Use the checked-out FlashRT repository directly; installation is optional.
+    sys.path.insert(0, str(repo))
+    from flash_rt.frontends.torch.pi05_rtx import Pi05TorchFrontendRtx
+
+    return Pi05TorchFrontendRtx
+
+
+def make_setup_fn(args: argparse.Namespace):
+    frontend_cls = import_flashrt_frontend(args.flashrt_root)
+
+    def setup_fn(batch_size: int) -> bnb.BenchSpec:
+        if batch_size != 1:
+            raise ValueError("FlashRT PI0.5 wrapper supports only batch_size=1")
+
+        if args.precision == "bf16":
+            # FlashRT uses this environment switch to force the BF16 PI0.5 RTX path.
+            os.environ["FVK_PI05_RTX_FORCE_BF16"] = "1"
+        else:
+            os.environ.pop("FVK_PI05_RTX_FORCE_BF16", None)
+
+        model = frontend_cls(
+            args.checkpoint,
+            num_views=args.num_views,
+            chunk_size=args.chunk_size,
+            cache_frames=1,
+            use_fp8=(args.precision == "fp8_bf16"),
+            hardware=args.hardware,
+        )
+        obs = make_observation(args.num_views, args.seed, args.prompt)
+
+        # set_prompt builds the prompt-specific pipeline; calibration captures the graph.
+        model.set_prompt(args.prompt)
+        model.calibrate_with_real_data([obs])
+        model.infer(obs)
+        torch.cuda.synchronize()
+
+        call_count = 0
+
+        def step() -> None:
+            nonlocal call_count
+            if call_count == args.n_warmup:
+                model.latency_records.clear()
+            model.infer(obs)
+            # FlashRT PI0.5 may use an internal CUDA stream, so synchronize inside the
+            # step to make the common runner's timing boundary conservative.
+            torch.cuda.synchronize()
+            call_count += 1
+
+        spec = bnb.BenchSpec(
+            name="flashrt_pi05",
+            step_callable=step,
+            teardown_callable=lambda: None,
+        )
+        spec.flashrt_internal_latency_ms = LazySummary(
+            lambda: [float(x) for x in getattr(model, "latency_records", [])]
+        )  # type: ignore[attr-defined]
+        return spec
+
+    return setup_fn
+
+
+def make_extras_fn(args: argparse.Namespace):
+    def extras_fn(batch_size: int, spec: bnb.BenchSpec) -> dict[str, Any]:
+        return {
+            "runtime": "FlashRT",
+            "checkpoint": str(args.checkpoint),
+            "flashrt_root": str(args.flashrt_root),
+            "precision": args.precision,
+            "hardware": args.hardware,
+            "batch_size_contract": 1,
+            "num_views": args.num_views,
+            "chunk_size": args.chunk_size,
+            "prompt": args.prompt,
+            "seed": args.seed,
+            "internal_latency_ms": getattr(spec, "flashrt_internal_latency_ms", {}),
+            "timing_scope": "FlashRT direct Pi05TorchFrontendRtx.infer hot path after set_prompt and first graph-building infer; common runner uses perf-counter wall time because FlashRT runs work on an internal CUDA stream",
+        }
+
+    return extras_fn
+
+
+def main() -> None:
+    flashrt_default = env_path("FLASHRT_ROOT")
+    parser = argparse.ArgumentParser(
+        description=__doc__,
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    parser.add_argument(
+        "--flashrt-root",
+        type=Path,
+        default=flashrt_default,
+        required=flashrt_default is None,
+        help="Path to the FlashRT repository clone. Can also be set with FLASHRT_ROOT.",
+    )
+    parser.add_argument(
+        "--checkpoint",
+        type=Path,
+        required=True,
+        help="PI0.5 checkpoint directory readable by FlashRT.",
+    )
+    parser.add_argument("--num-views", type=int, default=2)
+    parser.add_argument("--chunk-size", type=int, default=50)
+    parser.add_argument("--prompt", default="do something")
+    parser.add_argument("--seed", type=int, default=0)
+    parser.add_argument("--precision", choices=("bf16", "fp8_bf16"), default="bf16")
+    parser.add_argument("--hardware", default="auto")
+
+    bnb.add_bench_cli_args(parser)
+    parser.set_defaults(bench_name="flashrt_pi05", batch_sizes=[1])
+    add_profile_cli_args(parser)
+    args = parser.parse_args()
+
+    if not torch.cuda.is_available():
+        raise RuntimeError("CUDA is required for FlashRT PI0.5 benchmarking")
+
+    install_profiler(profile_config_from_args(args))
+    runner = bnb.NBatchBenchRunner(
+        setup_fn=make_setup_fn(args),
+        extras_fn=make_extras_fn(args),
+        # FlashRT may execute on an internal CUDA stream. Force the common
+        # runner's perf-counter path; step() synchronizes CUDA before returning.
+        device=torch.device("cpu"),
+        **bnb.bench_runner_kwargs_from_args(args),
+    )
+    runner.run()
+
+
+if __name__ == "__main__":
+    main()