changh95's picture
tt-model push diffusion-planner-p150 (container)
026560c verified
|
Raw History Blame Contribute Delete
23 kB

Serving Diffusion Planner v5.0 (Autoware diffusion_planner) on Blackhole with tt-model-manager

This repo is the authoring source of the tt-model container package changh95/diffusion-planner-p150 (kind: tt-dit-server, schema 5.1): tt-model.yaml is the manifest, code/ the port, and code/tt_diffusion_planner/server/app.py the ASGI app uvicorn runs. Weights are a pinned pointer, never baked into the image. The device open, the decoders, the Python-API contract and the HTTP contract come from code/tt_diffusion_planner/ttaw/, the shared package of the Autoware ports vendored into this repo (ttaw/VENDORED.json records its version, source commit and file hashes; it is not edited here).

item value
tt-metal tree <workspace>/tt-metal (main 44d66500520, v0.80.0-dev20261006-78-g44d6650052, + patches/tt-metal-eth-dispatch.patch + patches/tt-metal-reshape-rm-sys1419.patch; torch 2.11.0+cpu)
weights AutowareFoundation/diffusion_planner @ 423efde67f5414734da43a7ad856c17ceb8b51aa (tag v5.0; files diffusion_planner_encoder.onnx, diffusion_planner_decoder.onnx, diffusion_planner_turn_indicator.onnx, diffusion_planner.param.json; 58.9 MB; Apache-2.0; public, ungated; sha256 of every file checked at load)
Autoware reference planning/autoware_diffusion_planner @ autoware_universe 9ceaccf026c31ffc5319bc9eeb4bd7bede0af3fd (package 0.53.0), the node's multi_step mode
shared package ttaw 0.23.2 (common @ 60b6dd7)
app tt_diffusion_planner.server.app:app (uvicorn; the ASGI lifespan does weights -> device -> graph -> trace capture)
device recipe ttnn.open_device(device_id, dispatch_core_config=DispatchCoreConfig(ETH), l1_small_size=32768, trace_region_size=192 MiB, num_command_queues=1) (DEVICE_DEFAULTS in code/tt_diffusion_planner/device.py; the 6 traces hold 60.3 MB)
graph ONE metal trace per plan: the encoder (6 MLP-Mixer trunks, small encoders, 6 fusion blocks), the cross-attention K / V hoisted once, 11 DiT evaluations with the 10 fp32 DPM-Solver++(2M) updates and the prefix constraint, the turn-indicator head, one packed readback. Six traces are captured at startup, the full capacity and one per agent bucket (32 / 64 / 96 / 128 / 192 decoder rows); each request replays the smallest that holds its valid neighbour rows (exact). 1,418 device programs at r96 (6,282 in the first release); custom generic_op kernels in code/tt_diffusion_planner/tt/kernels/
hardware one Blackhole p150 (hardware: p150, mesh_device: P150, TT_MESH_SHAPE=1x1), 12×10 compute grid

1. Run on the HOST (hardware validation, no Docker)

The tt-metal python_env has ttnn and torch; bin/tt-env.sh puts the workspace's dev overlay (fastapi, uvicorn, onnx) on PYTHONPATH instead of installing into it. On the shared workspace box every command that opens the chip goes through the device lock (bin/devrun).

ROOT=<workspace>
source $ROOT/bin/tt-env.sh                     # TT_METAL_HOME, PYTHONPATH (+ the dev overlay: fastapi, uvicorn, onnx), python_env
cd $ROOT/bundles/diffusion-planner-p150
export PYTHONPATH=$PWD/code:$PYTHONPATH
export HF_MODEL=AutowareFoundation/diffusion_planner TT_MODEL_WEIGHTS_REVISION=423efde67f5414734da43a7ad856c17ceb8b51aa TT_MESH_SHAPE=1x1

# host tests (no device; the device tests are skipped; ttnn is the fake of common/tests/host: a real ttnn host
# tensor opens the chip, so only devrun jobs may build one)
PYTHONPATH=$ROOT/research/packaging/scripts:$PYTHONPATH TT_VISIBLE_DEVICES=none \
    python -m pytest -q -p no:cacheprovider -p fake_ttnn_plugin code/tt_diffusion_planner/tests
