# Serve pi-0.5 (`lerobot/pi05_base`) on 2x Blackhole p300 with tt-model-manager

This repo is a **tt-model container package** (`tt-model.yaml`, the port in `code/`, the OCI image in `image/`). The
server is `code/models/experimental/pi0_5/server/app.py` (FastAPI, `kind: tt-dit-server`). The model is
`code/models/experimental/pi0`: the multi-config pi0.5 megakernel of the p150 package plus the single-robot mesh, the
batched expert, the multi-chip runtime and the multi-robot scheduler of this package.

## Terms

- Preset: one combination of a camera count (1-4), a prompt bucket (32 / 64 / 128 / 224 tokens, plus the extra buckets
  c1 192, c2 160, c4 96) and an action-row bucket (32 rows for H ≤ 32, 64 rows for H 33-64). The device programs of a
  preset are compiled from these parameters.
- Slab: one request's 18 K and 18 V caches in one DRAM tensor (`KVSlab`). The PREFIX program writes it, the batched
  expert reads it; between chips the host copies it.
- Batched expert: one EXPERT program for `nb` requests (`program_b.BatchedExpertMegakernel`). With 32 action rows two
  requests share one 64-row slot ("packing").
- SP slice: the contiguous rows of the PREFIX that one chip of the single-robot mesh computes.
- Cool / thermal steady state: the first minutes of load vs after ≥ 12 min, when the hottest chips are clamped to
  800 MHz by the firmware's GDDR thermal limit (QB2 fans are fixed). Numbers are labelled with the state they were
  measured in.

## Profiles

`tt-model serve` exports the profile's env (each sets `TT_MESH_SHAPE=1x4`). The megakernel backend is selected with
`PI05_MEGAKERNEL=mc`:

```
multi-robot     PI05_MEGAKERNEL=mc TT_MESH_SHAPE=1x4 PI05_LAYOUT=dp PI05_SLO_MS=100 PI05_SLO_P99_MS=110 PI05_SLO_CLOCK_GUARD=1.69
single-robot    PI05_MEGAKERNEL=mc TT_MESH_SHAPE=1x4 PI05_LAYOUT=sp
disaggregated   PI05_MEGAKERNEL=mc TT_MESH_SHAPE=1x4 PI05_LAYOUT=pd PI05_MAX_BATCH=4
```

(and for all: `PI05_NUM_IMAGES=2 PI05_ACTION_HORIZON=50 PI05_NUM_STEPS=10 PI05_TOKEN_LEN=224 PI05_BATCH_WINDOW_MS=0`
and the hang watchdog `PI05_HANG_S=10 PI05_HANG_ACTION=degrade PI05_RESET_ON_START=auto
PI05_HANG_MARKER=/cache/pi05_hang_marker`; override any of them with `-e` / the profile env.)

- `single-robot` (`server/mc_backend.py` -> `PI05SingleRobotMesh.open`): the model enables `FABRIC_1D_RING`, opens the
  1x4 mesh itself and closes mesh and fabric at shutdown. `/info` `megakernel.sp_placement` shows the chip of every
  mesh coordinate, the hot chip(s), the EXPERT chip and every preset's slices / camera chips.
- `multi-robot` (`server/mcx_backend.py` -> `PI05MultiChip(layout="dp")`): every chip opened as its own device. With
  `PI05_SLO_MS` the server calibrates every chip after warm-up (logged as `SLO calibration k=...` lines) and binds robots
  to chips; `/info` `megakernel.slo` shows the limits, per-chip timing and speed factors, the router and the latency
  counters (device side and HTTP side). The default SLO (p50 100 / p99 110 ms, clock guard 1.69) fits one robot per
  chip at 1-2 cameras / H 50 on cool chips. At 3-4 cameras, or when started on hot chips (e.g. a restart right after
  sustained load), no robot fits and the server refuses to start: raise `PI05_SLO_MS` / `PI05_SLO_P99_MS`, lower
  `PI05_SLO_CLOCK_GUARD`, or start from cool chips.
- `disaggregated` (`PI05MultiChip(layout="pd")`): chips 0..2 prefix, chip 3 batched expert; requests are accepted at any time and the
  expert chip batches the ready prefixes (up to `PI05_MAX_BATCH`). No SLO or admission control; every request uses the
  server's config (per-request heterogeneous configs and join / leave LLM-style serving are future work).