# device tests (per-module PCC + end-to-end agreement; 99 scenes with the workspace goldens, the 2 shipped samples elsewhere)
$ROOT/bin/devrun -t 1800 -- python -m pytest -q -s code/tt_diffusion_planner/tests/test_pcc_device.py code/tt_diffusion_planner/tests/test_e2e_device.py
# the server, with the serve pins of tt-model.yaml, and the container smoke test against it
$ROOT/bin/devrun -t 1800 -- bash -c 'env DIFFUSION_PLANNER_DISPATCH=eth DIFFUSION_PLANNER_NUM_CQS=1 \
    python -m uvicorn --host 127.0.0.1 --port 20000 --lifespan on tt_diffusion_planner.server.app:app & \
    S=$!; python3 code/tt_diffusion_planner/server/smoke_test.py --url http://127.0.0.1:20000 --wait 1200; R=$?; kill -TERM $S; wait $S; exit $R'

Boot log landmarks (they drive tt-model serve's progress view, boot_progress.py TT_DIT_PHASES): Loading weights -> Opening device -> Warming up: capturing trace ... (once per trace variant) -> Warmup complete -> uvicorn Application startup complete. A cold JIT cache compiles every kernel first (352 s measured for from_pretrained); later boots take about 13 s to READY (measured on the host, warm JIT cache). The packaged container measured the same on 2026-10-11: tt-model serve READY after 6 min 0 s on its first start (server boot 358 s, empty JIT cache) and after 20.9 s on the next one (server boot 19.0 s). Startup failures raise and uvicorn exits non-zero (no CPU fallback). SIGTERM: the lifespan releases the traces and closes the chip (within tt-model stop's 120 s budget).

Offline override (no Hub): DIFFUSION_PLANNER_WEIGHTS_DIR=<dir holding diffusion_planner_encoder.onnx, diffusion_planner_decoder.onnx, diffusion_planner_turn_indicator.onnx, diffusion_planner.param.json>.

2. Package, serve, push (Docker)

Docker 29 on this box is rootful (/var/lib/docker), and the user's session may not carry the docker group: wrap every docker-using command in sg docker -c "..." (no docker-env.sh is needed here; the reference bundles' docker-env.sh was the authors' rootless-Docker box). Run every tt-model command from this directory (source.tt_metal and extra_code[].root: code are CWD-relative) and point --out outside it.

ROOT=<workspace>; cd $ROOT/bundles/diffusion-planner-p150
TTM=$HOME/.local/share/uv/tools/tt-model/bin/python

# offline validation (seconds; no docker, no device) -- must print VALID
$TTM -c "from tt_kernel.container_manifest import load_container_manifest as L; m=L('tt-model.yaml', check_sources=True); p=m.resolve_profile(); print('VALID', m.name, m.kind, p.hardware, p.mesh_device, m.weights_ref)"
python3 $ROOT/research/packaging/scripts/check_bundle.py .            # conventions (add --stage publish before a push)

# build (one at a time on the box: the absolute lock file): ~35 min cold, ~12 min after a code/ change (the C++
# build re-runs, ccache-warm, whenever code/ changes), ~1 min for a manifest- or lock-only change
flock $ROOT/.package.lock sg docker -c "tt-model package --container tt-model.yaml --out $ROOT/build"   # log: ~/.cache/tt-model/build/diffusion-planner-p150.log
# provenance: the base-image digests this build resolved (tt-model's FROM tags float) -> build_info.json (published)
sg docker -c "python3 $ROOT/research/packaging/scripts/record_build_info.py --staged $ROOT/build/diffusion-planner-p150 --bundle ."

# serve + smoke + stop, all inside ONE device-lock window (there is one serve profile, the default; serve returns
# once READY and leaves the container running; the script's trap saves the container log and always runs
# `tt-model stop`, also when devrun's timeout fires: -k 150 > the 120 s grace). No token reaches the public weights
# download. Evidence: logs/smoke/ (--log-dir).
$ROOT/bin/devrun -t 3600 -k 150 -- env -u HF_TOKEN -u HUGGING_FACE_HUB_TOKEN \
    sg docker -c "bash code/scripts/container_smoke.sh $ROOT/build/diffusion-planner-p150 20000"

# publish: fill the provenance from the staged manifest, run the publish gate, overlay the hand-written docs, the pip
# project and the build record onto the staged dir, then push (HF_TOKEN only from the environment)
python3 $ROOT/research/packaging/scripts/instantiate_bundle.py --values $ROOT/research/diffusion-planner/bundle_values.json \
    --fill . --manifest $ROOT/build/diffusion-planner-p150/tt_kernel_manifest.json
python3 $ROOT/research/packaging/scripts/check_bundle.py . --stage publish
rsync -a README.md SERVING.md OPT_BASELINE.md OPT_REPORT.md VERIFICATION_*.md tt-model.yaml pyproject.toml build_info.json \
    examples media patches $ROOT/build/diffusion-planner-p150/
cp code/PYTHON.md $ROOT/build/diffusion-planner-p150/PYTHON.md
tt-model push $ROOT/build/diffusion-planner-p150 --publish

The smoke test fails unless /info reports ETH dispatch, the 12x10 grid and what the staged package pins (DIFFUSION_PLANNER_DISPATCH, DIFFUSION_PLANNER_NUM_CQS, DIFFUSION_PLANNER_VARIANT, the weights revision), and it compares the served trajectory of the shipped sample with the stored CPU-reference output (code/tt_diffusion_planner/samples/kashiwanoha_dense.reference.json: average displacement <= 0.3 m, final displacement <= 1.0 m). On the host, the same smoke test against uvicorn with the serve pins passed: PASS diffusion-planner-p150: profile=- variant=default dispatch=eth grid=12x10 cqs=1 n=80 turn=DISABLE reference=kashiwanoha_dense.reference.json (trajectory ade=0.0109 fde=0.0173; predicted_agents max_abs=6.283) device_ms=109.402 total_ms=133.384 rtt_ms=139.0

(predicted_agents max_abs is reported, not gated. Here it is the yaw column at ±π: -3.1413 vs +3.1415 rad, the same heading. The neighbour paths are gated in the end-to-end tests, on the per-agent maximum displacement.)

serve publishes the first free port from 20000 and exports into the container: HF_MODEL=AutowareFoundation/diffusion_planner, MESH_DEVICE=P150, TT_MESH_SHAPE=1x1, TT_MODEL_WEIGHTS_REVISION=423efde67f5414734da43a7ad856c17ceb8b51aa, then serve.env. The HF cache is mounted at /hf (rw), the JIT cache at /cache (host ~/.cache/tt-model/diffusion-planner-p150/cache), and the container sees only its own chip (--device /dev/tenstorrent/<n>). tt-cli users: tt serve changh95/diffusion-planner-p150 / tt model stop changh95/diffusion-planner-p150 (point tt at this tt-model with tt config set tools.override.tt-model ~/.local/bin/tt-model or TT_TOOL_BIN_TT_MODEL).

What push does to the repo: code/ and image/ on the Hub become exactly the staged trees (extra_code.paths + models/common/lightweightmodule.py); every top-level file of the staged dir is uploaded (the overlaid README replaces the generated card; pyproject.toml, PYTHON.md and build_info.json arrive the same way, never through code/); top-level files already on the Hub and not in the staged dir are kept.

3. The request / response contract