- Without `TT_MESH_SHAPE` (or `1x1`) and `PI05_MEGAKERNEL=mc`: the single-chip `PI05MegakernelTTNN` on `TT_DEVICE_ID`
  (the p150 path).

## Environment

| Variable | Values | Default | Meaning |
|---|---|---|---|
| `PI05_NUM_IMAGES` | 1-4 | 2 | cameras; a request must send exactly this many images |
| `PI05_ACTION_HORIZON` | 1-64 | 50 | action chunk H (32-row bucket for H ≤ 32, else 64) |
| `PI05_NUM_STEPS` | 1-10 | 10 | flow-matching steps N |
| `PI05_TOKEN_LEN` | ≤ 224 | 224 | the tokenizer's buffer; a request runs in the smallest prompt bucket that holds its real tokens (or `prompt_bucket`) |
| `TT_MESH_SHAPE` | `1x1`, `1x4` | from the profile | `1x1`: one chip; `1x4`: the 4 chips (layout from `PI05_LAYOUT`) |
| `PI05_LAYOUT` | `sp`, `dp`, `pd` | `dp` on 1x4 | `sp`: single-robot mesh; `dp`: every chip the whole model; `pd`: prefix / expert disaggregation |
| **single-robot** | | | |
| `PI05_SP_PLACEMENT` | `auto`, `legacy` | `auto` | `auto`: thermal-aware placement (`tt/sp_thermal.py`); `legacy`: the pre-thermal layout (EXPERT on coordinate 0, even split), for A / B |
| `PI05_SP_HOT` | `<asic id or PCI BDF>,...`, `none` | measured box rank | the hot chips, hottest first (by sysfs `tt_asic_id`, stable across reboots; /dev order is not); `none`: no hot chip |
| **multi-robot** | | | |
| `PI05_SLO_MS` | ms | unset | p50 latency target: calibrate, derive robots per chip / batch cap / R_max, bind robots, 503 beyond R_max. Unset: per-request least-loaded routing, `PI05_MAX_BATCH` |
| `PI05_SLO_P99_MS` | ms ≥ SLO | `PI05_SLO_MS` | p99 target |
| `PI05_SLO_MARGIN_MS` | ≥ 0 | 5 | p99 margin over the calibrated cycle (at least the calibration jitter) |
| `PI05_SLO_CLOCK_GUARD` | ≥ 1 | 1 | factor on every chip's calibrated cycle; 1.69 = 1350 / 800 MHz, for sustained full load (the profile sets it) |
| `PI05_MR_GUARD` | unset / `none`, `hot`, `<asic id or BDF>:<factor>,...` | off | per-chip guards instead (`hot`: the SP hot rank's first chip 1.69). Under full multi-robot load 3 of 4 chips clamp, so `hot` alone is too optimistic; opt-in |
| `PI05_SLO_HOST_MS` | ms | measured at start | per-request host time (decode, tokenize, prepare) added to the predictions |
| `PI05_ROUTING` | `robot`, `least` | `robot` with an SLO | bind robots to chips / per-request least-loaded |
| `PI05_ROBOT_IDLE_S` | s | 5 | a robot without requests this long frees its slot |
| `PI05_JOIN` | `work`, `lateness` | `work` | batch join rule; `lateness` for robots with think time (needs the SLO) |
| `PI05_STEAL` | 0 / 1 | 0 | an idle chip takes requests that would miss their deadline (needs the SLO) |
| `PI05_MC_LOOP` | `v2`, `v1` | `v2` | chip loop: chained prefills / the previous loop |
| `PI05_MAX_BATCH` | 1-8 | 1 (`dp`), 4 (`pd`); with an SLO: the L1 maximum | requests per batched expert (refused above what L1 holds: 4-7 at 64 rows, 8 at 32 rows) |
| `PI05_EXPERT_PACK` | `auto`, `0`, `1` | `auto` | two H ≤ 32 requests per 64-row slot; `0` = exact mode (every request bit-identical to the single request) |
| `PI05_BATCH_WINDOW_MS` | ≥ 0 | 4 (profiles set 0) | `v1` loop's batch window |
| `PI05_HTTP_THREADS` | n | 128 (multi-chip) | request threads |
| **disaggregated** | | | |
| `PI05_PREFIX_CHIPS` | n-1 | n-1 | prefix chips (one expert chip) |
| **hang watchdog** (every layout; README "Stability on QuietBox 2") | | | |
| `PI05_HANG_S` | s, `0` = off | 10 | a device call stuck this long is a hang (must exceed `PI05_BATCH_WINDOW_MS` + 2 s). `dp`: the chip is quarantined (also at once on a tt-metal dispatch timeout), its requests get 503 + Retry-After, its robots are re-bound, `/health` = `degraded`. `sp`, `pd`, 1 chip: `/health` = `hung` and the server exits with `PI05_HANG_EXIT_CODE` |
| `PI05_HANG_ACTION` | `degrade`, `exit` | `degrade` | `dp` after a quarantine: keep serving on the healthy chips / drain for up to `PI05_HANG_DRAIN_S` (5 s) and exit with `PI05_HANG_EXIT_CODE` for a supervisor restart. With no healthy chip left it always exits |
| `PI05_HANG_EXIT_CODE` | code | 75 | exit code after a hang; a SIGTERM after a quarantine also exits with it, without closing the devices (a hung chip's close never returns) |
| `PI05_RESET_ON_START` | `auto`, `1`, `0` | `auto` | before opening the chips: `tt-smi -r` when the hang marker exists (`auto`) / always / never; then every chip's heartbeat and an open / close check, up to 3 attempts. Refused while another process holds the chips |
| `PI05_HANG_MARKER` | path | `~/.cache/pi05/hang_marker` (profiles: `/cache/pi05_hang_marker`) | written on a hang, removed after a good reset. In a container it must be on a persistent volume |
| **A / B and debugging** | | | |
| `PI05_PE_DEFINES` | `NAME=VALUE,...` | unset | extra VISION / PREFIX kernel defines, e.g. `PE_NO_XRES=1` (the fp32 residual through DRAM instead of L1: the A / B arm of the energy work; same outputs, slower), `PE_KV_SERIAL=0` (the previous release's concurrent K / V feed: same outputs, hangs within minutes under load; only for reproducing the hang) |

## Run on the host (no Docker)

```bash
TREE=<a tt-metal tree at f856a38a361, built (./build_metal.sh) with its python_env>
cp -r code/models/experimental/pi0 code/models/experimental/pi0_5 $TREE/models/experimental/
source $TREE/python_env/bin/activate
uv pip install fastapi uvicorn "pydantic>=2" pillow
export TT_METAL_HOME=$TREE PYTHONPATH=$TREE ARCH_NAME=blackhole
export PI05_WEIGHTS_DIR=~/.cache/huggingface/hub/models--lerobot--pi05_base/snapshots/b211f3d44c36b6acfcf7ae94a64e8e96f75a64ba
export PI05_MEGAKERNEL=mc TT_MESH_SHAPE=1x4 PI05_NUM_IMAGES=2 PI05_ACTION_HORIZON=50 PI05_NUM_STEPS=10 PI05_TOKEN_LEN=224
# single-robot (4-chip mesh)
PI05_LAYOUT=sp python -m uvicorn --host 0.0.0.0 --port 20000 --lifespan on models.experimental.pi0_5.server.app:app
# multi-robot with an SLO
PI05_LAYOUT=dp PI05_SLO_MS=100 PI05_SLO_P99_MS=110 PI05_SLO_CLOCK_GUARD=1.69 PI05_BATCH_WINDOW_MS=0 \
  python -m uvicorn --host 0.0.0.0 --port 20000 --lifespan on models.experimental.pi0_5.server.app:app
# other shell
python code/models/experimental/pi0_5/server/smoke_test.py --url http://127.0.0.1:20000 \
  --images media/sample_base.png media/sample_wrist.png
# closed-loop robots against the server (HTTP mode: robot_id per robot, 503 retires a robot)
python -m models.experimental.pi0.tests.perf.bench_multirobot --url http://127.0.0.1:20000 --robots 4,5 --duration 30 --repeats 1
```

Only one process may use the chips at a time. If a run hangs: `tt-smi -r`, then open and close the 4 chips once (the
server does this itself on its next start, `PI05_RESET_ON_START`). Never run `tt-smi`, tt-triage or another tool that
opens the chips while a server runs: read hwmon / sysfs instead (README "Stability on QuietBox 2").

Supervising the server: on a hang `single-robot` and `disaggregated` exit with code 75, and so does `multi-robot` with
`PI05_HANG_ACTION=exit` or once no healthy chip is left. A supervisor (systemd, a container restart policy; not part
of `tt-model serve`) should restart the server on that exit code, and should also give each start a timeout (~7 min;
a normal start is ready in ~1-3 min, including ~51 s for a reset): the model build has no watchdog of its own and
after a host reboot 3 of 13 builds on this box hung in the weight upload.

In Python:

```python
from models.experimental.pi0.common.weight_loader import PI0WeightLoader
from models.experimental.pi0.tt.ttnn_pi05_sp import PI05SingleRobotMesh
from models.experimental.pi0.tt.ttnn_pi05_multichip import PI05MultiChip

# one robot on all 4 chips (opens FABRIC_1D_RING + the 1x4 mesh; close() releases both)
model = PI05SingleRobotMesh.open(config, PI0WeightLoader(weights))
model.warmup()
actions = model.sample_actions(images, img_masks, lang_tokens, lang_masks=lang_masks, noise=noise)
model.close()

# many robots: every chip the whole model
mc = PI05MultiChip(config, PI0WeightLoader(weights), device_ids=[0, 1, 2, 3], layout="dp", max_batch=4)
ticket = mc.submit(images, lang_tokens, lang_masks, noise)      # thread-safe
actions = ticket.wait()                                          # [1, H, 32]; ticket.batched_as
mc.close()
```

## Tests

| Test | Device | What |
|---|---|---|
| `tests/pcc/test_pi05_megakernel_host.py`, `test_pi05_sp_host.py`, `test_sp_thermal_host.py`, `test_mr_sched_host.py`, `test_bench_multirobot_host.py`, `test_energy_profile_host.py`, `pi0_5/tests/test_server_mc.py` | none | host proofs and policy tests, including the hang watchdog (562 tests with `test_fused_host.py` and `test_server_masks.py`) |
| `tests/pcc/test_pcc_pi05_megakernel.py` | 1 chip | the single-chip megakernel against the fp32 reference (extra buckets included) |
| `tests/pcc/gen_sp_golden.py` + `test_pcc_pi05_sp.py` | 1 chip, then 4 | single-robot mesh `torch.equal` to the 1-chip golden (every preset, closed loops, rebuild, exchange timeout) |
| `tests/pcc/test_pcc_pi05_batched.py`, `check_batched_vs_single.py`, `check_multichip_vs_single.py` | 1 / 4 chips | batched and multi-chip requests vs the single request |
| `tests/perf/bench_pi05_sp.py` | 4 chips | single-robot mesh vs one chip, trace and host-to-host p50 / p99 |
| `tests/perf/bench_multirobot.py` | 4 chips | closed-loop robots under an SLO, in-process or against a server (`--url`) |
| `tests/perf/energy_profile.py` | 1-4 chips | time / power / energy per call, cool and at thermal steady state; interleaved A / B with CIs |
| `tests/perf/mr_hang_probe.py` | 4 chips | hang-rate probe of the 4-robot loop (hangs per load window, state capture on a hang) |

## Layout of `code/`

```
code/models/experimental/pi0/
  tt/ttnn_pi05_model.py          PI05MegakernelTTNN: the single-chip model
  tt/ttnn_pi05_sp.py             PI05SingleRobotMesh: the single-robot mesh (1x4 FABRIC_1D_RING)
  tt/sp_thermal.py               thermal-aware chip placement (hot rank, slices, EXPERT chip, per-chip MR guards)
  tt/megakernel/mesh_launch.py   the mesh launches (VISION, relay, SP PREFIX, EXPERT), SP split, K / V exchange plan
  tt/megakernel/kernels_mesh/    relay and no-op kernels
  tt/ttnn_pi05_batched.py        ModelShape / ExpertLayout, PrefixStage, ExpertStage, SlotIO, PI05BatchedTTNN
  tt/ttnn_pi05_multichip.py      PI05MultiChip: DPWorker, PrefixWorker, ExpertWorker (dp, pd)
  tt/mr_sched.py                 multi-robot policy: cycle model, calibration limits, robot router, admission
  tt/megakernel/program_b.py     BatchedExpertMegakernel, KVSlab
  tt/megakernel/README.md        the programs, the presets, every measurement
code/models/experimental/pi0_5/server/
  app.py, mc_backend.py          the server and the single-chip / single-robot-mesh backend
  mcx_backend.py                 the multi-chip backend (dp with the SLO policy, pd)
```