route purpose
GET /health {"status": "ok" | "starting" | "error", "model", "device": {"dispatch", "grid", "cores", "device_id"}, "error"}, always 200; ok only after warm-up
GET /v1/health same as /health (tt-model's hint for non-chat packages points here)
GET /info model, task, io, Autoware package + commit, weights {repo, tag, revision, path, license}, device (dispatch, grid, CQs, fallback), the input schema (the 15 tensors with shapes and dtypes), labels (turn-indicator logit order), variant, warm variants, runtime params + defaults, the numerics options and precision policy in effect, the trace (variants, persistent inputs, trace buffers), warm-up and boot times
GET /v1/models {"object": "list", "data": [{"id": "AutowareFoundation/diffusion_planner", "object": "model", "owned_by": "changh95"}]} -- a stub so OpenAI-shaped probes do not 404; NOT a chat API
POST /predict one plan -> the model output (below)

3.1 Request (application/json; unknown fields are a 422)

field type meaning
inputs object the planner tensors: {"format": "npz", "data": <base64 .npz>} or {"format": "json", "arrays": {name: nested list}} with exactly the 15 raw tensors of the node's create_input_data(): sampled_trajectories [1,321,81,4], ego_agent_past [1,31,4], ego_current_state [1,10], neighbor_agents_past [1,320,31,11], static_objects [1,5,10], lanes [1,140,20,33], lanes_speed_limit [1,140,1], route_lanes [1,25,20,33], route_lanes_speed_limit [1,25,1], polygons [1,10,40,3], line_strings [1,60,20,4], goal_pose [1,4], ego_shape [1,3], turn_indicators [1,31], delay [1,1]; float32, ego (base_link) frame, BEFORE normalization; names, shapes and finite values are checked (400)
params object per-request knobs (host post-processing only): velocity_smoothing_window (8, 1..79), stopping_threshold (0.3 m/s), turn_indicator_keep_offset (-1.25), return_denoising_steps (false)
output_format str json (default) or npz (adds the poses as a lossless base64 NPZ array)

Base64 may be standard or URL-safe, wrapped, with or without a data: prefix. python3 code/tt_diffusion_planner/server/client.py --inputs <scene.npz> [--param name=value] --out req.json builds the request with the standard library only (no numpy); --url http://127.0.0.1:20000 sends it.

3.2 Response (200)

The served body of the shipped sample (uvicorn on the host with the serve pins; trajectory cut to 3 of 80 rows and the predicted-agents array to its header):

{
 "model": "diffusion-planner-p150",
 "frame_id": "base_link",
 "meta": {"predicted_agent_columns": ["x", "y", "yaw", "cos", "sin"], "predicted_agent_rows": [0, 1, 2, "... 85 more"], "force_stop": false, "time_from_start_s": [0.1, 0.2, 0.3, "... 77 more"], "valid_counts": {"ego": 1, "neighbor": 88, "static": 0, "lane": 123, "route": 17, "polygon": 0, "line_string": 60, "goal": 1, "ego_shape": 1, "turn": 1}},
 "timing_ms": {"preprocess": 2.4, "device": 22.8, "postprocess": 1.2, "total": 36.9, "decode": 8.4, "model_call": 26.7},
 "num_poses": 80,
 "columns": ["x", "y", "yaw", "cos", "sin", "velocity", "acceleration"],
 "trajectory": [[0.3644, 0.01, -0.0169, 0.9964, -0.0169, 3.8033, -0.0392], [0.7626, -0.0021, -0.0412, 0.9968, -0.0412, 3.7994, -0.5496], [1.1517, -0.0281, -0.0674, 0.9942, -0.0672, 3.7444, -0.4683], "... 77 more rows"],
 "turn_indicator": {"command": 1, "command_name": "DISABLE", "keep_selected": true, "held": false, "logits": [-15.960197448730469, -5.119836807250977, -5.23490047454834, -0.7636525630950928, 5.502235412597656], "probabilities": [1.655269699085693e-09, 8.448460721410811e-05, 7.530191942350939e-05, 0.00658634165301919, 0.9931545257568359]},
 "predicted_agents": {"format": "npz", "key": "predicted_agents", "dtype": "float32", "shape": [88, 80, 5], "data": "<base64 npz>"}
}
  • trajectory: 80 points at 0.1-8.0 s in base_link (the ego frame of the input tensors), post-processed like the node's ~/output/trajectory: velocity from consecutive 3-D points / 0.1 s, forward moving average over velocity_smoothing_window points, force stop (poses frozen once the smoothed speed falls below stopping_threshold while the ego moves), acceleration by finite difference. yaw is what tf2::getYaw reads from the node's (unnormalised) quaternion of the raw network cos / sin, which are returned as they are.
  • turn_indicator: command 0 NO_COMMAND, 1 DISABLE, 2 ENABLE_LEFT, 3 ENABLE_RIGHT (when the head selects KEEP, the command repeats the last input report turn_indicators[0, 30]); the 5 raw logits (NONE, DISABLE, LEFT, RIGHT, KEEP) and the decision's probabilities. No hold window is applied (see 3.5).
  • predicted_agents: one 80-point path (x, y, yaw, cos, sin) per non-empty neighbour row, in input order (meta.predicted_agent_rows gives the rows).
  • meta: force_stop, time_from_start_s, valid_counts (entities the encoder saw), and with return_denoising_steps the ego row of the 11 solver iterates (denoising_steps, [11, 81, 4], the node's ~/debug/denoising_steps).

timing_ms: decode (base64 + npz parsing + the schema check), preprocess (the node's normalization and the encoder's host features), device (host tensors + H2D + one trace replay of the request's agent bucket + D2H), postprocess, model_call, total (server side, after the body arrived).

3.3 Errors

400 undecodable or malformed input (bad base64, a missing or unknown tensor, a wrong shape, a non-finite value, unknown or out-of-range params), 422 schema violation (unknown field, wrong type), 503 while starting (or when the boot failed), 500 inference failed: <ExceptionType>: <message> if the device call raises.

3.4 Environment the app reads (lifespan only, never at import)

var set by meaning / default
HF_MODEL launcher (weights.repo) weights repo id; default AutowareFoundation/diffusion_planner
TT_MODEL_WEIGHTS_REVISION launcher (weights.revision) pinned commit; TT_WEIGHTS_REVISION (serve.env) is the same for older tt-model
DIFFUSION_PLANNER_WEIGHTS_DIR you (offline) local weights directory; overrides the Hub
TT_MESH_SHAPE launcher (runtime.mesh_shape_env) 1x1; anything else -> RuntimeError at startup
TT_DEVICE_ID you chip to open, default 0
DIFFUSION_PLANNER_DISPATCH serve.env eth (default) | worker (A/B only, 11×10) | auto (ETH if the patch is present)
DIFFUSION_PLANNER_NUM_CQS serve.env 1 (at the first release 2 CQs replayed ~10 % slower on ETH, OPT_BASELINE.md; not re-measured since)
DIFFUSION_PLANNER_VARIANT serve.env default (the only v5.0 graph)
DIFFUSION_PLANNER_LN_FP32, _SPLIT_MATMUL, _ATTN_MATMUL, _HIDDEN_FP32, _ATTN_FP32_ACC serve.env the numerics knobs, pinned at their gated defaults (enc.mixer.*,dec.*; enc.island.*,enc.pre.*,dec.*; enc.fusion.attn,dec.*; empty; empty: tt/config.py KNOBS); changing one changes the numerics and needs the gates re-run
DIFFUSION_PLANNER_ENC_CH2D, _ATTN_FAST, _DEC_MMCFG, _LN_KERNEL, _LN_RESID, _SPLIT_KCAT, _ATTN_SMASK, _ATTN_SMSM, _ATTN_FUSED, _KCAT_EMIT, _LN_TR, _KCAT_L1, _ATTN_L1, _DEC_L1, _ENC_KCAT, _KCAT_ACT, _LN_SPLIT, _LN_SFPU_BCAST, _COMPACT, _AGENT_BUCKETS, _HOST_FAST, _INPUT_TRIM (and the rejected _KCAT_ACT_ONCE, _ENC_L1, _FUS_L1, _LIN_ACT = 0) serve.env the optimization knobs, pinned at the shipped defaults (each one an A/B of OPT_REPORT.md; tt/config.py KNOBS); the gates were run with exactly these values
DIFFUSION_PLANNER_PRECISION serve.env empty: the precision policy DEFAULT_PRECISION of tt/config.py (HiFi4 + fp32 accumulation everywhere); extra rules such as dec.*=HiFi2+fp32 are for experiments
DIFFUSION_PLANNER_WARMUP you JSON list of warm-up variants, default or none
DIFFUSION_PLANNER_TRACE_REGION, DIFFUSION_PLANNER_L1_SMALL, DIFFUSION_PLANNER_WORKER_L1_SIZE you device-open overrides (validated values: DEVICE_DEFAULTS in device.py)
DIFFUSION_PLANNER_MAX_BODY_MB you request size guard, default 256 (the shipped sample's request is 0.15 MB of base64 npz)
TT_METAL_VISIBLE_DEVICES, MESH_DEVICE, HF_HUB_DISABLE_IMPLICIT_TOKEN serve.env / launcher 0, P150, 1 (the weights are public: no token is ever sent)

Server-side pipeline: JSON -> tt_diffusion_planner.io decoders (schema check) -> DiffusionPlanner.__call__ (the node's host pre-processing ported from planning/autoware_diffusion_planner -> H2D into the persistent device inputs -> execute_trace of the plan's agent-bucket trace -> one D2H -> the node's host post-processing) under one lock -> Output.to_dict().

3.5 What the client keeps (the API is stateless)

Every request is one independent plan. The node keeps state between plans; a client that wants its behaviour over a sequence of plans keeps the same state and puts it into the request or applies it to the response:

  • The input tensors. Converting ROS messages and the Lanelet2 map into the 15 tensors is the client's: the per-UUID agent buffers with their 0.1 s resampling, the ego history, lane / route / polygon / line-string selection and encoding, traffic-light states, speed limits, the goal and the turn-indicator report history.
  • The turn-indicator hold window. The node holds its last non-KEEP command for turn_indicator_hold_duration (1.0 s in its YAML). The server evaluates every request with a fresh manager (no held command). To reproduce the node, keep the last non-KEEP command and its time, and return it instead while less than 1.0 s has passed; the node's manager ships in the package (tt_diffusion_planner.host.postprocess.TurnIndicatorManager, fed with turn_indicator.logits, the plan time and the last report).
  • The initial solver state x_T (sampled_trajectories, normalised space). Zeros is the node's default (temperature: [0.0]). With a temperature > 0 the node draws N(0, 1) x temperature for every element (a fresh std::random_device seed per call); send that noise in sampled_trajectories. With the RTC prefix (delay_step > 0, at most 40) the node also writes its previous plan into the ego row, slots t = 0 .. delay_step: the previous poses re-expressed in the current ego frame, x as (x - 10) / 20, y as y / 20, cos / sin as they are. Non-zero x_T was checked on the device against ONNX Runtime (noise of scale 0.5 and 1.0, an RTC prefix: within the gates, VERIFICATION_2026-10-08.md).
  • delay is accepted and ignored, as in the node's multi-step mode.
  • The map frame. The output is in base_link (the ego frame of the inputs); the node transforms it to map with the current ego pose before publishing.

4. Caveats

  • Batch 1, one plan per request; concurrent clients queue on the lock.
  • Fixed shapes of the v5.0 export (320 neighbours, 140 lanes, 25 route lanes, 10 polygons, 60 line strings, 31 history and 80 future steps) and 10 DPM-Solver steps (11 decoder evaluations) are compiled into the traces. The neighbour trunk and the decoder run on the smallest agent bucket that holds the scene (exact; one trace per bucket), so the device time depends on the neighbour count (17.4-26.7 ms per plan); the map entities are computed at full capacity.
  • The node's guidance services (start / stop / centerline guidance) are off (the node's default too); a guided mode would need the host in the solver loop.
  • Warm-up captures the traces inside the lifespan, so READY means warm; the first cold boot pays the ttnn JIT.
  • Weights are pinned by sha in weights.revision, exported by the launcher and repeated in serve.env; snapshot_download(..., revision=<sha>) is a cache hit after serve's pre-download and falls back to local_files_only=True if the Hub is unreachable. The sha256 of the four files and major_version == 5 are checked at load.
  • The runtime image has no host C/C++ compiler: device kernels JIT-compile (sfpi ships in the image); the host code is numpy.
  • tt-model curl and the ready card's /v1/models hint are OpenAI-shaped and are not this API; use the routes above.