tt-model push diffusion-planner-p150 (container)
Browse filesThis view is limited to 50 files because it contains too many changes. See raw diff
- .gitattributes +24 -0
- OPT_BASELINE.md +535 -0
- OPT_REPORT.md +142 -0
- PYTHON.md +179 -0
- README.md +195 -0
- SERVING.md +242 -0
- VERIFICATION_2026-10-08.md +210 -0
- build_info.json +48 -0
- code/PYTHON.md +179 -0
- code/conftest.py +59 -0
- code/models/common/lightweightmodule.py +12 -0
- code/scripts/README.md +16 -0
- code/scripts/bench.py +207 -0
- code/scripts/bringup_device.py +167 -0
- code/scripts/container_smoke.sh +137 -0
- code/scripts/fetch_samples.sh +24 -0
- code/scripts/precision_exp.py +147 -0
- code/scripts/profile_ops.py +250 -0
- code/scripts/ref_golden.py +169 -0
- code/scripts/split_error.py +124 -0
- code/tt_diffusion_planner/__init__.py +41 -0
- code/tt_diffusion_planner/api.py +125 -0
- code/tt_diffusion_planner/calib/README.md +6 -0
- code/tt_diffusion_planner/device.py +61 -0
- code/tt_diffusion_planner/host/__init__.py +27 -0
- code/tt_diffusion_planner/host/features.py +203 -0
- code/tt_diffusion_planner/host/normalize.py +67 -0
- code/tt_diffusion_planner/host/pipeline.py +107 -0
- code/tt_diffusion_planner/host/postprocess.py +222 -0
- code/tt_diffusion_planner/host/solver.py +273 -0
- code/tt_diffusion_planner/io.py +32 -0
- code/tt_diffusion_planner/reference/__init__.py +34 -0
- code/tt_diffusion_planner/reference/config.py +175 -0
- code/tt_diffusion_planner/reference/goldens.py +162 -0
- code/tt_diffusion_planner/reference/model.py +226 -0
- code/tt_diffusion_planner/reference/ort.py +137 -0
- code/tt_diffusion_planner/reference/pipeline.py +123 -0
- code/tt_diffusion_planner/reference/rewrites.py +183 -0
- code/tt_diffusion_planner/reference/weights.py +442 -0
- code/tt_diffusion_planner/samples/README.md +23 -0
- code/tt_diffusion_planner/samples/kashiwanoha_dense.npz +3 -0
- code/tt_diffusion_planner/samples/kashiwanoha_dense.reference.json +0 -0
- code/tt_diffusion_planner/samples/straight_road.npz +3 -0
- code/tt_diffusion_planner/samples/straight_road.reference.json +887 -0
- code/tt_diffusion_planner/server/__init__.py +5 -0
- code/tt_diffusion_planner/server/app.py +51 -0
- code/tt_diffusion_planner/server/client.py +22 -0
- code/tt_diffusion_planner/server/smoke_test.py +162 -0
- code/tt_diffusion_planner/tests/__init__.py +1 -0
- code/tt_diffusion_planner/tests/_research.py +57 -0
.gitattributes
CHANGED
|
@@ -33,3 +33,27 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
|
|
| 33 |
*.zip filter=lfs diff=lfs merge=lfs -text
|
| 34 |
*.zst filter=lfs diff=lfs merge=lfs -text
|
| 35 |
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 33 |
*.zip filter=lfs diff=lfs merge=lfs -text
|
| 34 |
*.zst filter=lfs diff=lfs merge=lfs -text
|
| 35 |
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
| 36 |
+
image/blobs/sha256/069395655fd57335d9426528f7f49731be6cf7ed963ca649e4ff7c9697d37e0a filter=lfs diff=lfs merge=lfs -text
|
| 37 |
+
image/blobs/sha256/0936b6c31255cfe6d029dab00fc251d8771f5e3d228ac7da23e395e1161c6992 filter=lfs diff=lfs merge=lfs -text
|
| 38 |
+
image/blobs/sha256/3ac9de4540f40bdf99f054641b2fb36a38f49253d1727fb08a73272d4208ce1b filter=lfs diff=lfs merge=lfs -text
|
| 39 |
+
image/blobs/sha256/46f21edbef18d35cb9e2f400eeb6d8f8fac5873f8620e8d96829a9c01447ed52 filter=lfs diff=lfs merge=lfs -text
|
| 40 |
+
image/blobs/sha256/98c4455a98982b35380ec62f90d5cf0edff02eeb4b3bb18595d871fdeb0c0c67 filter=lfs diff=lfs merge=lfs -text
|
| 41 |
+
image/blobs/sha256/a15921f6aeba449853d8dfc7215c466ad00cff30329183f8b839522d33bc6552 filter=lfs diff=lfs merge=lfs -text
|
| 42 |
+
image/blobs/sha256/be04a2af1cfe08908467229c0dd35fc629cb65ba98a497e7b627499f0c76821b filter=lfs diff=lfs merge=lfs -text
|
| 43 |
+
image/blobs/sha256/c006f8f8bd1ac6c419a231b5d94521d87d03b141b76873d7a4d7fb68a933e57a filter=lfs diff=lfs merge=lfs -text
|
| 44 |
+
image/blobs/sha256/c74a07e1fb0aaf18cdccf3f41385a21dd137109c155363830e2ac2c565d7ac43 filter=lfs diff=lfs merge=lfs -text
|
| 45 |
+
image/blobs/sha256/ca0cef561ac457e6f5fe6e064c167bfc523b6e975437121969a8f39630a0a424 filter=lfs diff=lfs merge=lfs -text
|
| 46 |
+
image/blobs/sha256/d447e3306f087f0a9b594fb8a92f46880560372c42349c9fc8c10e5bb634aadd filter=lfs diff=lfs merge=lfs -text
|
| 47 |
+
image/blobs/sha256/e1ff664083a6df132a0751cb98c17c425fd5908c3434270f38e2de56763f06ba filter=lfs diff=lfs merge=lfs -text
|
| 48 |
+
image/blobs/sha256/e324ba29d9a1eb6f52840139e2d929d488845bcadb738439b521f7900e1e22a3 filter=lfs diff=lfs merge=lfs -text
|
| 49 |
+
image/blobs/sha256/fc48d390b577ea739097ee841cf6e598acbbe40f0f42571c6cb450d57d7a09f2 filter=lfs diff=lfs merge=lfs -text
|
| 50 |
+
media/dp_kashiwanoha_dense_tt_vs_cpu.png filter=lfs diff=lfs merge=lfs -text
|
| 51 |
+
media/dp_nuscenes_scene-0061_bev_tt_NC.gif filter=lfs diff=lfs merge=lfs -text
|
| 52 |
+
media/dp_nuscenes_scene-0061_kf06_bev_tt_NC.png filter=lfs diff=lfs merge=lfs -text
|
| 53 |
+
media/dp_nuscenes_scene-0061_kf18_bev_tt_NC.png filter=lfs diff=lfs merge=lfs -text
|
| 54 |
+
media/dp_nuscenes_scene-0061_kf18_cam_front_tt_NC.jpg filter=lfs diff=lfs merge=lfs -text
|
| 55 |
+
media/dp_nuscenes_scene-0103_kf12_bev_tt_NC.png filter=lfs diff=lfs merge=lfs -text
|
| 56 |
+
media/dp_nuscenes_scene-0103_kf12_cam_front_tt_NC.jpg filter=lfs diff=lfs merge=lfs -text
|
| 57 |
+
media/dp_nuscenes_scene-0757_kf11_bev_tt_NC.png filter=lfs diff=lfs merge=lfs -text
|
| 58 |
+
media/dp_nuscenes_scene-0916_kf11_bev_tt_NC.png filter=lfs diff=lfs merge=lfs -text
|
| 59 |
+
media/dp_straight_road_tt_vs_cpu.png filter=lfs diff=lfs merge=lfs -text
|
OPT_BASELINE.md
ADDED
|
@@ -0,0 +1,535 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# diffusion-planner-p150 baseline on the p150 (before optimization)
|
| 2 |
+
|
| 3 |
+
Date 2026-10-08. Code is the baseline commit `5541833` (first correct port, no optimization): the port's last commit
|
| 4 |
+
`34b6049` plus ttaw re-vendored to 0.17.1 and two measurement scripts, `code/scripts/bench.py` (rewritten: the full
|
| 5 |
+
stage table) and `code/scripts/profile_ops.py` (new). `git diff 34b6049 5541833` touches no model code (`tt/`,
|
| 6 |
+
`host/`, `reference/`, `api.py`, `device.py`, `io.py`, `server/` are unchanged).
|
| 7 |
+
tt-metal `44d66500520` (v0.80.0-dev20261006-78) + the ETH-dispatch patch (workspace `patches/tt-metal-eth-dispatch.patch`,
|
| 8 |
+
sha256 `08d0ddf6…45cc`, 4 files; `eth_patch: true` in every log). ttaw 0.17.1 @ common `290f22b`
|
| 9 |
+
(`code/tt_diffusion_planner/ttaw/VENDORED.json`; `vendor.py --check`: 0 differences). Weights
|
| 10 |
+
`AutowareFoundation/diffusion_planner` @ `423efde67f5` (tag `v5.0`). Input: `code/tt_diffusion_planner/samples/kashiwanoha_dense.npz`
|
| 11 |
+
(one planning instant on the kashiwanoha test map: 88 neighbours, 123 lanes, 17 route lanes, 60 line strings, ego at
|
| 12 |
+
6 m/s), batch 1. The stage bench also runs `straight_road.npz` (shipped) and one public-data instant, nuScenes
|
| 13 |
+
v1.0-mini `scene-0103_kf14` (42 neighbours, 103 lanes, 7 route lanes, 10 polygons, 26 line strings; CC BY-NC-SA 4.0,
|
| 14 |
+
used locally, never shipped).
|
| 15 |
+
|
| 16 |
+
Configuration of every number unless a row says otherwise: ETH dispatch, 1 CQ, 12×10 grid
|
| 17 |
+
(`device.compute_with_storage_grid_size()` printed in each log), one p150b (KMD 2.10.0, firmware bundle 19.13.1.0).
|
| 18 |
+
AICLK 1350 MHz: sysfs `tt_aiclk` was sampled every 50 ms during every bench, median 1350 MHz and min 1343 MHz. Board
|
| 19 |
+
power had a median of 76 W and a max of 86 W, at 63-68 °C. The profiler CSV header also gives 1350 MHz. The host is
|
| 20 |
+
shared (AMD EPYC-Rome, 8 vCPUs, up to 8 other agents, load average 8-12 during these runs), so host-side stages are
|
| 21 |
+
quoted as p50 / p99 (min).
|
| 22 |
+
|
| 23 |
+
## Summary
|
| 24 |
+
|
| 25 |
+
| | baseline |
|
| 26 |
+
|---|---|
|
| 27 |
+
| device latency per plan (one trace replay + sync) | **102.13 ms** p50, 104.81 p99, 102.05 min |
|
| 28 |
+
| back-to-back replays (device time per plan) | **102.04 ms** = 9.80 plans/s; the same on all three scenes (dense capacities) |
|
| 29 |
+
| `model(inputs=arrays)` end to end | **117.90 ms** p50, 134.56 p99, 112.87 min (`.npz` path: 124.76 ms) |
|
| 30 |
+
| programs per plan (device profiler, traced replay) | **6,282** (309 unique programs; the fake-ttnn count was 6,408) |
|
| 31 |
+
| device kernel sum / op-to-op gaps / span (profiled replay) | 99.27 / 4.39 / 103.66 ms. Kernel-bound, not gap-bound: the median program runs 5.76 µs, and 4,126 programs under 10 µs add up to 19.2 ms |
|
| 32 |
+
| largest single sink | the encoder's channel-MLP and pre-projection matmuls on 4-D `[1, E, T, C]` activations, which run on **4-8 cores**: 90 programs, **25.8 ms** (the 10 slowest programs of the plan are all neighbour-mixer channel matmuls, ~970 µs each) |
|
| 33 |
+
| cost of the numerics defaults (PORT_LOG decision 11), re-measured | **+33.4 ms** over the round-1 defaults (68.64 ms b2b) and **+58.3 ms** over the fastest graph (43.74 ms, which fails the gates). Per knob: decoder split matmuls 20.2, fp32 LayerNorm 18.1, fp32 matmul attention 15.5, encoder split matmuls 6.0 ms |
|
| 34 |
+
| dispatch / CQ (D14) | ETH-1CQ stays. WORKER is equal (−0.17 ms on 11×10). 2 CQs cost +9.8 ms on ETH (±0 on WORKER) |
|
| 35 |
+
| accuracy | 44 device tests, every gate value and all 99 per-scene e2e numbers identical to the port's final run |
|
| 36 |
+
|
| 37 |
+
## How to run
|
| 38 |
+
|
| 39 |
+
```bash
|
| 40 |
+
ROOT=/home/ubuntu/experiments/tt-models; cd $ROOT/bundles/diffusion-planner-p150; source $ROOT/bin/tt-env.sh
|
| 41 |
+
export PYTHONPATH=$PWD/code:$PYTHONPATH HF_HUB_OFFLINE=1 # weights: $DIFFUSION_PLANNER_WEIGHTS_DIR > workspace assets > HF cache
|
| 42 |
+
# accuracy gates (per-module PCC + end-to-end agreement vs the fp32 CPU reference; 99 scenes with the workspace goldens)
|
| 43 |
+
$ROOT/bin/devrun -t 1800 -- env TTAW_GATES_READONLY=1 python -m pytest -q -s \
|
| 44 |
+
code/tt_diffusion_planner/tests/test_pcc_device.py code/tt_diffusion_planner/tests/test_e2e_device.py
|
| 45 |
+
# stage breakdown (load / host_pre / pack / host_in / H2D / trace / D2H / host_post / e2e / b2b: p50, p99, min; AICLK)
|
| 46 |
+
$ROOT/bin/devrun -t 900 -- python code/scripts/bench.py --iters 100 --json /tmp/bench.json \
|
| 47 |
+
--input code/tt_diffusion_planner/samples/kashiwanoha_dense.npz --input code/tt_diffusion_planner/samples/straight_road.npz \
|
| 48 |
+
--input $ROOT/research/diffusion-planner/public_data/inputs/nuscenes/scene-0103_kf14.npz
|
| 49 |
+
# dispatch / CQ matrix: the same per configuration, one process each: --dispatch eth|worker --num-cqs 1|2
|
| 50 |
+
# numerics cost: the same with DIFFUSION_PLANNER_SPLIT_MATMUL / _LN_FP32 / _ATTN_MATMUL set (timing only, see below)
|
| 51 |
+
# device profile: one eager plan (stage + layer-kind signposts) and one traced replay, between signposts
|
| 52 |
+
$ROOT/bin/devrun -t 3600 -- python -m tracy -r -p -v --op-support-count 16000 --no-web-server \
|
| 53 |
+
-o $ROOT/generated/profiler/diffusion-planner_baseline code/scripts/profile_ops.py
|
| 54 |
+
tt-perf-report <ops_perf_results_*.csv> --start-signpost trace --end-signpost trace_end --arch p150
|
| 55 |
+
python $ROOT/logs/diffusion-planner/baseline/scripts/analyze_profile.py <the same csv> out.json --md out.md
|
| 56 |
+
```
|
| 57 |
+
|
| 58 |
+
The job scripts of this baseline are in `logs/diffusion-planner/baseline/scripts/` (workspace):
|
| 59 |
+
|
| 60 |
+
- `windowA.sh` (one devrun window, 04:35-04:58 UTC): job 1, the device suite and the alloc-tracking run; then job 2,
|
| 61 |
+
the bench and the dispatch / CQ matrix.
|
| 62 |
+
- `windowB.sh` (05:54-06:15 UTC): job 3, the numerics ablation; job 4, the Tracy profile; job 5, the
|
| 63 |
+
ttnn-visualizer captures.
|
| 64 |
+
- `host_suite.sh`, `analyze_profile.py`, `profile_breakdown.py`, `tables.py`, `fake_flops.py`, `visualizer_*.py`.
|
| 65 |
+
|
| 66 |
+
The plan issues 6,282 programs, more than the profiler's default 1,000-program buffer. So the profile runs with
|
| 67 |
+
`--op-support-count 16000`: 48 B per program per RISC, about 77 MB of DRAM per channel. `ttnn.ReadDeviceProfiler` is
|
| 68 |
+
called before and after each section, and nothing was dropped (no full-buffer warning).
|
| 69 |
+
|
| 70 |
+
## Accuracy (baseline)
|
| 71 |
+
|
| 72 |
+
Re-run on this commit after re-vendoring ttaw 0.11.0 -> 0.17.1, with the gates read-only:
|
| 73 |
+
|
| 74 |
+
- device suite: **44 passed** in 48.6 s, and 44 passed again under `TT_METAL_TRACE_ALLOC_TRACKING=1`;
|
| 75 |
+
- every gate value and every per-scene number of the 99 scenes is identical to the port's final run (job 19, ttaw
|
| 76 |
+
0.11.0);
|
| 77 |
+
- host suite on the fake ttnn: 105 passed, 46 skipped (the 44 device tests, the ORT module, one dead test).
|
| 78 |
+
|
| 79 |
+
Logs: `logs/diffusion-planner/baseline/device_suite.log`, `alloc_tracking.log`, `e2e_device.json`, `host_suite.log`.
|
| 80 |
+
|
| 81 |
+
| check | gate | baseline |
|
| 82 |
+
|---|---|---|
|
| 83 |
+
| `enc.ego` / `enc.neighbor` / `enc.lane` / `enc.route` / `enc.polygon` / `enc.line_string` PCC vs the fp32 reference (valid rows, 2 samples) | ≥ 0.999 | 0.999972 / 0.999952 / 0.999992 / 0.999998 / 0.999994 / 0.999998 |
|
| 84 |
+
| `enc.goal` / `enc.ego_shape` / `enc.turn` PCC | ≥ 0.999 | 0.999991 / 1.000000 / 0.999995 |
|
| 85 |
+
| `enc.encoding` PCC (valid tokens) | ≥ 0.999 | 0.999978 |
|
| 86 |
+
| `dec.eval` PCC (teacher-forced decoder, min over the 11 evaluations x 2 samples) | ≥ 0.999 | ≥ 0.9999995 (prints 1) |
|
| 87 |
+
| end-to-end ego trajectory vs the reference, max over 99 scenes (2 samples, 5 research scenes, 92 nuScenes-mini instants): max / mean displacement | ≤ 1.0 m / ≤ 0.3 m | 0.313 m / 0.143 m (both on nuScenes scene-0103_kf14) |
|
| 88 |
+
| turn-indicator command identical | = 1.0 | 99 / 99 |
|
| 89 |
+
| neighbours: median per-agent max displacement, max over scenes | ≤ 1.5 m | 0.086 m |
|
| 90 |
+
| CPU reference vs ONNX Runtime on the shipped ONNX (checks the reference itself; research venv, PORT_LOG 5.1) | ≥ 0.9999 | final x0 PCC ≥ 0.9999999999982 (7 scenes), ≥ 0.99999999983 (92 nuScenes instants); 48 tests passed |
|
| 91 |
+
|
| 92 |
+
## Performance (warm, batch 1, 100 iterations per stage)
|
| 93 |
+
|
| 94 |
+
`code/scripts/bench.py --iters 100` (job 2, `matrix_eth-1cq.{log,json}`):
|
| 95 |
+
|
| 96 |
+
- `trace` is one replay of the whole plan plus a device sync, as served;
|
| 97 |
+
- `b2b` is the device time per plan over 50 back-to-back replays (3 rounds);
|
| 98 |
+
- `e2e` is `model(inputs=<decoded arrays>)`, and `e2e_path` is `model(inputs=<.npz path>)`.
|
| 99 |
+
|
| 100 |
+
Values are ms, p50 / p99 (min).
|
| 101 |
+
|
| 102 |
+
| stage | kashiwanoha_dense (shipped default) | straight_road (shipped) | nuScenes scene-0103_kf14 (public data) |
|
| 103 |
+
|---|---:|---:|---:|
|
| 104 |
+
| `load`: `.npz` decode + schema check (path inputs only) | 7.02 / 18.20 (5.79) | 6.75 / 16.81 (5.86) | |
|
| 105 |
+
| `host_pre`: the node's pre-processing | 4.76 / 12.55 (3.84) | 4.32 / 12.58 (3.57) | 4.25 / 6.16 (3.74) |
|
| 106 |
+
| `pack`: the 19 persistent trace inputs | 0.45 / 4.52 (0.36) | 0.44 / 1.06 (0.38) | 0.43 / 0.69 (0.34) |
|
| 107 |
+
| `host_in`: TILE host tensors | 2.41 / 6.80 (2.02) | 2.41 / 6.79 (2.00) | 2.31 / 3.35 (1.92) |
|
| 108 |
+
| H2D (19 tensors, 3.6 MB) + sync | 1.14 / 6.32 (0.90) | 1.09 / 5.43 (0.91) | 1.03 / 2.72 (0.90) |
|
| 109 |
+
| **device trace, one blocking plan** | **102.13 / 104.81 (102.05)** | 102.14 / 104.32 (102.07) | 102.13 / 102.19 (102.05) |
|
| 110 |
+
| **back-to-back trace replays, per plan** | **102.04** (rounds 102.09, 102.04, 102.04) = 9.80 plans/s | 102.04 | 102.04 |
|
| 111 |
+
| D2H (one packed read, 460 KB) | 0.58 / 4.24 (0.31) | 0.59 / 1.50 (0.36) | 0.56 / 0.75 (0.30) |
|
| 112 |
+
| `host_post`: the node's post-processing | 4.32 / 11.54 (3.57) | 3.91 / 9.47 (3.29) | 3.93 / 6.96 (3.50) |
|
| 113 |
+
| **e2e `model(inputs=arrays)`** | **117.90 / 134.56 (112.87)** | 126.58 / 144.70 (114.23) | 114.52 / 122.91 (113.05) |
|
| 114 |
+
| e2e `model(inputs=path)` | 124.76 / 147.74 (118.99) | 121.92 / 143.29 (118.41) | |
|
| 115 |
+
|
| 116 |
+
- **Timing breakdown.** `model(...).timing_ms` p50 (kashiwanoha):
|
| 117 |
+
- preprocess 4.77 ms;
|
| 118 |
+
- device 106.83 ms (`host_in` + H2D + replay + D2H + unpack);
|
| 119 |
+
- postprocess 4.36 ms;
|
| 120 |
+
- total 117.35 ms.
|
| 121 |
+
- **Start-up.** The first call after `from_pretrained` takes 118.2 ms (1.0x). `from_pretrained` takes 11.6 s with a
|
| 122 |
+
warm kernel cache: ~5.1 s to open the device, 0.75 s to build (read the ONNX weights and upload them), 5.8 s to warm
|
| 123 |
+
up and capture.
|
| 124 |
+
- **The device time does not depend on the scene.** All three scenes take 102.13-102.14 ms. Every plan computes the
|
| 125 |
+
full capacities whatever the scene holds (no compaction):
|
| 126 |
+
- 320 neighbours, 140 lanes, 25 route lanes, 10 polygons and 60 line strings in the encoder;
|
| 127 |
+
- 352 decoder rows (321 agents);
|
| 128 |
+
- 576 tokens (564 real).
|
| 129 |
+
- **Host overhead.** Around the trace the host adds ~15.8 ms per in-process request: pre 4.8 + pack 0.5 +
|
| 130 |
+
host_in 2.4 + H2D 1.1 + D2H 0.6 + post 4.3 ms, plus the schema check. Handing over an `.npz` path adds 7 ms.
|
| 131 |
+
- **Noise.** The p99 / mean spread of every host stage (and of `e2e`) comes from the shared host. Straight_road's
|
| 132 |
+
e2e p50 is higher than kashiwanoha's only because other agents loaded the CPU during that loop. The device numbers
|
| 133 |
+
(`trace` min, `b2b`) are stable to ±0.05 ms.
|
| 134 |
+
- **Memory.** 309 unique programs; the trace buffers hold 74.6 MB of DRAM; 48.1 MB of weights and constants are
|
| 135 |
+
uploaded.
|
| 136 |
+
|
| 137 |
+
## Dispatch / CQ matrix (D14)
|
| 138 |
+
|
| 139 |
+
One process per configuration, all in one lock window (job 2, `matrix_<tag>.{log,json}`); default sample, 100
|
| 140 |
+
iterations.
|
| 141 |
+
|
| 142 |
+
| config | grid | trace p50 / p99 (min) | b2b per plan | e2e p50 / p99 (min) | H2D p50 | D2H p50 | load (JIT cache) | first call |
|
| 143 |
+
|---|---|---:|---:|---:|---:|---:|---:|---:|
|
| 144 |
+
| **ETH-1CQ (default)** | 12x10 | **102.13** / 104.81 (102.05) | **102.04** | **117.90** / 134.56 (112.87) | 1.14 | 0.58 | 11.6 s (warm) | 118.2 |
|
| 145 |
+
| ETH-2CQ | 12x10 | 112.02 / 112.96 (111.94) | 111.87 | 125.64 / 141.03 (123.65) | 1.02 | 0.50 | 11.0 s (warm) | 125.0 |
|
| 146 |
+
| WORKER-2CQ | 11x10 | 102.03 / 106.28 (101.90) | 101.88 | 126.63 / 142.67 (114.26) | 0.94 | 0.52 | 417.6 s (cold) | 114.7 |
|
| 147 |
+
| WORKER-1CQ | 11x10 | 101.97 / 104.53 (101.89) | 101.87 | 116.10 / 135.84 (112.62) | 0.98 | 0.58 | 446.1 s (cold) | 113.6 |
|
| 148 |
+
| ETH-1CQ, repeated last | 12x10 | 102.14 / 102.37 (102.07) | 102.04 | 117.35 / 124.73 (114.13) | 1.14 | 0.63 | 9.4 s (warm) | 117.8 |
|
| 149 |
+
|
| 150 |
+
- **WORKER vs ETH: equal, so ETH stays the default.** WORKER dispatch is 0.17 ms (0.17 %) faster per plan, on 110
|
| 151 |
+
cores instead of 120. Two effects cancel:
|
| 152 |
+
- The 12th column hardly matters to this graph (Grid usage below). The 3,775 programs that do use all 120 cores
|
| 153 |
+
are short element-wise programs, bound by a fixed cost per program, not by the core count. The 25.8 ms of 4-8-core
|
| 154 |
+
matmuls do not use the grid at all.
|
| 155 |
+
- WORKER saves ~0.5 µs of dispatch per program (probe P1).
|
| 156 |
+
|
| 157 |
+
That is not "clearly better", so ETH stays the default (D14, `serve.env`). Re-measure once the op count drops,
|
| 158 |
+
when the 12th column should start to pay (DISPATCH.md).
|
| 159 |
+
- **CQ count: 1.**
|
| 160 |
+
- On ETH, a second CQ makes the replay itself +9.8 ms (+9.6 %) per plan, about +1.6 µs per program.
|
| 161 |
+
- On WORKER it costs nothing (101.88 vs 101.87 ms). So the penalty belongs to the patched ETH 2-CQ dispatch
|
| 162 |
+
topology. It was not profiled here; `--profile-dispatch-cores` would show it.
|
| 163 |
+
- What 2 CQs could hide is the upload of the next plan while the trace runs: ~1.1 ms of H2D per 117 ms request,
|
| 164 |
+
and only when requests are pipelined.
|
| 165 |
+
|
| 166 |
+
Decision: **1 CQ** (the D14 default); `DIFFUSION_PLANNER_NUM_CQS=1` stays pinned.
|
| 167 |
+
- A WORKER open recompiles every kernel: 7 min here, for both CQ counts. The served configuration never pays this.
|
| 168 |
+
- No drift over the 21-minute window: ETH-1CQ measured 102.13 ms first and 102.14 ms last (trace p50), and 102.04 /
|
| 169 |
+
102.04 ms (b2b).
|
| 170 |
+
|
| 171 |
+
## What the numerics defaults cost (precision ablation)
|
| 172 |
+
|
| 173 |
+
Job 3 (`precision_<tag>.{log,json}`) ran one process per configuration on ETH-1CQ with the default sample: 30
|
| 174 |
+
iterations, and b2b over 3 x 30 replays. The configurations differ only by the `DIFFUSION_PLANNER_*` knobs. **This
|
| 175 |
+
is timing only: no gate was run here, and only `shipped` is gate-clean.** The PORT_LOG evidence for three of the
|
| 176 |
+
others:
|
| 177 |
+
|
| 178 |
+
- the round-1 defaults fail `ego.mean_err_m` on nuScenes scene-0103_kf14 (0.347 m > 0.3, job 11);
|
| 179 |
+
- bf16 SDPA instead of the fp32 matmul attention (PORT_LOG's "decoder split alone", the same knobs) gives 0.74 /
|
| 180 |
+
0.33 m (max / mean) on scene-0103_kf14, over the 0.3 m mean gate (job 13, `logs/diffusion-planner/split_error_decsplit.json`);
|
| 181 |
+
- the fastest graph (fused LN, no split, SDPA) fails `enc.ego` (0.99856 / 0.99261, job 4).
|
| 182 |
+
|
| 183 |
+
The decoder-split, encoder-split and fused-LN removals were never gated.
|
| 184 |
+
|
| 185 |
+
| config | `SPLIT_MATMUL` | `LN_FP32` | `ATTN_MATMUL` | trace p50 | b2b | vs shipped | unique programs | trace buffers |
|
| 186 |
+
|---|---|---|---|---:|---:|---:|---:|---:|
|
| 187 |
+
| **shipped** (decision 11) | `enc.island.*,enc.pre.*,dec.*` | `enc.mixer.*,dec.*` | `enc.fusion.attn,dec.*` | 102.11 | 102.04 | | 309 | 74.6 MB |
|
| 188 |
+
| bf16 SDPA instead of fp32 matmul attention | same | same | none | 86.64 | 86.57 | −15.47 | 305 | 69.9 MB |
|
| 189 |
+
| no split matmuls in the decoder | `enc.island.*,enc.pre.*` | same | same | 81.93 | 81.85 | −20.19 | 283 | 50.4 MB |
|
| 190 |
+
| no split matmuls in the encoder | `dec.*` | same | same | 96.08 | 96.02 | −6.02 | 209 | 72.9 MB |
|
| 191 |
+
| fused `ttnn.layer_norm` everywhere | same | none | same | 84.03 | 83.94 | −18.10 | 304 | 43.9 MB |
|
| 192 |
+
| round-1 defaults (PORT_LOG job 11: 68.7 ms) | `enc.island.*,enc.pre.*,dec.preproj.*` | `enc.mixer.*,dec.*` | none | 68.72 | 68.64 | **−33.40** | 287 | 50.4 MB |
|
| 193 |
+
| fastest graph (PORT_LOG job 4: 43.7 ms) | none | none | none | 43.79 | 43.74 | −58.30 | 173 | 16.2 MB |
|
| 194 |
+
|
| 195 |
+
- **Precision cost of the shipped defaults.** It is **33.4 ms** over round 1, which reproduces PORT_LOG's 102.1 vs
|
| 196 |
+
68.7 ms. Over the fastest graph it is 58.3 ms.
|
| 197 |
+
- **The knobs add up almost exactly.** The four single-knob removals sum to 59.8 ms, against 58.3 ms with all of
|
| 198 |
+
them off.
|
| 199 |
+
- **On the card this cost is the first optimization target:** it is the price of the accuracy the nuScenes e2e gate
|
| 200 |
+
needs. Items 2-4 below recover it with fused kernels that keep the fp32-level accuracy. They do not relax precision:
|
| 201 |
+
precision policy §9.3, and every numerics change re-checked on all 99 scenes. Item 1 is a separate sink of similar
|
| 202 |
+
size that this profile found, unrelated to precision.
|
| 203 |
+
|
| 204 |
+
## Host vs device, and every host<->device transfer per plan
|
| 205 |
+
|
| 206 |
+
Host, per plan: numpy, the node's own pre- and post-processing (the PLAN.md 2.12 host rows). There is no host
|
| 207 |
+
fallback inside the plan.
|
| 208 |
+
|
| 209 |
+
- `load`: decodes the `.npz` and checks the 15 raw tensors against `INPUT_SCHEMA` (names, shapes, finite values).
|
| 210 |
+
Only `model(inputs=<path>)` does this; the server decodes its JSON / base64 body instead.
|
| 211 |
+
- `host_pre` (`host.prepare`):
|
| 212 |
+
- normalization (all-zero rows kept) and the speed masks;
|
| 213 |
+
- the encoder's pre-matmul input plumbing: history truncation, velocity zeroing, validity masks, position features
|
| 214 |
+
with the ONNX `atan2` decomposition and the pseudo-heading quirk, point deltas;
|
| 215 |
+
- the decoder masks and `x_T`.
|
| 216 |
+
- `pack` (`tt.inputs.plan_inputs`) builds the 19 persistent trace inputs: the input columns of the exact affine
|
| 217 |
+
rewrites, the 576-token arrays, the key-bias rows, and `cs` / `y0` on 352 rows.
|
| 218 |
+
- `host_in`: converts them to ttnn host tensors (TILE layout, host tilization).
|
| 219 |
+
- `host_post` (`host.make_output`): denormalization, trajectory velocity / force-stop / acceleration, the predicted
|
| 220 |
+
paths of the non-empty neighbour rows, the turn-indicator decision.
|
| 221 |
+
|
| 222 |
+
Device: everything else, as ONE metal trace (`plan`):
|
| 223 |
+
|
| 224 |
+
- the encoder: 6 mixer trunks, entity heads, small encoders, token assembly, 6 fusion blocks;
|
| 225 |
+
- the hoisted cross K / V;
|
| 226 |
+
- 11 x (DiT evaluation + fp32 DPM-Solver++(2M) update + prefix constraint);
|
| 227 |
+
- the turn head and the output pack.
|
| 228 |
+
|
| 229 |
+
**H2D per plan: 19 `copy_host_to_device_tensor` into the persistent inputs** (CQ0; with 2 CQs, CQ1 + events): 1.19 MB
|
| 230 |
+
logical, 3.60 MB as TILE-padded fp32 / bf16.
|
| 231 |
+
|
| 232 |
+
| input | shape | dtype | logical B | TILE-padded B |
|
| 233 |
+
|---|---|---|---:|---:|
|
| 234 |
+
| `neighbor_x` | [1, 320, 6, 9] | fp32 | 69,120 | 1,310,720 |
|
| 235 |
+
| `lane_x` | [1, 140, 20, 8] | fp32 | 89,600 | 573,440 |
|
| 236 |
+
| `cs` (current states in the t = 0 columns) | [1, 1, 352, 324] | fp32 | 456,192 | 495,616 |
|
| 237 |
+
| `y0` (`x_T * mask0`; zeros at temperature 0) | [1, 1, 352, 324] | fp32 | 456,192 | 495,616 |
|
| 238 |
+
| `line_string_x` | [1, 60, 20, 6] | fp32 | 28,800 | 245,760 |
|
| 239 |
+
| `route_x` | [1, 25, 20, 8] | fp32 | 16,000 | 102,400 |
|
| 240 |
+
| `polygon_x` | [1, 10, 40, 5] | fp32 | 8,000 | 81,920 |
|
| 241 |
+
| `token_valid` | [1, 1, 576, 1] | fp32 | 2,304 | 73,728 |
|
| 242 |
+
| `pos_aug` | [1, 1, 576, 15] | fp32 | 34,560 | 73,728 |
|
| 243 |
+
| `neighbor_aux` | [1, 1, 320, 4] | fp32 | 5,120 | 40,960 |
|
| 244 |
+
| `fusion_key_row` | [1, 1, 1, 576] | bf16 | 1,152 | 36,864 |
|
| 245 |
+
| `agent_key_row` | [1, 1, 1, 352] | bf16 | 704 | 22,528 |
|
| 246 |
+
| `lane_aux` | [1, 1, 140, 29] | fp32 | 16,240 | 20,480 |
|
| 247 |
+
| `ego_x`, `static_x`, `route_aux`, `goal_x`, `ego_shape_x`, `turn_x` | small | fp32 | 3,344 | 6 x 4,096 |
|
| 248 |
+
|
| 249 |
+
**D2H per plan: 1 packed read** (`pack_outputs`), 117,664 fp32 = 470,656 B in one row:
|
| 250 |
+
|
| 251 |
+
- `final_x0` [352, 324];
|
| 252 |
+
- the turn logits [5];
|
| 253 |
+
- the ego rows of the 11 published iterates [11, 324].
|
| 254 |
+
|
| 255 |
+
There is no mid-graph host round trip and no host fallback op.
|
| 256 |
+
|
| 257 |
+
## Grid usage
|
| 258 |
+
|
| 259 |
+
`tt/` sets no core grid or program config on any op. The only grid consumer, C20's SDPA config, reads
|
| 260 |
+
`device.compute_with_storage_grid_size()`, and the default fp32 matmul attention does not use it. So every op picks
|
| 261 |
+
its own core count:
|
| 262 |
+
|
| 263 |
+
| core count per program (traced replay) | programs | kernel ms | what |
|
| 264 |
+
|---|---:|---:|---|
|
| 265 |
+
| 120 | 3,775 | 39.4 | element-wise / typecast / reduce programs, spread over all cores whatever their size |
|
| 266 |
+
| 64-119 | 1,050 | 14.1 | decoder matmuls on 66-72 cores, softmax on 88, token-mixing matmuls on 100-117 |
|
| 267 |
+
| 16-63 | 708 | 11.9 | decoder matmuls on 48 cores (561 programs, 9.7 ms) |
|
| 268 |
+
| **2-15** | **712** | **33.6** | **matmuls on 4 / 8 cores: 148 programs, 26.0 ms, of which the encoder's batched `[1, E, T, C] @ [C, 128]` channel matmuls are 90 programs, 25.8 ms**; the attention P·V matmuls on 11 cores (66 programs, 4.6 ms); decoder LN reductions on 11 cores |
|
| 269 |
+
| 1 | 37 | 0.2 | slices, small reshapes |
|
| 270 |
+
|
| 271 |
+
Matmul efficiency is low everywhere. tt-perf-report gives a weighted mean of 7.2 % of the FLOP roofline over the
|
| 272 |
+
1,365 matmuls (max 12.9 %), and 6.8 % of the DRAM roofline (35 GB/s) over the modelled ops.
|
| 273 |
+
|
| 274 |
+
## What is already fused / traced
|
| 275 |
+
|
| 276 |
+
- **The whole plan is one trace:** encoder + 11 DiT evaluations + 10 solver updates + turn head + pack, replayed
|
| 277 |
+
with persistent inputs. Replay == eager bit for bit (`test_plan_replay_equals_eager`, and the profile run). The
|
| 278 |
+
device suite is green under `TT_METAL_TRACE_ALLOC_TRACKING=1`.
|
| 279 |
+
- **Exact rewrites already in the graph** (`tt/params.py`, proven on the host):
|
| 280 |
+
- the cross-attention K / V of the 3 DiT blocks is hoisted out of the solver loop (once per plan, not 11 times);
|
| 281 |
+
- the per-step adaLN is folded into the LayerNorm affine rows (no adaLN MLP on the device);
|
| 282 |
+
- the solver update is folded into `y' = A y - B m_k + Cm m_(k-1)`, with the prefix constraint as `x = y + cs` and
|
| 283 |
+
the t = 0 output columns zeroed in the last projection's weights;
|
| 284 |
+
- the ego / neighbour pre-projection runs as a pad-relative fp32 island;
|
| 285 |
+
- the small embeddings (neighbour type, lane speed / attributes, position) are one matmul over host-built columns;
|
| 286 |
+
- the turn head is `W_sel` / `W_pool / 564`, and the agent embedding is folded into bias rows.
|
| 287 |
+
- **Fused in stock ops:**
|
| 288 |
+
- bias + GELU / tanh-GELU in `ttnn.linear` (2-D inputs);
|
| 289 |
+
- Q | K | V in one linear and one `nlp_create_qkv_heads` (DiT self-attention);
|
| 290 |
+
- K | V in one linear + `split_q_kv` (fusion);
|
| 291 |
+
- one packed readback.
|
| 292 |
+
- **Not fused: the numerics defaults** (PORT_LOG decision 11), counted in programs per call:
|
| 293 |
+
- each split hi / lo matmul is **8.3 programs** on average instead of 1: typecast, typecast, subtract, 3
|
| 294 |
+
matmuls / linears, 2 adds, then GELU / typecast as needed;
|
| 295 |
+
- each fp32 LayerNorm is **9 programs** instead of 1: mean, subtract, multiply, mean, add, rsqrt, multiply, then
|
| 296 |
+
multiply / add for the affine;
|
| 297 |
+
- each fp32 attention is **5 programs** (matmul, scale, mask add, softmax, matmul), plus the head split / merge.
|
| 298 |
+
- **Not fused: glue.** Residual adds, adaLN gate multiplies, the mixer token-mixing transposes, the solver update
|
| 299 |
+
(3-5 element-wise programs per step) and the ego-row slices.
|
| 300 |
+
- **Not done at all:**
|
| 301 |
+
- every activation is DRAM-interleaved: no sharding, no L1 residency;
|
| 302 |
+
- no custom kernel, no `generic_op`, no megakernel;
|
| 303 |
+
- no compaction: all 321 agent rows and 564 tokens are computed whatever the scene holds.
|
| 304 |
+
|
| 305 |
+
## Device profile (traced replay)
|
| 306 |
+
|
| 307 |
+
The profile (job 4, `profile_ops.py`) has two sections:
|
| 308 |
+
|
| 309 |
+
- one eager plan with stage (`m:`) and layer-kind (`c:`) signposts: 778.6 ms, dispatch-bound at a median of 101 µs
|
| 310 |
+
per eager op;
|
| 311 |
+
- one traced replay of the same 6,282 programs. The op codes match the eager run one by one, so every replayed program
|
| 312 |
+
inherits its stage and kind.
|
| 313 |
+
|
| 314 |
+
The replay with the profiler on took 105.13 ms on the host, against 102.1 ms without it. The eager run after the
|
| 315 |
+
capture logs tt-metal's "allocating device buffers ... active trace" warning once; `run_eager` frees its buffers
|
| 316 |
+
before the replay, and the replay equals the eager plan bit for bit.
|
| 317 |
+
|
| 318 |
+
| | traced replay |
|
| 319 |
+
|---|---|
|
| 320 |
+
| programs per plan | **6,282** |
|
| 321 |
+
| device kernel sum | **99.27 ms** |
|
| 322 |
+
| op-to-op gaps (sum / median / p90 / max) | **4.39 ms** / 0.56 µs / 0.65 µs / 7.97 µs |
|
| 323 |
+
| span (first FW start to last FW end, 1350 MHz) | **103.66 ms** |
|
| 324 |
+
| kernel time per program, p10 / p50 / p90 / p99 / max | 2.7 / 5.76 / 32.2 / 96.6 / 984 µs |
|
| 325 |
+
| programs under 5 µs / under 10 µs | 2,673 / 4,126 (19.2 ms together) |
|
| 326 |
+
| math fidelity | HiFi4 on all 6,106 compute programs |
|
| 327 |
+
|
| 328 |
+
The plan is **kernel-bound**: gaps are 4 % of the span. The time goes to a few slow, badly parallelised programs and
|
| 329 |
+
to thousands of small ones, each costing ~5-8 µs whatever its size. The FLOPs are minor: the device graph does 165
|
| 330 |
+
GFLOP of matmul per plan with every split pass counted (132 GFLOP of it in the decoder; `fake_flops.log`). That is
|
| 331 |
+
~1 ms at the 120-core HiFi4 peak (~162 TFLOPS, TT_PLATFORM.md).
|
| 332 |
+
|
| 333 |
+
**By stage** (kernel + gaps, ms):
|
| 334 |
+
|
| 335 |
+
| stage | programs | ms | share |
|
| 336 |
+
|---|---:|---:|---:|
|
| 337 |
+
| encoder: 6 mixer trunks (pre-projection 8.8, 6 x 6 MixerBlocks 39.2, pool + heads 0.3) | 1,274 | 48.35 | 46.6 % |
|
| 338 |
+
| encoder: fusion (6 blocks + final LN), small encoders, tokens, masks | 149 | 3.23 | 3.1 % |
|
| 339 |
+
| decoder: 11 evaluations, 4.67 ms each incl. its solver update (433-443 programs each; e00-e09 4.38-4.40 ms kernel, e10 4.45) | 4,791 | 51.42 | 49.6 % |
|
| 340 |
+
| cross K / V hoist, turn head, output pack | 68 | 0.66 | 0.6 % |
|
| 341 |
+
|
| 342 |
+
**By layer kind and part** (kernel + gaps, ms; `kind_by_group.md`):
|
| 343 |
+
|
| 344 |
+
| part | split matmul | plain linear | fp32 LN | fp32 attention | glue | heads / mask / fused LN | total |
|
| 345 |
+
|---|---:|---:|---:|---:|---:|---:|---:|
|
| 346 |
+
| encoder mixer trunks | 8.31 (24 calls, 346 µs each) | **23.27** (159 calls) | 12.96 (72 calls, 180 µs each) | | 3.77 | 0.04 | 48.35 |
|
| 347 |
+
| encoder fusion + small | | 0.72 | | 1.96 (6 calls) | 0.21 | 0.33 | 3.23 |
|
| 348 |
+
| decoder x11 + solver | **26.09** (308 calls, 85 µs each) | | 8.59 (165 calls, 52 µs each) | **14.47** (66 calls, 219 µs each) | 1.34 | 0.93 | 51.42 |
|
| 349 |
+
| cross K / V, turn, pack | 0.38 | 0.01 | | | 0.24 | 0.02 | 0.66 |
|
| 350 |
+
| **all** | **34.78** | **24.00** | **21.56** | **16.43** | **5.56** | 1.33 | **103.66** |
|
| 351 |
+
|
| 352 |
+
**One DiT evaluation** (e5, 4.67 ms):
|
| 353 |
+
|
| 354 |
+
| | ms | share |
|
| 355 |
+
|---|---:|---:|
|
| 356 |
+
| split matmuls (pre-projection, self Q / K / V + out, cross Q + out, 2 MLPs, final) | 2.37 | 50.8 % |
|
| 357 |
+
| self- and cross-attention as fp32 matmuls (3 x 0.20 + 3 x 0.24) | 1.32 | 28.2 % |
|
| 358 |
+
| fp32 LayerNorms (12 x 48.5 µs + the final layer's 3 LNs 197 µs) | 0.78 | 16.7 % |
|
| 359 |
+
| glue, head split / merge, solver update | 0.20 | 4.3 % |
|
| 360 |
+
|
| 361 |
+
**By op code:**
|
| 362 |
+
|
| 363 |
+
| op | programs | kernel ms | share |
|
| 364 |
+
|---|---:|---:|---:|
|
| 365 |
+
| Matmul | 1,365 | 52.34 | 52.7 % |
|
| 366 |
+
| BinaryNg (element-wise binary) | 2,948 | 31.48 | 31.7 % |
|
| 367 |
+
| Reduce (LN means) | 480 | 3.69 | 3.7 % |
|
| 368 |
+
| Softmax | 72 | 3.47 | 3.5 % |
|
| 369 |
+
| Unary (GELU, rsqrt) | 428 | 2.81 | 2.8 % |
|
| 370 |
+
| Typecast | 676 | 2.41 | 2.4 % |
|
| 371 |
+
| Transpose (mixer token mixing) | 84 | 1.34 | 1.4 % |
|
| 372 |
+
| NlpCreateHeads / NLPConcatHeads | 150 | 1.00 | 1.0 % |
|
| 373 |
+
| other (reshape, untilize, LayerNorm, repeat, tilize, concat, slice, pad) | 79 | 0.73 | 0.7 % |
|
| 374 |
+
|
| 375 |
+
**Top-10 device programs** (traced replay, single programs). All ten are the neighbour mixer's channel-MLP matmuls
|
| 376 |
+
`[1, 320, 64, 128] @ [128, 128]` (6 MixerBlocks x 2 linears), each on **8 cores**:
|
| 377 |
+
|
| 378 |
+
| # | op | stage | in0 @ in1 | dtypes in0 / in1 / out | cores | kernel µs |
|
| 379 |
+
|---:|---|---|---|---|---:|---:|
|
| 380 |
+
| 1 | Matmul (`ch2`) | `enc.neighbor.mix` | [1,320,64,128] @ [1,1,128,128] | bf16 / bf16 / fp32 | 8 | 984.1 |
|
| 381 |
+
| 2 | Matmul (`ch2`) | `enc.neighbor.mix` | same | bf16 / bf16 / fp32 | 8 | 977.6 |
|
| 382 |
+
| 3 | Matmul (`ch2`) | `enc.neighbor.mix` | same | bf16 / bf16 / fp32 | 8 | 972.8 |
|
| 383 |
+
| 4 | Matmul (`ch2`) | `enc.neighbor.mix` | same | bf16 / bf16 / fp32 | 8 | 971.8 |
|
| 384 |
+
| 5 | Matmul (`ch2`) | `enc.neighbor.mix` | same | bf16 / bf16 / fp32 | 8 | 971.2 |
|
| 385 |
+
| 6 | Matmul (`ch2`) | `enc.neighbor.mix` | same | bf16 / bf16 / fp32 | 8 | 970.1 |
|
| 386 |
+
| 7 | Matmul (`ch1`, + GELU) | `enc.neighbor.mix` | [1,320,64,128] @ [1,1,128,128] | fp32 / bf16 / bf16 | 8 | 966.4 |
|
| 387 |
+
| 8 | Matmul (`ch1`) | `enc.neighbor.mix` | same | fp32 / bf16 / bf16 | 8 | 965.2 |
|
| 388 |
+
| 9 | Matmul (`ch1`) | `enc.neighbor.mix` | same | fp32 / bf16 / bf16 | 8 | 965.2 |
|
| 389 |
+
| 10 | Matmul (`ch1`) | `enc.neighbor.mix` | same | fp32 / bf16 / bf16 | 8 | 965.0 |
|
| 390 |
+
|
| 391 |
+
**Top op signatures by total time** (`profile_breakdown.md`). Only matmul shapes are shown: element-wise programs are
|
| 392 |
+
shared across shapes, and the profiler's op-info cache reports the first call's shapes for them.
|
| 393 |
+
|
| 394 |
+
| op | where | kind | in0 @ in1 | cores | programs | ms | mean µs |
|
| 395 |
+
|---|---|---|---|---:|---:|---:|---:|
|
| 396 |
+
| Matmul | `enc.neighbor.mix` channel MLP | linear | [1,320,64,128] @ [128,128] | 8 | 12 | 11.64 | 970 |
|
| 397 |
+
| BinaryNg | mixer blocks | fp32 LN | | 120 | 360 | 9.92 | 27.6 |
|
| 398 |
+
| Matmul | decoder MLP fc2 (mlp1 + mlp2, 3 split passes each) | split | [1,1,352,1024] @ [1024,256] | 48 | 198 | 5.71 | 28.8 |
|
| 399 |
+
| Matmul | `enc.lane.mix` channel MLP | linear | [1,140,64,128] @ [128,128] | 8 | 12 | 5.09 | 425 |
|
| 400 |
+
| BinaryNg | decoder attention scale / mask add | fp32 attention | | 120 | 132 | 4.72 | 35.8 |
|
| 401 |
+
| Matmul | decoder MLP fc1 (mlp1 + mlp2, 3 split passes each) | split | [1,1,352,256] @ [256,1024] | 66 | 198 | 3.39 | 17.1 |
|
| 402 |
+
| Matmul | `enc.neighbor.pre` (pad-relative island) | split | [1,320,6,128] @ [128,128], [1,320,6,9] @ [9,128] | 4 | 6 | 3.15 | 525 |
|
| 403 |
+
| Matmul | decoder self-attention P·V | fp32 attention | [1,8,352,352] @ [1,8,352,32] | 11 | 33 | 2.33 | 70.7 |
|
| 404 |
+
| Matmul | decoder cross-attention P·V | fp32 attention | [1,8,352,576] @ [1,8,576,32] | 11 | 33 | 2.26 | 68.4 |
|
| 405 |
+
| Matmul | `enc.line_string.mix` channel MLP | linear | [1,60,64,128] @ [128,128] | 8 | 12 | 2.19 | 183 |
|
| 406 |
+
| Softmax | decoder cross-attention | fp32 attention | | 88 | 33 | 1.83 | 55.4 |
|
| 407 |
+
| BinaryNg | mixer residual adds | glue | | 120 | 72 | 1.75 | 24.3 |
|
| 408 |
+
| Reduce | mixer LN means | fp32 LN | | 120 | 72 | 1.39 | 19.3 |
|
| 409 |
+
|
| 410 |
+
Compare the 2-D token-mixing matmul of the same trunk: `[1,1,40960,64] @ [64,64]` runs on 117 cores in 69 µs. The
|
| 411 |
+
4-D channel MLP with twice the FLOPs takes 970 µs on 8 cores: 14x the time, about 7x slower per FLOP.
|
| 412 |
+
|
| 413 |
+
Artifacts:
|
| 414 |
+
|
| 415 |
+
- **Profile:** `generated/profiler/diffusion-planner_baseline/reports/2026_10_08_06_13_03/`, holding
|
| 416 |
+
`ops_perf_results_*.csv`, the Tracy file and `profile_log_device.csv.zst` (7.3 GB raw; `zstd -d` it before loading
|
| 417 |
+
the folder in ttnn-visualizer), plus `.logs/cpp_device_perf_report.csv`.
|
| 418 |
+
- **ttnn-visualizer memory / graph reports** (not for timing): `generated/ttnn_visualizer/diffusion-planner_baseline/{decode_once,encoder_taps}/`:
|
| 419 |
+
- one decoder evaluation: 508 ops, 20.9 MB `db.sqlite`;
|
| 420 |
+
- the encoder: 1,483 ops, 63.5 MB.
|
| 421 |
+
- No whole-plan capture was made. At the measured ~1.3k buffer rows per op it would be ~8M rows (a ~2.5 GB capture
|
| 422 |
+
JSON) for 11 repeats of the same evaluation. The two captures cover every distinct layer of the plan except the
|
| 423 |
+
solver update, the turn head and the output pack.
|
| 424 |
+
- Open them with `ttnn-visualizer --profiler-path <that dir> --performance-path <the profile folder>`.
|
| 425 |
+
- Capture: `logs/diffusion-planner/baseline/scripts/visualizer_capture.py` (job 5). Import:
|
| 426 |
+
`visualizer_import.py`, on the CPU.
|
| 427 |
+
|
| 428 |
+
## Ranked optimization opportunities
|
| 429 |
+
|
| 430 |
+
Gains are estimated on the 102.0 ms replay and come from the profile above. They are **not additive**: each one
|
| 431 |
+
shrinks the base of the next. Every item must keep the frozen gates, re-checked on all 99 e2e scenes (PORT_LOG 8:
|
| 432 |
+
the plan's sensitivity to device numerics is chaotic per scene). The precision cost of decision 11 (33.4 ms over round
|
| 433 |
+
1) is recovered by items 2-4, not by relaxing precision.
|
| 434 |
+
|
| 435 |
+
1. **Run the encoder's channel MLPs on 2-D activations (≈ −20 to −22 ms; exact rewrite, low effort).**
|
| 436 |
+
- *What is slow.* The mixer `channels_mlp` (`ch1` / `ch2`) and the channel pre-projections run `ttnn.linear` on
|
| 437 |
+
4-D `[1, E, T, C]` activations. That is a batched matmul on **4-8 cores**: 90 programs, 25.8 ms.
|
| 438 |
+
- *MixerBlocks.* `[1, E, 64, 128]` -> `[1, 1, E·64, 128]` is a free view (64 rows = 2 whole tiles per entity).
|
| 439 |
+
Then the matmul can spread over the grid, as the token-mixing matmuls already do (69 µs on 117 cores for the
|
| 440 |
+
neighbours).
|
| 441 |
+
- *Pre-projections.* T = 6 / 20 / 40 rows do not fill a tile, so the view is not free. Two options:
|
| 442 |
+
- have `pack` lay the inputs out as `[1, 1, E·T, C]` and re-tile once before the token transpose;
|
| 443 |
+
- or give the 4-D matmul an explicit multi-core program config, with the HiFi4 compute config passed too (a
|
| 444 |
+
`program_config` alone falls back to LoFi, PLAN 0.2).
|
| 445 |
+
- *Estimate.* The channel matmul at 2x the measured token-MLP time of its trunk saves ~17 ms on the MixerBlocks
|
| 446 |
+
and ~5 ms on the pre-projections. The math is unchanged (the K = 128 reduction per output). The gates must still
|
| 447 |
+
be re-run, because accumulation order may differ.
|
| 448 |
+
- *Generalisation.* No `[1, E, ...]`-batched `ttnn.linear` anywhere (the PR:P12 rule extended from token mixing
|
| 449 |
+
to channel mixing).
|
| 450 |
+
2. **A fused split (bf16x3) matmul, one program per linear (≈ −17 to −19 ms).**
|
| 451 |
+
- *Today.* 338 split linears, 2,802 programs, 34.8 ms: decoder 26.1, encoder 8.3, cross K/V 0.4.
|
| 452 |
+
- *The kernel.* A `generic_op` matmul that reads the fp32 activation once, splits hi / lo on the fly (SFPU), and
|
| 453 |
+
accumulates `x_hi W_hi + x_hi W_lo + x_lo W_hi` in fp32 DEST.
|
| 454 |
+
- *What it removes.* All typecast / subtract / add programs and their gaps, 9.3 ms in the decoder. The three
|
| 455 |
+
passes also share operand reads: estimated at 1.5x one pass, 8 instead of 16 ms of decoder matmul.
|
| 456 |
+
- *Accuracy.* The same ~1e-5 relative error as today's 3-matmul form.
|
| 457 |
+
- *Cheaper first step.* Explicit program configs for the 352-row decoder matmuls, which run on 48-72 cores at
|
| 458 |
+
~7 % FLOP utilisation. Pass the compute config with them: a `program_config` without one silently falls back to
|
| 459 |
+
LoFi (PLAN 0.2).
|
| 460 |
+
3. **A fused fp32 LayerNorm, one program per LN (≈ −15 to −16 ms).**
|
| 461 |
+
- *Today.* 237 LNs, 2,133 programs, 21.6 ms: mixers 72 x 180 µs on `[E, 64, 128]`; decoder 165 x 52 µs on
|
| 462 |
+
`[352, 256]`.
|
| 463 |
+
- *The kernel.* One program per LN (mean / variance in fp32 by Welford or two passes in L1, rsqrt, affine with
|
| 464 |
+
the folded adaLN rows), at ~2 element-wise passes: ~45 µs in the mixers and ~12 µs in the decoder.
|
| 465 |
+
- *Limit.* The fused `ttnn.layer_norm` would save 18.1 ms (ablation), but it is ~2.5e-3 relative (PR:P10). On the
|
| 466 |
+
offset-dominated mixer rows it reaches rel-L2 0.03, and the teacher-forced neighbour mixer drops to PCC 0.9978
|
| 467 |
+
(PORT_LOG job 5).
|
| 468 |
+
4. **Attention (≈ −5.5 ms with stock ops, ≈ −12 to −13 ms with a fused kernel).**
|
| 469 |
+
- *Today.* 72 fp32 matmul attentions, 16.4 ms, ~228 µs per call. The bf16 SDPA would cost ~13 µs per call
|
| 470 |
+
(ablation: −15.5 ms), but fails the e2e mean gate.
|
| 471 |
+
- *(a) Stock ops.* Fold the 1/√32 scale into `W_q` / `b_q`: fp64 fold, then the hi / lo split; gate it. Use a
|
| 472 |
+
fused scale-mask-softmax for the mask add. Together this removes the 132 decoder (+ 12 fusion) scale / mask
|
| 473 |
+
programs: 4.7 + 0.7 ms.
|
| 474 |
+
- *Parallelism.* The P·V matmuls (`[1, 8, 352, Sk] @ [1, 8, Sk, 32]`) run on 11 cores (66 programs, 4.6 ms); a
|
| 475 |
+
head-parallel program config would spread them.
|
| 476 |
+
- *(b) A fused fp32-accurate flash-attention kernel* (`generic_op`): scores stay in L1 / DEST, `Q·Kᵀ` and `P·V`
|
| 477 |
+
use hi / lo operands, and the mask comes from the key-bias row. At ~3-4x the bf16 SDPA's cost (~50 µs per call)
|
| 478 |
+
it saves ~12.8 ms and supersedes (a).
|
| 479 |
+
5. **Glue (≈ −3 to −4 ms).**
|
| 480 |
+
- *Mixers: 3.8 ms.* Token-mixing transposes (84 programs, 1.3 ms) and residual adds (72 programs, 1.8 ms). The
|
| 481 |
+
entity-parallel mixer kernel of PLAN 5.3 (no transposes) or residual adds in the matmul epilogue remove them.
|
| 482 |
+
- *Decoder.* Fold the per-step adaLN gates exactly into per-step copies of the gated weights (`attn.out`,
|
| 483 |
+
`mlp1.fc2`): this removes 66 gate multiplies, at ~65 MB of extra DRAM for the 11 x 3 copies with hi / lo parts.
|
| 484 |
+
- *Solver.* The update `A y - B m_k + Cm m_(k-1)` becomes one ternary / `generic_op` program instead of 3-5 (83
|
| 485 |
+
programs, 0.45 ms).
|
| 486 |
+
6. **Megakernels (D19; after 1-4, gains measured against what is left).**
|
| 487 |
+
- **MK-D, the persistent DiT-evaluation kernel.**
|
| 488 |
+
- *Today.* One evaluation is 4.67 ms in ~435 programs. Its matmul work is 12.0 GFLOP, split passes included
|
| 489 |
+
(74 µs at the HiFi4 peak).
|
| 490 |
+
- *Fits in L1.* Its weights are 8.7 MB (DiT blocks) + 1.8 MB (pre-projection, final layer) as bf16 hi parts, plus
|
| 491 |
+
the fp32 lo parts; the activations are `[352, 256]` fp32 (360 KB); the hoisted cross K / V is 3.5 MB.
|
| 492 |
+
- *Parameters.* Compile-time: agent bucket (352), heads 8 x 32, MLP 1024, depth 3, NFE 11. The folded adaLN
|
| 493 |
+
rows and the solver coefficients move from 11 unrolled trace copies into an L1 table (RT-dev constants per
|
| 494 |
+
steps value).
|
| 495 |
+
- *Target.* 0.5-1 ms per evaluation would cut the decoder from 51 to ~6-11 ms.
|
| 496 |
+
- **MK-E (one mixer trunk) and MK-F (one fusion block)** for the encoder.
|
| 497 |
+
- **The single persistent megakernel for the whole plan** (PLAN 5.3) is plausible. The 48 MB of uploaded weights /
|
| 498 |
+
constants and the largest activation (neighbour mixer `[320, 64, 128]` fp32, 10.5 MB) fit the ~175 MB of
|
| 499 |
+
aggregate L1. The compute floor is ~1 ms.
|
| 500 |
+
- It ships only if it is faster with the gates intact; the attempt is recorded either way (D19).
|
| 501 |
+
7. **Compaction (LOAD-time buckets; 20-40 % of what remains after 1-4, to be measured).**
|
| 502 |
+
- The device time is scene-independent: kashiwanoha uses 88 / 320 neighbours, 123 / 140 lanes and 89 / 352
|
| 503 |
+
decoder rows.
|
| 504 |
+
- Valid-entity buckets need one trace per bucket and the exact compaction rewrites of PLAN 5.3:
|
| 505 |
+
- they shrink the bandwidth-bound encoder work in proportion;
|
| 506 |
+
- they shrink the attention scores quadratically (`[8, 352, 352]` -> `[8, 96, 96]` at 89 agents);
|
| 507 |
+
- they help the latency-bound decoder matmuls little.
|
| 508 |
+
8. **Host side, e2e only (≈ −4 to −6 ms per request).** 15.8 ms of host work and transfers surround the 102 ms replay.
|
| 509 |
+
- Pack `neighbor_x` as a 2-D `[1, 1, 1920, 9]` input: 1.31 MB of TILE padding becomes 245 KB, which also feeds
|
| 510 |
+
item 1.
|
| 511 |
+
- Skip the `y0` upload while `x_T` is unchanged (zeros at temperature 0, the node's default) and send `cs` as its
|
| 512 |
+
4 columns. Together: H2D 3.6 -> ~1.6 MB, and less host tilization (`host_in` 2.4 ms).
|
| 513 |
+
- Vectorise the per-entity loops of `host_pre` (4.8 ms) and `host_post` (4.3 ms: the predicted paths of up to 320
|
| 514 |
+
agents).
|
| 515 |
+
- The 7 ms `.npz` decode is the caller's choice (pass arrays).
|
| 516 |
+
9. **Dispatch configuration: no gain today; re-measure after 1-5.**
|
| 517 |
+
- Gaps are only 4.4 ms (0.56 µs median). What pays is removing programs: each removed short program saves its
|
| 518 |
+
~5-8 µs kernel plus ~0.6 µs.
|
| 519 |
+
- Re-run the ETH / WORKER and 1 / 2 CQ matrix once the op count drops below ~2,000. ETH's 12th column matters only
|
| 520 |
+
once ops become throughput-bound.
|
| 521 |
+
- The ETH-2CQ penalty (+1.6 µs per program) needs a dispatch-core profile before 2 CQs are reconsidered.
|
| 522 |
+
|
| 523 |
+
## Raw logs and CSVs
|
| 524 |
+
|
| 525 |
+
Under `logs/diffusion-planner/baseline/` (workspace) unless absolute:
|
| 526 |
+
|
| 527 |
+
| what | files |
|
| 528 |
+
|---|---|
|
| 529 |
+
| re-vendoring + suites | `host_suite.log` (105 passed / 46 skipped; again on the final tree with the new scripts: `host_suite_final.log`), `device_suite.log` (44 passed), `alloc_tracking.log` (44 passed), `e2e_device.json` (99 scenes), `windowA.log` |
|
| 530 |
+
| stage bench + dispatch / CQ matrix | `matrix_{eth-1cq,eth-2cq,worker-2cq,worker-1cq,eth-1cq-repeat}.{log,json}`, `matrix.rcs`, `tables.md` |
|
| 531 |
+
| numerics ablation | `precision_{shipped,no_attnmm,no_decsplit,no_encsplit,no_ln32,round1,fastest}.{log,json}`, `precision.rcs`, `windowB.log` |
|
| 532 |
+
| device profile | `profile_tracy.log`, `profile_run.json`; ops CSV `/home/ubuntu/experiments/tt-models/generated/profiler/diffusion-planner_baseline/reports/2026_10_08_06_13_03/ops_perf_results_2026_10_08_06_13_03.csv`; summaries `profile_summary.{json,md}` (`analyze_profile.py`), `profile_breakdown.md`, `kind_by_group.md`; tt-perf-report `tt_perf_report_trace.{txt,csv}`, `tt_perf_report_trace_summary.{csv,png}` |
|
| 533 |
+
| ttnn-visualizer | `visualizer_capture.log`, `visualizer_import.log`; reports in `/home/ubuntu/experiments/tt-models/generated/ttnn_visualizer/diffusion-planner_baseline/` |
|
| 534 |
+
| op counts / FLOPs on the fake ttnn (CPU) | `fake_signposts.log` (6,408 fake ops by stage / kind), `fake_flops.log` (165 GFLOP per plan); the port's per-configuration counts: `logs/diffusion-planner/opcount_fake_r2.log` |
|
| 535 |
+
| scripts | `scripts/` (`windowA.sh`, `windowB.sh`, `job{1..5}_*.sh`, `chain.sh`, `analyze_profile.py`, `profile_breakdown.py`, `tables.py`, `fake_*.py`, `visualizer_*.py`) |
|
OPT_REPORT.md
ADDED
|
@@ -0,0 +1,142 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# diffusion-planner-p150 optimization report (p150, ETH dispatch, 12×10 grid)
|
| 2 |
+
|
| 3 |
+
**Status: baseline port; optimization pending.** This is the first public release: the functional port, measured
|
| 4 |
+
once (`OPT_BASELINE.md`, baseline commit `5541833`, 2026-10-08) and not optimized yet. No optimization round has run,
|
| 5 |
+
so there is no step, no rejected attempt and no hang to report, and every number below is the baseline. The whole
|
| 6 |
+
plan (encoder, the 11 DiT evaluations with the DPM-Solver++(2M) updates, the turn-indicator head) runs as one metal
|
| 7 |
+
trace of 6,282 programs: 102.04 ms per back-to-back replay (9.80 plans/s of device throughput), 117.9 ms per
|
| 8 |
+
synchronous `model()` call on the shipped sample. The trace is kernel-bound (op-to-op gaps 4.4 ms of a 103.7 ms span).
|
| 9 |
+
Two sinks of similar size come first: the **numerics defaults** that the end-to-end gates need (split hi / lo matmuls,
|
| 10 |
+
fp32 LayerNorm, fp32 matmul attention) cost **33.4 ms** over the first device round, and the encoder's channel-MLP and
|
| 11 |
+
pre-projection matmuls run on only 4-8 cores (**25.8 ms**, unrelated to precision).
|
| 12 |
+
|
| 13 |
+
All numbers: `code/scripts/bench.py` (median of 100 warm iterations, batch 1,
|
| 14 |
+
`code/tt_diffusion_planner/samples/kashiwanoha_dense.npz`), ETH dispatch, 1 CQ, 12×10 grid, the pinned numerics
|
| 15 |
+
(`DIFFUSION_PLANNER_SPLIT_MATMUL=enc.island.*,enc.pre.*,dec.*`, `_LN_FP32=enc.mixer.*,dec.*`,
|
| 16 |
+
`_ATTN_MATMUL=enc.fusion.attn,dec.*`), unless marked otherwise. Device profile: `code/scripts/profile_ops.py` under the
|
| 17 |
+
device profiler (one traced replay between signposts). Accuracy gates: `OPT_BASELINE.md` "How to run".
|
| 18 |
+
|
| 19 |
+
## Summary
|
| 20 |
+
|
| 21 |
+
| | baseline `5541833` (2026-10-08) | **final (= baseline: no round yet)** |
|
| 22 |
+
|---|---|---|
|
| 23 |
+
| device trace, one blocking plan | 102.13 ms (p99 104.81) | same |
|
| 24 |
+
| back-to-back traces | 102.04 ms (9.80 plans/s) | same |
|
| 25 |
+
| e2e `model(inputs=arrays)` p50 / p99 | 117.90 / 134.56 ms | same |
|
| 26 |
+
| e2e `model(inputs=<.npz path>)` p50 | 124.76 ms | same |
|
| 27 |
+
| host pre-processing / pack / host tensors / H2D / D2H / post-processing | 4.76 / 0.45 / 2.41 / 1.14 / 0.58 / 4.32 ms | same |
|
| 28 |
+
| device programs per plan (unique programs) | 6,282 (309) | same |
|
| 29 |
+
| kernel sum / op-to-op gaps / span (profile) | 99.27 / 4.39 / 103.66 ms | same |
|
| 30 |
+
| encoder / 11 decoder evaluations (+ solver updates) | 51.6 / 51.4 ms (4.67 ms per evaluation) | same |
|
| 31 |
+
| accuracy gates (PCC / agreement vs the fp32 CPU reference) | 44 / 44 device tests: module PCC ≥ 0.999952, encoding 0.999978, decoder evaluation ≥ 0.9999995; 99 scenes: ego max 0.313 m / mean 0.143 m (gates 1.0 / 0.3 m), turn command 99 / 99, neighbours ≤ 0.086 m (gate 1.5 m) | same |
|
| 32 |
+
| served `/predict` `timing_ms.total`, median of 50 (uvicorn on the host, the shipped sample) | 123.1 ms, measured on `c0d84f9` (same device code; the release verification) | same |
|
| 33 |
+
| `from_pretrained` load: empty JIT cache / warm cache | 315 s / 8.6 s (`c0d84f9`) | same |
|
| 34 |
+
|
| 35 |
+
The host is shared with other agents' jobs, so the host-side rows move with its load (several ms); the device rows
|
| 36 |
+
repeat to ±0.05 ms. AICLK 1350 MHz (min 1343) during the stage bench. The release re-check on `c0d84f9`
|
| 37 |
+
(`VERIFICATION_2026-10-08.md`) reproduced the device rows: back-to-back replay 102.03 vs 102.04 ms, one plan 102.10 vs 102.13 ms; the host-bound end to end 115.6 vs 117.9 ms p50 moves with the host load.
|
| 38 |
+
|
| 39 |
+
## Steps (chronological; each row is one commit)
|
| 40 |
+
|
| 41 |
+
| step | commit | what | trace / b2b ms | e2e ms | accuracy | revert switch |
|
| 42 |
+
|---|---|---|---|---|---|---|
|
| 43 |
+
| baseline | `5541833` | first correct port: the whole plan in one trace, ETH 1CQ 12×10, the numerics defaults of PORT_LOG decision 11, host pre- and post-processing | 102.13 / 102.04 | 117.90 | all gates (above) | – |
|
| 44 |
+
|
| 45 |
+
Later commits up to the first release changed no model numerics and no device code: the baseline documentation
|
| 46 |
+
(`a6e5bbf`), the ttaw 0.19.0 re-vendor with the `serve.env` pins of the numerics knobs, the quickstart picture and two
|
| 47 |
+
test changes (`d521981`), the ttaw 0.20.0 re-vendor (`c0d84f9`), and the release docs and demo media. After the re-vendor the device suite was re-run with
|
| 48 |
+
identical gate values and per-scene numbers (`VERIFICATION_2026-10-08.md`).
|
| 49 |
+
|
| 50 |
+
## Round 1
|
| 51 |
+
|
| 52 |
+
Not started. The optimization phase follows the publication of every model of the collection (`research/PLAN.md`
|
| 53 |
+
§4.1, §5): each step is one commit with one `DIFFUSION_PLANNER_*` A/B knob, keeps the frozen gates and is re-checked on
|
| 54 |
+
all 99 end-to-end scenes (the plan's sensitivity to device numerics is chaotic per scene: PORT_LOG known issues). The
|
| 55 |
+
precision policy (PLAN.md §9.3) may be relaxed per module only with gate evidence; the precision cost is to be
|
| 56 |
+
recovered with fused kernels that keep fp32-level accuracy, not by dropping precision.
|
| 57 |
+
|
| 58 |
+
### Findings (measured)
|
| 59 |
+
|
| 60 |
+
The baseline profile (`OPT_BASELINE.md` "Device profile") is the starting point:
|
| 61 |
+
|
| 62 |
+
- The trace is **kernel-bound, not dispatch-bound**: op-to-op gaps total 4.39 ms of a 103.66 ms span (median gap
|
| 63 |
+
0.56 µs). Removing a short program saves its ~5-8 µs of kernel time plus ~0.6 µs.
|
| 64 |
+
- The plan does 165 GFLOP of matmul (split passes counted), ~1 ms at the 120-core HiFi4 peak: the time goes to badly
|
| 65 |
+
parallelised matmuls and thousands of small element-wise programs. Matmul efficiency is 7.2 % of the FLOP roofline
|
| 66 |
+
on average.
|
| 67 |
+
- By layer kind: split matmuls 34.8 ms, plain linears 24.0 ms (23.3 of them in the encoder's mixer trunks), fp32
|
| 68 |
+
LayerNorm 21.6 ms, fp32 matmul attention 16.4 ms, glue 5.6 ms.
|
| 69 |
+
- What each numerics default costs (device time per plan, timing only; only the shipped configuration is gate-clean):
|
| 70 |
+
decoder split matmuls 20.2 ms, fp32 LayerNorm 18.1 ms, fp32 matmul attention 15.5 ms, encoder split matmuls 6.0 ms;
|
| 71 |
+
all four together 58.3 ms (the fastest graph, 43.74 ms, fails `enc.ego`), the decision-11 additions over the first
|
| 72 |
+
device round 33.4 ms (round-1 defaults: 68.64 ms, which fail `ego.mean_err_m` on nuScenes scene-0103_kf14).
|
| 73 |
+
- Dispatch / CQ matrix (back-to-back): ETH-1CQ 102.04, ETH-2CQ 111.87, WORKER-2CQ 101.88, WORKER-1CQ 101.87 ms. WORKER
|
| 74 |
+
is equal (the 12th column hardly matters to this graph yet), so ETH stays the default (D14); 2 CQs cost +9.8 ms on ETH,
|
| 75 |
+
so 1 CQ stays pinned.
|
| 76 |
+
|
| 77 |
+
### Megakernel / fusion work
|
| 78 |
+
|
| 79 |
+
None yet. The candidates (D19) are in the backlog (item 6): the persistent DiT-evaluation kernel first.
|
| 80 |
+
|
| 81 |
+
## Known hangs (all rounds)
|
| 82 |
+
|
| 83 |
+
| when (UTC) | command | cause | status |
|
| 84 |
+
|---|---|---|---|
|
| 85 |
+
| – | – | none: no hang, timeout, reset or FAULT marker in any device job of the port, the baseline or the release docs | – |
|
| 86 |
+
|
| 87 |
+
## Rejected / not kept
|
| 88 |
+
|
| 89 |
+
None yet. Configurations that were measured and are **not** the default, for accuracy (PORT_LOG decision 11,
|
| 90 |
+
`OPT_BASELINE.md` "What the numerics defaults cost"):
|
| 91 |
+
|
| 92 |
+
- the first device round's defaults (split only for the mixer inputs and the decoder pre-projection, bf16 SDPA):
|
| 93 |
+
68.7 ms, but ego mean 0.347 m > 0.3 m on nuScenes scene-0103_kf14 (job 11);
|
| 94 |
+
- configuration B (A + split and fp32 LayerNorm in the whole encoder): 109 ms, worst ego 0.35 / 0.15 m;
|
| 95 |
+
- configuration C (split and fp32 LayerNorm everywhere): 169 ms, worst ego 0.22 / 0.10 m: about a third less error
|
| 96 |
+
than the shipped configuration A (0.31 / 0.14 m, 105 ms in the same job 15) at +61 % device time (PORT_LOG open
|
| 97 |
+
question 8);
|
| 98 |
+
- the fastest graph (fused LayerNorm, no split, bf16 SDPA): 43.7 ms, fails the encoder PCC gate (`enc.ego` 0.99856 /
|
| 99 |
+
0.99261).
|
| 100 |
+
|
| 101 |
+
## Profile at the end (trace replay)
|
| 102 |
+
|
| 103 |
+
The baseline profile (`OPT_BASELINE.md`): 6,282 programs, kernel sum 99.27 ms, span 103.66 ms; by op code Matmul
|
| 104 |
+
52.3 ms (1,365 programs), BinaryNg 31.5 ms (2,948), Reduce 3.7 ms, Softmax 3.5 ms, Unary 2.8 ms, Typecast 2.4 ms. The
|
| 105 |
+
ten slowest programs are the neighbour mixer's channel matmuls `[1, 320, 64, 128] @ [128, 128]`, ~970 µs each on 8
|
| 106 |
+
cores.
|
| 107 |
+
|
| 108 |
+
## Remaining backlog (gains on the 102.0 ms replay unless marked e2e; estimates, not measurements)
|
| 109 |
+
|
| 110 |
+
From `OPT_BASELINE.md` "Ranked optimization opportunities". The gains are not additive: each item shrinks the base of
|
| 111 |
+
the next.
|
| 112 |
+
|
| 113 |
+
1. **Encoder channel MLPs on 2-D activations (≈ −20 to −22 ms; exact rewrite, low effort).** The mixer `channels_mlp`
|
| 114 |
+
and the channel pre-projections run `ttnn.linear` on 4-D `[1, E, T, C]` activations, a batched matmul on 4-8 cores
|
| 115 |
+
(90 programs, 25.8 ms). A free `[1, 1, E·64, 128]` view for the MixerBlocks (as the token-mixing matmuls already do:
|
| 116 |
+
69 µs on 117 cores) and a re-tiled 2-D layout or an explicit multi-core program config (with the HiFi4 compute
|
| 117 |
+
config: a `program_config` alone falls back to LoFi) for the pre-projections.
|
| 118 |
+
2. **A fused split (bf16x3) matmul, one program per linear (≈ −17 to −19 ms).** 338 split linears are 2,802 programs
|
| 119 |
+
and 34.8 ms today; a `generic_op` matmul that splits hi / lo on the fly and accumulates the three products in fp32
|
| 120 |
+
DEST removes the typecast / subtract / add programs (9.3 ms in the decoder) and shares the operand reads.
|
| 121 |
+
3. **A fused fp32 LayerNorm, one program per LN (≈ −15 to −16 ms).** 237 LNs are 2,133 programs and 21.6 ms; the fused
|
| 122 |
+
`ttnn.layer_norm` is ~2.5e-3 relative and fails the mixers (PORT_LOG job 5), so the kernel must keep fp32
|
| 123 |
+
statistics.
|
| 124 |
+
4. **Attention (≈ −5.5 ms with stock ops, ≈ −12 to −13 ms with a fused kernel).** 72 fp32 matmul attentions, 16.4 ms:
|
| 125 |
+
fold the 1/√32 scale into the Q weights and fuse the mask into the softmax (stock ops), spread the P·V matmuls over
|
| 126 |
+
heads, or one fp32-accurate flash-attention `generic_op` (scores in L1, hi / lo operands).
|
| 127 |
+
5. **Glue (≈ −3 to −4 ms).** Mixer transposes and residual adds, the adaLN gate multiplies (exact fold into per-step
|
| 128 |
+
weight copies), the solver update as one program instead of 3-5.
|
| 129 |
+
6. **Megakernels (D19; after 1-4).** MK-D, the persistent DiT-evaluation kernel (today 4.67 ms in ~435 programs for
|
| 130 |
+
12.0 GFLOP; its weights fit in L1; target 0.5-1 ms per evaluation, the decoder from 51 to ~6-11 ms); MK-E / MK-F for
|
| 131 |
+
a mixer trunk / a fusion block; a single persistent megakernel for the whole plan is plausible (48 MB of weights and
|
| 132 |
+
constants, the largest activation 10.5 MB, ~175 MB of aggregate L1). It ships only if faster with the gates intact;
|
| 133 |
+
the attempt is recorded either way.
|
| 134 |
+
7. **Compaction (LOAD-time buckets; 20-40 % of what remains after 1-4, to be measured).** The device time is
|
| 135 |
+
scene-independent: kashiwanoha uses 88 / 320 neighbours and 89 / 352 decoder rows. Valid-entity buckets (one trace
|
| 136 |
+
per bucket, exact compaction rewrites) shrink the encoder in proportion and the attention scores quadratically.
|
| 137 |
+
8. **Host side (e2e ≈ −4 to −6 ms per request).** 15.8 ms of host work and transfers surround the replay: a 2-D
|
| 138 |
+
`neighbor_x` upload (1.31 MB of TILE padding -> 245 KB), skipping the `y0` upload while `x_T` is zero and sending
|
| 139 |
+
`cs` as its 4 columns (H2D 3.6 -> ~1.6 MB), vectorised per-entity loops in the pre- and post-processing.
|
| 140 |
+
9. **Dispatch configuration: no gain today; re-measure after 1-5.** Re-run the ETH / WORKER and 1 / 2 CQ matrix once
|
| 141 |
+
the program count drops below ~2,000; the ETH-2CQ penalty (+1.6 µs per program) needs a dispatch-core profile
|
| 142 |
+
first.
|
PYTHON.md
ADDED
|
@@ -0,0 +1,179 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Python API: Diffusion Planner v5.0 (Autoware diffusion_planner) on Blackhole
|
| 2 |
+
|
| 3 |
+
Use this API from Python code (a pipeline, a notebook, a ROS 2 node wrapper). You do not need the HTTP server: the
|
| 4 |
+
API and the server share the decoders, the device trace and the post-processing, so the outputs and the speed are
|
| 5 |
+
the same.
|
| 6 |
+
|
| 7 |
+
## Install
|
| 8 |
+
|
| 9 |
+
Install the package on top of an environment that already has `ttnn` (a tt-metal `python_env` at `44d66500520`
|
| 10 |
+
with `patches/tt-metal-eth-dispatch.patch`, or the tt-model container). From the root of the model repository (the
|
| 11 |
+
directory that holds `pyproject.toml`, `README.md` and `code/`):
|
| 12 |
+
|
| 13 |
+
```bash
|
| 14 |
+
pip install -e . # the Python API (numpy<2, pillow, pyyaml, onnx, huggingface_hub)
|
| 15 |
+
pip install -e ".[server,test]" # + the HTTP server and the tests
|
| 16 |
+
```
|
| 17 |
+
|
| 18 |
+
The pip project is the repository's top-level `pyproject.toml`; it installs the package from
|
| 19 |
+
`code/tt_diffusion_planner` (there is no `pyproject.toml` inside `code/`, because the container build copies `code/`
|
| 20 |
+
over the tt-metal tree). ttnn and torch come from tt-metal and are not declared.
|
| 21 |
+
|
| 22 |
+
The package carries `tt_diffusion_planner.ttaw`, the shared code of the Autoware ports to Blackhole (device open, trace
|
| 23 |
+
runner, decoders, model base class, HTTP app), vendored at the version recorded in
|
| 24 |
+
`code/tt_diffusion_planner/ttaw/VENDORED.json`.
|
| 25 |
+
|
| 26 |
+
| You want to run | Extras |
|
| 27 |
+
|---|---|
|
| 28 |
+
| the Python API | none |
|
| 29 |
+
| the HTTP server (`tt_diffusion_planner.server.app`, see `SERVING.md`) | `server` |
|
| 30 |
+
| host tests (no device; device tests are skipped): `TT_VISIBLE_DEVICES=none python -m pytest -q code/tt_diffusion_planner/tests` | `server,test` |
|
| 31 |
+
| device tests: `python -m pytest -q -s code/tt_diffusion_planner/tests/test_pcc_device.py code/tt_diffusion_planner/tests/test_e2e_device.py` | `test` |
|
| 32 |
+
|
| 33 |
+
## Quickstart
|
| 34 |
+
|
| 35 |
+
```python
|
| 36 |
+
from tt_diffusion_planner import DiffusionPlanner
|
| 37 |
+
|
| 38 |
+
with DiffusionPlanner.from_pretrained(device_id=0) as model:
|
| 39 |
+
out = model(inputs="code/tt_diffusion_planner/samples/kashiwanoha_dense.npz")
|
| 40 |
+
print(out.to_dict()) # the POST /predict body
|
| 41 |
+
```
|
| 42 |
+
|
| 43 |
+
`examples/quickstart.py` runs the same snippet, writes `quickstart.json` and a bird's-eye view of the input and the
|
| 44 |
+
plan (`quickstart_bev.png`).
|
| 45 |
+
|
| 46 |
+
## `DiffusionPlanner.from_pretrained(...)`
|
| 47 |
+
|
| 48 |
+
```python
|
| 49 |
+
DiffusionPlanner.from_pretrained(
|
| 50 |
+
model_id=None, # HF repo or a local directory with the weights files; default AutowareFoundation/diffusion_planner
|
| 51 |
+
*,
|
| 52 |
+
revision=None, # default for the default repo: the validated commit 423efde67f5 (tag v5.0)
|
| 53 |
+
variant=None, # "default" (the only v5.0 graph); default $DIFFUSION_PLANNER_VARIANT or "default"
|
| 54 |
+
device_id=None, # chip to open; default $TT_DEVICE_ID or 0
|
| 55 |
+
device=None, # an already-opened ttnn device (tt_diffusion_planner.device.open_device); close() does not close it
|
| 56 |
+
dispatch=None, # "eth" (p150 target, 12x10 grid) | "worker" (A/B only, 11x10) | "auto"; default $DIFFUSION_PLANNER_DISPATCH or "eth"
|
| 57 |
+
num_command_queues=None, # default $DIFFUSION_PLANNER_NUM_CQS or 1
|
| 58 |
+
weights_dir=None, # explicit local weights directory; no Hub access
|
| 59 |
+
warmup_variants="default", # trace variants to capture now; see "Warm-up"
|
| 60 |
+
verbose=False,
|
| 61 |
+
precision=None, # the only compile parameter: extra precision-policy rules, e.g. "dec.*=HiFi2+fp32" (experiments only)
|
| 62 |
+
) -> DiffusionPlanner
|
| 63 |
+
```
|
| 64 |
+
|
| 65 |
+
What it does: resolves the weights first, so a Hub problem never claims the chip (`weights_dir` >
|
| 66 |
+
`$DIFFUSION_PLANNER_WEIGHTS_DIR` > a local `model_id` directory > the HF snapshot at the pinned revision, restricted to
|
| 67 |
+
the three v5.0 ONNX files and `diffusion_planner.param.json`, with an offline fallback to the cache; the sha256 of every
|
| 68 |
+
file and the weights' `major_version == 5` are checked), opens the chip (ETH dispatch, 12×10, 1 CQ; the other open
|
| 69 |
+
parameters are `DEVICE_DEFAULTS` in `tt_diffusion_planner/device.py`, overridable with `DIFFUSION_PLANNER_*`), reads the
|
| 70 |
+
ONNX initializers as data, uploads the weights and constants (48.1 MB), builds the graph, then compiles and captures
|
| 71 |
+
the metal trace. If ETH dispatch cannot open (tt-metal without the patch), it warns and falls back to WORKER dispatch
|
| 72 |
+
(`model.info["device"]["fallback"]` names it). Any other keyword argument is a `TypeError`.
|
| 73 |
+
|
| 74 |
+
The numerics are not arguments: the published configuration is the default of the `DIFFUSION_PLANNER_LN_FP32`,
|
| 75 |
+
`_SPLIT_MATMUL`, `_ATTN_MATMUL`, `_HIDDEN_FP32` and `_ATTN_FP32_ACC` knobs (`tt_diffusion_planner.tt.config.KNOBS`,
|
| 76 |
+
pinned in `tt-model.yaml` `serve.env`). Setting one of them in the environment changes the graph and invalidates the
|
| 77 |
+
accuracy figures of the card until the gates are re-run.
|
| 78 |
+
|
| 79 |
+
## Warm-up
|
| 80 |
+
|
| 81 |
+
`from_pretrained` returns a warm model: it builds the graph, runs the plan once eagerly (this first run compiles every kernel into the JIT cache), then captures the whole plan as one metal trace (`warmup_variants="default"`: the variant `plan`) with program-cache misses forbidden, so no later call compiles anything. `model.warmup()` is idempotent; `warmup_variants="none"` defers the capture to `model.warmup()`.
|
| 82 |
+
|
| 83 |
+
Measured on the shipped sample (`model.info["warmup_ms"]`, 2026-10-09): the load takes 315 s with an empty JIT cache and 8.6 s with a warm one (build 0.53 s: the ONNX initializers read and 48.1 MB of weights and constants uploaded; warm-up and capture 3.7 s; the rest is the device open). The first call then takes 120 ms and the second 120 ms (the stage bench's steady state: 118 ms p50 for decoded arrays, 125 ms for an `.npz` path). The trace holds 74.6 MB of DRAM (`trace_region_size` 192 MiB).
|
| 84 |
+
|
| 85 |
+
## Call: `model(...)`
|
| 86 |
+
|
| 87 |
+
| Argument | Type | Description |
|
| 88 |
+
|---|---|---|
|
| 89 |
+
| `inputs` | mapping / `.npz` path / bytes / JSON envelope | the 15 raw planner tensors (see "Input types") |
|
| 90 |
+
| `velocity_smoothing_window` | int, 1..79, default 8 | forward moving average of the trajectory velocity, in points |
|
| 91 |
+
| `stopping_threshold` | float >= 0, default 0.3 | force stop below this smoothed speed (m/s), when the ego moves |
|
| 92 |
+
| `turn_indicator_keep_offset` | float, default -1.25 | added to the KEEP logit before the turn-indicator decision |
|
| 93 |
+
| `return_denoising_steps` | bool, default False | add the ego row of the 11 solver iterates (`out.meta["denoising_steps"]`, `[11, 81, 4]`, the node's `~/debug/denoising_steps`) |
|
| 94 |
+
|
| 95 |
+
### Input types
|
| 96 |
+
|
| 97 |
+
- `inputs=`: the 15 raw tensors of the Autoware node's `DiffusionPlannerCore::create_input_data()` (batch 1,
|
| 98 |
+
float32, ego `base_link` frame, BEFORE normalization; names and shapes in `tt_diffusion_planner.INPUT_SCHEMA`):
|
| 99 |
+
a `{name: array}` mapping (numpy or torch), an `.npz` path or its bytes, or the `/predict` envelope
|
| 100 |
+
`{"format": "npz", "data": <base64>}` / `{"format": "json", "arrays": {...}}`. Names, shapes and finite values are
|
| 101 |
+
checked (`InputError`). `tt_diffusion_planner.load_inputs(source)` is the same decoder.
|
| 102 |
+
- Any other input (`points`, `images`, `calibration`, ...) is refused (`InputError`).
|
| 103 |
+
|
| 104 |
+
### What the caller keeps (the API is stateless)
|
| 105 |
+
|
| 106 |
+
One call is one independent plan. The Autoware node keeps state between plans; to reproduce it over a sequence of
|
| 107 |
+
plans, the caller keeps the same state (SERVING.md 3.5 has the details):
|
| 108 |
+
|
| 109 |
+
- **The tensors.** The node's pre-processing from ROS messages and the Lanelet2 map (per-UUID agent buffers and their
|
| 110 |
+
0.1 s resampling, the ego history, lane / route / polygon / line-string selection and encoding, traffic lights,
|
| 111 |
+
speed limits, goal, turn-indicator report history) is not part of the bundle.
|
| 112 |
+
- **The turn-indicator hold window** (`turn_indicator_hold_duration`, 1.0 s in the node's YAML). Each call decides
|
| 113 |
+
with a fresh manager; apply the node's hold across calls with the node's own manager:
|
| 114 |
+
|
| 115 |
+
```python
|
| 116 |
+
from tt_diffusion_planner.host.postprocess import TurnIndicatorManager
|
| 117 |
+
|
| 118 |
+
manager = TurnIndicatorManager() # hold 1.0 s, KEEP offset -1.25 (the node's YAML)
|
| 119 |
+
out = model(inputs=tensors)
|
| 120 |
+
decision = manager.evaluate(out.turn_indicator["logits"], stamp_s=now_s, prev_report=int(tensors["turn_indicators"][0, 30]))
|
| 121 |
+
command = decision.command # the held command while less than 1.0 s has passed
|
| 122 |
+
```
|
| 123 |
+
|
| 124 |
+
- **The initial solver state** `sampled_trajectories` (`x_T`, normalised space): zeros is the node's default
|
| 125 |
+
(`temperature: [0.0]`); for a temperature > 0 send N(0, 1) x temperature; for the RTC prefix (`delay_step` > 0) put
|
| 126 |
+
the previous plan into the ego row, slots t = 0 .. delay_step (x as (x - 10) / 20, y as y / 20, cos / sin as they are,
|
| 127 |
+
in the current ego frame). `delay` is accepted and ignored (the node's multi-step mode never reads it).
|
| 128 |
+
- **The map frame.** Outputs are in `base_link`; the node transforms them with the current ego pose to `map`.
|
| 129 |
+
|
| 130 |
+
## Output
|
| 131 |
+
|
| 132 |
+
`model(...)` returns a `tt_diffusion_planner.Output` (= `ttaw.outputs.Trajectory`); `out.to_dict()` is exactly the `POST /predict` body (SERVING.md section 3.2).
|
| 133 |
+
|
| 134 |
+
| field | type | meaning |
|
| 135 |
+
|---|---|---|
|
| 136 |
+
| `poses` | float32 `[80, 7]` | the ego trajectory at 0.1-8.0 s in `base_link`: x, y, yaw, cos, sin, velocity, acceleration (`out.columns`), post-processed like the node's `~/output/trajectory` |
|
| 137 |
+
| `turn_indicator` | dict | `command` (0 NO_COMMAND, 1 DISABLE, 2 ENABLE_LEFT, 3 ENABLE_RIGHT), `command_name`, `keep_selected`, `held` (always false: no hold window), the 5 raw `logits` (NONE, DISABLE, LEFT, RIGHT, KEEP), the decision's `probabilities` |
|
| 138 |
+
| `predicted_agents` | float32 `[N, 80, 5]` | x, y, yaw, cos, sin of each non-empty neighbour row, in input order |
|
| 139 |
+
| `meta` | dict | `predicted_agent_rows` (the rows of `predicted_agents`), `predicted_agent_columns`, `force_stop`, `time_from_start_s`, `valid_counts` (the entities the encoder saw), and with `return_denoising_steps` the encoded `denoising_steps` |
|
| 140 |
+
| `timing_ms` | dict | `preprocess`, `device` (host tensors + H2D + replay + D2H), `postprocess`, `total` |
|
| 141 |
+
|
| 142 |
+
`out.to_dicts()` gives one `{x, y, yaw, cos, sin, velocity, acceleration}` dict per trajectory point; `out.to_dict("npz")` adds the poses as a lossless base64 NPZ array.
|
| 143 |
+
|
| 144 |
+
## Lifetime and information
|
| 145 |
+
|
| 146 |
+
- `model.close()` releases the trace and the persistent device tensors and closes the chip if the model opened it;
|
| 147 |
+
idempotent. `with` calls it for you; an unclosed model is closed when Python exits.
|
| 148 |
+
- `model.info`: weights (repo, tag, revision, path), device (dispatch, grid, CQs, fallback), variant, warm variants,
|
| 149 |
+
warm-up times, runtime parameter defaults, the input schema, the numerics options and precision policy in effect, the
|
| 150 |
+
trace (variants, persistent inputs, trace buffers in MB).
|
| 151 |
+
- Calls from several threads are safe: the device calls are serialised. One model per process per chip.
|
| 152 |
+
|
| 153 |
+
## Speed
|
| 154 |
+
|
| 155 |
+
Warm calls, batch 1, ETH dispatch, 1 CQ, 12×10, the pinned numerics (`code/scripts/bench.py`, 100 iterations; the numbers of `OPT_BASELINE.md`, 2026-10-08, on a shared host; p50, with p99 in brackets):
|
| 156 |
+
|
| 157 |
+
| stage | shipped sample `kashiwanoha_dense` |
|
| 158 |
+
|---|---:|
|
| 159 |
+
| `.npz` decode + schema check (path inputs only) | 7.02 (18.20) ms |
|
| 160 |
+
| host pre-processing (the node's normalization, the encoder's host features, masks, `x_T`) | 4.75 (12.55) ms |
|
| 161 |
+
| packing the 19 persistent trace inputs · ttnn host tensors · H2D | 0.45 · 2.41 · 1.14 ms |
|
| 162 |
+
| **device trace, one blocking plan** | **102.13** (104.81) ms |
|
| 163 |
+
| D2H (one packed read, 460 KB) | 0.58 ms |
|
| 164 |
+
| host post-processing (trajectory, predicted paths, turn decision) | 4.32 (11.54) ms |
|
| 165 |
+
| **`model(inputs=arrays)` end to end** | **117.90** (134.56) ms |
|
| 166 |
+
| `model(inputs=<.npz path>)` | 124.76 (147.74) ms |
|
| 167 |
+
| back-to-back replays (device time per plan) | 102.04 ms = 9.80 plans/s |
|
| 168 |
+
|
| 169 |
+
The device time does not depend on the scene: every plan computes the full capacities (re-checked on c0d84f9: kashiwanoha_dense 102.10, straight_road 102.11, a nuScenes instant 102.09 ms). Throughput above one plan per ~118 ms needs pipelining of the host work of neighbouring requests (not implemented); 2 CQs do not help a synchronous request (`OPT_BASELINE.md`). Where the time goes and what comes next: `OPT_REPORT.md`.
|
| 170 |
+
|
| 171 |
+
## Limits
|
| 172 |
+
|
| 173 |
+
- Batch 1 on the chip; one model per process.
|
| 174 |
+
- Fixed shapes of the v5.0 export (320 neighbours, 140 lanes, 25 route lanes, 10 polygons, 60 line strings, 31 history
|
| 175 |
+
and 80 future steps) and 10 DPM-Solver steps (11 decoder evaluations) are compiled into the trace; every plan computes
|
| 176 |
+
the full capacities, so the device time does not depend on the scene.
|
| 177 |
+
- The node's guidance services (start / stop / centerline guidance) are not available (the node's default is off).
|
| 178 |
+
- Accuracy is agreement with the fp32 CPU reference of the same network (README "Demo & Performances"); the planner's
|
| 179 |
+
driving quality is the weights' (trained by TIER IV on data that is not public).
|
README.md
ADDED
|
@@ -0,0 +1,195 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
---
|
| 2 |
+
tags:
|
| 3 |
+
- blackhole
|
| 4 |
+
- p150
|
| 5 |
+
- tt-dit-server
|
| 6 |
+
- tt-model-cache
|
| 7 |
+
- tt-model-container
|
| 8 |
+
- tenstorrent
|
| 9 |
+
- ttnn
|
| 10 |
+
- tt-metal
|
| 11 |
+
- tt-nn
|
| 12 |
+
- autoware
|
| 13 |
+
- autonomous-driving
|
| 14 |
+
- planning
|
| 15 |
+
- trajectory-generation
|
| 16 |
+
- motion-prediction
|
| 17 |
+
- diffusion
|
| 18 |
+
- diffusion-planner
|
| 19 |
+
- tt-model-catalog
|
| 20 |
+
pipeline_tag: robotics
|
| 21 |
+
license: apache-2.0
|
| 22 |
+
license_link: https://huggingface.co/AutowareFoundation/diffusion_planner
|
| 23 |
+
base_model:
|
| 24 |
+
- AutowareFoundation/diffusion_planner
|
| 25 |
+
---
|
| 26 |
+
|
| 27 |
+
# diffusion-planner-p150
|
| 28 |
+
|
| 29 |
+
Diffusion Planner v5.0 (Autoware diffusion_planner): the network Autoware deploys in `autoware_diffusion_planner`, ported to one Tenstorrent Blackhole p150 with tt-nn. The whole plan runs on the chip as one metal trace: the scene encoder, the 11 DiT decoder evaluations of the DPM-Solver++(2M) loop with their solver updates, and the turn-indicator head. The Autoware planner tensors in (ego and neighbour histories, lanes, route, polygons, line strings, goal, ego shape, turn-indicator history); an 8 s ego trajectory, the predicted 8 s paths of the neighbours and a turn-indicator command out, with the node's exact pre- and post-processing.
|
| 30 |
+
Weights: [AutowareFoundation/diffusion_planner `v5.0`](https://huggingface.co/AutowareFoundation/diffusion_planner/tree/423efde67f5414734da43a7ad856c17ceb8b51aa) · Paper: [arXiv:2501.15564](https://arxiv.org/abs/2501.15564) · Autoware package: [autoware_diffusion_planner](https://github.com/autowarefoundation/autoware_universe/tree/9ceaccf026c31ffc5319bc9eeb4bd7bede0af3fd/planning/autoware_diffusion_planner) · Training code: [tier4/Diffusion-Planner (TIER IV training fork)](https://github.com/tier4/Diffusion-Planner) · Port: [`code/`](https://huggingface.co/changh95/diffusion-planner-p150/tree/main/code)
|
| 31 |
+
|
| 32 |
+
Runs on **p150** (mesh `P150`). Configuration: dispatch on the ETH cores, 1 command queue, 12×10 compute grid. Numerics (the default and the `serve.env` pins): fp32 residual streams and solver state, HiFi4 with fp32 accumulation, split hi / lo matmuls for the mixer inputs and every decoder linear, fp32 LayerNorm in the mixers and the decoder, fp32 matmul attention in the fusion encoder and the decoder: the configuration the end-to-end accuracy gates need (Caveats). All numbers on this card were measured in this configuration.
|
| 33 |
+
|
| 34 |
+
Packaged and published with [tt-model-manager](https://github.com/tenstorrent/tt-model-manager) 0.1.0 (manifest schema 5.1).
|
| 35 |
+
|
| 36 |
+
## Quickstart (Python)
|
| 37 |
+
|
| 38 |
+
Prerequisite: a tt-metal / ttnn environment at tt-metal [`44d66500520`](https://github.com/tenstorrent/tt-metal/commit/44d66500520fda9f2c7060c0f6b41ec48f7ab37e) with [`patches/tt-metal-eth-dispatch.patch`](patches/tt-metal-eth-dispatch.patch) applied. ttnn is not on PyPI.
|
| 39 |
+
|
| 40 |
+
```bash
|
| 41 |
+
hf download changh95/diffusion-planner-p150 --exclude "image/*" --local-dir diffusion-planner-p150 && cd diffusion-planner-p150
|
| 42 |
+
pip install -e . # adds numpy<2, pillow, pyyaml, onnx, huggingface_hub; ttnn and torch come from tt-metal
|
| 43 |
+
pip install -e ".[server,test]" # optional: the HTTP server and the tests
|
| 44 |
+
```
|
| 45 |
+
|
| 46 |
+
Run the snippet from the model repo root: `code/tt_diffusion_planner/samples/kashiwanoha_dense.npz` is a path relative to it.
|
| 47 |
+
|
| 48 |
+
```python
|
| 49 |
+
from tt_diffusion_planner import DiffusionPlanner
|
| 50 |
+
|
| 51 |
+
with DiffusionPlanner.from_pretrained(device_id=0) as model: # weights -> your HF cache, trace captured
|
| 52 |
+
out = model(inputs="code/tt_diffusion_planner/samples/kashiwanoha_dense.npz") # the 15 raw planner tensors: .npz path, its bytes, or {name: array}
|
| 53 |
+
|
| 54 |
+
print(out.columns) # x, y, yaw, cos, sin, velocity, acceleration (base_link, 0.1-8.0 s)
|
| 55 |
+
print(out.poses[:5])
|
| 56 |
+
print(out.turn_indicator["command_name"], out.predicted_agents.shape)
|
| 57 |
+
```
|
| 58 |
+
|
| 59 |
+
- `from_pretrained` downloads the three v5.0 ONNX files and `diffusion_planner.param.json` (58.9 MB) of [`AutowareFoundation/diffusion_planner`](https://huggingface.co/AutowareFoundation/diffusion_planner) at the pinned commit `423efde67f5` (tag `v5.0`) to your HF cache (no token needed), checks their sha256, opens the chip, builds the graph and captures the metal trace. The first load compiles the kernels (315 s with an empty JIT cache, firmware and every kernel compiled); later loads take about 9 s (device open, reading and uploading the weights, warm-up and trace capture).
|
| 60 |
+
- The trace is captured during the load, so the first call is as fast as the later calls and no call compiles anything.
|
| 61 |
+
- The `with` block releases the trace and closes the chip. Without `with`, call `model.close()`.
|
| 62 |
+
|
| 63 |
+
| | |
|
| 64 |
+
|---|---|
|
| 65 |
+
| **Input** | `inputs=`: the 15 raw tensors of the node's `create_input_data()` (ego frame, before normalization; batch 1): an `.npz` path or its bytes, a `{name: array}` mapping, or the `/predict` JSON envelope. Checked against `DiffusionPlanner.INPUT_SCHEMA` (names, shapes, finite values). Converting ROS messages and the Lanelet2 map into these tensors stays with the client. |
|
| 66 |
+
| **Options** | `velocity_smoothing_window=8`, `stopping_threshold=0.3`, `turn_indicator_keep_offset=-1.25`, `return_denoising_steps=False`. `from_pretrained(device_id=0, dispatch="eth", weights_dir=None, device=None)`. |
|
| 67 |
+
| **Output** | `Trajectory`: `poses` float32 [80, 7] (x, y, yaw, cos, sin, velocity, acceleration; base_link), `turn_indicator` (command, logits, probabilities), `predicted_agents` float32 [N, 80, 5], `timing_ms`, `meta` (the neighbour rows, force stop, the solver iterates on request). |
|
| 68 |
+
| **Methods** | `out.to_dict()` gives the `/predict` JSON. `out.to_dicts()` gives one dict per trajectory point. |
|
| 69 |
+
|
| 70 |
+
- The API gives the same output as the HTTP server `/predict`: both share the decoders, the device trace and the host post-processing (checked on the device by `test_api_equals_server`).
|
| 71 |
+
- The API is stateless: one call is one plan. What the Autoware node keeps between plans stays with the caller: the turn-indicator hold window (1.0 s; the node's manager ships as `tt_diffusion_planner.host.postprocess.TurnIndicatorManager`), the initial solver state `sampled_trajectories` (zeros = the node's default temperature 0; noise for a temperature > 0; the previous plan for the RTC prefix) and the agent and ego histories. Details: [`code/PYTHON.md`](code/PYTHON.md) "What the caller keeps".
|
| 72 |
+
- One model uses one chip; calls from several threads are serialised.
|
| 73 |
+
- Full reference: [`code/PYTHON.md`](code/PYTHON.md). Runnable example: [`examples/quickstart.py`](examples/quickstart.py) (also writes `quickstart_bev.png`, the input tensors and the plan from above).
|
| 74 |
+
|
| 75 |
+
## Serving (HTTP)
|
| 76 |
+
|
| 77 |
+
```bash
|
| 78 |
+
tt-model pull changh95/diffusion-planner-p150 --with-weights
|
| 79 |
+
tt-model serve changh95/diffusion-planner-p150 # or with tt-cli: tt serve changh95/diffusion-planner-p150
|
| 80 |
+
python3 code/tt_diffusion_planner/server/client.py --inputs code/tt_diffusion_planner/samples/kashiwanoha_dense.npz --out req.json
|
| 81 |
+
curl -s localhost:20000/predict -H 'Content-Type: application/json' -d @req.json
|
| 82 |
+
tt model stop changh95/diffusion-planner-p150
|
| 83 |
+
```
|
| 84 |
+
|
| 85 |
+
- The image does not contain the weights. `--with-weights` puts them in your HF cache.
|
| 86 |
+
- The server uses port 20000 (or the next free port). It is ready when the log shows `Application startup complete`.
|
| 87 |
+
- One serve profile, the default (`tt-model profiles changh95/diffusion-planner-p150`). `serve.env` pins ETH dispatch, 1 CQ and the numerics knobs.
|
| 88 |
+
- Run the request lines from the model repo root (the `hf download` of the Quickstart): the client and the sample are files of this repo. `code/tt_diffusion_planner/server/client.py` needs only the Python standard library (no numpy); add `--url http://127.0.0.1:20000` to send the request.
|
| 89 |
+
- `POST /predict`: `inputs` (the 15 raw tensors of the Autoware node's `create_input_data()` in the ego frame, before normalization: a base64 `.npz` or `{"format": "json", "arrays": {...}}`); optional `params` (`velocity_smoothing_window` 8, `stopping_threshold` 0.3, `turn_indicator_keep_offset` -1.25, `return_denoising_steps` false), `output_format`. Also `GET /health`, `GET /info`, `GET /v1/models` (stub). Contract: [`SERVING.md`](SERVING.md) section 3.
|
| 90 |
+
|
| 91 |
+
The response for the shipped sample (served on the p150; trajectory cut to 3 of 80 rows):
|
| 92 |
+
|
| 93 |
+
```json
|
| 94 |
+
{
|
| 95 |
+
"model": "diffusion-planner-p150",
|
| 96 |
+
"frame_id": "base_link",
|
| 97 |
+
"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}},
|
| 98 |
+
"timing_ms": {"preprocess": 3.8, "device": 105.3, "postprocess": 3.2, "total": 122.8, "decode": 8.5, "model_call": 112.5},
|
| 99 |
+
"num_poses": 80,
|
| 100 |
+
"columns": ["x", "y", "yaw", "cos", "sin", "velocity", "acceleration"],
|
| 101 |
+
"trajectory": [[0.365, 0.01, -0.017, 0.9962, -0.0169, 3.8059, -0.0431], [0.7631, -0.0022, -0.0414, 0.9966, -0.0413, 3.8016, -0.5485], [1.1524, -0.0282, -0.0675, 0.9942, -0.0674, 3.7468, -0.4708], "... 77 more rows"],
|
| 102 |
+
"turn_indicator": {"command": 1, "command_name": "DISABLE", "keep_selected": true, "held": false, "logits": [-15.960084915161133, -5.125061511993408, -5.230083465576172, -0.7658059597015381, 5.504596710205078], "probabilities": [1.651601411190029e-09, 8.38487030705437e-05, 7.548934809165075e-05, 0.006556871347129345, 0.993184506893158]},
|
| 103 |
+
"predicted_agents": {"format": "npz", "key": "predicted_agents", "dtype": "float32", "shape": [88, 80, 5], "data": "<base64 npz>"}
|
| 104 |
+
}
|
| 105 |
+
```
|
| 106 |
+
|
| 107 |
+
- `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 points, 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; `cos` / `sin` are the raw network outputs.
|
| 108 |
+
- `turn_indicator.command`: 0 NO_COMMAND, 1 DISABLE, 2 ENABLE_LEFT, 3 ENABLE_RIGHT (KEEP repeats the last input report); the node's 1 s hold window needs state across calls and is not applied (see above).
|
| 109 |
+
- `predicted_agents`: one 80-point path (x, y, yaw, cos, sin) per non-empty neighbour row, in input order (`meta.predicted_agent_rows`).
|
| 110 |
+
|
| 111 |
+
## Demo
|
| 112 |
+
|
| 113 |
+
| Shipped sample `kashiwanoha_dense.npz`: the p150 plan next to the fp32 CPU reference | The 11 solver iterates of the same plan (ego row), p150 and CPU |
|
| 114 |
+
|:---:|:---:|
|
| 115 |
+
|  |  |
|
| 116 |
+
| **Shipped sample `straight_road.npz`** | **Agreement with the fp32 CPU reference on all 99 gated scenes** |
|
| 117 |
+
|  |  |
|
| 118 |
+
|
| 119 |
+
On nuScenes v1.0-mini planning instants (p150 outputs; the scenes are converted from the dataset into the planner tensors and are not in this repository). **Non-commercial, CC BY-NC-SA 4.0.** The grey path with hollow dots is the logged drive; the model never saw nuScenes (see Caveats).
|
| 120 |
+
|
| 121 |
+
| scene-0061, 2 Hz sequence (33 plans): following a van, then a left turn | scene-0061 key-frame 18, in the turn: CAM_FRONT with the plan as a vehicle-width ribbon |
|
| 122 |
+
|:---:|:---:|
|
| 123 |
+
|  |  |
|
| 124 |
+
| **scene-0757 key-frame 11: the turn head commands RIGHT before the driver's blinker** | **scene-0757 key-frame 11, CAM_FRONT** |
|
| 125 |
+
|  |  |
|
| 126 |
+
| **scene-0103 key-frame 12 (mini_val, Boston)** | **scene-0103 key-frame 12, CAM_FRONT** |
|
| 127 |
+
|  |  |
|
| 128 |
+
| **scene-0061 key-frame 6: approach behind the van** | **scene-0916 key-frame 11 (mini_val): parking-lot right turn** |
|
| 129 |
+
|  |  |
|
| 130 |
+
|
| 131 |
+
nuScenes renders: rendered from the nuScenes dataset (v1.0-mini, CAN bus expansion and map expansion v1.3), © Motional AD Inc., CC BY-NC-SA 4.0 and the nuScenes Terms of Use; non-commercial use only; Motional does not endorse this work. The kashiwanoha sample is derived from the Apache-2.0 map AutowareFoundation/map-carla-kashiwanoha. Sources, changes and the full attributions: [`media/ATTRIBUTION.md`](media/ATTRIBUTION.md).
|
| 132 |
+
|
| 133 |
+
## Demo & Performances
|
| 134 |
+
|
| 135 |
+
Warm, batch 1; the stage bench 2026-10-08, the served rows, load times and the re-check 2026-10-09. Latency: the stage bench of [`OPT_BASELINE.md`](OPT_BASELINE.md) (`code/scripts/bench.py`, 100 iterations per stage) on the shipped sample `kashiwanoha_dense.npz` (88 neighbours, 123 lanes, 17 route lanes, 60 line strings; the device time is the same for every scene: every plan computes the full capacities); the served rows from uvicorn on the host (the app the container runs, with its serve pins) and a loopback client, 50 requests of the shipped sample. The host is shared with other jobs, so the host stages move with its load by several ms; the device rows repeat to ±0.05 ms. Accuracy: the p150 output against the fp32 CPU reference of the same network on 99 scenes: the 2 shipped samples, 5 research scenes and 92 nuScenes v1.0-mini planning instants (the frozen end-to-end gates), and against the research pipeline's ONNX Runtime outputs as an independent oracle.
|
| 136 |
+
|
| 137 |
+
| Metric | Performance |
|
| 138 |
+
|---|---:|
|
| 139 |
+
| Agreement with the fp32 CPU reference, shipped sample `kashiwanoha_dense` (8 s ego plan, 88 neighbours) | ego max **2.3 cm** / mean 1.1 cm; turn command identical; neighbours: median per-agent max 8.6 cm |
|
| 140 |
+
| Agreement, shipped sample `straight_road` | ego max 4.9 cm / mean 1.8 cm; turn command identical; neighbours 6.3 cm |
|
| 141 |
+
| Agreement, all 99 gated scenes (2 samples, 5 research scenes, 92 nuScenes-mini instants) | worst ego max **0.313 m** / mean **0.143 m** (gates 1.0 / 0.3 m; nuscenes/scene-0103_kf14); ego mean: median 1.3 cm, 95th percentile 6.2 cm; turn command identical **99 / 99**; neighbours: median per-agent max ≤ 0.086 m (gate 1.5 m) |
|
| 142 |
+
| Agreement with an independent oracle: the research pipeline's ONNX Runtime outputs (raw x0), 92 nuScenes instants + the 33-plan scene-0061 sequence | worst ego max 0.313 m / mean 0.143 m, turn command identical 92 / 92; sequence: worst ego max 0.118 m / mean 0.049 m, turn 33 / 33; every instant within the gates |
|
| 143 |
+
| Module PCC vs the fp32 reference (encoder categories, encoding, teacher-forced decoder evaluation; gate 0.999) | ≥ 0.999952 / 0.999978 / ≥ 0.9999995 |
|
| 144 |
+
| Open-loop vs the nuScenes log, 92 instants (a sanity check against one human driver, not a planning metric) | p150 ADE / FDE at 8 s 4.75 / 11.87 m (CPU reference 4.74 / 11.87; constant velocity 4.99 / 13.06); turn command = logged blinker 88 / 92 (CPU 88 / 92) |
|
| 145 |
+
| Python `model()` call, shipped sample (host pre-processing, H2D, trace, D2H, host post-processing) | **117.9 ms p50** (p99 134.6) · 8.5 plans/s |
|
| 146 |
+
| Served `/predict` `timing_ms.total` (uvicorn on the host, the shipped sample) | **123.1 ms median** (min 122.5; of which decode 8.6) |
|
| 147 |
+
| Served client round trip, loopback (base64 `.npz` request, 0.15 MB) | 127.5 ms median |
|
| 148 |
+
| Device trace, one blocking plan (encoder + 11 DiT evaluations + 10 solver updates + turn head) | **102.13 ms** |
|
| 149 |
+
| Back-to-back trace replays | **102.04 ms per plan** · 9.80 plans/s |
|
| 150 |
+
| Host pre-processing · pack · host tensors · H2D · D2H · host post-processing | 4.75 · 0.45 · 2.41 · 1.14 · 0.58 · 4.32 ms |
|
| 151 |
+
| `from_pretrained` load: empty JIT cache / warm cache | 315 s / 8.6 s |
|
| 152 |
+
|
| 153 |
+
All numbers in this table were measured with dispatch on the ETH cores, 1 command queue and a 12×10 compute grid on one p150, with the pinned numerics of `serve.env`. Accuracy is agreement with the fp32 CPU reference of the same Autoware network (same weights, same pre- and post-processing); no dataset-level accuracy is claimed: the paper's benchmark is nuPlan closed loop (an account-gated dataset and simulator, not run here), and the deployed v5.0 weights were trained by TIER IV on data that is not public, so no public benchmark is in-domain. Details: [`VERIFICATION_2026-10-08.md`](VERIFICATION_2026-10-08.md), [`OPT_BASELINE.md`](OPT_BASELINE.md), [`OPT_REPORT.md`](OPT_REPORT.md).
|
| 154 |
+
|
| 155 |
+
No GPU comparison: no GPU was available on the host where this port was built and measured, so this card makes no GPU speed claim. The reference rows are the port's own fp32 CPU reference on the same host (a correctness baseline, not a speed target). Autoware's CHANGELOG quotes 5.13 ms mean (300 runs) for an older single-step engine of this planner on an RTX PRO 6000 Blackwell with TensorRT (precision not stated); it is not like-for-like with this v5.0 multi-step port. p150 power was not measured, so no efficiency comparison is made.
|
| 156 |
+
|
| 157 |
+
## Caveats
|
| 158 |
+
|
| 159 |
+
- First release: **baseline port, optimization pending.** The plan is one metal trace and kernel-bound (6,282 programs; op-to-op gaps 4.4 ms of a 103.7 ms span). The **first optimization target is the precision cost**: the numerics defaults that the end-to-end gates need (split hi / lo matmuls, fp32 LayerNorm, fp32 matmul attention) cost 33.4 ms per plan (102.0 ms vs 68.6 ms with the first device round's defaults, which fail the gates); fused kernels are to recover it without dropping precision. A second sink of similar size is unrelated to precision: the encoder's channel-MLP and pre-projection matmuls run on 4-8 cores (25.8 ms). [`OPT_REPORT.md`](OPT_REPORT.md) ranks what comes next.
|
| 160 |
+
- Deployment status in Autoware: `autoware_diffusion_planner` is an alternative to the default rule-based planning stack, selected with `planning_setting:=diffusion_planner` (package README); it is aimed at Autoware's proposed new planning framework. This bundle is not a ROS 2 node (Python API and HTTP) and not a certified Autoware component; do not use it for safety-critical driving decisions or closed-loop vehicle control.
|
| 161 |
+
- Stateless API: the node's state between plans (the turn-indicator hold window, the RTC prefix and temperature of the initial solver state, the agent buffers and the ego history) is the client's ("What the caller keeps" in [`code/PYTHON.md`](code/PYTHON.md)). The node's guidance services (start / stop / centerline guidance) are off, as in the node's default.
|
| 162 |
+
- Precision policy of this release: fp32 residual streams and solver state; HiFi4 with fp32 accumulation for every matmul; the ego / neighbour pre-projection as a pad-relative fp32 island; split hi / lo matmuls (bf16 hi + fp32 lo parts, ~1e-5 relative, because a device fp32 matmul rounds its operands like TF32) for the mixer inputs and every decoder linear; an fp32 LayerNorm decomposition in the mixers and the decoder (the fused `ttnn.layer_norm` loses the per-entity signal on the mixers' offset-dominated rows); fp32 matmul attention in the fusion encoder and the decoder; the turn head in fp32; the other weights and the hidden MLP activations of the mixers and the fusion encoder in bf16. The plan's sensitivity to these choices is chaotic per scene: with the first device round's numerics, one nuScenes instant (scene-0103_kf14) moved 0.35 m on average while the module PCCs differed only in the 5th decimal. So every numerics change is re-checked on all 99 scenes.
|
| 163 |
+
- Validation scope: the p150 output agrees with the fp32 CPU reference on 99 scenes (above). The two shipped samples and the five research scenes are synthetic-but-faithful scenes built with a Python port of the node's tensor construction; the 92 nuScenes instants are converted from nuScenes v1.0-mini, a domain the model never saw (Singapore and Boston, right-hand traffic in Boston, oracle tracks from 2 Hz annotations, no traffic-light states, no speed limits, stop areas instead of stop lines). On them the plans are plausible but conservative (moving plans about 17 % shorter than the logged drive); the open-loop numbers above are a sanity check against one human driver, not a planning metric.
|
| 164 |
+
- Neighbour predictions: the gate is the median over agents of each agent's max displacement, so single agents can differ more. On the shipped sample `kashiwanoha_dense` (container smoke, served p150 output vs the stored CPU reference) the per-agent max displacement is 8.6 cm median and 0.26 m at the 90th percentile; the worst of the 88 agents differs by 4.31 m. The ego plan is gated on its max and mean; neighbour paths only on that median.
|
| 165 |
+
- Documented deviations from the node: the output stays in `base_link` (the node transforms it to `map` with the ego pose); `yaw` is what `tf2::getYaw` reads from the node's quaternion of the unnormalised cos / sin rotation (it differs from atan2(sin, cos) when |(cos, sin)| ≠ 1, as in the node); the speed masks follow the node's TensorRT path (`> FLT_EPSILON`); the turn indicator is decided without the hold window; `delay` is accepted and ignored (the node's multi-step mode never reads it).
|
| 166 |
+
- Does not scale to multiple p150 in a mesh configuration. The build uses a 12×10 compute grid of Tensix cores: the dispatch functions move from one Tensix column to the ETH cores (`patches/tt-metal-eth-dispatch.patch`), so this build assumes that you do not need chip-to-chip ethernet communication.
|
| 167 |
+
- `dispatch="worker"` (server: `DIFFUSION_PLANNER_DISPATCH=worker`) is an A/B opt-in. On a p150 it gives an 11×10 grid (101.87 ms per replay, the same as ETH within 0.2 %); if ETH dispatch is not available (tt-metal without the patch), the model falls back to it with a warning. The numbers on this card do not apply to that mode.
|
| 168 |
+
- Batch 1, one plan per request; requests are serialised on the chip. The 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 are compiled into the trace, and every plan computes them in full, so the device time does not depend on the scene.
|
| 169 |
+
- Not an OpenAI-compatible API; `GET /v1/models` is a stub so the tt-model ready card does not 404.
|
| 170 |
+
- p150 power was not measured, so no efficiency comparison is made.
|
| 171 |
+
|
| 172 |
+
## Licensing
|
| 173 |
+
|
| 174 |
+
- Weights: [AutowareFoundation/diffusion_planner](https://huggingface.co/AutowareFoundation/diffusion_planner) at tag `v5.0` (commit `423efde67f5414734da43a7ad856c17ceb8b51aa`), Apache-2.0 per its model card. Not redistributed here: the package only points to them. The upstream card states that TIER IV trained the models on TIER IV synthetic and real driving data; the dataset composition is not publicly documented.
|
| 175 |
+
- Pre- and post-processing ported from autoware_universe [`planning/autoware_diffusion_planner`](https://github.com/autowarefoundation/autoware_universe/tree/9ceaccf026c31ffc5319bc9eeb4bd7bede0af3fd/planning/autoware_diffusion_planner) (Apache-2.0).
|
| 176 |
+
- Port and serving code (`code/`): Apache-2.0. `patches/tt-metal-eth-dispatch.patch` modifies tt-metal (Apache-2.0).
|
| 177 |
+
- Sample data: `code/tt_diffusion_planner/samples/kashiwanoha_dense.npz` is derived from the Apache-2.0 Lanelet2 map [AutowareFoundation/map-carla-kashiwanoha](https://huggingface.co/datasets/AutowareFoundation/map-carla-kashiwanoha) (0.2.0) with a scripted ego, route and agents; `straight_road.npz` is a procedural scene generated by this repo. Both Apache-2.0, each with its stored CPU-reference output (`*.reference.json`). Only these redistributable samples ship; the nuScenes-derived planning instants of the accuracy tables are not in this repository.
|
| 178 |
+
- Demo media (`media/`, sources and changes in [`media/ATTRIBUTION.md`](media/ATTRIBUTION.md)):
|
| 179 |
+
- nuScenes renders (`media/*_NC.*`), **non-commercial, CC BY-NC-SA 4.0**: Rendered from the nuScenes dataset (v1.0-mini, CAN bus expansion and map expansion v1.3), © Motional AD Inc., CC BY-NC-SA 4.0 and the nuScenes Terms of Use (https://www.nuscenes.org/terms-of-use). Non-commercial use only; adaptations under the same license. Motional does not endorse this work. Cite: H. Caesar et al., *nuScenes: A Multimodal Dataset for Autonomous Driving*, CVPR 2020.
|
| 180 |
+
- Renders of the shipped samples: Apache-2.0 (kashiwanoha map: AutowareFoundation/map-carla-kashiwanoha@0.2.0, Apache-2.0).
|
| 181 |
+
|
| 182 |
+
## Provenance
|
| 183 |
+
|
| 184 |
+
These are the exact sources the container image was built from:
|
| 185 |
+
|
| 186 |
+
| component | built from |
|
| 187 |
+
| --- | --- |
|
| 188 |
+
| tt-metal | [`44d66500520fda9f2c7060c0f6b41ec48f7ab37e`](https://github.com/tenstorrent/tt-metal/commit/44d66500520fda9f2c7060c0f6b41ec48f7ab37e) + [`patches/tt-metal-eth-dispatch.patch`](patches/tt-metal-eth-dispatch.patch) (sha256 `08d0ddf6…45cc`, 4 files; dirty tree: the image includes the patch) |
|
| 189 |
+
| weights | [`AutowareFoundation/diffusion_planner@423efde67f5414734da43a7ad856c17ceb8b51aa`](https://huggingface.co/AutowareFoundation/diffusion_planner/tree/423efde67f5414734da43a7ad856c17ceb8b51aa) (tag `v5.0`), files `diffusion_planner_encoder.onnx, diffusion_planner_decoder.onnx, diffusion_planner_turn_indicator.onnx, diffusion_planner.param.json` (sha256 `2856886a…ca49`, `eb30c0c0…57ca`, `07acfb58…a732`, `ee3145b6…a268`, checked at load) |
|
| 190 |
+
| Autoware reference | autoware_universe [`9ceaccf026c31ffc5319bc9eeb4bd7bede0af3fd`](https://github.com/autowarefoundation/autoware_universe/tree/9ceaccf026c31ffc5319bc9eeb4bd7bede0af3fd) (`planning/autoware_diffusion_planner`, package 0.53.0, `multi_step` mode) |
|
| 191 |
+
| shared package | `ttaw` 0.20.0, vendored as `code/tt_diffusion_planner/ttaw` from the Autoware ports' shared `common` repository at commit `89dec49` (`code/tt_diffusion_planner/ttaw/VENDORED.json`: version, commit and per-file sha256) |
|
| 192 |
+
| `code/` digest (image) | `c0e7abb7888098a9` (sha256, first 16 hex digits; `built.code_sha256` of `tt_kernel_manifest.json`) |
|
| 193 |
+
| image | `tt-model/diffusion-planner-p150:3b96d8ea7190` (`sha256:3b96d8ea71902fe6a00f1792dd41290839b7758070431464cd613cde3f6bf909`) |
|
| 194 |
+
| base images | build stage `ghcr.io/tenstorrent/tt-metal/tt-metalium/ubuntu-22.04-dev-amd64:latest` @ `sha256:df9d279c7f85c17c6fad982d196802682d669cca1b7ced9cbaad8181339cd5fc`; runtime stage `docker.io/library/ubuntu:22.04` @ `sha256:5ec03bb3441e8b0bf3b4f9cd4629a1ae763010dc3035bb8da3ae6cf026486401` (tt-model's `FROM` tags float; these are the digests this build resolved, see [`build_info.json`](build_info.json)) |
|
| 195 |
+
| built | 2026-10-09T04:37:19+00:00 by tt-model 0.1.0 |
|
SERVING.md
ADDED
|
@@ -0,0 +1,242 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Serving Diffusion Planner v5.0 (Autoware diffusion_planner) on Blackhole with tt-model-manager
|
| 2 |
+
|
| 3 |
+
This repo is the authoring source of the **tt-model container package** `changh95/diffusion-planner-p150`
|
| 4 |
+
(`kind: tt-dit-server`, schema 5.1): `tt-model.yaml` is the manifest, `code/` the port, and
|
| 5 |
+
`code/tt_diffusion_planner/server/app.py` the ASGI app uvicorn runs. Weights are a pinned pointer, never baked into the
|
| 6 |
+
image. The device open, the decoders, the Python-API contract and the HTTP contract come from
|
| 7 |
+
`code/tt_diffusion_planner/ttaw/`, the shared package of the Autoware ports vendored into this repo
|
| 8 |
+
(`ttaw/VENDORED.json` records its version, source commit and file hashes; it is not edited here).
|
| 9 |
+
|
| 10 |
+
| item | value |
|
| 11 |
+
|---|---|
|
| 12 |
+
| tt-metal tree | `/home/ubuntu/experiments/tt-models/tt-metal` (main `44d66500520`, `v0.80.0-dev20261006-78-g44d6650052`, + `patches/tt-metal-eth-dispatch.patch`; torch 2.11.0+cpu) |
|
| 13 |
+
| 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) |
|
| 14 |
+
| Autoware reference | `planning/autoware_diffusion_planner` @ autoware_universe `9ceaccf026c31ffc5319bc9eeb4bd7bede0af3fd` (package 0.53.0), the node's `multi_step` mode |
|
| 15 |
+
| shared package | `ttaw` 0.20.0 (`common` @ `89dec49`) |
|
| 16 |
+
| app | `tt_diffusion_planner.server.app:app` (uvicorn; the ASGI lifespan does weights -> device -> graph -> trace capture) |
|
| 17 |
+
| 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 plan's trace holds 74.6 MB) |
|
| 18 |
+
| 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; 6,282 device programs |
|
| 19 |
+
| hardware | one Blackhole p150 (`hardware: p150`, `mesh_device: P150`, `TT_MESH_SHAPE=1x1`), 12×10 compute grid |
|
| 20 |
+
|
| 21 |
+
## 1. Run on the HOST (hardware validation, no Docker)
|
| 22 |
+
|
| 23 |
+
The tt-metal `python_env` has ttnn and torch; `bin/tt-env.sh` puts the workspace's dev overlay (fastapi, uvicorn, onnx)
|
| 24 |
+
on `PYTHONPATH` instead of installing into it. On the shared workspace box every command that opens the chip goes
|
| 25 |
+
through the device lock (`bin/devrun`).
|
| 26 |
+
|
| 27 |
+
```bash
|
| 28 |
+
ROOT=/home/ubuntu/experiments/tt-models
|
| 29 |
+
source $ROOT/bin/tt-env.sh # TT_METAL_HOME, PYTHONPATH (+ the dev overlay: fastapi, uvicorn, onnx), python_env
|
| 30 |
+
cd $ROOT/bundles/diffusion-planner-p150
|
| 31 |
+
export PYTHONPATH=$PWD/code:$PYTHONPATH
|
| 32 |
+
export HF_MODEL=AutowareFoundation/diffusion_planner TT_MODEL_WEIGHTS_REVISION=423efde67f5414734da43a7ad856c17ceb8b51aa TT_MESH_SHAPE=1x1
|
| 33 |
+
|
| 34 |
+
# host tests (no device; the device tests are skipped; ttnn is the fake of common/tests/host: a real ttnn host
|
| 35 |
+
# tensor opens the chip, so only devrun jobs may build one)
|
| 36 |
+
PYTHONPATH=$ROOT/research/packaging/scripts:$PYTHONPATH TT_VISIBLE_DEVICES=none \
|
| 37 |
+
python -m pytest -q -p no:cacheprovider -p fake_ttnn_plugin code/tt_diffusion_planner/tests
|
| 38 |
+
# device tests (per-module PCC + end-to-end agreement; 99 scenes with the workspace goldens, the 2 shipped samples elsewhere)
|
| 39 |
+
$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
|
| 40 |
+
# the server, with the serve pins of tt-model.yaml, and the container smoke test against it
|
| 41 |
+
$ROOT/bin/devrun -t 1800 -- bash -c 'env DIFFUSION_PLANNER_DISPATCH=eth DIFFUSION_PLANNER_NUM_CQS=1 \
|
| 42 |
+
python -m uvicorn --host 127.0.0.1 --port 20000 --lifespan on tt_diffusion_planner.server.app:app & \
|
| 43 |
+
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'
|
| 44 |
+
```
|
| 45 |
+
|
| 46 |
+
Boot log landmarks (they drive `tt-model serve`'s progress view, `boot_progress.py` TT_DIT_PHASES):
|
| 47 |
+
`Loading weights` -> `Opening device` -> `Warming up: capturing trace ...` -> `Warmup complete` -> uvicorn
|
| 48 |
+
`Application startup complete`. A cold JIT cache compiles every kernel first (315 s measured for
|
| 49 |
+
`from_pretrained`); later boots take about 10 s to READY (measured on the host, warm JIT cache). Startup failures raise and uvicorn exits non-zero (no CPU
|
| 50 |
+
fallback). SIGTERM: the lifespan releases the trace and closes the chip (within `tt-model stop`'s 120 s budget).
|
| 51 |
+
|
| 52 |
+
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>`.
|
| 53 |
+
|
| 54 |
+
## 2. Package, serve, push (Docker)
|
| 55 |
+
|
| 56 |
+
Docker 29 on this box is **rootful** (`/var/lib/docker`), and the user's session may not carry the `docker` group:
|
| 57 |
+
wrap every docker-using command in `sg docker -c "..."` (no `docker-env.sh` is needed here; the reference bundles'
|
| 58 |
+
`docker-env.sh` was the authors' rootless-Docker box). Run every `tt-model` command **from this directory**
|
| 59 |
+
(`source.tt_metal` and `extra_code[].root: code` are CWD-relative) and point `--out` outside it.
|
| 60 |
+
|
| 61 |
+
```bash
|
| 62 |
+
ROOT=/home/ubuntu/experiments/tt-models; cd $ROOT/bundles/diffusion-planner-p150
|
| 63 |
+
TTM=$HOME/.local/share/uv/tools/tt-model/bin/python
|
| 64 |
+
|
| 65 |
+
# offline validation (seconds; no docker, no device) -- must print VALID
|
| 66 |
+
$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)"
|
| 67 |
+
python3 $ROOT/research/packaging/scripts/check_bundle.py . # conventions (add --stage publish before a push)
|
| 68 |
+
|
| 69 |
+
# build (one at a time on the box: the absolute lock file): ~35 min cold, ~12 min after a code/ change (the C++
|
| 70 |
+
# build re-runs, ccache-warm, whenever code/ changes), ~1 min for a manifest- or lock-only change
|
| 71 |
+
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
|
| 72 |
+
# provenance: the base-image digests this build resolved (tt-model's FROM tags float) -> build_info.json (published)
|
| 73 |
+
sg docker -c "python3 $ROOT/research/packaging/scripts/record_build_info.py --staged $ROOT/build/diffusion-planner-p150 --bundle ."
|
| 74 |
+
|
| 75 |
+
# serve + smoke + stop, all inside ONE device-lock window (there is one serve profile, the default; serve returns
|
| 76 |
+
# once READY and leaves the container running; the script's trap saves the container log and always runs
|
| 77 |
+
# `tt-model stop`, also when devrun's timeout fires: -k 150 > the 120 s grace). No token reaches the public weights
|
| 78 |
+
# download. Evidence: logs/smoke/ (--log-dir).
|
| 79 |
+
$ROOT/bin/devrun -t 3600 -k 150 -- env -u HF_TOKEN -u HUGGING_FACE_HUB_TOKEN \
|
| 80 |
+
sg docker -c "bash code/scripts/container_smoke.sh $ROOT/build/diffusion-planner-p150 20000"
|
| 81 |
+
|
| 82 |
+
# publish: fill the provenance from the staged manifest, run the publish gate, overlay the hand-written docs, the pip
|
| 83 |
+
# project and the build record onto the staged dir, then push (HF_TOKEN only from the environment)
|
| 84 |
+
python3 $ROOT/research/packaging/scripts/instantiate_bundle.py --values $ROOT/research/diffusion-planner/bundle_values.json \
|
| 85 |
+
--fill . --manifest $ROOT/build/diffusion-planner-p150/tt_kernel_manifest.json
|
| 86 |
+
python3 $ROOT/research/packaging/scripts/check_bundle.py . --stage publish
|
| 87 |
+
rsync -a README.md SERVING.md OPT_BASELINE.md OPT_REPORT.md VERIFICATION_*.md tt-model.yaml pyproject.toml build_info.json \
|
| 88 |
+
examples media patches $ROOT/build/diffusion-planner-p150/
|
| 89 |
+
cp code/PYTHON.md $ROOT/build/diffusion-planner-p150/PYTHON.md
|
| 90 |
+
tt-model push $ROOT/build/diffusion-planner-p150 --publish
|
| 91 |
+
```
|
| 92 |
+
|
| 93 |
+
The smoke test fails unless `/info` reports ETH dispatch, the 12x10 grid and what the staged package pins
|
| 94 |
+
(`DIFFUSION_PLANNER_DISPATCH`, `DIFFUSION_PLANNER_NUM_CQS`, `DIFFUSION_PLANNER_VARIANT`, the weights revision), and it
|
| 95 |
+
compares the served trajectory of the shipped sample with the stored CPU-reference output
|
| 96 |
+
(`code/tt_diffusion_planner/samples/kashiwanoha_dense.reference.json`: average displacement <= 0.3 m, final
|
| 97 |
+
displacement <= 1.0 m). On the host, the same smoke test against uvicorn with the serve pins passed:
|
| 98 |
+
`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`
|
| 99 |
+
|
| 100 |
+
(`predicted_agents max_abs` is reported, not gated. Here it is the yaw column at ±π: -3.1413 vs +3.1415
|
| 101 |
+
rad, the same heading. The neighbour paths are gated in the end-to-end tests, on the per-agent maximum displacement.)
|
| 102 |
+
|
| 103 |
+
`serve` publishes the first free port from 20000 and exports into the container: `HF_MODEL=AutowareFoundation/diffusion_planner`,
|
| 104 |
+
`MESH_DEVICE=P150`, `TT_MESH_SHAPE=1x1`, `TT_MODEL_WEIGHTS_REVISION=423efde67f5414734da43a7ad856c17ceb8b51aa`, then `serve.env`. The HF cache is
|
| 105 |
+
mounted at `/hf` (rw), the JIT cache at `/cache` (host `~/.cache/tt-model/diffusion-planner-p150/cache`), and the container
|
| 106 |
+
sees only its own chip (`--device /dev/tenstorrent/<n>`). tt-cli users: `tt serve changh95/diffusion-planner-p150` /
|
| 107 |
+
`tt model stop changh95/diffusion-planner-p150` (point tt at this tt-model with
|
| 108 |
+
`tt config set tools.override.tt-model ~/.local/bin/tt-model` or `TT_TOOL_BIN_TT_MODEL`).
|
| 109 |
+
|
| 110 |
+
What `push` does to the repo: `code/` and `image/` on the Hub become exactly the staged trees (`extra_code.paths` +
|
| 111 |
+
`models/common/lightweightmodule.py`); every top-level file of the staged dir is uploaded (the overlaid README
|
| 112 |
+
replaces the generated card; `pyproject.toml`, `PYTHON.md` and `build_info.json` arrive the same way, never through
|
| 113 |
+
`code/`); top-level files already on the Hub and not in the staged dir are kept.
|
| 114 |
+
|
| 115 |
+
## 3. The request / response contract
|
| 116 |
+
|
| 117 |
+
| route | purpose |
|
| 118 |
+
|---|---|
|
| 119 |
+
| `GET /health` | `{"status": "ok" \| "starting" \| "error", "model", "device": {"dispatch", "grid", "cores", "device_id"}, "error"}`, always 200; `ok` only after warm-up |
|
| 120 |
+
| `GET /v1/health` | same as `/health` (tt-model's hint for non-chat packages points here) |
|
| 121 |
+
| `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 |
|
| 122 |
+
| `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 |
|
| 123 |
+
| `POST /predict` | one plan -> the model output (below) |
|
| 124 |
+
|
| 125 |
+
### 3.1 Request (`application/json`; unknown fields are a 422)
|
| 126 |
+
|
| 127 |
+
| field | type | meaning |
|
| 128 |
+
|---|---|---|
|
| 129 |
+
| `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) |
|
| 130 |
+
| `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) |
|
| 131 |
+
| `output_format` | str | `json` (default) or `npz` (adds the poses as a lossless base64 NPZ array) |
|
| 132 |
+
|
| 133 |
+
Base64 may be standard or URL-safe, wrapped, with or without a `data:` prefix. `python3
|
| 134 |
+
code/tt_diffusion_planner/server/client.py --inputs <scene.npz> [--param name=value] --out req.json` builds the request
|
| 135 |
+
with the standard library only (no numpy); `--url http://127.0.0.1:20000` sends it.
|
| 136 |
+
|
| 137 |
+
### 3.2 Response (200)
|
| 138 |
+
|
| 139 |
+
The served body of the shipped sample (uvicorn on the host with the serve pins; trajectory cut to 3 of 80 rows and
|
| 140 |
+
the predicted-agents array to its header):
|
| 141 |
+
|
| 142 |
+
```json
|
| 143 |
+
{
|
| 144 |
+
"model": "diffusion-planner-p150",
|
| 145 |
+
"frame_id": "base_link",
|
| 146 |
+
"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}},
|
| 147 |
+
"timing_ms": {"preprocess": 3.8, "device": 105.3, "postprocess": 3.2, "total": 122.8, "decode": 8.5, "model_call": 112.5},
|
| 148 |
+
"num_poses": 80,
|
| 149 |
+
"columns": ["x", "y", "yaw", "cos", "sin", "velocity", "acceleration"],
|
| 150 |
+
"trajectory": [[0.365, 0.01, -0.017, 0.9962, -0.0169, 3.8059, -0.0431], [0.7631, -0.0022, -0.0414, 0.9966, -0.0413, 3.8016, -0.5485], [1.1524, -0.0282, -0.0675, 0.9942, -0.0674, 3.7468, -0.4708], "... 77 more rows"],
|
| 151 |
+
"turn_indicator": {"command": 1, "command_name": "DISABLE", "keep_selected": true, "held": false, "logits": [-15.960084915161133, -5.125061511993408, -5.230083465576172, -0.7658059597015381, 5.504596710205078], "probabilities": [1.651601411190029e-09, 8.38487030705437e-05, 7.548934809165075e-05, 0.006556871347129345, 0.993184506893158]},
|
| 152 |
+
"predicted_agents": {"format": "npz", "key": "predicted_agents", "dtype": "float32", "shape": [88, 80, 5], "data": "<base64 npz>"}
|
| 153 |
+
}
|
| 154 |
+
```
|
| 155 |
+
|
| 156 |
+
- `trajectory`: 80 points at 0.1-8.0 s in `base_link` (the ego frame of the input tensors), post-processed like the
|
| 157 |
+
node's `~/output/trajectory`: velocity from consecutive 3-D points / 0.1 s, forward moving average over
|
| 158 |
+
`velocity_smoothing_window` points, force stop (poses frozen once the smoothed speed falls below `stopping_threshold`
|
| 159 |
+
while the ego moves), acceleration by finite difference. `yaw` is what `tf2::getYaw` reads from the node's
|
| 160 |
+
(unnormalised) quaternion of the raw network `cos` / `sin`, which are returned as they are.
|
| 161 |
+
- `turn_indicator`: `command` 0 NO_COMMAND, 1 DISABLE, 2 ENABLE_LEFT, 3 ENABLE_RIGHT (when the head selects KEEP, the
|
| 162 |
+
command repeats the last input report `turn_indicators[0, 30]`); the 5 raw `logits` (NONE, DISABLE, LEFT, RIGHT, KEEP)
|
| 163 |
+
and the decision's `probabilities`. No hold window is applied (see 3.5).
|
| 164 |
+
- `predicted_agents`: one 80-point path (x, y, yaw, cos, sin) per non-empty neighbour row, in input order
|
| 165 |
+
(`meta.predicted_agent_rows` gives the rows).
|
| 166 |
+
- `meta`: `force_stop`, `time_from_start_s`, `valid_counts` (entities the encoder saw), and with
|
| 167 |
+
`return_denoising_steps` the ego row of the 11 solver iterates (`denoising_steps`, `[11, 81, 4]`, the node's
|
| 168 |
+
`~/debug/denoising_steps`).
|
| 169 |
+
|
| 170 |
+
`timing_ms`: `decode` (base64 + npz parsing + the schema check), `preprocess` (the node's normalization and the encoder's
|
| 171 |
+
host features), `device` (host tensors + H2D + one trace replay + D2H), `postprocess`, `model_call`, `total` (server
|
| 172 |
+
side, after the body arrived).
|
| 173 |
+
|
| 174 |
+
### 3.3 Errors
|
| 175 |
+
|
| 176 |
+
**400** undecodable or malformed input (bad base64, a missing or unknown tensor, a wrong shape, a non-finite value,
|
| 177 |
+
unknown or out-of-range `params`), **422** schema violation (unknown field, wrong type), **503** while starting (or
|
| 178 |
+
when the boot failed), **500** `inference failed: <ExceptionType>: <message>` if the device call raises.
|
| 179 |
+
|
| 180 |
+
### 3.4 Environment the app reads (lifespan only, never at import)
|
| 181 |
+
|
| 182 |
+
| var | set by | meaning / default |
|
| 183 |
+
|---|---|---|
|
| 184 |
+
| `HF_MODEL` | launcher (`weights.repo`) | weights repo id; default `AutowareFoundation/diffusion_planner` |
|
| 185 |
+
| `TT_MODEL_WEIGHTS_REVISION` | launcher (`weights.revision`) | pinned commit; `TT_WEIGHTS_REVISION` (`serve.env`) is the same for older tt-model |
|
| 186 |
+
| `DIFFUSION_PLANNER_WEIGHTS_DIR` | you (offline) | local weights directory; overrides the Hub |
|
| 187 |
+
| `TT_MESH_SHAPE` | launcher (`runtime.mesh_shape_env`) | `1x1`; anything else -> RuntimeError at startup |
|
| 188 |
+
| `TT_DEVICE_ID` | you | chip to open, default 0 |
|
| 189 |
+
| `DIFFUSION_PLANNER_DISPATCH` | `serve.env` | `eth` (default) \| `worker` (A/B only, 11×10) \| `auto` (ETH if the patch is present) |
|
| 190 |
+
| `DIFFUSION_PLANNER_NUM_CQS` | `serve.env` | `1` (2 CQs replay ~10 % slower on ETH, OPT_BASELINE.md) |
|
| 191 |
+
| `DIFFUSION_PLANNER_VARIANT` | `serve.env` | `default` (the only v5.0 graph) |
|
| 192 |
+
| `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 |
|
| 193 |
+
| `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 |
|
| 194 |
+
| `DIFFUSION_PLANNER_WARMUP` | you | JSON list of warm-up variants, `default` or `none` |
|
| 195 |
+
| `DIFFUSION_PLANNER_TRACE_REGION`, `DIFFUSION_PLANNER_L1_SMALL`, `DIFFUSION_PLANNER_WORKER_L1_SIZE` | you | device-open overrides (validated values: `DEVICE_DEFAULTS` in `device.py`) |
|
| 196 |
+
| `DIFFUSION_PLANNER_MAX_BODY_MB` | you | request size guard, default 256 (the shipped sample's request is 0.15 MB of base64 npz) |
|
| 197 |
+
| `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) |
|
| 198 |
+
|
| 199 |
+
Server-side pipeline: JSON -> `tt_diffusion_planner.io` decoders (schema check) -> `DiffusionPlanner.__call__` (the node's
|
| 200 |
+
host pre-processing ported from `planning/autoware_diffusion_planner` -> H2D into 19 persistent device inputs ->
|
| 201 |
+
`execute_trace` of the whole plan -> one D2H -> the node's host post-processing) under one lock -> `Output.to_dict()`.
|
| 202 |
+
|
| 203 |
+
### 3.5 What the client keeps (the API is stateless)
|
| 204 |
+
|
| 205 |
+
Every request is one independent plan. The node keeps state between plans; a client that wants its behaviour over a
|
| 206 |
+
sequence of plans keeps the same state and puts it into the request or applies it to the response:
|
| 207 |
+
|
| 208 |
+
- **The input tensors.** Converting ROS messages and the Lanelet2 map into the 15 tensors is the client's: the per-UUID
|
| 209 |
+
agent buffers with their 0.1 s resampling, the ego history, lane / route / polygon / line-string selection and
|
| 210 |
+
encoding, traffic-light states, speed limits, the goal and the turn-indicator report history.
|
| 211 |
+
- **The turn-indicator hold window.** The node holds its last non-KEEP command for `turn_indicator_hold_duration`
|
| 212 |
+
(1.0 s in its YAML). The server evaluates every request with a fresh manager (no held command). To reproduce the node,
|
| 213 |
+
keep the last non-KEEP command and its time, and return it instead while less than 1.0 s has passed; the node's
|
| 214 |
+
manager ships in the package (`tt_diffusion_planner.host.postprocess.TurnIndicatorManager`, fed with
|
| 215 |
+
`turn_indicator.logits`, the plan time and the last report).
|
| 216 |
+
- **The initial solver state `x_T` (`sampled_trajectories`, normalised space).** Zeros is the node's default
|
| 217 |
+
(`temperature: [0.0]`). With a temperature > 0 the node draws N(0, 1) x temperature for every element (a fresh
|
| 218 |
+
`std::random_device` seed per call); send that noise in `sampled_trajectories`. With the RTC prefix (`delay_step` > 0,
|
| 219 |
+
at most 40) the node also writes its previous plan into the ego row, slots t = 0 .. delay_step: the previous poses
|
| 220 |
+
re-expressed in the current ego frame, `x` as (x - 10) / 20, `y` as y / 20, `cos` / `sin` as they are. Non-zero
|
| 221 |
+
`x_T` was checked on the device against ONNX Runtime (noise of scale 0.5 and 1.0, an RTC prefix: within the gates,
|
| 222 |
+
`VERIFICATION_2026-10-08.md`).
|
| 223 |
+
- **`delay`** is accepted and ignored, as in the node's multi-step mode.
|
| 224 |
+
- **The map frame.** The output is in `base_link` (the ego frame of the inputs); the node transforms it to `map` with the
|
| 225 |
+
current ego pose before publishing.
|
| 226 |
+
|
| 227 |
+
## 4. Caveats
|
| 228 |
+
|
| 229 |
+
- Batch 1, one plan per request; concurrent clients queue on the lock.
|
| 230 |
+
- Fixed shapes of the v5.0 export (320 neighbours, 140 lanes, 25 route lanes, 10 polygons, 60 line strings, 31 history
|
| 231 |
+
and 80 future steps) and 10 DPM-Solver steps (11 decoder evaluations) are compiled into the trace; every plan computes
|
| 232 |
+
the full capacities whatever the scene holds, so the device time does not depend on the scene.
|
| 233 |
+
- The node's guidance services (start / stop / centerline guidance) are off (the node's default too); a guided mode
|
| 234 |
+
would need the host in the solver loop.
|
| 235 |
+
- Warm-up captures the trace inside the lifespan, so READY means warm; the first cold boot pays the ttnn JIT.
|
| 236 |
+
- Weights are pinned by sha in `weights.revision`, exported by the launcher and repeated in `serve.env`;
|
| 237 |
+
`snapshot_download(..., revision=<sha>)` is a cache hit after `serve`'s pre-download and falls back to
|
| 238 |
+
`local_files_only=True` if the Hub is unreachable. The sha256 of the four files and `major_version == 5` are checked
|
| 239 |
+
at load.
|
| 240 |
+
- The runtime image has no host C/C++ compiler: device kernels JIT-compile (sfpi ships in the image); the host code is
|
| 241 |
+
numpy.
|
| 242 |
+
- `tt-model curl` and the ready card's `/v1/models` hint are OpenAI-shaped and are not this API; use the routes above.
|
VERIFICATION_2026-10-08.md
ADDED
|
@@ -0,0 +1,210 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# diffusion-planner-p150: independent verification, 2026-10-08
|
| 2 |
+
|
| 3 |
+
## Verdict
|
| 4 |
+
|
| 5 |
+
**PASS.** This is the light verification pass of the first publish (`research/PLAN.md` §6.3): the frozen gates re-run,
|
| 6 |
+
the baseline numbers re-checked, and every device-open path audited for ETH dispatch. It was made by the agent that
|
| 7 |
+
wrote the release docs, which is neither the porter nor the baseline agent, on commit `c0d84f9` (the `code/` of this
|
| 8 |
+
repo: ttaw 0.20.0 vendored, the numerics knobs pinned in `serve.env`). Against the baseline commit `5541833`:
|
| 9 |
+
|
| 10 |
+
- **Accuracy:** the device suite passes (44 passed; again under alloc tracking), every gate value and all 99 per-scene end-to-end numbers identical to the baseline run and to the port's final run: worst ego 0.313 / 0.143 m (max / mean; gates 1.0 / 0.3 m), turn command 99 / 99, neighbours ≤ 0.086 m (gate 1.5 m).
|
| 11 |
+
- **Independent oracle:** against the research pipeline's ONNX Runtime goldens, which the bundle's reference did not produce, all 92 / 92 nuScenes instants and 33 / 33 instants of the scene-0061 sequence pass the same gates (worst ego max 0.313 m).
|
| 12 |
+
- **Speed:** the stage bench reproduces `OPT_BASELINE.md`: back-to-back replay 102.03 vs 102.04 ms, one plan 102.10 vs 102.13 ms; the host-bound end to end 115.6 vs 117.9 ms p50.
|
| 13 |
+
- **Audit:** ETH dispatch, 1 CQ and the 12×10 grid on every device-open path; no fallback in any run.
|
| 14 |
+
- **No numerics change** since the port: the diff `5541833..c0d84f9` is the ttaw re-vendor (0.17.1 -> 0.19.0 -> 0.20.0: on this
|
| 15 |
+
bundle's path only `trace.py`'s `pack_outputs`, whose single-row layout for the planner's 117,664-element readback is
|
| 16 |
+
unchanged), the `serve.env` pins of the numerics knobs (equal to the code defaults), two host-test changes, the
|
| 17 |
+
quickstart picture and documentation. No gate threshold moved.
|
| 18 |
+
|
| 19 |
+
The full adversarial verification of the port is `VERIFY_PORT.md` of the porting workspace (round 1, 2026-10-08, PASS
|
| 20 |
+
with 9 low findings; not part of this repository): it re-ran every suite from a clean shell, audited the request path,
|
| 21 |
+
the gates, the oracle chain and the Autoware fidelity line by line, compared the device with ONNX Runtime goldens that
|
| 22 |
+
the bundle's own reference did not produce, and ran five adversarial scenes. Its results that this card relies on are
|
| 23 |
+
quoted below, and its findings are resolved as listed in "Review findings".
|
| 24 |
+
|
| 25 |
+
## Setup
|
| 26 |
+
|
| 27 |
+
| item | value |
|
| 28 |
+
|---|---|
|
| 29 |
+
| Hardware | one Blackhole p150b, ETH dispatch, 1 CQ, 12×10 compute grid (120 cores), as printed by every run (`dispatch eth`, `fallback null`, `eth_patch true`) |
|
| 30 |
+
| tt-metal | `44d66500520` with `patches/tt-metal-eth-dispatch.patch` (sha256 `08d0ddf6…45cc`; the tree's `git diff` against `44d6650` is byte-identical to the patch: 4 files) |
|
| 31 |
+
| Weights | `AutowareFoundation/diffusion_planner` @ `423efde67f5414734da43a7ad856c17ceb8b51aa` (tag `v5.0`): the Python API, the quickstart and the server read them from the HF cache at the pinned revision (downloaded without a token, sha256 equal to the workspace copy, checked again at every load); the device tests from the workspace copy |
|
| 32 |
+
| Verified commit | `c0d84f9` (ttaw 0.20.0, `common` `89dec49`, vendored from the committed tree; `vendor.py --check`: 0 differences) |
|
| 33 |
+
| Baseline commit | `5541833` (`OPT_BASELINE.md`) |
|
| 34 |
+
| Workload | `code/tt_diffusion_planner/samples/kashiwanoha_dense.npz` (+ `straight_road.npz` and nuScenes scene-0103_kf14 for the stage bench), batch 1, warm; 99 scenes for the accuracy gates and 125 nuScenes instants for the oracle check (local only) |
|
| 35 |
+
| Bench method | `code/scripts/bench.py --iters 100 --warmup 5 --b2b-iters 50 --b2b-rounds 3`, one run, the three scenes of the baseline |
|
| 36 |
+
| Served method | uvicorn on the host with every `serve.env` pin of `tt-model.yaml` (the container's app and environment), 5 warm-up + 50 requests of the shipped sample from a loopback client |
|
| 37 |
+
| Accuracy reference | the fp32 CPU reference (`code/tt_diffusion_planner/reference/`), itself equal to ONNX Runtime on the shipped ONNX (48 tests, PORT_LOG 5.1); independently, the research pipeline's ONNX Runtime goldens |
|
| 38 |
+
| Host | AMD EPYC VM, 8 vCPUs, shared with other agents' jobs (load average 5.42, 7.45 during the runs); host-side stages move with it |
|
| 39 |
+
| Logs | `logs/diffusion-planner/docs/` of the workspace (job scripts in `scripts/`) |
|
| 40 |
+
|
| 41 |
+
## Measured baseline vs re-check
|
| 42 |
+
|
| 43 |
+
| measurement (shipped sample `kashiwanoha_dense` unless noted) | baseline `5541833` (`OPT_BASELINE.md`) | re-check `c0d84f9` | change |
|
| 44 |
+
|---|---:|---:|---:|
|
| 45 |
+
| back-to-back trace replays, per plan (p50) | 102.04 ms | 102.03 ms | -0.01 ms |
|
| 46 |
+
| device trace, one blocking plan (p50) | 102.13 ms | 102.10 ms | -0.04 ms |
|
| 47 |
+
| device trace, one blocking plan (min) | 102.05 ms | 102.03 ms | -0.02 ms |
|
| 48 |
+
| e2e `model(inputs=arrays)` p50 | 117.90 ms | 115.61 ms | -1.9 % |
|
| 49 |
+
| e2e `model(inputs=arrays)` p99 | 134.56 ms | 121.79 ms | -9.5 % |
|
| 50 |
+
| e2e `model(inputs=<.npz path>)` p50 | 124.76 ms | 118.77 ms | -4.8 % |
|
| 51 |
+
| host pre-processing | 4.75 ms | 3.91 ms | -17.8 % |
|
| 52 |
+
| pack (19 persistent inputs) | 0.45 ms | 0.40 ms | -11.8 % |
|
| 53 |
+
| host tensors (TILE) | 2.41 ms | 2.05 ms | -14.9 % |
|
| 54 |
+
| H2D | 1.14 ms | 0.94 ms | -17.1 % |
|
| 55 |
+
| D2H (one packed read) | 0.58 ms | 0.37 ms | -36.3 % |
|
| 56 |
+
| host post-processing | 4.32 ms | 3.60 ms | -16.7 % |
|
| 57 |
+
| back-to-back replays, `straight_road` | 102.04 ms | 102.04 ms | -0.00 ms |
|
| 58 |
+
| back-to-back replays, nuScenes scene-0103_kf14 | 102.04 ms | 102.04 ms | -0.00 ms |
|
| 59 |
+
| first call after `from_pretrained` | 118.20 ms | 114.79 ms | -2.9 % |
|
| 60 |
+
| `from_pretrained` load (warm JIT cache) | 11.64 s | 10.38 s | -10.8 % |
|
| 61 |
+
|
| 62 |
+
AICLK during the re-check: median 1350 MHz (min 1343, max 1350; 760 sysfs samples, sampled during the timed loops); the device rows reproduce the baseline to the 0.05 ms the baseline quotes, the host rows move with the load of the shared host.
|
| 63 |
+
|
| 64 |
+
Not in the baseline, measured for the card on `c0d84f9`:
|
| 65 |
+
|
| 66 |
+
| measurement | value |
|
| 67 |
+
|---|---:|
|
| 68 |
+
| `from_pretrained` load with an EMPTY JIT cache (firmware + every kernel compiled); first / second call | 314.6 s; 121 / 119 ms |
|
| 69 |
+
| `from_pretrained` load with that cache warm (new process); first / second call | 8.6 s; 120 / 120 ms |
|
| 70 |
+
| served `/predict` `timing_ms`, median of 50: decode · preprocess · device · postprocess · total | 8.6 · 3.8 · 105.3 · 3.3 · **123.1 ms** (min total 122.5) |
|
| 71 |
+
| served client round trip, loopback (median / min; request 0.15 MB) | 127.5 / 127.0 ms |
|
| 72 |
+
| server boot to ready (uvicorn, warm JIT cache) | 10.0 s |
|
| 73 |
+
| the bundle's `server/smoke_test.py` against it (the container smoke's check) | `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` |
|
| 74 |
+
| served body vs the stored CPU reference (the smoke gates) | PASS ({"trajectory": {"n": 80, "ade": 0.01089988605897666, "fde": 0.017330320250936008}, "predicted_agents": {"max_abs_err": 6.282742738723755, "n": 35200}}); `predicted_agents` is reported, not gated: its largest difference is the yaw column at ±π (-3.1413 vs +3.1415 rad, the same heading) |
|
| 75 |
+
| the README's request recipe (`client.py --out req.json` + `curl`) | HTTP body with 80 poses, turn DISABLE; same trajectory as the client: True, same predicted agents: True |
|
| 76 |
+
| `examples/quickstart.py` as a user runs it (weights from the HF cache, offline, no token) | rc 0; turn DISABLE; `quickstart.json` and `quickstart_bev.png` written |
|
| 77 |
+
|
| 78 |
+
## Accuracy gate results
|
| 79 |
+
|
| 80 |
+
`test_pcc_device.py` + `test_e2e_device.py` on `c0d84f9`, one devrun job, `TTAW_GATES_READONLY=1`: 44 passed in 42.29s.
|
| 81 |
+
Every value equals the baseline run (`OPT_BASELINE.md`, ttaw 0.17.1) and the port's final run (PORT_LOG job 19, ttaw
|
| 82 |
+
0.11.0) to the printed precision, and all 99 per-scene end-to-end numbers are identical to both
|
| 83 |
+
(`logs/diffusion-planner/docs/compare_runs.json`).
|
| 84 |
+
|
| 85 |
+
| gate (PLAN §2.12; `tests/*.gates.json`) | threshold | value |
|
| 86 |
+
|---|---:|---:|
|
| 87 |
+
| `enc.ego` PCC vs the fp32 reference (valid rows, 2 samples) | ≥ 0.999 | 0.999972 |
|
| 88 |
+
| `enc.neighbor` PCC vs the fp32 reference (valid rows, 2 samples) | ≥ 0.999 | 0.999952 |
|
| 89 |
+
| `enc.lane` PCC vs the fp32 reference (valid rows, 2 samples) | ≥ 0.999 | 0.999992 |
|
| 90 |
+
| `enc.route` PCC vs the fp32 reference (valid rows, 2 samples) | ≥ 0.999 | 0.999998 |
|
| 91 |
+
| `enc.polygon` PCC vs the fp32 reference (valid rows, 2 samples) | ≥ 0.999 | 0.999994 |
|
| 92 |
+
| `enc.line_string` PCC vs the fp32 reference (valid rows, 2 samples) | ≥ 0.999 | 0.999998 |
|
| 93 |
+
| `enc.goal` PCC vs the fp32 reference (valid rows, 2 samples) | ≥ 0.999 | 0.999991 |
|
| 94 |
+
| `enc.ego_shape` PCC vs the fp32 reference (valid rows, 2 samples) | ≥ 0.999 | 1 |
|
| 95 |
+
| `enc.turn` PCC vs the fp32 reference (valid rows, 2 samples) | ≥ 0.999 | 0.999995 |
|
| 96 |
+
| `enc.encoding` PCC (valid tokens) | ≥ 0.999 | 0.999978 |
|
| 97 |
+
| `dec.eval` PCC (teacher-forced decoder, min over 11 evaluations x 2 samples) | ≥ 0.999 | 1 (≥ 0.9999995) |
|
| 98 |
+
| `ego.max_err_m`, max over 99 scenes | ≤ 1.0 m | 0.313 m (nuscenes/scene-0103_kf14) |
|
| 99 |
+
| `ego.mean_err_m`, max over scenes | ≤ 0.3 m | 0.143 m (nuscenes/scene-0103_kf14) |
|
| 100 |
+
| `turn.command_agreement` | = 1.0 | 1.000 (99 / 99) |
|
| 101 |
+
| `neighbors.median_max_err_m`, max over scenes | ≤ 1.5 m | 0.086 m (kashiwanoha_dense) |
|
| 102 |
+
| diagnostics, .kashiwanoha_dense: `final_x0` valid-agent PCC; turn logits max abs | reported | 0.999994; 0.0398 |
|
| 103 |
+
| diagnostics, straight_road: `final_x0` valid-agent PCC; turn logits max abs | reported | 1.000000; 0.0372 |
|
| 104 |
+
| replay == eager (bit for bit); prefix constraint exact; API == `/predict` | equal | equal |
|
| 105 |
+
|
| 106 |
+
Distribution over the 99 scenes: ego mean displacement median 1.3 cm, 95th percentile 6.2 cm. Under `TT_METAL_TRACE_ALLOC_TRACKING=1` (no program compiled after the first capture): 44 passed in 40.50s.
|
| 107 |
+
|
| 108 |
+
Also re-run, CPU only: the host suite on the fake ttnn, 107 passed and 45 skipped (the 44 device tests and the ONNX
|
| 109 |
+
Runtime module, which runs in the research venv), after the re-vendor and the test changes of `d521981` (0.20.0 re-vendored in `c0d84f9`).
|
| 110 |
+
|
| 111 |
+
### Against an independent oracle
|
| 112 |
+
|
| 113 |
+
The p150 outputs of the 92 nuScenes v1.0-mini planning instants and of the 33-plan scene-0061 sequence (2 Hz), produced
|
| 114 |
+
through the public API (`model(inputs=...)`), compared with the research pipeline's ONNX Runtime goldens
|
| 115 |
+
(`research/diffusion-planner/public_data/goldens/`: ORT 1.30 on the shipped ONNX files with an independent host
|
| 116 |
+
normalization and solver port), with the PLAN §2.12 gates on the raw denormalised x0
|
| 117 |
+
(`research/diffusion-planner/public_data/scripts/dp_public_metrics.py gates`):
|
| 118 |
+
|
| 119 |
+
| set | instants | all four gates pass | worst ego max | worst ego mean | worst neighbour median-max | turn command identical | logits max abs |
|
| 120 |
+
|---|---:|---:|---:|---:|---:|---:|---:|
|
| 121 |
+
| nuScenes v1.0-mini planning instants (1 Hz, all 10 scenes) | 92 | 92 / 92 | 0.313 m (scene-0103_kf14) | 0.143 m (scene-0103_kf14) | 0.082 m | 92 / 92 | 0.095 |
|
| 122 |
+
| scene-0061 sequence (2 Hz, the demo GIF) | 33 | 33 / 33 | 0.118 m (scene-0061_kf36) | 0.049 m (scene-0061_kf09) | 0.027 m | 33 / 33 | 0.074 |
|
| 123 |
+
|
| 124 |
+
`VERIFY_PORT.md` (round 1) also ran five adversarial scenes on the device against ONNX Runtime, which the stored goldens
|
| 125 |
+
do not cover (every golden has `x_T = 0` and at most 88 agents): full capacity (320 neighbours, 140 lanes, 25 route
|
| 126 |
+
lanes, 10 polygons): ego 0.042 / 0.025 m (max / mean); ego only (every neighbour removed): 0.062 / 0.031 m; `x_T` = 0.5
|
| 127 |
+
N(0, 1): 0.022 / 0.007 m; `x_T` = 1.0 N(0, 1): 0.067 / 0.031 m; an RTC prefix in the ego row of `x_T`: 0.024 / 0.012
|
| 128 |
+
m; turn command equal and outputs finite on all five.
|
| 129 |
+
|
| 130 |
+
### Open-loop sanity metrics (not a planning metric)
|
| 131 |
+
|
| 132 |
+
The p150 plans of the 92 nuScenes instants against the logged drive, with the research metric script
|
| 133 |
+
(`dp_public_metrics.py eval`), next to the CPU reference's numbers (`public_data/metrics/open_loop_cpu_fp32.json`). The
|
| 134 |
+
model never saw nuScenes and the inputs have known gaps (no traffic-light states, no speed limits), so this is a sanity
|
| 135 |
+
check against one human driver, not a paper metric:
|
| 136 |
+
|
| 137 |
+
| group (instants) | ADE@3s p150 / CPU | FDE@3s | ADE@8s | FDE@8s | constant velocity ADE@8s / FDE@8s | plan / logged length (m), p150 / CPU |
|
| 138 |
+
|---|---:|---:|---:|---:|---:|---:|
|
| 139 |
+
| all (92) | 1.12 / 1.12 | 2.60 / 2.59 | 4.75 / 4.74 | 11.87 / 11.87 | 4.99 / 13.06 | 37.5 / 37.5 vs 44.2 |
|
| 140 |
+
| moving (v0 >= 1 m/s) (69) | 1.42 / 1.41 | 3.27 / 3.26 | 5.88 / 5.88 | 14.50 / 14.51 | 6.56 / 17.21 | 48.5 / 48.5 vs 58.8 |
|
| 141 |
+
| mini_val (18) | 1.44 / 1.44 | 3.42 / 3.42 | 6.42 / 6.43 | 15.86 / 15.89 | 9.47 / 24.11 | 35.6 / 35.6 vs 38.0 |
|
| 142 |
+
| mini_train (74) | 1.04 / 1.04 | 2.40 / 2.39 | 4.34 / 4.33 | 10.90 / 10.89 | 3.90 / 10.37 | 38.0 / 38.0 vs 45.7 |
|
| 143 |
+
|
| 144 |
+
Turn-indicator command equal to the logged (debounced) blinker: p150 88 / 92, CPU 88 / 92. Moving plans (v0 ≥ 1 m/s) are 17 % shorter than the logged drive.
|
| 145 |
+
|
| 146 |
+
## Disclosed numerics changes
|
| 147 |
+
|
| 148 |
+
- None since the port. The configuration is decision 11 of `PORT_LOG.md`: fp32 residual streams and solver state;
|
| 149 |
+
HiFi4 + fp32 accumulation for every matmul; the ego / neighbour pre-projection as a pad-relative fp32 island; split
|
| 150 |
+
hi / lo matmuls (`DIFFUSION_PLANNER_SPLIT_MATMUL=enc.island.*,enc.pre.*,dec.*`); the fp32 LayerNorm decomposition
|
| 151 |
+
(`_LN_FP32=enc.mixer.*,dec.*`); fp32 matmul attention (`_ATTN_MATMUL=enc.fusion.attn,dec.*`); the turn head in fp32;
|
| 152 |
+
the other weights and the mixer / fusion hidden activations in bf16. This is the precision policy of PLAN.md §9.3
|
| 153 |
+
(HiFi4, fp32 accumulation, two-term or fp32 weights where the gates need them).
|
| 154 |
+
- What the configuration costs (timing only, `OPT_BASELINE.md`): 33.4 ms per plan over the first device round's
|
| 155 |
+
defaults, which fail `ego.mean_err_m` (0.347 m on nuScenes scene-0103_kf14, PORT_LOG job 11); 58.3 ms over the
|
| 156 |
+
fastest graph, which fails `enc.ego`. The card names this cost as the first optimization target.
|
| 157 |
+
- The knobs are now pinned in `serve.env` at these values, and a host test (`test_serve_env_pins_the_numerics`) checks
|
| 158 |
+
the pins against `KNOBS.serve_env()`; the served `/info` reported the same options (null).
|
| 159 |
+
|
| 160 |
+
## Review findings
|
| 161 |
+
|
| 162 |
+
- **No gate loosened.** `tests/test_pcc_device.gates.json` and `tests/test_e2e_device.gates.json` are unchanged since
|
| 163 |
+
`3223604` (the port's job 18 freeze); the thresholds equal PLAN §2.12; `GateRegistry` refuses a looser declaration and
|
| 164 |
+
the runs used `TTAW_GATES_READONLY=1`.
|
| 165 |
+
- **No cached or replayed outputs.** The gates compare trace replays with the fp32 CPU reference's goldens, the oracle
|
| 166 |
+
check with ONNX Runtime goldens written by another pipeline; the only TT-vs-TT check is replay == eager (a trace
|
| 167 |
+
correctness check, not an accuracy gate).
|
| 168 |
+
- **The diff `5541833..c0d84f9`** touches no file of `tt_diffusion_planner/tt/`, `host/`, `reference/`, `api.py`,
|
| 169 |
+
`device.py`, `io.py` or `server/`: the ttaw re-vendor, `tt-model.yaml` (`serve.env` pins), `examples/quickstart.py`
|
| 170 |
+
and two host tests (`test_api_host.py`: the dead `test_build_says_what_is_missing` replaced by
|
| 171 |
+
`test_build_refuses_unknown_compile_params`; `test_bundle_host.py`: `test_serve_env_pins_the_numerics`). The docs
|
| 172 |
+
commit on top changes only documentation, media and `tt-model.yaml` comments, card text and one `verify:` line.
|
| 173 |
+
- **No hang, timeout, reset or FAULT marker** in the device jobs of this pass.
|
| 174 |
+
|
| 175 |
+
Status of the `VERIFY_PORT.md` findings:
|
| 176 |
+
|
| 177 |
+
| # | finding | status |
|
| 178 |
+
|---|---|---|
|
| 179 |
+
| L1 | vendored ttaw outdated (0.11.0 vs common 0.15.1) | resolved: re-vendored 0.17.1 at the baseline step (`5541833`) then 0.19.0 (`d521981`) and 0.20.0 (`c0d84f9`) in this pass; the device suite after the first and the last gave identical gate values and per-scene numbers |
|
| 180 |
+
| L2 | numerics knobs not pinned in `serve.env`, no host test | resolved (`d521981`): the five knobs and an empty `DIFFUSION_PLANNER_PRECISION` pinned; `test_serve_env_pins_the_numerics` |
|
| 181 |
+
| L3 | publish documents outstanding; template text wrong for a planner | resolved: README, SERVING, PYTHON, OPT_BASELINE, OPT_REPORT, this file, the patch copy and the demo media; the card's risks rewritten; packaging, the container smoke, `build_info.json` and the push remain for the packaging step |
|
| 182 |
+
| L4 | dead test `test_build_says_what_is_missing` | resolved (`d521981`): replaced by a check that `_build` refuses unknown compile parameters |
|
| 183 |
+
| L5 | "6,408 programs" was a fake-ttnn op count | resolved: `OPT_BASELINE.md` quotes 6,282 programs per plan from the device profiler |
|
| 184 |
+
| L6 | op-count evidence only in a session scratchpad | resolved: `logs/diffusion-planner/opcount_fake_r2.log` |
|
| 185 |
+
| L7 | stateless API: hold window and `x_T` handling left to the client | resolved by documentation: README "Quickstart", `code/PYTHON.md` "What the caller keeps" (with the node's `TurnIndicatorManager` applied across calls), `SERVING.md` 3.5 |
|
| 186 |
+
| L8 | the bf16 SDPA path (`DIFFUSION_PLANNER_ATTN_MATMUL=none`) has no gate run since decision 11 | unchanged: it is not a shipped option (`serve.env` pins the matmul attention); PORT_LOG job 13 shows it fails the mean gate on scene-0103_kf14 together with the decoder split (0.74 / 0.33 m), and `OPT_BASELINE.md` times it only |
|
| 187 |
+
| L9 | C20 PCC range quoted as 0.99966-0.99983 | resolved in `PORT_LOG.md` (0.99961-0.99983 over all suite shapes, 0.99972-0.99983 at the planner's shapes); `ttaw/API.md` section 17 (common) keeps the old text until the next functional C20 change, as PORT_LOG section 9 says; the gate (0.9995) is unaffected |
|
| 188 |
+
|
| 189 |
+
## p150 ETH-dispatch compliance
|
| 190 |
+
|
| 191 |
+
| path | dispatch / CQs / grid |
|
| 192 |
+
|---|---|
|
| 193 |
+
| Python API `DiffusionPlanner.from_pretrained()` (`ttaw.api_base` → `tt_diffusion_planner.device.open_device`) | **ETH / 1 / 12×10** by default (`DeviceConfig.dispatch = "eth"`, `DEVICE_DEFAULTS["num_command_queues"] = 1`; `DIFFUSION_PLANNER_DISPATCH` / `_NUM_CQS` override). A failing ETH open falls back to WORKER with a `RuntimeWarning`, and `model.info["device"]["fallback"]` names it. Measured in this pass (quickstart, load times, demo runs): `dispatch eth`, `grid 12x10`, `fallback null` |
|
| 194 |
+
| HTTP server (`tt_diffusion_planner.server.app`, `config_from_env`) | ETH unless `DIFFUSION_PLANNER_DISPATCH` says otherwise; `/health` and `/info` report the device actually opened. Measured (served bench): `/info` `dispatch eth`, `grid 12x10`, `num_command_queues 1`, `fallback None` |
|
| 195 |
+
| `tt-model.yaml` serve env (one serve profile, the default) | `DIFFUSION_PLANNER_DISPATCH=eth`, `_NUM_CQS=1`, `_VARIANT=default` and the numerics pins; `check_bundle.py` refuses a profile whose `*_DISPATCH` is not `eth`; host tests check the pins against `device.py` and `tt/config.py` |
|
| 196 |
+
| image `verify:` | asserts the ETH-patch marker `single_chip_arch_1cq_no_dispatch_s` in `/opt/tt-metal/tt_metal/impl/dispatch/topology.cpp` and that ttaw sees the patch (`eth_dispatch_patch_present`) |
|
| 197 |
+
| container smoke (the one serve profile) | `server/smoke_test.py` asserts `/info` = `eth` / `12x10` (constants, not options) and the profile's pins (dispatch, CQs, variant, weights revision), and compares the served plan with the stored CPU reference; run against uvicorn on the host in this pass (`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`); to be run on the built image by the packaging step |
|
| 198 |
+
| gate tests (`code/conftest.py`) | `open_device(..., allow_fallback=False)`: an ETH open that fails is an error, never a silent WORKER run |
|
| 199 |
+
| `code/scripts/bench.py`, `profile_ops.py` | `from_pretrained(dispatch=--dispatch or DIFFUSION_PLANNER_DISPATCH or eth)`; the JSON records the device actually opened (`fallback null` in every run); the WORKER rows of `OPT_BASELINE.md` were explicit A/B runs |
|
| 200 |
+
| `code/scripts/bringup_device.py`, `precision_exp.py`, `split_error.py` | `open_device(allow_fallback=False)` (ETH) |
|
| 201 |
+
| `examples/quickstart.py` | `from_pretrained(device_id=...)`: ETH / 1 / 12×10 (run in this pass) |
|
| 202 |
+
| vendored `ttaw` | 0.20.0 @ `89dec49`, `source_dirty: false`; `common/tools/vendor.py --check`: 0 differences; `check_bundle.py`: no drift warning |
|
| 203 |
+
| hard-coded grids | none: `tt/` sets no core grid or program config; the only grid consumer, C20's SDPA program config, reads `compute_with_storage_grid_size()` and is not on the default path (matmul attention) |
|
| 204 |
+
|
| 205 |
+
## Not covered by this pass
|
| 206 |
+
|
| 207 |
+
The container build, its smoke test on the built image, `build_info.json`, the HF pre-flight and the push (the
|
| 208 |
+
packaging step); the Python install through `pip install -e .` (the quickstart ran with `code/` on `PYTHONPATH`, the
|
| 209 |
+
same import path as the device tests); the dispatch / CQ matrix and the device profile (taken once, in
|
| 210 |
+
`OPT_BASELINE.md`, on the same device code).
|
build_info.json
ADDED
|
@@ -0,0 +1,48 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"schema": "ttaw-build-info/1",
|
| 3 |
+
"bundle": "diffusion-planner-p150",
|
| 4 |
+
"recorded_at": "2026-10-09T04:48:03Z",
|
| 5 |
+
"image": {
|
| 6 |
+
"tag": "tt-model/diffusion-planner-p150:3b96d8ea7190",
|
| 7 |
+
"digest": "sha256:3b96d8ea71902fe6a00f1792dd41290839b7758070431464cd613cde3f6bf909",
|
| 8 |
+
"built_at": "2026-10-09T04:37:19+00:00",
|
| 9 |
+
"code_sha256": "c0e7abb7888098a9319ab5c66a10c4fd4009fc1542326b26507554c04f05a481",
|
| 10 |
+
"size_bytes": 4069025450,
|
| 11 |
+
"layers": 23
|
| 12 |
+
},
|
| 13 |
+
"base_images": {
|
| 14 |
+
"build": {
|
| 15 |
+
"stage": "prep",
|
| 16 |
+
"ref": "ghcr.io/tenstorrent/tt-metal/tt-metalium/ubuntu-22.04-dev-amd64:latest",
|
| 17 |
+
"digest": "sha256:df9d279c7f85c17c6fad982d196802682d669cca1b7ced9cbaad8181339cd5fc",
|
| 18 |
+
"source": "docker build log",
|
| 19 |
+
"local_tag_matches": false,
|
| 20 |
+
"local_id": "sha256:3fd1e6013e658c65df1bc6b084543e7174d854c15af97373e493df0c1787c393",
|
| 21 |
+
"created": "2026-10-07T01:17:29.186036377Z",
|
| 22 |
+
"note": "the local tag moved after this build; the digest above is the one the image was built from"
|
| 23 |
+
},
|
| 24 |
+
"runtime": {
|
| 25 |
+
"stage": "runtime",
|
| 26 |
+
"ref": "docker.io/library/ubuntu:22.04",
|
| 27 |
+
"digest": "sha256:5ec03bb3441e8b0bf3b4f9cd4629a1ae763010dc3035bb8da3ae6cf026486401",
|
| 28 |
+
"source": "docker build log",
|
| 29 |
+
"local_tag_matches": null,
|
| 30 |
+
"note": "not in the local image store (BuildKit pulled it into its cache)"
|
| 31 |
+
}
|
| 32 |
+
},
|
| 33 |
+
"tt_metal": {
|
| 34 |
+
"sha": "44d66500520fda9f2c7060c0f6b41ec48f7ab37e",
|
| 35 |
+
"describe": "v0.80.0-dev20261006-78-g44d6650052-dirty",
|
| 36 |
+
"dirty": true,
|
| 37 |
+
"scm_version": "0.65.2.dev11169+g44d66500520",
|
| 38 |
+
"mode": "local",
|
| 39 |
+
"remote": "https://github.com/tenstorrent/tt-metal.git",
|
| 40 |
+
"branch": "main",
|
| 41 |
+
"pushed": true
|
| 42 |
+
},
|
| 43 |
+
"tools": {
|
| 44 |
+
"tt_model": "0.1.0",
|
| 45 |
+
"docker": "29.8.2"
|
| 46 |
+
},
|
| 47 |
+
"notes": "tt-model's FROM tags float (build: tt-metalium dev image :latest, runtime: ubuntu:<version>); these are the digests this image was built from. A moved build base costs a cold C++ build."
|
| 48 |
+
}
|
code/PYTHON.md
ADDED
|
@@ -0,0 +1,179 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Python API: Diffusion Planner v5.0 (Autoware diffusion_planner) on Blackhole
|
| 2 |
+
|
| 3 |
+
Use this API from Python code (a pipeline, a notebook, a ROS 2 node wrapper). You do not need the HTTP server: the
|
| 4 |
+
API and the server share the decoders, the device trace and the post-processing, so the outputs and the speed are
|
| 5 |
+
the same.
|
| 6 |
+
|
| 7 |
+
## Install
|
| 8 |
+
|
| 9 |
+
Install the package on top of an environment that already has `ttnn` (a tt-metal `python_env` at `44d66500520`
|
| 10 |
+
with `patches/tt-metal-eth-dispatch.patch`, or the tt-model container). From the root of the model repository (the
|
| 11 |
+
directory that holds `pyproject.toml`, `README.md` and `code/`):
|
| 12 |
+
|
| 13 |
+
```bash
|
| 14 |
+
pip install -e . # the Python API (numpy<2, pillow, pyyaml, onnx, huggingface_hub)
|
| 15 |
+
pip install -e ".[server,test]" # + the HTTP server and the tests
|
| 16 |
+
```
|
| 17 |
+
|
| 18 |
+
The pip project is the repository's top-level `pyproject.toml`; it installs the package from
|
| 19 |
+
`code/tt_diffusion_planner` (there is no `pyproject.toml` inside `code/`, because the container build copies `code/`
|
| 20 |
+
over the tt-metal tree). ttnn and torch come from tt-metal and are not declared.
|
| 21 |
+
|
| 22 |
+
The package carries `tt_diffusion_planner.ttaw`, the shared code of the Autoware ports to Blackhole (device open, trace
|
| 23 |
+
runner, decoders, model base class, HTTP app), vendored at the version recorded in
|
| 24 |
+
`code/tt_diffusion_planner/ttaw/VENDORED.json`.
|
| 25 |
+
|
| 26 |
+
| You want to run | Extras |
|
| 27 |
+
|---|---|
|
| 28 |
+
| the Python API | none |
|
| 29 |
+
| the HTTP server (`tt_diffusion_planner.server.app`, see `SERVING.md`) | `server` |
|
| 30 |
+
| host tests (no device; device tests are skipped): `TT_VISIBLE_DEVICES=none python -m pytest -q code/tt_diffusion_planner/tests` | `server,test` |
|
| 31 |
+
| device tests: `python -m pytest -q -s code/tt_diffusion_planner/tests/test_pcc_device.py code/tt_diffusion_planner/tests/test_e2e_device.py` | `test` |
|
| 32 |
+
|
| 33 |
+
## Quickstart
|
| 34 |
+
|
| 35 |
+
```python
|
| 36 |
+
from tt_diffusion_planner import DiffusionPlanner
|
| 37 |
+
|
| 38 |
+
with DiffusionPlanner.from_pretrained(device_id=0) as model:
|
| 39 |
+
out = model(inputs="code/tt_diffusion_planner/samples/kashiwanoha_dense.npz")
|
| 40 |
+
print(out.to_dict()) # the POST /predict body
|
| 41 |
+
```
|
| 42 |
+
|
| 43 |
+
`examples/quickstart.py` runs the same snippet, writes `quickstart.json` and a bird's-eye view of the input and the
|
| 44 |
+
plan (`quickstart_bev.png`).
|
| 45 |
+
|
| 46 |
+
## `DiffusionPlanner.from_pretrained(...)`
|
| 47 |
+
|
| 48 |
+
```python
|
| 49 |
+
DiffusionPlanner.from_pretrained(
|
| 50 |
+
model_id=None, # HF repo or a local directory with the weights files; default AutowareFoundation/diffusion_planner
|
| 51 |
+
*,
|
| 52 |
+
revision=None, # default for the default repo: the validated commit 423efde67f5 (tag v5.0)
|
| 53 |
+
variant=None, # "default" (the only v5.0 graph); default $DIFFUSION_PLANNER_VARIANT or "default"
|
| 54 |
+
device_id=None, # chip to open; default $TT_DEVICE_ID or 0
|
| 55 |
+
device=None, # an already-opened ttnn device (tt_diffusion_planner.device.open_device); close() does not close it
|
| 56 |
+
dispatch=None, # "eth" (p150 target, 12x10 grid) | "worker" (A/B only, 11x10) | "auto"; default $DIFFUSION_PLANNER_DISPATCH or "eth"
|
| 57 |
+
num_command_queues=None, # default $DIFFUSION_PLANNER_NUM_CQS or 1
|
| 58 |
+
weights_dir=None, # explicit local weights directory; no Hub access
|
| 59 |
+
warmup_variants="default", # trace variants to capture now; see "Warm-up"
|
| 60 |
+
verbose=False,
|
| 61 |
+
precision=None, # the only compile parameter: extra precision-policy rules, e.g. "dec.*=HiFi2+fp32" (experiments only)
|
| 62 |
+
) -> DiffusionPlanner
|
| 63 |
+
```
|
| 64 |
+
|
| 65 |
+
What it does: resolves the weights first, so a Hub problem never claims the chip (`weights_dir` >
|
| 66 |
+
`$DIFFUSION_PLANNER_WEIGHTS_DIR` > a local `model_id` directory > the HF snapshot at the pinned revision, restricted to
|
| 67 |
+
the three v5.0 ONNX files and `diffusion_planner.param.json`, with an offline fallback to the cache; the sha256 of every
|
| 68 |
+
file and the weights' `major_version == 5` are checked), opens the chip (ETH dispatch, 12×10, 1 CQ; the other open
|
| 69 |
+
parameters are `DEVICE_DEFAULTS` in `tt_diffusion_planner/device.py`, overridable with `DIFFUSION_PLANNER_*`), reads the
|
| 70 |
+
ONNX initializers as data, uploads the weights and constants (48.1 MB), builds the graph, then compiles and captures
|
| 71 |
+
the metal trace. If ETH dispatch cannot open (tt-metal without the patch), it warns and falls back to WORKER dispatch
|
| 72 |
+
(`model.info["device"]["fallback"]` names it). Any other keyword argument is a `TypeError`.
|
| 73 |
+
|
| 74 |
+
The numerics are not arguments: the published configuration is the default of the `DIFFUSION_PLANNER_LN_FP32`,
|
| 75 |
+
`_SPLIT_MATMUL`, `_ATTN_MATMUL`, `_HIDDEN_FP32` and `_ATTN_FP32_ACC` knobs (`tt_diffusion_planner.tt.config.KNOBS`,
|
| 76 |
+
pinned in `tt-model.yaml` `serve.env`). Setting one of them in the environment changes the graph and invalidates the
|
| 77 |
+
accuracy figures of the card until the gates are re-run.
|
| 78 |
+
|
| 79 |
+
## Warm-up
|
| 80 |
+
|
| 81 |
+
`from_pretrained` returns a warm model: it builds the graph, runs the plan once eagerly (this first run compiles every kernel into the JIT cache), then captures the whole plan as one metal trace (`warmup_variants="default"`: the variant `plan`) with program-cache misses forbidden, so no later call compiles anything. `model.warmup()` is idempotent; `warmup_variants="none"` defers the capture to `model.warmup()`.
|
| 82 |
+
|
| 83 |
+
Measured on the shipped sample (`model.info["warmup_ms"]`, 2026-10-09): the load takes 315 s with an empty JIT cache and 8.6 s with a warm one (build 0.53 s: the ONNX initializers read and 48.1 MB of weights and constants uploaded; warm-up and capture 3.7 s; the rest is the device open). The first call then takes 120 ms and the second 120 ms (the stage bench's steady state: 118 ms p50 for decoded arrays, 125 ms for an `.npz` path). The trace holds 74.6 MB of DRAM (`trace_region_size` 192 MiB).
|
| 84 |
+
|
| 85 |
+
## Call: `model(...)`
|
| 86 |
+
|
| 87 |
+
| Argument | Type | Description |
|
| 88 |
+
|---|---|---|
|
| 89 |
+
| `inputs` | mapping / `.npz` path / bytes / JSON envelope | the 15 raw planner tensors (see "Input types") |
|
| 90 |
+
| `velocity_smoothing_window` | int, 1..79, default 8 | forward moving average of the trajectory velocity, in points |
|
| 91 |
+
| `stopping_threshold` | float >= 0, default 0.3 | force stop below this smoothed speed (m/s), when the ego moves |
|
| 92 |
+
| `turn_indicator_keep_offset` | float, default -1.25 | added to the KEEP logit before the turn-indicator decision |
|
| 93 |
+
| `return_denoising_steps` | bool, default False | add the ego row of the 11 solver iterates (`out.meta["denoising_steps"]`, `[11, 81, 4]`, the node's `~/debug/denoising_steps`) |
|
| 94 |
+
|
| 95 |
+
### Input types
|
| 96 |
+
|
| 97 |
+
- `inputs=`: the 15 raw tensors of the Autoware node's `DiffusionPlannerCore::create_input_data()` (batch 1,
|
| 98 |
+
float32, ego `base_link` frame, BEFORE normalization; names and shapes in `tt_diffusion_planner.INPUT_SCHEMA`):
|
| 99 |
+
a `{name: array}` mapping (numpy or torch), an `.npz` path or its bytes, or the `/predict` envelope
|
| 100 |
+
`{"format": "npz", "data": <base64>}` / `{"format": "json", "arrays": {...}}`. Names, shapes and finite values are
|
| 101 |
+
checked (`InputError`). `tt_diffusion_planner.load_inputs(source)` is the same decoder.
|
| 102 |
+
- Any other input (`points`, `images`, `calibration`, ...) is refused (`InputError`).
|
| 103 |
+
|
| 104 |
+
### What the caller keeps (the API is stateless)
|
| 105 |
+
|
| 106 |
+
One call is one independent plan. The Autoware node keeps state between plans; to reproduce it over a sequence of
|
| 107 |
+
plans, the caller keeps the same state (SERVING.md 3.5 has the details):
|
| 108 |
+
|
| 109 |
+
- **The tensors.** The node's pre-processing from ROS messages and the Lanelet2 map (per-UUID agent buffers and their
|
| 110 |
+
0.1 s resampling, the ego history, lane / route / polygon / line-string selection and encoding, traffic lights,
|
| 111 |
+
speed limits, goal, turn-indicator report history) is not part of the bundle.
|
| 112 |
+
- **The turn-indicator hold window** (`turn_indicator_hold_duration`, 1.0 s in the node's YAML). Each call decides
|
| 113 |
+
with a fresh manager; apply the node's hold across calls with the node's own manager:
|
| 114 |
+
|
| 115 |
+
```python
|
| 116 |
+
from tt_diffusion_planner.host.postprocess import TurnIndicatorManager
|
| 117 |
+
|
| 118 |
+
manager = TurnIndicatorManager() # hold 1.0 s, KEEP offset -1.25 (the node's YAML)
|
| 119 |
+
out = model(inputs=tensors)
|
| 120 |
+
decision = manager.evaluate(out.turn_indicator["logits"], stamp_s=now_s, prev_report=int(tensors["turn_indicators"][0, 30]))
|
| 121 |
+
command = decision.command # the held command while less than 1.0 s has passed
|
| 122 |
+
```
|
| 123 |
+
|
| 124 |
+
- **The initial solver state** `sampled_trajectories` (`x_T`, normalised space): zeros is the node's default
|
| 125 |
+
(`temperature: [0.0]`); for a temperature > 0 send N(0, 1) x temperature; for the RTC prefix (`delay_step` > 0) put
|
| 126 |
+
the previous plan into the ego row, slots t = 0 .. delay_step (x as (x - 10) / 20, y as y / 20, cos / sin as they are,
|
| 127 |
+
in the current ego frame). `delay` is accepted and ignored (the node's multi-step mode never reads it).
|
| 128 |
+
- **The map frame.** Outputs are in `base_link`; the node transforms them with the current ego pose to `map`.
|
| 129 |
+
|
| 130 |
+
## Output
|
| 131 |
+
|
| 132 |
+
`model(...)` returns a `tt_diffusion_planner.Output` (= `ttaw.outputs.Trajectory`); `out.to_dict()` is exactly the `POST /predict` body (SERVING.md section 3.2).
|
| 133 |
+
|
| 134 |
+
| field | type | meaning |
|
| 135 |
+
|---|---|---|
|
| 136 |
+
| `poses` | float32 `[80, 7]` | the ego trajectory at 0.1-8.0 s in `base_link`: x, y, yaw, cos, sin, velocity, acceleration (`out.columns`), post-processed like the node's `~/output/trajectory` |
|
| 137 |
+
| `turn_indicator` | dict | `command` (0 NO_COMMAND, 1 DISABLE, 2 ENABLE_LEFT, 3 ENABLE_RIGHT), `command_name`, `keep_selected`, `held` (always false: no hold window), the 5 raw `logits` (NONE, DISABLE, LEFT, RIGHT, KEEP), the decision's `probabilities` |
|
| 138 |
+
| `predicted_agents` | float32 `[N, 80, 5]` | x, y, yaw, cos, sin of each non-empty neighbour row, in input order |
|
| 139 |
+
| `meta` | dict | `predicted_agent_rows` (the rows of `predicted_agents`), `predicted_agent_columns`, `force_stop`, `time_from_start_s`, `valid_counts` (the entities the encoder saw), and with `return_denoising_steps` the encoded `denoising_steps` |
|
| 140 |
+
| `timing_ms` | dict | `preprocess`, `device` (host tensors + H2D + replay + D2H), `postprocess`, `total` |
|
| 141 |
+
|
| 142 |
+
`out.to_dicts()` gives one `{x, y, yaw, cos, sin, velocity, acceleration}` dict per trajectory point; `out.to_dict("npz")` adds the poses as a lossless base64 NPZ array.
|
| 143 |
+
|
| 144 |
+
## Lifetime and information
|
| 145 |
+
|
| 146 |
+
- `model.close()` releases the trace and the persistent device tensors and closes the chip if the model opened it;
|
| 147 |
+
idempotent. `with` calls it for you; an unclosed model is closed when Python exits.
|
| 148 |
+
- `model.info`: weights (repo, tag, revision, path), device (dispatch, grid, CQs, fallback), variant, warm variants,
|
| 149 |
+
warm-up times, runtime parameter defaults, the input schema, the numerics options and precision policy in effect, the
|
| 150 |
+
trace (variants, persistent inputs, trace buffers in MB).
|
| 151 |
+
- Calls from several threads are safe: the device calls are serialised. One model per process per chip.
|
| 152 |
+
|
| 153 |
+
## Speed
|
| 154 |
+
|
| 155 |
+
Warm calls, batch 1, ETH dispatch, 1 CQ, 12×10, the pinned numerics (`code/scripts/bench.py`, 100 iterations; the numbers of `OPT_BASELINE.md`, 2026-10-08, on a shared host; p50, with p99 in brackets):
|
| 156 |
+
|
| 157 |
+
| stage | shipped sample `kashiwanoha_dense` |
|
| 158 |
+
|---|---:|
|
| 159 |
+
| `.npz` decode + schema check (path inputs only) | 7.02 (18.20) ms |
|
| 160 |
+
| host pre-processing (the node's normalization, the encoder's host features, masks, `x_T`) | 4.75 (12.55) ms |
|
| 161 |
+
| packing the 19 persistent trace inputs · ttnn host tensors · H2D | 0.45 · 2.41 · 1.14 ms |
|
| 162 |
+
| **device trace, one blocking plan** | **102.13** (104.81) ms |
|
| 163 |
+
| D2H (one packed read, 460 KB) | 0.58 ms |
|
| 164 |
+
| host post-processing (trajectory, predicted paths, turn decision) | 4.32 (11.54) ms |
|
| 165 |
+
| **`model(inputs=arrays)` end to end** | **117.90** (134.56) ms |
|
| 166 |
+
| `model(inputs=<.npz path>)` | 124.76 (147.74) ms |
|
| 167 |
+
| back-to-back replays (device time per plan) | 102.04 ms = 9.80 plans/s |
|
| 168 |
+
|
| 169 |
+
The device time does not depend on the scene: every plan computes the full capacities (re-checked on c0d84f9: kashiwanoha_dense 102.10, straight_road 102.11, a nuScenes instant 102.09 ms). Throughput above one plan per ~118 ms needs pipelining of the host work of neighbouring requests (not implemented); 2 CQs do not help a synchronous request (`OPT_BASELINE.md`). Where the time goes and what comes next: `OPT_REPORT.md`.
|
| 170 |
+
|
| 171 |
+
## Limits
|
| 172 |
+
|
| 173 |
+
- Batch 1 on the chip; one model per process.
|
| 174 |
+
- Fixed shapes of the v5.0 export (320 neighbours, 140 lanes, 25 route lanes, 10 polygons, 60 line strings, 31 history
|
| 175 |
+
and 80 future steps) and 10 DPM-Solver steps (11 decoder evaluations) are compiled into the trace; every plan computes
|
| 176 |
+
the full capacities, so the device time does not depend on the scene.
|
| 177 |
+
- The node's guidance services (start / stop / centerline guidance) are not available (the node's default is off).
|
| 178 |
+
- Accuracy is agreement with the fp32 CPU reference of the same network (README "Demo & Performances"); the planner's
|
| 179 |
+
driving quality is the weights' (trained by TIER IV on data that is not public).
|
code/conftest.py
ADDED
|
@@ -0,0 +1,59 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-License-Identifier: Apache-2.0
|
| 2 |
+
"""pytest fixtures of diffusion-planner-p150 (no dependency on tt-metal's own conftest).
|
| 3 |
+
|
| 4 |
+
- ``device`` (session): one chip opened like the published numbers (``tt_diffusion_planner.device.open_device``:
|
| 5 |
+
ETH dispatch, 12x10 grid, the port's validated sizes and CQs). ``--device-id N`` or ``TT_DEVICE_ID`` selects
|
| 6 |
+
the chip (default 0); ``DIFFUSION_PLANNER_DISPATCH=worker`` is the A/B opt-in, and the other ``DIFFUSION_PLANNER_*``
|
| 7 |
+
variables apply as for the server. A failing ETH open is an error here, never a silent WORKER fallback: gates
|
| 8 |
+
are only valid on the published setup.
|
| 9 |
+
- Tests marked ``device`` are skipped when ttnn is missing or ``TT_VISIBLE_DEVICES=none`` (host-only runs).
|
| 10 |
+
On the shared workspace box run them through the lock and name the test files:
|
| 11 |
+
``bin/devrun -t 1800 -- python -m pytest -q -s code/tt_diffusion_planner/tests/test_pcc_device.py``.
|
| 12 |
+
"""
|
| 13 |
+
from __future__ import annotations
|
| 14 |
+
|
| 15 |
+
import gc
|
| 16 |
+
import importlib.util
|
| 17 |
+
import os
|
| 18 |
+
|
| 19 |
+
import pytest
|
| 20 |
+
|
| 21 |
+
|
| 22 |
+
def pytest_addoption(parser):
|
| 23 |
+
parser.addoption("--device-id", action="store", default=None, help="chip id (default $TT_DEVICE_ID or 0)")
|
| 24 |
+
|
| 25 |
+
|
| 26 |
+
def pytest_configure(config):
|
| 27 |
+
config.addinivalue_line("markers", "device: needs a Tenstorrent chip (skipped on host-only runs)")
|
| 28 |
+
|
| 29 |
+
|
| 30 |
+
def _no_device_reason():
|
| 31 |
+
if os.environ.get("TT_VISIBLE_DEVICES", "").lower() == "none":
|
| 32 |
+
return "TT_VISIBLE_DEVICES=none (host-only run)"
|
| 33 |
+
if importlib.util.find_spec("ttnn") is None:
|
| 34 |
+
return "ttnn is not installed"
|
| 35 |
+
return None
|
| 36 |
+
|
| 37 |
+
|
| 38 |
+
def pytest_collection_modifyitems(config, items):
|
| 39 |
+
reason = _no_device_reason()
|
| 40 |
+
if reason:
|
| 41 |
+
skip = pytest.mark.skip(reason=reason)
|
| 42 |
+
for item in items:
|
| 43 |
+
if "device" in item.keywords:
|
| 44 |
+
item.add_marker(skip)
|
| 45 |
+
|
| 46 |
+
|
| 47 |
+
@pytest.fixture(autouse=True)
|
| 48 |
+
def _gc_between_tests():
|
| 49 |
+
gc.collect()
|
| 50 |
+
|
| 51 |
+
|
| 52 |
+
@pytest.fixture(scope="session")
|
| 53 |
+
def device(request):
|
| 54 |
+
from tt_diffusion_planner.device import close_device, open_device
|
| 55 |
+
|
| 56 |
+
cli = request.config.getoption("--device-id")
|
| 57 |
+
dev = open_device(int(cli) if cli is not None else None, allow_fallback=False)
|
| 58 |
+
yield dev
|
| 59 |
+
close_device(dev)
|
code/models/common/lightweightmodule.py
ADDED
|
@@ -0,0 +1,12 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-FileCopyrightText: © 2023 Tenstorrent USA, Inc.
|
| 2 |
+
|
| 3 |
+
# SPDX-License-Identifier: Apache-2.0
|
| 4 |
+
|
| 5 |
+
|
| 6 |
+
class LightweightModule:
|
| 7 |
+
"""Torch modules add a surprising amount of host overhead for attribute
|
| 8 |
+
access and method calls. This class is a lightweight alternative that
|
| 9 |
+
just wraps a forward function for now."""
|
| 10 |
+
|
| 11 |
+
def __call__(self, *args, **kwargs):
|
| 12 |
+
return self.forward(*args, **kwargs)
|
code/scripts/README.md
ADDED
|
@@ -0,0 +1,16 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# code/scripts
|
| 2 |
+
|
| 3 |
+
| script | purpose | device |
|
| 4 |
+
|---|---|---|
|
| 5 |
+
| `bench.py` | stage breakdown of warm plans (load / host_pre / pack / host_in / H2D / trace / D2H / host_post / e2e / b2b, p50 / p99 / min, AICLK) on one or more scene `.npz` files, `--dispatch` / `--num-cqs` for the dispatch / CQ matrix (the card and `OPT_BASELINE.md` numbers) | yes, via `bin/devrun` |
|
| 6 |
+
| `profile_ops.py` | one eager plan (stage `m:` and layer-kind `c:` signposts) + one traced replay between signposts under the device profiler (`python -m tracy -r -p -v --op-support-count 16000 ...`: a plan is 6,282 programs, over the default 1,000-program buffer) | yes |
|
| 7 |
+
| `bringup_device.py`, `precision_exp.py`, `split_error.py` | the port's numerics experiments (module PCC per configuration, LayerNorm / split-matmul / attention variants, where the plan error of a scene comes from: device encoder vs device decoder); their logs and findings are in `PORT_LOG.md` (workspace) | yes |
|
| 8 |
+
| `ref_golden.py` | goldens of the fp32 CPU reference (per-module taps, final outputs, the stored `/predict` references next to the samples) for the device tests; research venv (onnxruntime): `tools/research-venv/bin/python code/scripts/ref_golden.py` | no |
|
| 9 |
+
| `container_smoke.sh` | serve ONE serve profile of the built package (`--profile NAME`; default profile otherwise), run `server/smoke_test.py` against it (asserts ETH dispatch, the 12x10 grid and the profile's pins; compares with the stored CPU reference), keep the evidence (container log, `/info`, the `/predict` output and the result, in `logs/smoke/` of the repo or `--log-dir DIR`), always stop it; one `bin/devrun -t 3600 -k 150` window per profile | yes |
|
| 10 |
+
| `fetch_samples.sh` | downloads public sample data that may not be redistributed (sha256-checked); diffusion-planner-p150 has none to fetch: both shipped samples are Apache-2.0, and the nuScenes-derived planning instants of the accuracy tables live in the tt-models workspace only | no |
|
| 11 |
+
|
| 12 |
+
Everything here ships in `code/` on the Hub and inside the image (`source.extra_code` lists `scripts`).
|
| 13 |
+
|
| 14 |
+
The card's demo media (`media/`) were rendered from p150 outputs by workspace scripts that are not shipped, because
|
| 15 |
+
most of the data they read (nuScenes) may not be redistributed (`media/ATTRIBUTION.md`). `examples/quickstart.py`
|
| 16 |
+
writes a simple render (`quickstart_bev.png`: the input tensors and the plan, bird's-eye view) for any scene `.npz`.
|
code/scripts/bench.py
ADDED
|
@@ -0,0 +1,207 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#!/usr/bin/env python3
|
| 2 |
+
# SPDX-License-Identifier: Apache-2.0
|
| 3 |
+
"""Stage breakdown of warm plans -- the numbers OPT_BASELINE.md / OPT_REPORT.md / the card quote.
|
| 4 |
+
|
| 5 |
+
bin/devrun -t 900 -- python code/scripts/bench.py --iters 100 --json out.json
|
| 6 |
+
bin/devrun -t 900 -- python code/scripts/bench.py --dispatch worker --num-cqs 2 --iters 100 --json out.json
|
| 7 |
+
bin/devrun -t 900 -- python code/scripts/bench.py --input <a.npz> --input <b.npz> ...
|
| 8 |
+
|
| 9 |
+
``--input`` takes the planner tensors as an ``.npz`` with the 15 ``INPUT_SCHEMA`` names (the shipped samples) or with
|
| 10 |
+
``raw/<name>`` keys (the research / public-data scene files); default: the shipped ``kashiwanoha_dense.npz``.
|
| 11 |
+
Per input, ``ttaw.profiling.StageBench`` collects ``--iters`` warm iterations of each stage (p50 / p99 / mean / min):
|
| 12 |
+
|
| 13 |
+
- ``load``: decoding the ``.npz`` + the ``INPUT_SCHEMA`` check (``model(inputs=<path>)`` does it outside timing_ms);
|
| 14 |
+
- ``host_pre``: the node's pre-processing (``host.prepare``: normalization, speed masks, encoder host features,
|
| 15 |
+
decoder masks, the solver's initial state);
|
| 16 |
+
- ``pack``: the 19 persistent trace inputs (``tt.inputs.plan_inputs``);
|
| 17 |
+
- ``host_in``: their ttnn host tensors (fp32 / bf16 TILE, ``ttaw.tensors.to_host_tensor``);
|
| 18 |
+
- ``h2d``: the upload into the persistent device inputs + device sync;
|
| 19 |
+
- ``trace``: one replay of the ``plan`` trace + device sync (the device latency of one plan);
|
| 20 |
+
- ``d2h``: the one packed readback (``final_x0`` + logits + the ego rows of the 11 iterates);
|
| 21 |
+
- ``host_post``: the node's post-processing (``host.make_output``);
|
| 22 |
+
- ``e2e``: ``model(inputs=<decoded arrays>)``, the in-process API call (schema check included);
|
| 23 |
+
- ``e2e_path``: ``model(inputs=<path>)`` (adds ``load``; schema-named ``.npz`` files only);
|
| 24 |
+
- ``b2b``: back-to-back replays with no host work in between (device time per plan), ``--b2b-rounds`` rounds of
|
| 25 |
+
``--b2b-iters`` replays; ``plans_per_s`` = 1000 / median b2b.
|
| 26 |
+
|
| 27 |
+
Also: ``model(...).timing_ms`` (preprocess / device / postprocess / total), the first call after ``from_pretrained``
|
| 28 |
+
and the load time (weights, build, warm-up + capture), AICLK / power / temperature sampled from sysfs during the
|
| 29 |
+
timed loops (``ttaw.profiling.AiclkSampler``), the staged path checked bit for bit against ``model()``, the device
|
| 30 |
+
configuration (dispatch, CQs, grid) and the numerics options in effect (``DIFFUSION_PLANNER_*`` knobs). Always quote
|
| 31 |
+
the configuration line with the numbers.
|
| 32 |
+
"""
|
| 33 |
+
from __future__ import annotations
|
| 34 |
+
|
| 35 |
+
import argparse
|
| 36 |
+
import json
|
| 37 |
+
import statistics
|
| 38 |
+
import time
|
| 39 |
+
from pathlib import Path
|
| 40 |
+
from typing import Any, Dict, Optional
|
| 41 |
+
|
| 42 |
+
import numpy as np
|
| 43 |
+
|
| 44 |
+
from tt_diffusion_planner import DiffusionPlanner
|
| 45 |
+
|
| 46 |
+
SAMPLE = Path(__file__).resolve().parents[1] / "tt_diffusion_planner" / "samples" / "kashiwanoha_dense.npz"
|
| 47 |
+
|
| 48 |
+
|
| 49 |
+
def load_scene(path: str) -> Dict[str, Any]:
|
| 50 |
+
"""``{"name", "path", "arrays", "schema_npz"}``: the 15 raw tensors of a schema-named or ``raw/``-prefixed npz."""
|
| 51 |
+
from tt_diffusion_planner.reference import config as C
|
| 52 |
+
|
| 53 |
+
p = Path(path)
|
| 54 |
+
with np.load(p, allow_pickle=False) as z:
|
| 55 |
+
files = set(z.files)
|
| 56 |
+
if all(k in files for k in C.INPUT_NAMES):
|
| 57 |
+
arrays, schema_npz = {k: np.array(z[k]) for k in C.INPUT_NAMES}, True
|
| 58 |
+
elif all(f"raw/{k}" in files for k in C.INPUT_NAMES):
|
| 59 |
+
arrays, schema_npz = {k: np.array(z[f"raw/{k}"]) for k in C.INPUT_NAMES}, False
|
| 60 |
+
else:
|
| 61 |
+
raise SystemExit(f"{p}: neither the INPUT_SCHEMA names nor raw/<name> keys")
|
| 62 |
+
name = p.stem[len("golden_"):] if p.stem.startswith("golden_") else p.stem
|
| 63 |
+
if p.parent.name not in ("samples", "ort"):
|
| 64 |
+
name = f"{p.parent.name}/{name}"
|
| 65 |
+
return {"name": name, "path": str(p), "arrays": arrays, "schema_npz": schema_npz}
|
| 66 |
+
|
| 67 |
+
|
| 68 |
+
def counts(arrays: Dict[str, np.ndarray]) -> Dict[str, int]:
|
| 69 |
+
"""Valid entities of a scene (non-empty rows), for the report."""
|
| 70 |
+
def rows(a, axis):
|
| 71 |
+
return int(np.any(np.abs(a) > 0, axis=axis).sum())
|
| 72 |
+
return {"neighbors": rows(arrays["neighbor_agents_past"][0], (1, 2)), "lanes": rows(arrays["lanes"][0], (1, 2)),
|
| 73 |
+
"route_lanes": rows(arrays["route_lanes"][0], (1, 2)), "polygons": rows(arrays["polygons"][0], (1, 2)),
|
| 74 |
+
"line_strings": rows(arrays["line_strings"][0], (1, 2))}
|
| 75 |
+
|
| 76 |
+
|
| 77 |
+
def unpack(out: Dict[str, np.ndarray]) -> Dict[str, Any]:
|
| 78 |
+
"""The packed readback -> the raw outputs of ``TtDiffusionPlanner.forward`` (same reshapes)."""
|
| 79 |
+
from tt_diffusion_planner.reference import config as C
|
| 80 |
+
from tt_diffusion_planner.tt import config as T
|
| 81 |
+
|
| 82 |
+
final = out["final_x0"].reshape(T.AGENTS, T.STATE_COLS)[:C.MAX_NUM_AGENTS]
|
| 83 |
+
steps = out["ego_steps"].reshape(-1, T.STATE_COLS)
|
| 84 |
+
return {"final_x0": final.reshape(C.MAX_NUM_AGENTS, C.OUTPUT_T + 1, C.POSE_DIM).astype(np.float32),
|
| 85 |
+
"logit": out["logit"].reshape(-1)[:C.TURN_INDICATOR_OUTPUT_DIM].astype(np.float32),
|
| 86 |
+
"denoising_steps": [s.reshape(1, C.OUTPUT_T + 1, C.POSE_DIM).astype(np.float32) for s in steps]}
|
| 87 |
+
|
| 88 |
+
|
| 89 |
+
def bench_scene(model, scene: Dict[str, Any], a: argparse.Namespace) -> Dict[str, Any]:
|
| 90 |
+
import ttnn
|
| 91 |
+
|
| 92 |
+
from tt_diffusion_planner.host import pipeline as hp
|
| 93 |
+
from tt_diffusion_planner.reference import config as C
|
| 94 |
+
from tt_diffusion_planner.tt import inputs as I
|
| 95 |
+
from tt_diffusion_planner.ttaw.io import load_named_arrays
|
| 96 |
+
from tt_diffusion_planner.ttaw.profiling import AiclkSampler, StageBench, time_b2b
|
| 97 |
+
from tt_diffusion_planner.ttaw.tensors import to_host_tensor
|
| 98 |
+
|
| 99 |
+
runner, dev = model.runner, model.device
|
| 100 |
+
arrays, params = scene["arrays"], model.validate_params({})
|
| 101 |
+
obs = model.normalization.observation
|
| 102 |
+
bench = StageBench(f"diffusion-planner {scene['name']}")
|
| 103 |
+
timing: Dict[str, list] = {}
|
| 104 |
+
sync = lambda: ttnn.synchronize_device(dev) # noqa: E731
|
| 105 |
+
for _ in range(a.warmup):
|
| 106 |
+
ref = model(inputs=arrays)
|
| 107 |
+
with AiclkSampler(chip=a.chip, interval_s=0.05) as clk:
|
| 108 |
+
for _ in range(a.iters): # the in-process API call
|
| 109 |
+
with bench.stage("e2e"):
|
| 110 |
+
ref = model(inputs=arrays)
|
| 111 |
+
for k, v in ref.timing_ms.items():
|
| 112 |
+
timing.setdefault(k, []).append(v)
|
| 113 |
+
slots = {k: runner._slot(k, "input") for k in I.INPUT_SPECS}
|
| 114 |
+
for _ in range(a.iters): # the same path, stage by stage
|
| 115 |
+
if scene["schema_npz"]:
|
| 116 |
+
with bench.stage("load"):
|
| 117 |
+
raw = load_named_arrays(scene["path"], C.INPUT_SCHEMA)
|
| 118 |
+
else:
|
| 119 |
+
raw = load_named_arrays(arrays, C.INPUT_SCHEMA)
|
| 120 |
+
with bench.stage("host_pre"):
|
| 121 |
+
prep = hp.prepare(raw, obs)
|
| 122 |
+
with bench.stage("pack"):
|
| 123 |
+
packed = I.plan_inputs(prep)
|
| 124 |
+
with bench.stage("host_in"):
|
| 125 |
+
host = {k: to_host_tensor(v, slots[k].dtype, slots[k].layout, shape=slots[k].shape)
|
| 126 |
+
for k, v in packed.items()}
|
| 127 |
+
with bench.stage("h2d"):
|
| 128 |
+
runner.upload(host)
|
| 129 |
+
sync()
|
| 130 |
+
with bench.stage("trace"):
|
| 131 |
+
runner.replay("plan")
|
| 132 |
+
sync()
|
| 133 |
+
with bench.stage("d2h"):
|
| 134 |
+
out = runner.read("plan")
|
| 135 |
+
with bench.stage("host_post"):
|
| 136 |
+
res = model._postprocess(unpack(out), prep, params)
|
| 137 |
+
if scene["schema_npz"]:
|
| 138 |
+
for _ in range(a.iters):
|
| 139 |
+
with bench.stage("e2e_path"):
|
| 140 |
+
model(inputs=scene["path"])
|
| 141 |
+
rounds = [time_b2b(lambda: runner.replay("plan"), sync, n=a.b2b_iters, warmup=3)
|
| 142 |
+
for _ in range(a.b2b_rounds)]
|
| 143 |
+
for r in rounds:
|
| 144 |
+
bench.add("b2b", r)
|
| 145 |
+
same = bool(np.array_equal(res.poses, ref.poses) and np.array_equal(res.predicted_agents, ref.predicted_agents)
|
| 146 |
+
and res.turn_indicator["command"] == ref.turn_indicator["command"])
|
| 147 |
+
summary = bench.summary()
|
| 148 |
+
b2b = statistics.median(rounds)
|
| 149 |
+
return {"name": scene["name"], "path": scene["path"], "valid": counts(arrays), "stages_ms": summary,
|
| 150 |
+
"timing_ms": {k: {"p50": float(np.percentile(v, 50)), "p99": float(np.percentile(v, 99)),
|
| 151 |
+
"min": float(min(v))} for k, v in timing.items()},
|
| 152 |
+
"b2b_rounds_ms": rounds, "plans_per_s_b2b": 1000.0 / b2b, "aiclk": clk.summary(),
|
| 153 |
+
"staged_equals_model": same, "turn_command": int(ref.turn_indicator["command"]),
|
| 154 |
+
"table": bench.table()}
|
| 155 |
+
|
| 156 |
+
|
| 157 |
+
def main() -> None:
|
| 158 |
+
ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
|
| 159 |
+
ap.add_argument("--iters", type=int, default=100)
|
| 160 |
+
ap.add_argument("--warmup", type=int, default=5)
|
| 161 |
+
ap.add_argument("--b2b-iters", type=int, default=50)
|
| 162 |
+
ap.add_argument("--b2b-rounds", type=int, default=3)
|
| 163 |
+
ap.add_argument("--input", action="append", default=None, help="scene .npz (repeatable)")
|
| 164 |
+
ap.add_argument("--dispatch", default=None, choices=["eth", "worker"])
|
| 165 |
+
ap.add_argument("--num-cqs", type=int, default=None, choices=[1, 2])
|
| 166 |
+
ap.add_argument("--chip", type=int, default=0, help="sysfs chip index for the AICLK sampler")
|
| 167 |
+
ap.add_argument("--tag", default="")
|
| 168 |
+
ap.add_argument("--json")
|
| 169 |
+
a = ap.parse_args()
|
| 170 |
+
from tt_diffusion_planner.reference.weights import find_weights_dir
|
| 171 |
+
|
| 172 |
+
scenes = [load_scene(p) for p in (a.input or [str(SAMPLE)])]
|
| 173 |
+
wd = find_weights_dir() # None: from_pretrained resolves the pinned HF snapshot
|
| 174 |
+
t0 = time.perf_counter()
|
| 175 |
+
with DiffusionPlanner.from_pretrained(dispatch=a.dispatch, num_command_queues=a.num_cqs,
|
| 176 |
+
weights_dir=str(wd) if wd else None) as model:
|
| 177 |
+
load_s = time.perf_counter() - t0
|
| 178 |
+
t1 = time.perf_counter()
|
| 179 |
+
model(inputs=scenes[0]["arrays"]) # first call after from_pretrained (traces captured)
|
| 180 |
+
first_ms = (time.perf_counter() - t1) * 1e3
|
| 181 |
+
info = model.tt.describe()
|
| 182 |
+
res: Dict[str, Any] = {
|
| 183 |
+
"tag": a.tag, "config": model.device_info, "iters": a.iters, "load_s": round(load_s, 2),
|
| 184 |
+
"warmup_ms": {k: round(v, 1) for k, v in model.warmup_ms.items()}, "first_call_ms": round(first_ms, 2),
|
| 185 |
+
"options": info["options"], "precision": info["precision"], "uploaded_mb": info["uploaded_mb"],
|
| 186 |
+
"trace_buffers_mb": info["trace_buffers_mb"],
|
| 187 |
+
"program_cache_entries": info["trace"].get("program_cache_entries"),
|
| 188 |
+
"num_command_queues": info["trace"]["num_command_queues"], "scenes": {}}
|
| 189 |
+
for scene in scenes:
|
| 190 |
+
r = bench_scene(model, scene, a)
|
| 191 |
+
res["scenes"][r["name"]] = r
|
| 192 |
+
cfg = model.device_info
|
| 193 |
+
print(f"\n## {r['name']} [{cfg.get('dispatch')} {res['num_command_queues']}CQ {cfg.get('grid')}] "
|
| 194 |
+
f"valid {r['valid']}")
|
| 195 |
+
print(r["table"])
|
| 196 |
+
print("timing_ms p50:", {k: round(v["p50"], 3) for k, v in r["timing_ms"].items()},
|
| 197 |
+
f"| b2b rounds {[round(x, 3) for x in r['b2b_rounds_ms']]} ms -> {r['plans_per_s_b2b']:.2f} plans/s",
|
| 198 |
+
f"| aiclk {r['aiclk'].get('aiclk_mhz')}", f"| check: staged == model() {r['staged_equals_model']}",
|
| 199 |
+
flush=True)
|
| 200 |
+
print(json.dumps({k: v for k, v in res.items() if k != "scenes"}, default=str))
|
| 201 |
+
if a.json:
|
| 202 |
+
Path(a.json).parent.mkdir(parents=True, exist_ok=True)
|
| 203 |
+
Path(a.json).write_text(json.dumps(res, indent=1, default=str) + "\n")
|
| 204 |
+
|
| 205 |
+
|
| 206 |
+
if __name__ == "__main__":
|
| 207 |
+
main()
|
code/scripts/bringup_device.py
ADDED
|
@@ -0,0 +1,167 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-License-Identifier: Apache-2.0
|
| 2 |
+
"""Device bring-up of the planner graph (development tool; needs the p150, run through ``bin/devrun``).
|
| 3 |
+
|
| 4 |
+
bin/devrun -t 1800 -- python code/scripts/bringup_device.py [--scenes kashiwanoha_dense straight_road]
|
| 5 |
+
[--no-capture] [--precision "dec.*=HiFi2+fp32"] [--json logs/diffusion-planner/bringup.json]
|
| 6 |
+
|
| 7 |
+
Protocol (PLAN.md 4.4): build ``TtDiffusionPlanner(debug=True)``; per scene run every variant EAGERLY twice
|
| 8 |
+
(``encoder_taps``, ``decode_once`` at evaluations 0 / 5 / 10, ``plan``), check the two eager runs are bit-identical and
|
| 9 |
+
compare them with the research goldens (``research/diffusion-planner/goldens/<scene>.npz``, the fp32 CPU reference:
|
| 10 |
+
PCC on valid rows, max abs; final x0 on valid agents; logits); then capture all variants (strict, no program
|
| 11 |
+
compiled after the capture) and check replay == eager bit for bit on the same inputs. Prints a table and writes a
|
| 12 |
+
JSON report. This is a diagnostic: the frozen gates live in ``tests/test_pcc_device.py`` / ``test_e2e_device.py``.
|
| 13 |
+
"""
|
| 14 |
+
from __future__ import annotations
|
| 15 |
+
|
| 16 |
+
import argparse
|
| 17 |
+
import json
|
| 18 |
+
import sys
|
| 19 |
+
import time
|
| 20 |
+
from pathlib import Path
|
| 21 |
+
|
| 22 |
+
import numpy as np
|
| 23 |
+
|
| 24 |
+
HERE = Path(__file__).resolve()
|
| 25 |
+
sys.path.insert(0, str(HERE.parents[1]))
|
| 26 |
+
|
| 27 |
+
from tt_diffusion_planner.host import pipeline as hp # noqa: E402
|
| 28 |
+
from tt_diffusion_planner.reference import config as C # noqa: E402
|
| 29 |
+
from tt_diffusion_planner.ttaw.metrics import error_stats # noqa: E402
|
| 30 |
+
|
| 31 |
+
GOLDENS = HERE.parents[4] / "research" / "diffusion-planner" / "goldens"
|
| 32 |
+
|
| 33 |
+
|
| 34 |
+
def stats(dev, ref):
|
| 35 |
+
s = error_stats(np.asarray(dev, np.float64), np.asarray(ref, np.float64))
|
| 36 |
+
return {"pcc": round(s["pcc"], 7), "max_abs": float(f"{s['max_abs']:.4g}"), "rel_l2": float(f"{s['rel_l2']:.4g}")}
|
| 37 |
+
|
| 38 |
+
|
| 39 |
+
def main() -> int:
|
| 40 |
+
ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
|
| 41 |
+
ap.add_argument("--scenes", nargs="+", default=["kashiwanoha_dense", "straight_road"])
|
| 42 |
+
ap.add_argument("--no-capture", action="store_true")
|
| 43 |
+
ap.add_argument("--precision", default=None)
|
| 44 |
+
ap.add_argument("--ln-fp32", default=None, help="module globs (comma-separated) for the fp32 LayerNorm")
|
| 45 |
+
ap.add_argument("--hidden-fp32", default=None, help="module globs (comma-separated) for fp32 hidden activations")
|
| 46 |
+
ap.add_argument("--split", default=None, help="module globs (comma-separated) for split (hi/lo) fp32 matmuls")
|
| 47 |
+
ap.add_argument("--attn-fp32-acc", default=None, help="module globs for SDPA with fp32 accumulation")
|
| 48 |
+
ap.add_argument("--attn-matmul", default=None, help="module globs for the fp32 matmul attention")
|
| 49 |
+
ap.add_argument("--plan-only", action="store_true", help="skip the encoder taps and decoder evaluations")
|
| 50 |
+
ap.add_argument("--evals", nargs="+", type=int, default=[0, 5, 10])
|
| 51 |
+
ap.add_argument("--json", default=None)
|
| 52 |
+
args = ap.parse_args()
|
| 53 |
+
|
| 54 |
+
import torch
|
| 55 |
+
|
| 56 |
+
torch.set_num_threads(4)
|
| 57 |
+
from tt_diffusion_planner.device import close_device, open_device
|
| 58 |
+
from tt_diffusion_planner.reference.weights import find_weights_dir, load_weights
|
| 59 |
+
from tt_diffusion_planner.tt.model import TtDiffusionPlanner
|
| 60 |
+
|
| 61 |
+
weights = load_weights(find_weights_dir())
|
| 62 |
+
dev = open_device(allow_fallback=False)
|
| 63 |
+
g = dev.compute_with_storage_grid_size()
|
| 64 |
+
report = {"grid": f"{g.x}x{g.y}", "precision": args.precision, "scenes": {}}
|
| 65 |
+
t0 = time.perf_counter()
|
| 66 |
+
from tt_diffusion_planner.tt.config import globs
|
| 67 |
+
|
| 68 |
+
opts = {k: globs(v) for k, v in (("ln_fp32", args.ln_fp32), ("hidden_fp32", args.hidden_fp32),
|
| 69 |
+
("split", args.split), ("attn_fp32_acc", args.attn_fp32_acc),
|
| 70 |
+
("attn_matmul", args.attn_matmul)) if v is not None}
|
| 71 |
+
tt = TtDiffusionPlanner(dev, weights, debug=True, precision=args.precision, **opts)
|
| 72 |
+
report["options"] = tt.build.options()
|
| 73 |
+
report["build_s"] = round(time.perf_counter() - t0, 2)
|
| 74 |
+
print(f"build {report['build_s']} s, uploaded {tt.build.uploaded_bytes / 2**20:.1f} MiB", flush=True)
|
| 75 |
+
runner = tt.runner
|
| 76 |
+
from tt_diffusion_planner.tt import inputs as I
|
| 77 |
+
|
| 78 |
+
eager_cache = {}
|
| 79 |
+
try:
|
| 80 |
+
for scene in args.scenes:
|
| 81 |
+
with np.load(GOLDENS / f"{scene}.npz", allow_pickle=False) as z:
|
| 82 |
+
gold = {k: z[k] for k in z.files if k != "__meta__"}
|
| 83 |
+
raw = {k: gold[f"in.{k}"] for k in C.INPUT_NAMES}
|
| 84 |
+
prep = hp.prepare(raw, weights.normalization.observation)
|
| 85 |
+
res = {}
|
| 86 |
+
if args.plan_only:
|
| 87 |
+
args.evals = []
|
| 88 |
+
# encoder taps (eager x2)
|
| 89 |
+
t1 = time.perf_counter()
|
| 90 |
+
taps = tt.encoder_taps(prep, eager=True)
|
| 91 |
+
taps2 = taps if args.plan_only else tt.encoder_taps(prep, eager=True)
|
| 92 |
+
res["encoder_eager_s"] = round(time.perf_counter() - t1, 2)
|
| 93 |
+
res["encoder_deterministic"] = all(np.array_equal(taps[k], taps2[k]) for k in taps)
|
| 94 |
+
for name, _ in C.TOKEN_LAYOUT:
|
| 95 |
+
rows = np.flatnonzero(gold[f"host.valid.{name}"])
|
| 96 |
+
if rows.size:
|
| 97 |
+
res[f"enc.{name}"] = stats(taps[f"enc.{name}"][rows], gold[f"enc.{name}"][rows])
|
| 98 |
+
for c in ("ego", "neighbor", "lane", "route", "polygon", "line_string"):
|
| 99 |
+
rows = gold[f"enc.{c}.pre.rows"]
|
| 100 |
+
if rows.size:
|
| 101 |
+
for part in ("pre", "mixer"):
|
| 102 |
+
res[f"enc.{c}.{part}"] = stats(taps[f"enc.{c}.{part}"][rows], gold[f"enc.{c}.{part}"])
|
| 103 |
+
tok = np.flatnonzero(gold["host.token_valid"])
|
| 104 |
+
for name in ["enc.tokens"] + [f"enc.fusion.{i}" for i in range(6)] + ["enc.encoding"]:
|
| 105 |
+
res[name] = stats(taps[name][tok], gold[name][tok])
|
| 106 |
+
# decode_once, teacher forced
|
| 107 |
+
arows = gold["dec.rows"]
|
| 108 |
+
for k in args.evals:
|
| 109 |
+
outs = [tt_decode(tt, runner, prep, gold, k) for _ in range(2)]
|
| 110 |
+
res[f"dec.eval{k}.deterministic"] = bool(np.array_equal(outs[0], outs[1]))
|
| 111 |
+
res[f"dec.eval{k}"] = stats(outs[0][arows][:, 1:], gold["dec.out"][k][:, 1:])
|
| 112 |
+
# plan (eager x2)
|
| 113 |
+
t1 = time.perf_counter()
|
| 114 |
+
p1 = runner.run_eager("plan", inputs=I.plan_inputs(prep))
|
| 115 |
+
p2 = runner.run_eager("plan", inputs=I.plan_inputs(prep))
|
| 116 |
+
res["plan_eager_s"] = round(time.perf_counter() - t1, 2)
|
| 117 |
+
res["plan_deterministic"] = all(np.array_equal(p1[k], p2[k]) for k in p1)
|
| 118 |
+
eager_cache[scene] = (prep, p1)
|
| 119 |
+
final = p1["final_x0"].reshape(352, 324)[:321].reshape(321, 81, 4)
|
| 120 |
+
res["plan.final_x0"] = stats(final[arows], gold["final_x0"][arows])
|
| 121 |
+
res["plan.logit"] = {"dev": [round(float(v), 4) for v in p1["logit"].reshape(-1)[:5]],
|
| 122 |
+
"ref": [round(float(v), 4) for v in gold["turn.logit"]]}
|
| 123 |
+
out = hp.make_output(final, p1["logit"].reshape(-1)[:5], prep, weights.normalization,
|
| 124 |
+
{k: v[3] for k, v in hp.RUNTIME_PARAMS.items()})
|
| 125 |
+
dpos = np.hypot(*(out.poses[:, :2] - gold["out.trajectory"][:, :2]).T)
|
| 126 |
+
res["plan.ego_max_m"], res["plan.ego_mean_m"] = round(float(dpos.max()), 4), round(float(dpos.mean()), 4)
|
| 127 |
+
res["plan.turn_equal"] = int(out.turn_indicator["command"]) == int(gold["out.turn_command"])
|
| 128 |
+
if gold["out.predicted_agents"].shape[0]:
|
| 129 |
+
dxy = out.predicted_agents[..., :2] - gold["out.predicted_agents"][..., :2]
|
| 130 |
+
per = np.hypot(*dxy.transpose(2, 0, 1))
|
| 131 |
+
res["plan.nb_median_max_m"] = round(float(np.median(per.max(axis=1))), 4)
|
| 132 |
+
report["scenes"][scene] = res
|
| 133 |
+
print(json.dumps({scene: res}, indent=1), flush=True)
|
| 134 |
+
if not args.no_capture:
|
| 135 |
+
t1 = time.perf_counter()
|
| 136 |
+
runner.capture()
|
| 137 |
+
report["capture_s"] = round(time.perf_counter() - t1, 2)
|
| 138 |
+
report["timings_ms"] = runner.timings_ms
|
| 139 |
+
for scene, (prep, p1) in eager_cache.items():
|
| 140 |
+
rep = runner("plan", inputs=I.plan_inputs(prep))
|
| 141 |
+
report["scenes"][scene]["plan_replay_equals_eager"] = all(np.array_equal(rep[k], p1[k]) for k in p1)
|
| 142 |
+
import ttnn
|
| 143 |
+
|
| 144 |
+
ttnn.synchronize_device(dev)
|
| 145 |
+
t2 = time.perf_counter()
|
| 146 |
+
runner.replay("plan", 5)
|
| 147 |
+
ttnn.synchronize_device(dev)
|
| 148 |
+
report["scenes"][scene]["plan_replay_ms"] = round((time.perf_counter() - t2) / 5 * 1e3, 3)
|
| 149 |
+
print(json.dumps({k: v for k, v in report.items() if k != "scenes"}, indent=1, default=str), flush=True)
|
| 150 |
+
print(json.dumps({s: {k: v for k, v in r.items() if k.startswith("plan_")}
|
| 151 |
+
for s, r in report["scenes"].items()}, indent=1), flush=True)
|
| 152 |
+
finally:
|
| 153 |
+
tt.release()
|
| 154 |
+
close_device(dev)
|
| 155 |
+
if args.json:
|
| 156 |
+
Path(args.json).parent.mkdir(parents=True, exist_ok=True)
|
| 157 |
+
Path(args.json).write_text(json.dumps(report, indent=1, default=str) + "\n")
|
| 158 |
+
return 0
|
| 159 |
+
|
| 160 |
+
|
| 161 |
+
def tt_decode(tt, runner, prep, gold, k):
|
| 162 |
+
return tt.decode_once(prep, gold["dec.x_in"][k], float(gold["dec.t"][k]), encoding=gold["enc.encoding"],
|
| 163 |
+
eager=True)
|
| 164 |
+
|
| 165 |
+
|
| 166 |
+
if __name__ == "__main__":
|
| 167 |
+
sys.exit(main())
|
code/scripts/container_smoke.sh
ADDED
|
@@ -0,0 +1,137 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#!/usr/bin/env bash
|
| 2 |
+
# SPDX-License-Identifier: Apache-2.0
|
| 3 |
+
# Serve ONE serve profile of the BUILT container package, run the smoke test against it, keep the evidence, and always
|
| 4 |
+
# stop it -- in one device-lock window. Run it once per serve profile (`tt-model profiles
|
| 5 |
+
# <staged dir>/tt_kernel_manifest.json` lists them; without --profile the package's default profile is served):
|
| 6 |
+
#
|
| 7 |
+
# ROOT=/home/ubuntu/experiments/tt-models
|
| 8 |
+
# $ROOT/bin/devrun -t 3600 -k 150 -- env -u HF_TOKEN -u HUGGING_FACE_HUB_TOKEN sg docker -c \
|
| 9 |
+
# "bash code/scripts/container_smoke.sh $ROOT/build/diffusion-planner-p150 [port] [--profile NAME] [--log-dir DIR]"
|
| 10 |
+
#
|
| 11 |
+
# The smoke test FAILS unless /info reports ETH dispatch and the 12x10 grid and runs what the package pins for the
|
| 12 |
+
# profile (dispatch, CQs, variant, weights revision); it also compares the output with the stored CPU reference of the
|
| 13 |
+
# sample when one exists (code/tt_diffusion_planner/server/smoke_test.py). Exit code: the smoke test's (0 = PASS); 1 when
|
| 14 |
+
# serve fails, 2 on a usage error.
|
| 15 |
+
#
|
| 16 |
+
# Evidence, kept whatever the outcome (--log-dir, default: logs/smoke/ of this repo, next to code/), named
|
| 17 |
+
# <name>[-<profile>]-<UTC time>.*:
|
| 18 |
+
# .container.log the container's whole log (boot, requests, shutdown), followed while the container stops
|
| 19 |
+
# .info.json GET /info as soon as the server is READY (device: dispatch, grid, cores; pins; versions)
|
| 20 |
+
# .smoke.json the smoke test's /predict output (the SMOKE_OUT environment variable overrides this path)
|
| 21 |
+
# .result.json profile, port, exit code, start / end times
|
| 22 |
+
#
|
| 23 |
+
# `tt-model serve` returns once the server is READY and leaves the container running, so the stop must happen before
|
| 24 |
+
# the lock is released, also when devrun's timeout TERMs this script: the EXIT trap saves the log, then stops the
|
| 25 |
+
# container cleanly with SIGTERM (120 s grace, hence devrun -k 150), never `docker kill`, which would leave the chip
|
| 26 |
+
# dirty. Needs docker access (sg docker) and python3.
|
| 27 |
+
set -u
|
| 28 |
+
usage() { echo "usage: container_smoke.sh <staged package dir> [port] [--profile NAME] [--log-dir DIR]" >&2; }
|
| 29 |
+
PROFILE=""
|
| 30 |
+
LOG_DIR=""
|
| 31 |
+
POSITIONAL=()
|
| 32 |
+
while [ $# -gt 0 ]; do
|
| 33 |
+
case "$1" in
|
| 34 |
+
--profile) [ $# -ge 2 ] || { usage; exit 2; }; PROFILE="$2"; shift 2 ;;
|
| 35 |
+
--profile=*) PROFILE="${1#--profile=}"; shift ;;
|
| 36 |
+
--log-dir) [ $# -ge 2 ] || { usage; exit 2; }; LOG_DIR="$2"; shift 2 ;;
|
| 37 |
+
--log-dir=*) LOG_DIR="${1#--log-dir=}"; shift ;;
|
| 38 |
+
-h|--help) usage; exit 0 ;;
|
| 39 |
+
-*) echo "container_smoke.sh: unknown option $1" >&2; usage; exit 2 ;;
|
| 40 |
+
*) POSITIONAL+=("$1"); shift ;;
|
| 41 |
+
esac
|
| 42 |
+
done
|
| 43 |
+
[ ${#POSITIONAL[@]} -ge 1 ] && [ ${#POSITIONAL[@]} -le 2 ] || { usage; exit 2; }
|
| 44 |
+
STAGED="${POSITIONAL[0]}"
|
| 45 |
+
PORT="${POSITIONAL[1]:-20000}"
|
| 46 |
+
MANIFEST="$STAGED/tt_kernel_manifest.json"
|
| 47 |
+
[ -f "$MANIFEST" ] || { echo "container_smoke.sh: $MANIFEST not found (run tt-model package first)" >&2; exit 2; }
|
| 48 |
+
HERE="$(cd "$(dirname "$0")" && pwd)"
|
| 49 |
+
PROFILE_ARGS=()
|
| 50 |
+
[ -n "$PROFILE" ] && PROFILE_ARGS=(--profile "$PROFILE")
|
| 51 |
+
|
| 52 |
+
# The package name and the served profile (tt-model's rule: --profile, else default_profile, else the first serve
|
| 53 |
+
# profile), hence the container name tt-model gives it: tt-model-<name>-<profile> (tt_kernel/container.py).
|
| 54 |
+
read -r NAME PROFILE_NAME < <(python3 - "$MANIFEST" "$PROFILE" <<'EOF'
|
| 55 |
+
import json, sys
|
| 56 |
+
m = json.load(open(sys.argv[1]))
|
| 57 |
+
c = m.get("container") or {}
|
| 58 |
+
profiles = c.get("serve_profiles") or [{}]
|
| 59 |
+
print(m.get("name") or "model", sys.argv[2] or c.get("default_profile") or profiles[0].get("name") or "default")
|
| 60 |
+
EOF
|
| 61 |
+
)
|
| 62 |
+
[ -n "${NAME:-}" ] || { echo "container_smoke.sh: cannot read the package name from $MANIFEST" >&2; exit 2; }
|
| 63 |
+
CONTAINER="tt-model-$NAME-$PROFILE_NAME"
|
| 64 |
+
|
| 65 |
+
[ -n "$LOG_DIR" ] || LOG_DIR="$(cd "$HERE/../.." && pwd)/logs/smoke"
|
| 66 |
+
if ! mkdir -p "$LOG_DIR" 2>/dev/null || [ ! -w "$LOG_DIR" ]; then
|
| 67 |
+
echo "container_smoke.sh: cannot write $LOG_DIR; keeping the evidence in ${TMPDIR:-/tmp}" >&2
|
| 68 |
+
LOG_DIR="${TMPDIR:-/tmp}"
|
| 69 |
+
fi
|
| 70 |
+
LOG_DIR="$(cd "$LOG_DIR" && pwd)"
|
| 71 |
+
STARTED="$(date -u +%Y-%m-%dT%H:%M:%SZ)"
|
| 72 |
+
STEM="$LOG_DIR/$NAME${PROFILE:+-$PROFILE}-$(date -u +%Y%m%dT%H%M%SZ)"
|
| 73 |
+
SMOKE_JSON="${SMOKE_OUT:-$STEM.smoke.json}"
|
| 74 |
+
|
| 75 |
+
fetch() { # fetch URL FILE: one GET with a 30 s timeout; the body goes to FILE
|
| 76 |
+
python3 - "$1" "$2" <<'EOF'
|
| 77 |
+
import sys, urllib.request
|
| 78 |
+
try:
|
| 79 |
+
with urllib.request.urlopen(sys.argv[1], timeout=30) as r:
|
| 80 |
+
body = r.read()
|
| 81 |
+
except Exception as e: # noqa: BLE001 -- best effort: the smoke test reports the server's state
|
| 82 |
+
sys.exit(f"container_smoke.sh: GET {sys.argv[1]} failed: {e}")
|
| 83 |
+
open(sys.argv[2], "wb").write(body)
|
| 84 |
+
EOF
|
| 85 |
+
}
|
| 86 |
+
|
| 87 |
+
CHILD=""
|
| 88 |
+
STOPPED=0
|
| 89 |
+
cleanup() {
|
| 90 |
+
local rc=$?
|
| 91 |
+
[ "$STOPPED" = 1 ] && return
|
| 92 |
+
STOPPED=1
|
| 93 |
+
if [ -n "$CHILD" ]; then kill -TERM "$CHILD" 2>/dev/null; wait "$CHILD" 2>/dev/null; fi
|
| 94 |
+
# The log is lost with the container: follow it (whole history, then the shutdown lines) while it stops.
|
| 95 |
+
local logger=""
|
| 96 |
+
if docker inspect "$CONTAINER" >/dev/null 2>&1; then
|
| 97 |
+
docker logs --follow "$CONTAINER" > "$STEM.container.log" 2>&1 &
|
| 98 |
+
logger=$!
|
| 99 |
+
else
|
| 100 |
+
tt-model logs "${PROFILE_ARGS[@]}" "$MANIFEST" > "$STEM.container.log" 2>&1 || true
|
| 101 |
+
fi
|
| 102 |
+
tt-model stop "${PROFILE_ARGS[@]}" "$MANIFEST" || true
|
| 103 |
+
if [ -n "$logger" ]; then
|
| 104 |
+
for _ in $(seq 1 30); do kill -0 "$logger" 2>/dev/null || break; sleep 1; done
|
| 105 |
+
kill "$logger" 2>/dev/null
|
| 106 |
+
wait "$logger" 2>/dev/null
|
| 107 |
+
fi
|
| 108 |
+
python3 - "$STEM.result.json" "$NAME" "$PROFILE_NAME" "$PORT" "$rc" "$STARTED" "$CONTAINER" "$MANIFEST" <<'EOF'
|
| 109 |
+
import datetime, json, sys
|
| 110 |
+
out, name, profile, port, rc, started, container, manifest = sys.argv[1:]
|
| 111 |
+
ended = datetime.datetime.now(datetime.timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
|
| 112 |
+
json.dump({"bundle": name, "profile": profile, "port": int(port) if port.isdigit() else port, "rc": int(rc),
|
| 113 |
+
"result": "PASS" if rc == "0" else "FAIL",
|
| 114 |
+
"started": started, "ended": ended, "container": container, "manifest": manifest}, open(out, "w"), indent=1)
|
| 115 |
+
EOF
|
| 116 |
+
echo "container_smoke.sh: rc=$rc; evidence in $STEM.*"
|
| 117 |
+
}
|
| 118 |
+
trap cleanup EXIT
|
| 119 |
+
trap 'exit 143' TERM
|
| 120 |
+
trap 'exit 130' INT
|
| 121 |
+
|
| 122 |
+
# Each step runs in the background and is waited for: bash defers a trap while a FOREGROUND command runs, so a TERM
|
| 123 |
+
# sent to this script alone (devrun's timeout signals the whole process group) would otherwise wait for the step.
|
| 124 |
+
step() {
|
| 125 |
+
"$@" &
|
| 126 |
+
CHILD=$!
|
| 127 |
+
wait "$CHILD"
|
| 128 |
+
local rc=$?
|
| 129 |
+
CHILD=""
|
| 130 |
+
return "$rc"
|
| 131 |
+
}
|
| 132 |
+
|
| 133 |
+
# --port and --profile BEFORE the target: options after it are passed through to the container (tt-model cli rule)
|
| 134 |
+
step tt-model serve --port "$PORT" "${PROFILE_ARGS[@]}" "$MANIFEST" || exit 1
|
| 135 |
+
step fetch "http://127.0.0.1:$PORT/info" "$STEM.info.json"
|
| 136 |
+
step python3 "$HERE/../tt_diffusion_planner/server/smoke_test.py" --url "http://127.0.0.1:$PORT" --wait 600 \
|
| 137 |
+
--manifest "$MANIFEST" "${PROFILE_ARGS[@]}" --out "$SMOKE_JSON"
|
code/scripts/fetch_samples.sh
ADDED
|
@@ -0,0 +1,24 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#!/usr/bin/env bash
|
| 2 |
+
# SPDX-License-Identifier: Apache-2.0
|
| 3 |
+
# Download the public sample data that diffusion-planner-p150 may NOT redistribute (license unstated or non-commercial)
|
| 4 |
+
# into ${DIFFUSION_PLANNER_SAMPLES:-$HOME/.cache/tt_diffusion_planner/samples}, sha256-checked. Tests and benchmarks read it from there.
|
| 5 |
+
# Nothing downloaded here is executed; parse it with numpy / sqlite3 only.
|
| 6 |
+
set -euo pipefail
|
| 7 |
+
DEST="${DIFFUSION_PLANNER_SAMPLES:-$HOME/.cache/tt_diffusion_planner/samples}"
|
| 8 |
+
mkdir -p "$DEST"
|
| 9 |
+
|
| 10 |
+
fetch() { # fetch <url> <sha256> <file name>
|
| 11 |
+
local url="$1" sum="$2" name="$3"
|
| 12 |
+
if [ -f "$DEST/$name" ] && echo "$sum $DEST/$name" | sha256sum -c --status; then
|
| 13 |
+
echo "ok $name"; return
|
| 14 |
+
fi
|
| 15 |
+
curl -fL --retry 3 -o "$DEST/$name.part" "$url"
|
| 16 |
+
echo "$sum $DEST/$name.part" | sha256sum -c --status || { echo "sha256 mismatch: $name" >&2; exit 1; }
|
| 17 |
+
mv "$DEST/$name.part" "$DEST/$name"; echo "fetched $name"
|
| 18 |
+
}
|
| 19 |
+
|
| 20 |
+
# Example (Autoware demo rosbag; pinned in autoware/ansible/roles/demo_artifacts/tasks/main.yaml:48-53):
|
| 21 |
+
# fetch https://autoware-files.s3.us-west-2.amazonaws.com/recordings/bags/demos/sample-rosbag.zip \
|
| 22 |
+
# 5f9d36353393b3d249212153c19049822b1298db56512aa045b4f7f6fc37cf88 sample-rosbag.zip
|
| 23 |
+
# nothing to fetch: both shipped samples are redistributable (Apache-2.0)
|
| 24 |
+
echo "samples in $DEST"
|
code/scripts/precision_exp.py
ADDED
|
@@ -0,0 +1,147 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-License-Identifier: Apache-2.0
|
| 2 |
+
"""Precision experiment for the MLP-Mixer blocks on the device (development tool; run under ``bin/devrun``).
|
| 3 |
+
|
| 4 |
+
bin/devrun -t 1200 -- python code/scripts/precision_exp.py [--json logs/diffusion-planner/precision_exp.json]
|
| 5 |
+
|
| 6 |
+
Part A: ``ttnn.layer_norm`` vs :func:`tt.layers.layer_norm_fp32` on the real offset-dominated mixer rows (the
|
| 7 |
+
reference ``enc.neighbor.pre`` of the samples, block 0 ``norm1``) against float64. Part B: the 6 mixer blocks of ego /
|
| 8 |
+
neighbour / lane run alone on the reference's own block input (``enc.<cat>.pre``, teacher forcing) under several
|
| 9 |
+
option sets, compared with the reference ``enc.<cat>.mixer`` on the valid entities. Eager (no trace).
|
| 10 |
+
"""
|
| 11 |
+
from __future__ import annotations
|
| 12 |
+
|
| 13 |
+
import argparse
|
| 14 |
+
import json
|
| 15 |
+
import sys
|
| 16 |
+
from pathlib import Path
|
| 17 |
+
|
| 18 |
+
import numpy as np
|
| 19 |
+
|
| 20 |
+
HERE = Path(__file__).resolve()
|
| 21 |
+
sys.path.insert(0, str(HERE.parents[1]))
|
| 22 |
+
|
| 23 |
+
from tt_diffusion_planner.reference import config as C # noqa: E402
|
| 24 |
+
from tt_diffusion_planner.ttaw.metrics import error_stats # noqa: E402
|
| 25 |
+
|
| 26 |
+
GOLDENS = HERE.parents[4] / "research" / "diffusion-planner" / "goldens"
|
| 27 |
+
CONFIGS = {
|
| 28 |
+
"default": dict(),
|
| 29 |
+
"ln_fp32": dict(ln_fp32=("enc.mixer.*",)),
|
| 30 |
+
"hidden_fp32": dict(hidden_fp32=("enc.mixer.*",)),
|
| 31 |
+
"ln_hidden_fp32": dict(ln_fp32=("enc.mixer.*",), hidden_fp32=("enc.mixer.*",)),
|
| 32 |
+
"all_fp32_w": dict(ln_fp32=("enc.mixer.*",), hidden_fp32=("enc.mixer.*",),
|
| 33 |
+
precision="enc.mixer.*=HiFi4+fp32:w=fp32:a=fp32"),
|
| 34 |
+
"stream_bf16": dict(precision="enc.mixer.*=HiFi4+fp32:w=bf16:a=bf16"),
|
| 35 |
+
}
|
| 36 |
+
|
| 37 |
+
|
| 38 |
+
def st(dev, ref):
|
| 39 |
+
s = error_stats(np.asarray(dev, np.float64), np.asarray(ref, np.float64))
|
| 40 |
+
return {"pcc": round(s["pcc"], 7), "rel_l2": float(f"{s['rel_l2']:.4g}"), "max_abs": float(f"{s['max_abs']:.4g}")}
|
| 41 |
+
|
| 42 |
+
|
| 43 |
+
def main() -> int:
|
| 44 |
+
ap = argparse.ArgumentParser()
|
| 45 |
+
ap.add_argument("--scenes", nargs="+", default=["kashiwanoha_dense", "straight_road"])
|
| 46 |
+
ap.add_argument("--cats", nargs="+", default=["ego", "neighbor", "lane"])
|
| 47 |
+
ap.add_argument("--configs", nargs="*", default=list(CONFIGS))
|
| 48 |
+
ap.add_argument("--json", default=None)
|
| 49 |
+
args = ap.parse_args()
|
| 50 |
+
import torch
|
| 51 |
+
import ttnn
|
| 52 |
+
|
| 53 |
+
torch.set_num_threads(4)
|
| 54 |
+
from tt_diffusion_planner.device import close_device, open_device
|
| 55 |
+
from tt_diffusion_planner.reference.weights import find_weights_dir, load_weights
|
| 56 |
+
from tt_diffusion_planner.tt.encoder import ENTITIES, MixerTrunk
|
| 57 |
+
from tt_diffusion_planner.tt.layers import Build, layer_norm_fp32, policy
|
| 58 |
+
from tt_diffusion_planner.ttaw.precision import compute_kernel_config
|
| 59 |
+
from tt_diffusion_planner.ttaw.tensors import to_device, to_numpy
|
| 60 |
+
|
| 61 |
+
p = load_weights(find_weights_dir()).params
|
| 62 |
+
gold = {}
|
| 63 |
+
for s in args.scenes:
|
| 64 |
+
with np.load(GOLDENS / f"{s}.npz", allow_pickle=False) as z:
|
| 65 |
+
gold[s] = {k: z[k] for k in z.files if k.startswith(("enc.", "host.valid"))}
|
| 66 |
+
dev = open_device(allow_fallback=False)
|
| 67 |
+
report = {"ln": {}, "mixer": {}}
|
| 68 |
+
try:
|
| 69 |
+
# ---- A: LayerNorm accuracy on real rows
|
| 70 |
+
g = p["encoder.neighbor_encoder.blocks.0.norm1.gamma"]
|
| 71 |
+
b = p["encoder.neighbor_encoder.blocks.0.norm1.beta"]
|
| 72 |
+
gd, bd = (to_device(v.reshape(1, 1, 1, -1), dev, "float32") for v in (g, b))
|
| 73 |
+
for s in args.scenes:
|
| 74 |
+
x = gold[s]["enc.neighbor.pre"].reshape(1, 1, -1, C.MIXER_CHANNELS).astype(np.float32)
|
| 75 |
+
xx = x.astype(np.float64)
|
| 76 |
+
ref = (xx - xx.mean(-1, keepdims=True)) / np.sqrt(xx.var(-1, keepdims=True) + 1e-5) * g + b
|
| 77 |
+
tx = to_device(x, dev, "float32")
|
| 78 |
+
fused = to_numpy(ttnn.layer_norm(tx, epsilon=1e-5, weight=gd, bias=bd,
|
| 79 |
+
compute_kernel_config=compute_kernel_config("HiFi4", fp32_acc=True)))
|
| 80 |
+
dec = to_numpy(layer_norm_fp32(tx, gd, bd))
|
| 81 |
+
report["ln"][s] = {"fused_fp32": st(fused, ref), "decomposed_fp32": st(dec, ref),
|
| 82 |
+
"row_offset_ratio": float(np.abs(xx.mean(-1)).mean() / xx.std(-1).mean())}
|
| 83 |
+
print("LN", s, json.dumps(report["ln"][s]), flush=True)
|
| 84 |
+
# ---- B: mixer blocks alone, teacher-forced
|
| 85 |
+
for name in args.configs:
|
| 86 |
+
cfg = dict(CONFIGS[name])
|
| 87 |
+
build = Build(dev, policy(spec=cfg.pop("precision", None)), **cfg)
|
| 88 |
+
for cat in args.cats:
|
| 89 |
+
trunk = MixerTrunk(build, p, cat)
|
| 90 |
+
for s in args.scenes:
|
| 91 |
+
rows = gold[s][f"enc.{cat}.pre.rows"]
|
| 92 |
+
if rows.size == 0:
|
| 93 |
+
continue
|
| 94 |
+
x0 = np.zeros((1, ENTITIES[cat], C.MIXER_TOKENS, C.MIXER_CHANNELS), np.float32)
|
| 95 |
+
x0[0, rows] = gold[s][f"enc.{cat}.pre"]
|
| 96 |
+
out = to_numpy(trunk.mix(to_device(x0, dev, trunk.stream)))[0, rows]
|
| 97 |
+
r = st(out, gold[s][f"enc.{cat}.mixer"])
|
| 98 |
+
report["mixer"].setdefault(name, {}).setdefault(cat, {})[s] = r
|
| 99 |
+
print(f"MIX {name:15s} {cat:9s} {s:18s} {json.dumps(r)}", flush=True)
|
| 100 |
+
del trunk
|
| 101 |
+
# ---- C: the island on the device (TF32-like vs split matmuls) -> pre error, and the float64 CPU mixer on
|
| 102 |
+
# the device pre (separates the input error from the device mixer)
|
| 103 |
+
from tt_diffusion_planner.host import pipeline as hp
|
| 104 |
+
from tt_diffusion_planner.reference.model import Encoder, torch_params
|
| 105 |
+
from tt_diffusion_planner.tt import inputs as I
|
| 106 |
+
|
| 107 |
+
ref_enc = Encoder(torch_params(p, torch.float64))
|
| 108 |
+
weights = load_weights(find_weights_dir())
|
| 109 |
+
|
| 110 |
+
def cpu_mix(x, cat):
|
| 111 |
+
N = f"encoder.{cat}_encoder"
|
| 112 |
+
x = torch.from_numpy(np.asarray(x, np.float64))
|
| 113 |
+
for i in range(C.MIXER_DEPTH):
|
| 114 |
+
B = f"{N}.blocks.{i}"
|
| 115 |
+
x = x + ref_enc.mlp(ref_enc.ln(x, f"{B}.norm1").transpose(1, 2), f"{B}.tokens_mlp").transpose(1, 2)
|
| 116 |
+
x = x + ref_enc.mlp(ref_enc.ln(x, f"{B}.norm2"), f"{B}.channels_mlp")
|
| 117 |
+
return x.numpy()
|
| 118 |
+
|
| 119 |
+
preps = {}
|
| 120 |
+
for s in args.scenes:
|
| 121 |
+
with np.load(GOLDENS / f"{s}.npz", allow_pickle=False) as z:
|
| 122 |
+
raw = {k: z[f"in.{k}"] for k in C.INPUT_NAMES}
|
| 123 |
+
preps[s] = I.plan_inputs(hp.prepare(raw, weights.normalization.observation))
|
| 124 |
+
for mode, split in (("tf32", ()), ("split", ("enc.island.*",))):
|
| 125 |
+
build = Build(dev, policy(), ln_fp32=("enc.mixer.*",), split=split)
|
| 126 |
+
for cat in ("ego", "neighbor"):
|
| 127 |
+
trunk = MixerTrunk(build, p, cat)
|
| 128 |
+
for s in args.scenes:
|
| 129 |
+
rows = gold[s][f"enc.{cat}.pre.rows"]
|
| 130 |
+
pre_t = trunk.pre(to_device(preps[s][f"{cat}_x"], dev, "float32"))
|
| 131 |
+
pre = to_numpy(pre_t)[0, rows]
|
| 132 |
+
mix = to_numpy(trunk.mix(pre_t))[0, rows]
|
| 133 |
+
r = {"pre": st(pre, gold[s][f"enc.{cat}.pre"]),
|
| 134 |
+
"mix_device": st(mix, gold[s][f"enc.{cat}.mixer"]),
|
| 135 |
+
"mix_cpu64_on_device_pre": st(cpu_mix(pre, cat), gold[s][f"enc.{cat}.mixer"])}
|
| 136 |
+
report.setdefault("island", {}).setdefault(mode, {}).setdefault(cat, {})[s] = r
|
| 137 |
+
print(f"ISL {mode:6s} {cat:9s} {s:18s} {json.dumps(r)}", flush=True)
|
| 138 |
+
del trunk
|
| 139 |
+
finally:
|
| 140 |
+
close_device(dev)
|
| 141 |
+
if args.json:
|
| 142 |
+
Path(args.json).write_text(json.dumps(report, indent=1) + "\n")
|
| 143 |
+
return 0
|
| 144 |
+
|
| 145 |
+
|
| 146 |
+
if __name__ == "__main__":
|
| 147 |
+
sys.exit(main())
|
code/scripts/profile_ops.py
ADDED
|
@@ -0,0 +1,250 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#!/usr/bin/env python3
|
| 2 |
+
# SPDX-License-Identifier: Apache-2.0
|
| 3 |
+
"""Device profile of diffusion-planner-p150: one eager plan and one traced replay between Tracy signposts.
|
| 4 |
+
|
| 5 |
+
ROOT=/home/ubuntu/experiments/tt-models
|
| 6 |
+
$ROOT/bin/devrun -t 3600 -- python -m tracy -r -p -v --op-support-count 16000 --no-web-server \\
|
| 7 |
+
-o $ROOT/generated/profiler/diffusion-planner_baseline code/scripts/profile_ops.py
|
| 8 |
+
tt-perf-report <ops_perf_results_*.csv> --start-signpost trace --end-signpost trace_end
|
| 9 |
+
python -m tt_diffusion_planner.ttaw.profiling <the same csv or its directory> --start trace --end trace_end
|
| 10 |
+
|
| 11 |
+
Sections (signposts): ``eager`` .. ``eager_end``, one eager run of the ``plan`` variant after the warm-up
|
| 12 |
+
(``TraceRunner.run_eager``: the same graph and programs as the trace, nothing compiled), and ``trace`` ..
|
| 13 |
+
``trace_end``, ``--replays`` replays of the captured ``plan`` trace (the served device path) on the uploaded sample.
|
| 14 |
+
The device profiler buffer is flushed (``ttnn.ReadDeviceProfiler``) before and after each section.
|
| 15 |
+
|
| 16 |
+
Inside the eager section two levels of signposts attribute every op (and, through the identical op order, every
|
| 17 |
+
replayed op):
|
| 18 |
+
|
| 19 |
+
- ``m:<stage>``: the stage of the plan: ``enc.<category>.pre`` / ``.mix`` / ``.head`` (the six mixer trunks, then
|
| 20 |
+
pool + entity head), ``enc.static``, ``enc.<goal|ego_shape|turn>``, ``enc.tokens`` (concat, validity, position
|
| 21 |
+
embedding), ``mask`` (key-bias expansion), ``enc.fusion<i>.attn`` / ``.mlp``, ``enc.final_ln``, ``dec.cross_kv``
|
| 22 |
+
(hoisted cross K / V), ``dec.e<k>.preproj``, ``dec.e<k>.b<i>.attn`` / ``.mlp1`` / ``.cross`` / ``.mlp2``,
|
| 23 |
+
``dec.e<k>.final`` (evaluation k = 0..10, DiT block i), ``dec.e<k>.solver`` (the DPM-Solver++(2M) update, the
|
| 24 |
+
prefix constraint and the ego-row slice of the next iterate), ``turn``, ``pack``;
|
| 25 |
+
- ``c:<kind>``: the layer kind: ``split`` (split hi / lo matmul, ``SplitLinear``), ``linear`` (one ``ttnn.linear``),
|
| 26 |
+
``ln32`` (fp32 LayerNorm decomposition), ``ln`` (``ttnn.layer_norm``), ``attn_mm`` (fp32 matmul attention),
|
| 27 |
+
``sdpa``, ``heads`` (head split / merge), ``mask``, and ``glue`` for every op outside a layer (residual adds,
|
| 28 |
+
adaLN gates, mixer transposes, solver updates, slices, concats).
|
| 29 |
+
|
| 30 |
+
``--no-layer-signposts`` turns both levels off. One plan issues several thousand programs, more than the profiler's
|
| 31 |
+
default 1000-program buffer, so run it with ``--op-support-count`` above the program count of the largest section
|
| 32 |
+
(the warm-up before the first flush included). The output folder must be absolute and outside the bundle
|
| 33 |
+
(PLAN.md 5.1).
|
| 34 |
+
"""
|
| 35 |
+
from __future__ import annotations
|
| 36 |
+
|
| 37 |
+
import argparse
|
| 38 |
+
import functools
|
| 39 |
+
import json
|
| 40 |
+
import time
|
| 41 |
+
from pathlib import Path
|
| 42 |
+
from typing import Any, Callable, Dict, List
|
| 43 |
+
|
| 44 |
+
import numpy as np
|
| 45 |
+
|
| 46 |
+
from tt_diffusion_planner import DiffusionPlanner
|
| 47 |
+
from tt_diffusion_planner.ttaw.profiling import read_device_profiler, signpost, signposted
|
| 48 |
+
|
| 49 |
+
SAMPLE = Path(__file__).resolve().parents[1] / "tt_diffusion_planner" / "samples" / "kashiwanoha_dense.npz"
|
| 50 |
+
|
| 51 |
+
|
| 52 |
+
class _Proxy:
|
| 53 |
+
"""A callable stand-in for a layer object that emits a signpost (``label()``), then calls the layer."""
|
| 54 |
+
|
| 55 |
+
def __init__(self, inner: Any, label: Callable[[], str]):
|
| 56 |
+
self.inner, self._label = inner, label
|
| 57 |
+
|
| 58 |
+
def __call__(self, *args, **kwargs):
|
| 59 |
+
signpost(self._label())
|
| 60 |
+
return self.inner(*args, **kwargs)
|
| 61 |
+
|
| 62 |
+
def __getattr__(self, name: str):
|
| 63 |
+
return getattr(self.inner, name)
|
| 64 |
+
|
| 65 |
+
|
| 66 |
+
def install_signposts(tt) -> Callable[[], None]:
|
| 67 |
+
"""Signposts on a ``tt.model.TtDiffusionPlanner`` (instance attributes, layer classes and the attention
|
| 68 |
+
helpers); returns the function that removes them again."""
|
| 69 |
+
from tt_diffusion_planner.tt import layers as L
|
| 70 |
+
from tt_diffusion_planner.tt import model as M
|
| 71 |
+
from tt_diffusion_planner.ttaw.ops import attention as A
|
| 72 |
+
|
| 73 |
+
undo: List[Callable[[], None]] = []
|
| 74 |
+
state = {"k": -1, "ln": {}}
|
| 75 |
+
stack: List[str] = []
|
| 76 |
+
|
| 77 |
+
def set_attr(obj, name, value):
|
| 78 |
+
had = name in vars(obj)
|
| 79 |
+
old = vars(obj).get(name)
|
| 80 |
+
setattr(obj, name, value)
|
| 81 |
+
undo.append(lambda: setattr(obj, name, old) if had else delattr(obj, name))
|
| 82 |
+
|
| 83 |
+
def set_item(d, key, value):
|
| 84 |
+
old = d[key]
|
| 85 |
+
d[key] = value
|
| 86 |
+
undo.append(lambda: d.__setitem__(key, old))
|
| 87 |
+
|
| 88 |
+
# ---- layer kinds (class level, with a stack so nested layers restore the outer kind) ---------------------
|
| 89 |
+
def kind_wrap(owner, name, kind_of):
|
| 90 |
+
fn = getattr(owner, name)
|
| 91 |
+
|
| 92 |
+
@functools.wraps(fn)
|
| 93 |
+
def inner(*args, **kwargs):
|
| 94 |
+
kind = kind_of(*args)
|
| 95 |
+
stack.append(kind)
|
| 96 |
+
signpost(f"c:{kind}")
|
| 97 |
+
try:
|
| 98 |
+
return fn(*args, **kwargs)
|
| 99 |
+
finally:
|
| 100 |
+
stack.pop()
|
| 101 |
+
signpost(f"c:{stack[-1] if stack else 'glue'}")
|
| 102 |
+
|
| 103 |
+
setattr(owner, name, inner)
|
| 104 |
+
undo.append(lambda: setattr(owner, name, fn))
|
| 105 |
+
|
| 106 |
+
kind_wrap(L.Linear, "__call__", lambda *a: "linear")
|
| 107 |
+
kind_wrap(L.SplitLinear, "__call__", lambda *a: "split")
|
| 108 |
+
kind_wrap(L.LayerNorm, "__call__", lambda self, *a: "ln32" if self.mode == "fp32" else "ln")
|
| 109 |
+
kind_wrap(A, "attention_matmul", lambda *a: "attn_mm")
|
| 110 |
+
kind_wrap(A, "sdpa", lambda *a: "sdpa")
|
| 111 |
+
for name in ("split_qkv", "split_q_kv", "split_heads", "merge_heads"):
|
| 112 |
+
kind_wrap(A, name, lambda *a: "heads")
|
| 113 |
+
kind_wrap(A, "expand_key_bias", lambda *a: "mask")
|
| 114 |
+
|
| 115 |
+
# ---- stages ------------------------------------------------------------------------------------------------
|
| 116 |
+
def method_wrap(obj, name, before=None, after=None):
|
| 117 |
+
fn = getattr(obj, name)
|
| 118 |
+
|
| 119 |
+
@functools.wraps(fn)
|
| 120 |
+
def inner(*args, **kwargs):
|
| 121 |
+
if before:
|
| 122 |
+
signpost(before())
|
| 123 |
+
out = fn(*args, **kwargs)
|
| 124 |
+
if after:
|
| 125 |
+
signpost(after())
|
| 126 |
+
return out
|
| 127 |
+
|
| 128 |
+
set_attr(obj, name, inner)
|
| 129 |
+
|
| 130 |
+
enc, dec = tt.encoder, tt.decoder
|
| 131 |
+
for cat, trunk in enc.trunks.items():
|
| 132 |
+
method_wrap(trunk, "pre", before=lambda c=cat: f"m:enc.{c}.pre")
|
| 133 |
+
method_wrap(trunk, "mix", before=lambda c=cat: f"m:enc.{c}.mix")
|
| 134 |
+
method_wrap(trunk, "pool", before=lambda c=cat: f"m:enc.{c}.head")
|
| 135 |
+
set_attr(enc, "static1", _Proxy(enc.static1, lambda: "m:enc.static"))
|
| 136 |
+
for cat in list(enc.small):
|
| 137 |
+
set_item(enc.small, cat, _Proxy(enc.small[cat], lambda c=cat: f"m:enc.{c}"))
|
| 138 |
+
set_attr(enc, "pad_tokens", _Proxy(enc.pad_tokens, lambda: "m:enc.tokens"))
|
| 139 |
+
for i, blk in enumerate(enc.blocks):
|
| 140 |
+
set_item(blk, "kv", _Proxy(blk["kv"], lambda i=i: f"m:enc.fusion{i}.attn"))
|
| 141 |
+
set_item(blk, "n2", _Proxy(blk["n2"], lambda i=i: f"m:enc.fusion{i}.mlp"))
|
| 142 |
+
set_attr(enc, "final_norm", _Proxy(enc.final_norm, lambda: "m:enc.final_ln"))
|
| 143 |
+
method_wrap(dec, "cross_kv", before=lambda: "m:dec.cross_kv")
|
| 144 |
+
method_wrap(dec, "solve", before=lambda: "m:dec.solve")
|
| 145 |
+
|
| 146 |
+
def next_eval():
|
| 147 |
+
state["k"] += 1
|
| 148 |
+
state["ln"] = {}
|
| 149 |
+
return f"m:dec.e{state['k']}.preproj"
|
| 150 |
+
|
| 151 |
+
method_wrap(dec, "evaluate", before=next_eval, after=lambda: f"m:dec.e{state['k']}.solver")
|
| 152 |
+
for i, blk in enumerate(dec.blocks):
|
| 153 |
+
def ln_label(i=i):
|
| 154 |
+
n = state["ln"][i] = state["ln"].get(i, 0) + 1
|
| 155 |
+
return f"m:dec.e{state['k']}.b{i}.{'attn' if n % 2 else 'mlp1'}"
|
| 156 |
+
|
| 157 |
+
set_item(blk, "ln", _Proxy(blk["ln"], ln_label))
|
| 158 |
+
set_item(blk, "n3", _Proxy(blk["n3"], lambda i=i: f"m:dec.e{state['k']}.b{i}.cross"))
|
| 159 |
+
set_item(blk, "n4", _Proxy(blk["n4"], lambda i=i: f"m:dec.e{state['k']}.b{i}.mlp2"))
|
| 160 |
+
set_attr(dec, "fin_ln", _Proxy(dec.fin_ln, lambda: f"m:dec.e{state['k']}.final"))
|
| 161 |
+
set_attr(tt, "turn", _Proxy(tt.turn, lambda: "m:turn"))
|
| 162 |
+
pack = M.pack_outputs
|
| 163 |
+
|
| 164 |
+
def pack_wrap(*args, **kwargs):
|
| 165 |
+
signpost("m:pack")
|
| 166 |
+
return pack(*args, **kwargs)
|
| 167 |
+
|
| 168 |
+
M.pack_outputs = pack_wrap
|
| 169 |
+
undo.append(lambda: setattr(M, "pack_outputs", pack))
|
| 170 |
+
# the mask expansion is attributed to its own stage (encoder fusion mask, decoder self-attention mask)
|
| 171 |
+
expand = A.expand_key_bias
|
| 172 |
+
|
| 173 |
+
def expand_wrap(*args, **kwargs):
|
| 174 |
+
signpost("m:mask")
|
| 175 |
+
return expand(*args, **kwargs)
|
| 176 |
+
|
| 177 |
+
A.expand_key_bias = expand_wrap
|
| 178 |
+
undo.append(lambda: setattr(A, "expand_key_bias", expand))
|
| 179 |
+
|
| 180 |
+
def remove() -> None:
|
| 181 |
+
for fn in reversed(undo):
|
| 182 |
+
fn()
|
| 183 |
+
|
| 184 |
+
return remove
|
| 185 |
+
|
| 186 |
+
|
| 187 |
+
def main() -> None:
|
| 188 |
+
ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
|
| 189 |
+
ap.add_argument("--input", default=str(SAMPLE))
|
| 190 |
+
ap.add_argument("--dispatch", default=None, choices=["eth", "worker"])
|
| 191 |
+
ap.add_argument("--num-cqs", type=int, default=None, choices=[1, 2])
|
| 192 |
+
ap.add_argument("--replays", type=int, default=1, help="replays inside the trace section")
|
| 193 |
+
ap.add_argument("--no-eager", dest="eager", action="store_false")
|
| 194 |
+
ap.add_argument("--no-layer-signposts", dest="layers", action="store_false")
|
| 195 |
+
ap.add_argument("--json", default=None, help="write the run description (device, trace, timings) here")
|
| 196 |
+
a = ap.parse_args()
|
| 197 |
+
import ttnn
|
| 198 |
+
|
| 199 |
+
from tt_diffusion_planner.host import pipeline as hp
|
| 200 |
+
from tt_diffusion_planner.reference import config as C
|
| 201 |
+
from tt_diffusion_planner.reference.weights import find_weights_dir
|
| 202 |
+
from tt_diffusion_planner.tt import inputs as I
|
| 203 |
+
from tt_diffusion_planner.ttaw.io import load_named_arrays
|
| 204 |
+
|
| 205 |
+
wd = find_weights_dir()
|
| 206 |
+
t0 = time.perf_counter()
|
| 207 |
+
with DiffusionPlanner.from_pretrained(dispatch=a.dispatch, num_command_queues=a.num_cqs,
|
| 208 |
+
weights_dir=str(wd) if wd else None) as model:
|
| 209 |
+
tt, runner, dev = model.tt, model.runner, model.device
|
| 210 |
+
print("loaded in %.1f s:" % (time.perf_counter() - t0), json.dumps(model.device_info), flush=True)
|
| 211 |
+
raw = load_named_arrays(a.input, C.INPUT_SCHEMA)
|
| 212 |
+
inputs = I.plan_inputs(hp.prepare(raw, model.normalization.observation))
|
| 213 |
+
served = model(inputs=raw) # one served plan: upload + replay + readback
|
| 214 |
+
ttnn.synchronize_device(dev)
|
| 215 |
+
read_device_profiler(dev) # warm-up / capture / first plan out of the buffer
|
| 216 |
+
timings: Dict[str, float] = {}
|
| 217 |
+
if a.eager:
|
| 218 |
+
remove = install_signposts(tt) if a.layers else (lambda: None)
|
| 219 |
+
try:
|
| 220 |
+
t1 = time.perf_counter()
|
| 221 |
+
with signposted("eager"):
|
| 222 |
+
eager = runner.run_eager("plan", inputs=inputs)
|
| 223 |
+
ttnn.synchronize_device(dev)
|
| 224 |
+
timings["eager_ms"] = (time.perf_counter() - t1) * 1e3
|
| 225 |
+
finally:
|
| 226 |
+
remove()
|
| 227 |
+
read_device_profiler(dev)
|
| 228 |
+
final = np.asarray(eager["final_x0"], np.float32)
|
| 229 |
+
print("eager final_x0 finite:", bool(np.isfinite(final).all()), flush=True)
|
| 230 |
+
runner.upload(inputs)
|
| 231 |
+
ttnn.synchronize_device(dev)
|
| 232 |
+
read_device_profiler(dev)
|
| 233 |
+
t1 = time.perf_counter()
|
| 234 |
+
with signposted("trace"):
|
| 235 |
+
runner.replay("plan", n=a.replays)
|
| 236 |
+
ttnn.synchronize_device(dev)
|
| 237 |
+
timings["trace_ms"] = (time.perf_counter() - t1) * 1e3 / a.replays
|
| 238 |
+
read_device_profiler(dev)
|
| 239 |
+
out = runner.read("plan")
|
| 240 |
+
same = bool(np.array_equal(out["final_x0"], eager["final_x0"])) if a.eager else None
|
| 241 |
+
desc = {"device": model.device_info, "timings_ms": timings, "replays": a.replays,
|
| 242 |
+
"replay_equals_eager": same, "turn_command": int(served.turn_indicator["command"]),
|
| 243 |
+
"options": tt.build.options(), "trace": runner.describe()}
|
| 244 |
+
print("profile run:", json.dumps(desc, default=str), flush=True)
|
| 245 |
+
if a.json:
|
| 246 |
+
Path(a.json).write_text(json.dumps(desc, indent=1, default=str) + "\n")
|
| 247 |
+
|
| 248 |
+
|
| 249 |
+
if __name__ == "__main__":
|
| 250 |
+
main()
|
code/scripts/ref_golden.py
ADDED
|
@@ -0,0 +1,169 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#!/usr/bin/env python3
|
| 2 |
+
# SPDX-License-Identifier: Apache-2.0
|
| 3 |
+
"""Goldens of the fp32 CPU reference (no device): per-module taps and final outputs per scene.
|
| 4 |
+
|
| 5 |
+
# research venv (torch + onnx; onnxruntime only for --ort), from the bundle root:
|
| 6 |
+
PYTHONPATH=code python code/scripts/ref_golden.py --threads 4 [--ort] \
|
| 7 |
+
[--full-dir ../../research/diffusion-planner/goldens] [--scenes kashiwanoha_dense straight_road ...]
|
| 8 |
+
|
| 9 |
+
Scenes: the shipped samples (``code/tt_diffusion_planner/samples/*.npz``: small goldens -> ``tests/goldens``, and the
|
| 10 |
+
stored ``/predict`` body of the reference -> ``samples/<stem>.reference.json``) and, when the workspace research
|
| 11 |
+
directory is present, the seven ORT golden scenes of ``research/diffusion-planner/ort/golden_<scene>.npz`` and any
|
| 12 |
+
public-dataset scene in ``research/diffusion-planner/public_data/*.npz`` (full goldens only). ``--ort`` also runs ONNX
|
| 13 |
+
Runtime on every scene and writes ``<full-dir>/<scene>.ort_agreement.json`` (reference vs ORT on the deployed ONNX).
|
| 14 |
+
"""
|
| 15 |
+
from __future__ import annotations
|
| 16 |
+
|
| 17 |
+
import argparse
|
| 18 |
+
import hashlib
|
| 19 |
+
import json
|
| 20 |
+
import sys
|
| 21 |
+
import time
|
| 22 |
+
from pathlib import Path
|
| 23 |
+
|
| 24 |
+
import numpy as np
|
| 25 |
+
|
| 26 |
+
CODE = Path(__file__).resolve().parents[1]
|
| 27 |
+
if str(CODE) not in sys.path:
|
| 28 |
+
sys.path.insert(0, str(CODE))
|
| 29 |
+
|
| 30 |
+
from tt_diffusion_planner.reference import config as C # noqa: E402
|
| 31 |
+
from tt_diffusion_planner.reference.goldens import write_scene # noqa: E402
|
| 32 |
+
from tt_diffusion_planner.reference.pipeline import ReferencePlanner # noqa: E402
|
| 33 |
+
from tt_diffusion_planner.ttaw.metrics import pcc # noqa: E402
|
| 34 |
+
|
| 35 |
+
PKG = CODE / "tt_diffusion_planner"
|
| 36 |
+
SAMPLES = PKG / "samples"
|
| 37 |
+
SMALL_DIR = PKG / "tests" / "goldens"
|
| 38 |
+
RESEARCH = CODE.parents[2] / "research" / "diffusion-planner"
|
| 39 |
+
|
| 40 |
+
|
| 41 |
+
def public_ids(dataset: str, which: str):
|
| 42 |
+
"""Instant ids of ``public_data/inputs/<dataset>``: ``all``, or ``core`` (the ids the public-data goldens keep
|
| 43 |
+
solver iterates for: ``public_data/goldens/<dataset>/core_ids.json``)."""
|
| 44 |
+
root = RESEARCH / "public_data"
|
| 45 |
+
ids = sorted(p.stem for p in (root / "inputs" / dataset).glob("*.npz"))
|
| 46 |
+
if which == "core":
|
| 47 |
+
core = root / "goldens" / dataset / "core_ids.json"
|
| 48 |
+
keep = set(json.loads(core.read_text()).get("core_ids_with_denoising_steps", [])) if core.is_file() else set()
|
| 49 |
+
ids = [i for i in ids if i in keep]
|
| 50 |
+
return ids
|
| 51 |
+
|
| 52 |
+
|
| 53 |
+
def scene_sources(selected, public=None, public_which="core"):
|
| 54 |
+
"""``{scene: (raw source, kind)}``: shipped samples first, then research ORT scenes, then (``--public``) the
|
| 55 |
+
public-data instants of one dataset (nuScenes-derived: CC BY-NC-SA, local goldens only)."""
|
| 56 |
+
out = {}
|
| 57 |
+
if public is None:
|
| 58 |
+
for p in sorted(SAMPLES.glob("*.npz")):
|
| 59 |
+
out[p.stem] = (p, "sample")
|
| 60 |
+
for p in sorted((RESEARCH / "ort").glob("golden_*.npz")):
|
| 61 |
+
out.setdefault(p.stem[len("golden_"):], (p, "research"))
|
| 62 |
+
else:
|
| 63 |
+
for i in public_ids(public, public_which):
|
| 64 |
+
out[f"public/{public}/{i}"] = (RESEARCH / "public_data" / "inputs" / public / f"{i}.npz", "public")
|
| 65 |
+
if selected:
|
| 66 |
+
missing = sorted(set(selected) - set(out))
|
| 67 |
+
if missing:
|
| 68 |
+
raise SystemExit(f"unknown scenes {missing}; have {sorted(out)}")
|
| 69 |
+
out = {k: v for k, v in out.items() if k in selected}
|
| 70 |
+
return out
|
| 71 |
+
|
| 72 |
+
|
| 73 |
+
def load_raw(path: Path, kind: str):
|
| 74 |
+
with np.load(path, allow_pickle=False) as z:
|
| 75 |
+
if kind in ("research", "public"): # dp_reference.py / nuscenes_dp.py layout: raw/<name>
|
| 76 |
+
return {k: np.asarray(z["raw/" + k], np.float32) for k in C.INPUT_NAMES}
|
| 77 |
+
return {k: np.asarray(z[k], np.float32) for k in C.INPUT_NAMES}
|
| 78 |
+
|
| 79 |
+
|
| 80 |
+
def ort_agreement(weights_dir: Path, raw, g, threads: int):
|
| 81 |
+
from tt_diffusion_planner.reference.ort import OrtPlanner
|
| 82 |
+
|
| 83 |
+
o = OrtPlanner(weights_dir, threads=threads).run(raw)
|
| 84 |
+
rows = g["dec.rows"]
|
| 85 |
+
enc_rows = np.flatnonzero(g["host.token_valid"])
|
| 86 |
+
return {"encoding_pcc_valid": pcc(g["enc.encoding"][enc_rows], o.encoding[0][enc_rows]),
|
| 87 |
+
"encoding_max_abs": float(np.abs(g["enc.encoding"] - o.encoding[0]).max()),
|
| 88 |
+
"final_x0_pcc_valid": pcc(g["final_x0"][rows], o.final_x0[0][rows]),
|
| 89 |
+
"final_x0_max_abs_valid": float(np.abs(g["final_x0"][rows] - o.final_x0[0][rows]).max()),
|
| 90 |
+
"logit_max_abs": float(np.abs(g["turn.logit"] - o.logit[0]).max()),
|
| 91 |
+
"turn_logit_ref": g["turn.logit"].tolist(), "turn_logit_ort": o.logit[0].tolist()}
|
| 92 |
+
|
| 93 |
+
|
| 94 |
+
def stored_ort_agreement(g, golden: Path):
|
| 95 |
+
"""Reference vs a stored ORT golden in the ``dp_reference.py`` layout (``encoding``, ``final_x_normalized``,
|
| 96 |
+
``logit_multi``): the public-data goldens of ``research/diffusion-planner/public_data/goldens/<dataset>``."""
|
| 97 |
+
with np.load(golden, allow_pickle=False) as z:
|
| 98 |
+
enc, fx, lg = z["encoding"][0], z["final_x_normalized"][0], z["logit_multi"][0]
|
| 99 |
+
rows = g["dec.rows"]
|
| 100 |
+
tok = np.flatnonzero(g["host.token_valid"])
|
| 101 |
+
return {"ort_golden": str(golden),
|
| 102 |
+
"encoding_pcc_valid": pcc(g["enc.encoding"][tok], enc[tok]),
|
| 103 |
+
"encoding_max_abs": float(np.abs(g["enc.encoding"][tok] - enc[tok]).max()),
|
| 104 |
+
"final_x0_pcc_valid": pcc(g["final_x0"][rows], fx[rows]),
|
| 105 |
+
"final_x0_max_abs_valid": float(np.abs(g["final_x0"][rows] - fx[rows]).max()),
|
| 106 |
+
"logit_max_abs": float(np.abs(g["turn.logit"] - lg).max())}
|
| 107 |
+
|
| 108 |
+
|
| 109 |
+
def main(argv=None) -> int:
|
| 110 |
+
ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
| 111 |
+
ap.add_argument("--weights-dir", default=None)
|
| 112 |
+
ap.add_argument("--full-dir", type=Path, default=RESEARCH / "goldens",
|
| 113 |
+
help="where the large per-scene goldens go (never inside the bundle)")
|
| 114 |
+
ap.add_argument("--scenes", nargs="*", default=None)
|
| 115 |
+
ap.add_argument("--threads", type=int, default=4)
|
| 116 |
+
ap.add_argument("--ort", action="store_true", help="also compare with ONNX Runtime (research venv)")
|
| 117 |
+
ap.add_argument("--no-reference-json", action="store_true")
|
| 118 |
+
ap.add_argument("--public", default=None, help="a public_data/inputs/<dataset> (e.g. nuscenes): its instants "
|
| 119 |
+
"instead of the samples / research scenes, as lite goldens")
|
| 120 |
+
ap.add_argument("--public-ids", default="core", choices=["core", "all"])
|
| 121 |
+
a = ap.parse_args(argv)
|
| 122 |
+
if a.full_dir.resolve().is_relative_to(CODE.parent.resolve()):
|
| 123 |
+
raise SystemExit("--full-dir must be outside the bundle: the full goldens are 10-25 MB per scene")
|
| 124 |
+
ref = ReferencePlanner(a.weights_dir, threads=a.threads)
|
| 125 |
+
report, seen = {}, {}
|
| 126 |
+
for scene, (path, kind) in scene_sources(a.scenes, a.public, a.public_ids).items():
|
| 127 |
+
t0 = time.perf_counter()
|
| 128 |
+
raw = load_raw(path, kind)
|
| 129 |
+
digest = hashlib.sha256(b"".join(raw[k].tobytes() for k in C.INPUT_NAMES)).hexdigest()
|
| 130 |
+
if digest in seen: # e.g. research golden_straight == the shipped straight_road sample
|
| 131 |
+
report[scene] = {"same_inputs_as": seen[digest]}
|
| 132 |
+
print(scene, "skipped: same inputs as", seen[digest], flush=True)
|
| 133 |
+
continue
|
| 134 |
+
seen[digest] = scene
|
| 135 |
+
small = SMALL_DIR if kind == "sample" else None
|
| 136 |
+
r = write_scene(ref, raw, scene, a.full_dir, small, meta={"source": str(path), "kind": kind},
|
| 137 |
+
lite=(kind == "public"))
|
| 138 |
+
entry = {"paths": r["paths"], "valid_counts": r["info"]["valid_counts"]}
|
| 139 |
+
if kind == "sample" and not a.no_reference_json:
|
| 140 |
+
body = ref(inputs=raw).to_dict()
|
| 141 |
+
body["timing_ms"] = {}
|
| 142 |
+
ref_path = SAMPLES / f"{scene}.reference.json"
|
| 143 |
+
ref_path.write_text(json.dumps(body, indent=1) + "\n")
|
| 144 |
+
entry["reference_json"] = str(ref_path)
|
| 145 |
+
stored = None
|
| 146 |
+
if kind == "public":
|
| 147 |
+
stored = RESEARCH / "public_data" / "goldens" / a.public / f"golden_{path.stem}.npz"
|
| 148 |
+
if stored is not None and stored.is_file(): # the dataset work already ran ORT on this instant
|
| 149 |
+
agree = stored_ort_agreement(r["goldens"], stored)
|
| 150 |
+
(a.full_dir / f"{scene}.ort_agreement.json").write_text(json.dumps(agree, indent=1) + "\n")
|
| 151 |
+
entry["ort"] = {k: v for k, v in agree.items() if k != "ort_golden"}
|
| 152 |
+
elif a.ort:
|
| 153 |
+
agree = ort_agreement(ref.weights.path, raw, r["goldens"], a.threads)
|
| 154 |
+
a.full_dir.mkdir(parents=True, exist_ok=True)
|
| 155 |
+
(a.full_dir / f"{scene}.ort_agreement.json").write_text(json.dumps(agree, indent=1) + "\n")
|
| 156 |
+
entry["ort"] = {k: v for k, v in agree.items() if not k.startswith("turn_logit")}
|
| 157 |
+
entry["seconds"] = round(time.perf_counter() - t0, 1)
|
| 158 |
+
report[scene] = entry
|
| 159 |
+
print(scene, json.dumps(entry), flush=True)
|
| 160 |
+
a.full_dir.mkdir(parents=True, exist_ok=True)
|
| 161 |
+
index = a.full_dir / "index.json"
|
| 162 |
+
merged = json.loads(index.read_text()) if index.is_file() else {}
|
| 163 |
+
merged.update(report)
|
| 164 |
+
index.write_text(json.dumps(merged, indent=1, sort_keys=True) + "\n")
|
| 165 |
+
return 0
|
| 166 |
+
|
| 167 |
+
|
| 168 |
+
if __name__ == "__main__":
|
| 169 |
+
sys.exit(main())
|
code/scripts/split_error.py
ADDED
|
@@ -0,0 +1,124 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-License-Identifier: Apache-2.0
|
| 2 |
+
"""Where does the end-to-end error of a scene come from? (development tool; device, run under ``bin/devrun``)
|
| 3 |
+
|
| 4 |
+
bin/devrun -t 1200 -- python code/scripts/split_error.py --public scene-0103_kf14 --scenes straight_road
|
| 5 |
+
|
| 6 |
+
For each scene the ego / neighbour displacement vs the fp32 CPU reference of:
|
| 7 |
+
|
| 8 |
+
- ``device``: the full device plan;
|
| 9 |
+
- ``enc_only``: the device encoding (``encoder_taps``) + the CPU fp32 decoder and solver (encoder error alone);
|
| 10 |
+
- ``dec_only``: the CPU encoding + the device decoder (``decode_once`` replays) driven by the host solver (decoder
|
| 11 |
+
error alone, including its compounding over the 11 evaluations).
|
| 12 |
+
"""
|
| 13 |
+
from __future__ import annotations
|
| 14 |
+
|
| 15 |
+
import argparse
|
| 16 |
+
import json
|
| 17 |
+
import sys
|
| 18 |
+
from pathlib import Path
|
| 19 |
+
|
| 20 |
+
import numpy as np
|
| 21 |
+
|
| 22 |
+
HERE = Path(__file__).resolve()
|
| 23 |
+
sys.path.insert(0, str(HERE.parents[1]))
|
| 24 |
+
|
| 25 |
+
from tt_diffusion_planner.host import pipeline as hp # noqa: E402
|
| 26 |
+
from tt_diffusion_planner.host.solver import apply_prefix_constraint, dpm_solver_sample # noqa: E402
|
| 27 |
+
from tt_diffusion_planner.reference import config as C # noqa: E402
|
| 28 |
+
|
| 29 |
+
RES = HERE.parents[4] / "research" / "diffusion-planner"
|
| 30 |
+
|
| 31 |
+
|
| 32 |
+
def load_raw(scene: str, public: bool):
|
| 33 |
+
if public:
|
| 34 |
+
with np.load(RES / "public_data" / "inputs" / "nuscenes" / f"{scene}.npz", allow_pickle=False) as z:
|
| 35 |
+
return {k: np.array(z[f"raw/{k}"]) for k in C.INPUT_NAMES}
|
| 36 |
+
with np.load(RES / "goldens" / f"{scene}.npz", allow_pickle=False) as z:
|
| 37 |
+
return {k: np.array(z[f"in.{k}"]) for k in C.INPUT_NAMES}
|
| 38 |
+
|
| 39 |
+
|
| 40 |
+
def main() -> int:
|
| 41 |
+
ap = argparse.ArgumentParser()
|
| 42 |
+
ap.add_argument("--public", nargs="*", default=[])
|
| 43 |
+
ap.add_argument("--scenes", nargs="*", default=[])
|
| 44 |
+
ap.add_argument("--json", default=None)
|
| 45 |
+
for opt in ("ln-fp32", "hidden-fp32", "split", "attn-fp32-acc", "attn-matmul"):
|
| 46 |
+
ap.add_argument(f"--{opt}", default=None, help="module globs (comma-separated); default: the knob")
|
| 47 |
+
args = ap.parse_args()
|
| 48 |
+
from tt_diffusion_planner.tt.config import globs
|
| 49 |
+
|
| 50 |
+
opts = {k: globs(getattr(args, k)) for k in ("ln_fp32", "hidden_fp32", "split", "attn_fp32_acc", "attn_matmul")
|
| 51 |
+
if getattr(args, k) is not None}
|
| 52 |
+
import torch
|
| 53 |
+
|
| 54 |
+
torch.set_num_threads(4)
|
| 55 |
+
from tt_diffusion_planner.device import close_device, open_device
|
| 56 |
+
from tt_diffusion_planner.reference.model import Decoder, Encoder, torch_params
|
| 57 |
+
from tt_diffusion_planner.reference.weights import find_weights_dir, load_weights
|
| 58 |
+
from tt_diffusion_planner.tt.model import TtDiffusionPlanner
|
| 59 |
+
|
| 60 |
+
w = load_weights(find_weights_dir())
|
| 61 |
+
params = {k: v[3] for k, v in hp.RUNTIME_PARAMS.items()}
|
| 62 |
+
P = torch_params(w.params)
|
| 63 |
+
enc_ref, dec_ref = Encoder(P), Decoder(P)
|
| 64 |
+
dev = open_device(allow_fallback=False)
|
| 65 |
+
report = {}
|
| 66 |
+
try:
|
| 67 |
+
tt = TtDiffusionPlanner(dev, w, debug=True, **opts)
|
| 68 |
+
report["options"] = tt.build.options()
|
| 69 |
+
tt.capture()
|
| 70 |
+
for scene, public in [(s, True) for s in args.public] + [(s, False) for s in args.scenes]:
|
| 71 |
+
prep = hp.prepare(load_raw(scene, public), w.normalization.observation)
|
| 72 |
+
cs = prep.decoder.current_states
|
| 73 |
+
|
| 74 |
+
def solve(model_fn):
|
| 75 |
+
res = dpm_solver_sample(prep.x_T, model_fn, lambda x: apply_prefix_constraint(x, cs))
|
| 76 |
+
return res.final_x
|
| 77 |
+
|
| 78 |
+
def cpu_plan(encoding):
|
| 79 |
+
kv = dec_ref.cross_kv(torch.from_numpy(np.asarray(encoding, np.float32)))
|
| 80 |
+
with torch.no_grad():
|
| 81 |
+
return solve(lambda x, t: dec_ref.forward(x, t, kv, prep.decoder.agent_valid).numpy())
|
| 82 |
+
|
| 83 |
+
def out(final_x0):
|
| 84 |
+
o = hp.make_output(final_x0, np.zeros(5, np.float32), prep, w.normalization, params)
|
| 85 |
+
return o.poses[:, :2], o.predicted_agents[..., :2]
|
| 86 |
+
|
| 87 |
+
with torch.no_grad():
|
| 88 |
+
enc = enc_ref.forward(prep.features).numpy()
|
| 89 |
+
ref_ego, ref_nb = out(cpu_plan(enc))
|
| 90 |
+
dev_enc = tt.encoder_taps(prep)["enc.encoding"]
|
| 91 |
+
cases = {
|
| 92 |
+
"device": tt.forward(prep)["final_x0"],
|
| 93 |
+
"enc_only": cpu_plan(dev_enc),
|
| 94 |
+
"dec_only": solve(lambda x, t: _decode(tt, prep, x, t, enc)),
|
| 95 |
+
}
|
| 96 |
+
res = {}
|
| 97 |
+
for name, fx in cases.items():
|
| 98 |
+
ego, nb = out(fx)
|
| 99 |
+
d = np.hypot(*(ego - ref_ego).T)
|
| 100 |
+
r = {"ego_max_m": round(float(d.max()), 4), "ego_mean_m": round(float(d.mean()), 4)}
|
| 101 |
+
if nb.shape[0]:
|
| 102 |
+
per_agent = np.hypot(*(nb - ref_nb).transpose(2, 0, 1)).max(1)
|
| 103 |
+
r["nb_median_max_m"] = round(float(np.median(per_agent)), 4)
|
| 104 |
+
res[name] = r
|
| 105 |
+
res["encoding_pcc"] = float(np.corrcoef(dev_enc[prep.features.token_valid].ravel(),
|
| 106 |
+
enc[prep.features.token_valid].ravel())[0, 1])
|
| 107 |
+
report[scene] = res
|
| 108 |
+
print(scene, json.dumps(res), flush=True)
|
| 109 |
+
tt.release()
|
| 110 |
+
finally:
|
| 111 |
+
close_device(dev)
|
| 112 |
+
if args.json:
|
| 113 |
+
Path(args.json).write_text(json.dumps(report, indent=1) + "\n")
|
| 114 |
+
return 0
|
| 115 |
+
|
| 116 |
+
|
| 117 |
+
def _decode(tt, prep, x, t, enc):
|
| 118 |
+
out = tt.decode_once(prep, x, float(t), encoding=enc)
|
| 119 |
+
out[:, 0] = 0.0 # the t = 0 slot (masked on the device) is overwritten by the prefix constraint anyway
|
| 120 |
+
return out
|
| 121 |
+
|
| 122 |
+
|
| 123 |
+
if __name__ == "__main__":
|
| 124 |
+
sys.exit(main())
|
code/tt_diffusion_planner/__init__.py
ADDED
|
@@ -0,0 +1,41 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-License-Identifier: Apache-2.0
|
| 2 |
+
"""Diffusion Planner v5.0 (Autoware diffusion_planner) on a Tenstorrent Blackhole p150 (tt-nn), packaged as
|
| 3 |
+
changh95/diffusion-planner-p150.
|
| 4 |
+
|
| 5 |
+
Python API (see PYTHON.md)::
|
| 6 |
+
|
| 7 |
+
from tt_diffusion_planner import DiffusionPlanner
|
| 8 |
+
|
| 9 |
+
with DiffusionPlanner.from_pretrained(device_id=0) as model:
|
| 10 |
+
out = model(inputs="code/tt_diffusion_planner/samples/kashiwanoha_dense.npz")
|
| 11 |
+
|
| 12 |
+
Importing the package has no side effects (no device, no ttnn / torch import, no network); the names below load
|
| 13 |
+
on first use. The HTTP server is ``tt_diffusion_planner.server.app:app`` (SERVING.md). ``tt_diffusion_planner.ttaw``
|
| 14 |
+
is the vendored shared package of the Autoware ports (``ttaw/VENDORED.json`` records its version and file hashes;
|
| 15 |
+
never edit it here). ``tt_diffusion_planner.reference`` is the fp32 CPU reference (torch) and
|
| 16 |
+
``tt_diffusion_planner.host`` the node's pre- and post-processing (numpy).
|
| 17 |
+
"""
|
| 18 |
+
|
| 19 |
+
__version__ = "0.1.0"
|
| 20 |
+
__all__ = ["DiffusionPlanner", "Output", "open_device", "load_inputs", "INPUT_SCHEMA", "__version__"]
|
| 21 |
+
|
| 22 |
+
_LAZY = {
|
| 23 |
+
"DiffusionPlanner": (".api", "DiffusionPlanner"),
|
| 24 |
+
"Output": (".api", "Output"),
|
| 25 |
+
"open_device": (".device", "open_device"),
|
| 26 |
+
"load_inputs": (".io", "load_inputs"),
|
| 27 |
+
"INPUT_SCHEMA": (".reference.config", "INPUT_SCHEMA"),
|
| 28 |
+
}
|
| 29 |
+
|
| 30 |
+
|
| 31 |
+
def __getattr__(name):
|
| 32 |
+
if name in _LAZY:
|
| 33 |
+
import importlib
|
| 34 |
+
|
| 35 |
+
module, attr = _LAZY[name]
|
| 36 |
+
return getattr(importlib.import_module(module, __name__), attr)
|
| 37 |
+
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
| 38 |
+
|
| 39 |
+
|
| 40 |
+
def __dir__():
|
| 41 |
+
return sorted(list(globals()) + __all__)
|
code/tt_diffusion_planner/api.py
ADDED
|
@@ -0,0 +1,125 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-License-Identifier: Apache-2.0
|
| 2 |
+
"""Python API: Diffusion Planner v5.0 (Autoware diffusion_planner) on one Tenstorrent Blackhole p150.
|
| 3 |
+
|
| 4 |
+
from tt_diffusion_planner import DiffusionPlanner
|
| 5 |
+
|
| 6 |
+
with DiffusionPlanner.from_pretrained(device_id=0) as model: # weights -> HF cache, device open, traces captured
|
| 7 |
+
out = model(inputs="code/tt_diffusion_planner/samples/kashiwanoha_dense.npz")
|
| 8 |
+
print(out.to_dict()) # the same JSON as POST /predict
|
| 9 |
+
|
| 10 |
+
The contract shared by every bundle of the Autoware collection is the vendored ``ttaw.api_base.ModelBase``
|
| 11 |
+
(BUNDLE_CONVENTIONS.md section 8): ``from_pretrained`` resolves the pinned weights before it claims the chip, opens it
|
| 12 |
+
(ETH dispatch, 12x10), builds the graph and captures every trace variant in ``warmup_variants``, so the first call is
|
| 13 |
+
as fast as the later ones; calls are serialised by a lock (one chip, batch 1) and fill ``timing_ms``; ``close()`` is
|
| 14 |
+
idempotent, also runs at interpreter exit, and closes the chip only if the model opened it. The HTTP server
|
| 15 |
+
(``tt_diffusion_planner.server.app``) calls this class, so ``/predict`` and ``model(...)`` agree bit for bit.
|
| 16 |
+
|
| 17 |
+
Input: ``inputs=`` holds the 15 raw tensors of the node's ``DiffusionPlannerCore::create_input_data`` (ego frame,
|
| 18 |
+
before normalization, batch 1; ``INPUT_SCHEMA``). Converting ROS messages and the Lanelet2 map into them, and the
|
| 19 |
+
node's temporal state (agent buffers, ego history, RTC prefix of ``sampled_trajectories``), stay with the client.
|
| 20 |
+
The hooks delegate to ``tt_diffusion_planner.host`` (the node's pre- and post-processing, numpy) around
|
| 21 |
+
``tt_diffusion_planner.tt`` (the ttnn graph: encoder + 11 DiT evaluations with the DPM-Solver++(2M) update + turn head,
|
| 22 |
+
run by a ``ttaw.trace.TraceRunner``). Importing this module has no side effects (ttnn / torch only inside hooks).
|
| 23 |
+
"""
|
| 24 |
+
from __future__ import annotations
|
| 25 |
+
|
| 26 |
+
from typing import Any, Dict
|
| 27 |
+
|
| 28 |
+
from . import io as tio
|
| 29 |
+
from .device import DEVICE_DEFAULTS
|
| 30 |
+
from .host import pipeline as hp
|
| 31 |
+
from .reference import config as C
|
| 32 |
+
from .ttaw.api_base import ModelBase
|
| 33 |
+
from .ttaw.outputs import Trajectory
|
| 34 |
+
|
| 35 |
+
__all__ = ["DiffusionPlanner", "Output"]
|
| 36 |
+
|
| 37 |
+
# The result class of this model (ttaw.outputs: Detections3D, Detections2D, Segmentation3D, Mask2D, Trajectory).
|
| 38 |
+
Output = Trajectory
|
| 39 |
+
|
| 40 |
+
|
| 41 |
+
class DiffusionPlanner(ModelBase):
|
| 42 |
+
"""Diffusion Planner v5.0 (Autoware diffusion_planner) on one Blackhole p150. Create it with
|
| 43 |
+
:meth:`from_pretrained`."""
|
| 44 |
+
|
| 45 |
+
MODEL_NAME = "diffusion-planner-p150"
|
| 46 |
+
ENV_PREFIX = "DIFFUSION_PLANNER" # prefix of the environment knobs (SERVING.md section 3.4)
|
| 47 |
+
DEFAULT_REPO = "AutowareFoundation/diffusion_planner"
|
| 48 |
+
DEFAULT_TAG = "v5.0" # the Autoware ansible artifacts pin; the node loads only major version 5
|
| 49 |
+
# the commit DEFAULT_TAG points to (pinned: tags can move)
|
| 50 |
+
DEFAULT_REVISION = "423efde67f5414734da43a7ad856c17ceb8b51aa"
|
| 51 |
+
ALLOW_PATTERNS = ["diffusion_planner_encoder.onnx", "diffusion_planner_decoder.onnx",
|
| 52 |
+
"diffusion_planner_turn_indicator.onnx", "diffusion_planner.param.json"]
|
| 53 |
+
VARIANTS = ["default"] # load-time: the multi-step graph with dpm_solver_steps = 10
|
| 54 |
+
DEFAULT_VARIANT = "default"
|
| 55 |
+
INPUT_KIND = "planner" # lidar | camera | multicam | lidar+multicam | planner
|
| 56 |
+
CAMERA_ORDER = ()
|
| 57 |
+
POINT_FIELDS = tio.DEFAULT_POINT_FIELDS # () : no point cloud
|
| 58 |
+
# turn-indicator logit order (dimensions.hpp:74-79); the published command is the index for 0..3
|
| 59 |
+
LABELS = C.TURN_INDICATOR_LABELS
|
| 60 |
+
# Per-request knobs: name -> (type, min, max, default); host-side post-processing only (the node's YAML defaults).
|
| 61 |
+
RUNTIME_PARAMS = hp.RUNTIME_PARAMS
|
| 62 |
+
EXTRA_INPUTS = ()
|
| 63 |
+
# the ONNX-named raw tensors (name -> (shape, dtype)), decoded and checked on every call (API and server alike)
|
| 64 |
+
INPUT_SCHEMA = C.INPUT_SCHEMA
|
| 65 |
+
DEVICE_DEFAULTS = DEVICE_DEFAULTS # validated open parameters (device.py)
|
| 66 |
+
|
| 67 |
+
# ---- port-specific hooks (called by ModelBase; keep host work out of _forward) ---------------------------
|
| 68 |
+
def _build(self) -> None:
|
| 69 |
+
"""Weights (the three ONNX files + param JSON, read as data by ``reference.weights``) -> the ttnn graph of
|
| 70 |
+
``tt_diffusion_planner.tt`` registered as the variants of a ``ttaw.trace.TraceRunner`` (persistent inputs,
|
| 71 |
+
RT-dev solver / adaLN tables and states allocated here, before any capture). No capture here."""
|
| 72 |
+
from .reference.weights import load_weights
|
| 73 |
+
from .tt.model import TtDiffusionPlanner
|
| 74 |
+
|
| 75 |
+
unknown = sorted(set(self.compile_params) - {"precision"})
|
| 76 |
+
if unknown:
|
| 77 |
+
raise TypeError(f"unknown compile parameter(s) {unknown}; allowed: precision (extra precision rules, "
|
| 78 |
+
"e.g. 'dec.*=HiFi2+fp32'); the LN_FP32 / HIDDEN_FP32 options are DIFFUSION_PLANNER_* knobs")
|
| 79 |
+
self.planner_weights = load_weights(self.weights_path)
|
| 80 |
+
self.normalization = self.planner_weights.normalization
|
| 81 |
+
self.tt = TtDiffusionPlanner(self.device, self.planner_weights, precision=self.compile_params.get("precision"))
|
| 82 |
+
self.runner = self.tt.runner
|
| 83 |
+
|
| 84 |
+
def _warm_one(self, variant: Dict[str, Any]) -> None:
|
| 85 |
+
"""``TraceRunner.capture`` warms every pending variant eagerly (kernel JIT, program cache) before any capture,
|
| 86 |
+
then captures with program-cache misses forbidden; idempotent."""
|
| 87 |
+
self.runner.capture()
|
| 88 |
+
|
| 89 |
+
def _prepare(self, points: Any = None, inputs: Any = None, **other: Any) -> hp.Prepared:
|
| 90 |
+
"""The node's host pre-processing (``host.prepare``): normalization (all-zero rows kept), speed masks, the
|
| 91 |
+
encoder's host features, the decoder masks and the solver's initial state."""
|
| 92 |
+
given = sorted(k for k, v in {"points": points, **other}.items() if v is not None)
|
| 93 |
+
if given:
|
| 94 |
+
raise tio.InputError(f"this model takes only `inputs` (the planner tensors), not {given}")
|
| 95 |
+
if inputs is None:
|
| 96 |
+
raise tio.InputError("this model needs `inputs`: the 15 planner tensors of INPUT_SCHEMA")
|
| 97 |
+
return hp.prepare(inputs, self.normalization.observation)
|
| 98 |
+
|
| 99 |
+
def _forward(self, prepared: hp.Prepared) -> Dict[str, Any]:
|
| 100 |
+
"""Upload the host features into the persistent device inputs, replay the plan's trace(s) and read back
|
| 101 |
+
``final_x0`` (normalised [321, 81, 4]), the turn logits and, when asked, the solver iterates."""
|
| 102 |
+
return self.tt.forward(prepared)
|
| 103 |
+
|
| 104 |
+
def _postprocess(self, raw: Dict[str, Any], prepared: hp.Prepared, params: Dict[str, Any]) -> Output:
|
| 105 |
+
"""The node's post-processing (``host.make_output``): denormalisation, trajectory velocity / force-stop /
|
| 106 |
+
acceleration, predicted neighbour paths, turn-indicator decision."""
|
| 107 |
+
return hp.make_output(raw["final_x0"], raw["logit"], prepared, self.normalization, params,
|
| 108 |
+
model=self.MODEL_NAME, denoising_steps=raw.get("denoising_steps"))
|
| 109 |
+
|
| 110 |
+
def _release(self) -> None:
|
| 111 |
+
"""Release the traces and persistent device tensors; also called when ``from_pretrained`` fails half-way."""
|
| 112 |
+
runner = getattr(self, "runner", None)
|
| 113 |
+
if runner is not None:
|
| 114 |
+
runner.release()
|
| 115 |
+
|
| 116 |
+
def extra_info(self) -> Dict[str, Any]:
|
| 117 |
+
"""Additions to ``model.info`` and ``/info``: the trace variants, CQs and persistent tensors of the runner."""
|
| 118 |
+
runner = getattr(self, "runner", None)
|
| 119 |
+
info: Dict[str, Any] = {"dpm_solver_steps": C.DPM_SOLVER_STEPS, "input_names": list(C.INPUT_NAMES)}
|
| 120 |
+
tt = getattr(self, "tt", None)
|
| 121 |
+
if tt is not None:
|
| 122 |
+
info.update(tt.describe())
|
| 123 |
+
elif runner is not None:
|
| 124 |
+
info["trace"] = runner.describe()
|
| 125 |
+
return info
|
code/tt_diffusion_planner/calib/README.md
ADDED
|
@@ -0,0 +1,6 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Calibration presets
|
| 2 |
+
|
| 3 |
+
None: the planner takes no calibration. Its inputs are tensors in the ego (`base_link`) frame that the client builds
|
| 4 |
+
from ROS messages and the Lanelet2 map (README "Quickstart", SERVING.md 3.1), so there is no camera or LiDAR extrinsic to
|
| 5 |
+
send. The directory exists because the shared server of the Autoware ports lists `*.json` presets from it in `/info`
|
| 6 |
+
(`calibration_presets`, empty here); a request that carries `calibration` is refused (HTTP 400).
|
code/tt_diffusion_planner/device.py
ADDED
|
@@ -0,0 +1,61 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-License-Identifier: Apache-2.0
|
| 2 |
+
"""Opening the Blackhole p150 the way every published number of diffusion-planner-p150 was measured.
|
| 3 |
+
|
| 4 |
+
The implementation is the vendored ``ttaw.device`` (C01): dispatch on the idle ETH cores
|
| 5 |
+
(``ttnn.DispatchCoreConfig(ttnn.DispatchCoreType.ETH)``), which gives a 12x10 = 120-core compute grid on a p150
|
| 6 |
+
(11x10 with WORKER dispatch, the A/B switch). ETH dispatch needs ``patches/tt-metal-eth-dispatch.patch`` on tt-metal
|
| 7 |
+
44d6650 (the container image is built from a patched tree); if the ETH open fails, a ``RuntimeWarning`` is issued and
|
| 8 |
+
WORKER dispatch is used, unless ``allow_fallback=False``. Never hard-code the grid: use :func:`compute_grid` or
|
| 9 |
+
``device.compute_with_storage_grid_size()``.
|
| 10 |
+
|
| 11 |
+
This module binds it to the validated open parameters of this port, :data:`DEVICE_DEFAULTS` (the model class uses
|
| 12 |
+
the same dict), which the ``DIFFUSION_PLANNER_DISPATCH``, ``_NUM_CQS``, ``_L1_SMALL``, ``_TRACE_REGION`` and
|
| 13 |
+
``_WORKER_L1_SIZE`` variables and ``TT_DEVICE_ID`` override per process (SERVING.md section 3.4). Importing it has
|
| 14 |
+
no side effects (``ttnn`` is imported when a device is opened).
|
| 15 |
+
"""
|
| 16 |
+
from __future__ import annotations
|
| 17 |
+
|
| 18 |
+
import contextlib
|
| 19 |
+
import dataclasses
|
| 20 |
+
from typing import Any, Dict, Iterator, Optional
|
| 21 |
+
|
| 22 |
+
from .ttaw.device import DeviceConfig, close_device, compute_grid, core_grid, describe_device, full_core_range_set
|
| 23 |
+
|
| 24 |
+
__all__ = ["ENV_PREFIX", "DEVICE_DEFAULTS", "DeviceConfig", "device_config", "open_device", "device_session",
|
| 25 |
+
"close_device", "describe_device", "compute_grid", "core_grid", "full_core_range_set"]
|
| 26 |
+
|
| 27 |
+
ENV_PREFIX = "DIFFUSION_PLANNER"
|
| 28 |
+
|
| 29 |
+
# Validated device-open parameters of this port (fill per model; the card's numbers are measured with them).
|
| 30 |
+
DEVICE_DEFAULTS: Dict[str, Any] = {
|
| 31 |
+
"num_command_queues": 1, # 1, or 2 when the input upload (CQ1) overlaps the trace (CQ0)
|
| 32 |
+
"l1_small_size": 32768, # L1_SMALL bytes per core (conv / pool config tensors)
|
| 33 |
+
# DRAM bytes for the traces: measured 74.6 MB for the plan, 96.9 MB with the tests' debug variants (PORT_LOG 5.3)
|
| 34 |
+
"trace_region_size": 192 << 20,
|
| 35 |
+
}
|
| 36 |
+
|
| 37 |
+
|
| 38 |
+
def device_config(**overrides: Any) -> DeviceConfig:
|
| 39 |
+
""":data:`DEVICE_DEFAULTS` < the ``DIFFUSION_PLANNER_*`` / ``TT_DEVICE_ID`` environment < explicit non-None
|
| 40 |
+
``overrides``
|
| 41 |
+
(``device_id``, ``dispatch``, ``num_command_queues``, ``l1_small_size``, ``trace_region_size``,
|
| 42 |
+
``worker_l1_size``, ``allow_fallback``): the same resolution as ``DiffusionPlanner.from_pretrained``."""
|
| 43 |
+
config = DeviceConfig.from_env(ENV_PREFIX, **DEVICE_DEFAULTS)
|
| 44 |
+
return dataclasses.replace(config, **{k: v for k, v in overrides.items() if v is not None})
|
| 45 |
+
|
| 46 |
+
|
| 47 |
+
def open_device(device_id: Optional[int] = None, *, dispatch: Optional[str] = None, **overrides: Any):
|
| 48 |
+
"""Open one chip like the published numbers: ETH dispatch (``dispatch="worker"`` is the A/B opt-in) and this
|
| 49 |
+
port's sizes. Close it with :func:`close_device`, or use :func:`device_session`."""
|
| 50 |
+
return device_config(device_id=device_id, dispatch=dispatch, **overrides).open()
|
| 51 |
+
|
| 52 |
+
|
| 53 |
+
@contextlib.contextmanager
|
| 54 |
+
def device_session(device_id: Optional[int] = None, **overrides: Any) -> Iterator[Any]:
|
| 55 |
+
"""``with device_session() as dev:`` opens with :func:`open_device` and always closes, also when the body
|
| 56 |
+
raises."""
|
| 57 |
+
device = open_device(device_id, **overrides)
|
| 58 |
+
try:
|
| 59 |
+
yield device
|
| 60 |
+
finally:
|
| 61 |
+
close_device(device)
|
code/tt_diffusion_planner/host/__init__.py
ADDED
|
@@ -0,0 +1,27 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-License-Identifier: Apache-2.0
|
| 2 |
+
"""Host pre- and post-processing of the Autoware Diffusion Planner (numpy only, no ttnn, no torch).
|
| 3 |
+
|
| 4 |
+
Exact ports of the node's host code (autoware_universe @ 9ceaccf ``planning/autoware_diffusion_planner``) and of
|
| 5 |
+
the in-graph pre-processing of the encoder that the TT port moves to the host (PLAN.md section 2.12 "Host fallbacks":
|
| 6 |
+
normalization, position features, masks, post-processing):
|
| 7 |
+
|
| 8 |
+
- :mod:`.normalize` input normalization that keeps all-zero rows at zero, speed-limit masks
|
| 9 |
+
(``preprocessing_utils.cpp:34-84``, ``inference/utils.hpp:112-123``);
|
| 10 |
+
- :mod:`.features` what the encoder graph computes before its first matmul: history truncation, validity masks,
|
| 11 |
+
the neighbour velocity zeroing and valid-step flag, lane attributes and speed selection, polygon / line-string
|
| 12 |
+
deltas, the 14-dim position features (ONNX ``atan2`` decomposition and the polygon / line-string pseudo-heading
|
| 13 |
+
quirk), the fusion key mask; decoder agent mask and current states (SPEC 3.8, 4.3);
|
| 14 |
+
- :mod:`.solver` DPM-Solver++(2M) with denoise-to-zero, its float32 scalar schedule computed with the C library
|
| 15 |
+
like ``dpm_solver.cpp``, the prefix constraint (``multi_step_inference.cpp:300-340``);
|
| 16 |
+
- :mod:`.postprocess` denormalization, poses with Eigen's quaternion of the unnormalised rotation, the trajectory
|
| 17 |
+
velocity / force-stop / acceleration rules, predicted neighbour paths, the turn-indicator decision
|
| 18 |
+
(``postprocessing_utils.cpp``, ``turn_indicator_manager.cpp``);
|
| 19 |
+
- :mod:`.pipeline` ``prepare(raw) -> Prepared`` and ``make_output(...) -> Trajectory``: the two halves the Python
|
| 20 |
+
API, the HTTP server and the CPU reference share around the network.
|
| 21 |
+
|
| 22 |
+
Importing this package has no side effects.
|
| 23 |
+
"""
|
| 24 |
+
from .normalize import normalize_inputs, speed_masks # noqa: F401
|
| 25 |
+
from .pipeline import Prepared, make_output, prepare # noqa: F401
|
| 26 |
+
|
| 27 |
+
__all__ = ["normalize_inputs", "speed_masks", "Prepared", "prepare", "make_output"]
|
code/tt_diffusion_planner/host/features.py
ADDED
|
@@ -0,0 +1,203 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-License-Identifier: Apache-2.0
|
| 2 |
+
"""The encoder's in-graph pre-processing and the decoder's masks, computed on the host from normalised inputs.
|
| 3 |
+
|
| 4 |
+
Everything the v5.0 encoder graph does before its first matmul is tensor plumbing (slices, ``!= 0`` tests, OneHot,
|
| 5 |
+
``Atan`` with quadrant ``Where`` s): the TT port computes it here once per plan and uploads the results (PLAN.md 2.12;
|
| 6 |
+
SPEC 3.8, 8.2.6: padded lanes produce ``0/0`` headings that only a ``Where`` discards, which an arithmetic device
|
| 7 |
+
``where`` would propagate). The CPU reference consumes the same arrays, so the reference-vs-ONNX Runtime test also
|
| 8 |
+
proves this module against the graph (``/encoder/*/Where*``, ``/encoder/Concat_4`` / ``Concat_5`` taps).
|
| 9 |
+
|
| 10 |
+
Semantics (T4M ``model/module/encoder.py``; checked against the exported graph):
|
| 11 |
+
|
| 12 |
+
- ego: only the 6 OLDEST history rows are kept (rows 6..30 zeroed); the ego token is always valid and its position
|
| 13 |
+
feature is the last row of the truncated history, i.e. zeros;
|
| 14 |
+
- neighbours: rows 0..24 zeroed; a step is valid if any of dims 0..7 is non-zero, an agent if any step is; the type
|
| 15 |
+
one-hot (dims 8..10) and the position feature come from the last step; velocities (dims 4, 5) are zeroed AFTER the
|
| 16 |
+
validity test and a valid-step flag is appended (9 channels); invalid agents are zeroed;
|
| 17 |
+
- static objects: valid if any of the 10 values is non-zero (always zeros in Autoware, so never valid);
|
| 18 |
+
- lanes / route: dims 0..7 per point (zeroed for invalid lanes), attributes = dims 8..32 of point 0, speed limit and
|
| 19 |
+
its mask, position feature = point 10 with heading ``atan2(dy, dx)`` exported as ``Atan(dy / dx)`` plus quadrant
|
| 20 |
+
``Where`` s;
|
| 21 |
+
- polygons / line strings: ``[x, y, type one-hot..., dx, dy]`` with dx, dy = next point minus point (0 for the last
|
| 22 |
+
point), valid if any of the first FOUR columns is non-zero, position feature = point 20 / 10 with the NON-geometric
|
| 23 |
+
heading ``atan2(col3, col2)``: atan2(dx, is_intersection_area) for polygons, atan2(is_road_border, is_stop_line)
|
| 24 |
+
for line strings (SPEC 3.8: port the quirk exactly);
|
| 25 |
+
- goal / ego shape / turn indicators: always valid; positions (goal) and (0, 0, 1, 0); turn indicators drop the
|
| 26 |
+
current report (``[:, :-1]``, 30 values); the turn token reuses the ego-shape class id 8;
|
| 27 |
+
- fusion key mask: invalid tokens are masked with -inf, the ego key is forced valid (``encoder.py:833``);
|
| 28 |
+
- decoder: a neighbour is a valid attention key if its current state (``neighbor_agents_past[:, 30, :4]``) is not
|
| 29 |
+
all zero; current states = (ego_current_state[:4], neighbour current states) for the prefix constraint.
|
| 30 |
+
"""
|
| 31 |
+
from __future__ import annotations
|
| 32 |
+
|
| 33 |
+
from dataclasses import dataclass
|
| 34 |
+
from typing import Dict, Mapping
|
| 35 |
+
|
| 36 |
+
import numpy as np
|
| 37 |
+
|
| 38 |
+
from ..reference import config as C
|
| 39 |
+
|
| 40 |
+
__all__ = ["EncoderFeatures", "DecoderMasks", "encoder_features", "decoder_masks", "atan2_onnx", "line_features"]
|
| 41 |
+
|
| 42 |
+
PI_F32 = np.float32(3.1415927) # the ONNX constant of the exported atan2 (``/encoder/lane_encoder/Constant_13``)
|
| 43 |
+
|
| 44 |
+
|
| 45 |
+
def atan2_onnx(y: np.ndarray, x: np.ndarray) -> np.ndarray:
|
| 46 |
+
"""``torch.atan2`` as exported to ONNX opset 20: ``a = Atan(y / x)``; ``Where(x < 0, Where(y > 0, a + pi, a - pi),
|
| 47 |
+
a)`` (float32). Differs from IEEE atan2 only on signed zeros (x = -0, y = 0 -> -pi) and 0/0 (NaN, masked later)."""
|
| 48 |
+
y = np.asarray(y, np.float32)
|
| 49 |
+
x = np.asarray(x, np.float32)
|
| 50 |
+
with np.errstate(divide="ignore", invalid="ignore"):
|
| 51 |
+
a = np.arctan((y / x).astype(np.float32)).astype(np.float32)
|
| 52 |
+
alt = np.where(y > 0, a + PI_F32, a - PI_F32).astype(np.float32)
|
| 53 |
+
return np.where(x < 0, alt, a).astype(np.float32)
|
| 54 |
+
|
| 55 |
+
|
| 56 |
+
def _heading_pos(xy: np.ndarray, y_col: np.ndarray, x_col: np.ndarray) -> np.ndarray:
|
| 57 |
+
"""``[x, y, cos(h), sin(h)]`` with ``h = atan2_onnx(y_col, x_col)`` (float32)."""
|
| 58 |
+
h = atan2_onnx(y_col, x_col)
|
| 59 |
+
with np.errstate(invalid="ignore"):
|
| 60 |
+
return np.stack([xy[..., 0], xy[..., 1], np.cos(h), np.sin(h)], axis=-1).astype(np.float32)
|
| 61 |
+
|
| 62 |
+
|
| 63 |
+
def _onehot_pos(pos4: np.ndarray, cls: int) -> np.ndarray:
|
| 64 |
+
"""Append the 10-way class one-hot (``add_class_type``)."""
|
| 65 |
+
onehot = np.zeros(pos4.shape[:-1] + (C.POS_CLASS_NUM,), np.float32)
|
| 66 |
+
onehot[..., cls] = 1.0
|
| 67 |
+
return np.concatenate([pos4.astype(np.float32), onehot], axis=-1)
|
| 68 |
+
|
| 69 |
+
|
| 70 |
+
def line_features(x: np.ndarray) -> np.ndarray:
|
| 71 |
+
"""``LineEncoder``: ``[points..., dx, dy]`` with dx / dy = next point minus point (0 at the last point)."""
|
| 72 |
+
x = np.asarray(x, np.float32)
|
| 73 |
+
d = np.zeros(x.shape[:-1] + (2,), np.float32)
|
| 74 |
+
d[..., :-1, 0] = x[..., 1:, 0] - x[..., :-1, 0]
|
| 75 |
+
d[..., :-1, 1] = x[..., 1:, 1] - x[..., :-1, 1]
|
| 76 |
+
return np.concatenate([x, d], axis=-1)
|
| 77 |
+
|
| 78 |
+
|
| 79 |
+
@dataclass
|
| 80 |
+
class EncoderFeatures:
|
| 81 |
+
"""Host-side inputs of the encoder network (batch 1, no batch dim). ``valid[cat]`` is True for valid entities;
|
| 82 |
+
``token_valid`` (564) gates the positional embedding, ``key_valid`` (564) is the fusion key mask (ego forced
|
| 83 |
+
valid); ``pos`` (564 x 14) are the position features, with invalid rows set to 0 (the graph's NaN rows there are
|
| 84 |
+
discarded by a ``Where``)."""
|
| 85 |
+
|
| 86 |
+
ego: np.ndarray # [31, 4] truncated ego history
|
| 87 |
+
neighbor: np.ndarray # [320, 31, 9]
|
| 88 |
+
neighbor_type: np.ndarray # [320, 3]
|
| 89 |
+
static: np.ndarray # [5, 10]
|
| 90 |
+
lane: np.ndarray # [140, 20, 8]
|
| 91 |
+
lane_attr: np.ndarray # [140, 25]
|
| 92 |
+
lane_speed: np.ndarray # [140, 1]
|
| 93 |
+
lane_has_speed: np.ndarray # [140, 1] bool
|
| 94 |
+
route: np.ndarray # [25, 20, 8]
|
| 95 |
+
route_attr: np.ndarray # [25, 25]
|
| 96 |
+
route_speed: np.ndarray # [25, 1]
|
| 97 |
+
route_has_speed: np.ndarray # [25, 1] bool
|
| 98 |
+
polygon: np.ndarray # [10, 40, 5]
|
| 99 |
+
line_string: np.ndarray # [60, 20, 6]
|
| 100 |
+
goal: np.ndarray # [4]
|
| 101 |
+
ego_shape: np.ndarray # [3]
|
| 102 |
+
turn: np.ndarray # [30]
|
| 103 |
+
valid: Dict[str, np.ndarray]
|
| 104 |
+
token_valid: np.ndarray # [564] bool
|
| 105 |
+
key_valid: np.ndarray # [564] bool
|
| 106 |
+
pos: np.ndarray # [564, 14]
|
| 107 |
+
|
| 108 |
+
def counts(self) -> Dict[str, int]:
|
| 109 |
+
return {k: int(v.sum()) for k, v in self.valid.items()}
|
| 110 |
+
|
| 111 |
+
|
| 112 |
+
@dataclass
|
| 113 |
+
class DecoderMasks:
|
| 114 |
+
agent_valid: np.ndarray # [321] bool: ego + neighbours with a non-zero current state (self-attention keys)
|
| 115 |
+
current_states: np.ndarray # [321, 4] normalised (prefix constraint)
|
| 116 |
+
|
| 117 |
+
|
| 118 |
+
def _any_nonzero(a: np.ndarray, axes) -> np.ndarray:
|
| 119 |
+
return np.any(a != 0, axis=axes)
|
| 120 |
+
|
| 121 |
+
|
| 122 |
+
def encoder_features(norm: Mapping[str, np.ndarray], masks: Mapping[str, np.ndarray]) -> EncoderFeatures:
|
| 123 |
+
"""From the normalised inputs (``normalize_inputs``) and the speed masks (``speed_masks``), batch 1."""
|
| 124 |
+
f32 = lambda k: np.asarray(norm[k], np.float32)[0] # noqa: E731 (drop the batch dim)
|
| 125 |
+
# ego: keep the 6 oldest rows (encoder.py:170-175); the token is always valid, its position row is all zero
|
| 126 |
+
ego = np.zeros((C.INPUT_T + 1, C.POSE_DIM), np.float32)
|
| 127 |
+
ego[C.EGO_HISTORY_KEEP] = f32("ego_agent_past")[C.EGO_HISTORY_KEEP]
|
| 128 |
+
pos = {"ego": _onehot_pos(ego[-1:], C.POS_CLASS["ego"])}
|
| 129 |
+
|
| 130 |
+
# neighbours: keep the 6 newest rows (encoder.py:176-181, 441-451)
|
| 131 |
+
nb_raw = f32("neighbor_agents_past")
|
| 132 |
+
nb = np.zeros_like(nb_raw)
|
| 133 |
+
nb[:, C.NEIGHBOR_HISTORY_KEEP] = nb_raw[:, C.NEIGHBOR_HISTORY_KEEP]
|
| 134 |
+
nb_type = nb[:, -1, 8:11].copy()
|
| 135 |
+
x8 = nb[..., :8]
|
| 136 |
+
step_valid = _any_nonzero(x8, -1) # [320, 31]
|
| 137 |
+
nb_valid = step_valid.any(axis=-1) # [320]
|
| 138 |
+
feat = np.concatenate([x8, step_valid[..., None].astype(np.float32)], axis=-1)
|
| 139 |
+
feat[..., 4:6] = 0.0 # velocities zeroed after the validity test
|
| 140 |
+
feat[~nb_valid] = 0.0
|
| 141 |
+
pos["neighbor"] = _onehot_pos(x8[:, -1, :4], C.POS_CLASS["neighbor"])
|
| 142 |
+
|
| 143 |
+
static = f32("static_objects")
|
| 144 |
+
st_valid = _any_nonzero(static[..., :10], -1)
|
| 145 |
+
static = np.where(st_valid[:, None], static, 0.0).astype(np.float32)
|
| 146 |
+
pos["static"] = _onehot_pos(f32("static_objects")[:, :4], C.POS_CLASS["static"])
|
| 147 |
+
|
| 148 |
+
def lanes(key: str, speed_key: str, mask_key: str, cat: str):
|
| 149 |
+
x = f32(key)
|
| 150 |
+
attr = x[:, 0, C.LANE_FEATURE_DIM:].copy()
|
| 151 |
+
x8 = x[..., :C.LANE_FEATURE_DIM]
|
| 152 |
+
valid = _any_nonzero(x8, (-1, -2))
|
| 153 |
+
p = x8[:, C.LANE_POS_INDEX, :4]
|
| 154 |
+
pos[cat] = _onehot_pos(_heading_pos(p, p[:, 3], p[:, 2]), C.POS_CLASS[cat])
|
| 155 |
+
x8 = np.where(valid[:, None, None], x8, 0.0).astype(np.float32)
|
| 156 |
+
speed = f32(speed_key).reshape(-1, 1)
|
| 157 |
+
has = np.asarray(masks[mask_key])[0].reshape(-1, 1).astype(bool)
|
| 158 |
+
return x8, attr, speed, has, valid
|
| 159 |
+
|
| 160 |
+
lane, lane_attr, lane_speed, lane_has, lane_valid = lanes("lanes", "lanes_speed_limit",
|
| 161 |
+
"lanes_has_speed_limit", "lane")
|
| 162 |
+
route, route_attr, route_speed, route_has, route_valid = lanes("route_lanes", "route_lanes_speed_limit",
|
| 163 |
+
"route_lanes_has_speed_limit", "route")
|
| 164 |
+
|
| 165 |
+
def lines(key: str, cat: str, pos_index: int):
|
| 166 |
+
x = line_features(f32(key))
|
| 167 |
+
valid = _any_nonzero(x[..., :4], (-1, -2))
|
| 168 |
+
p = x[:, pos_index, :4]
|
| 169 |
+
pos[cat] = _onehot_pos(_heading_pos(p, p[:, 3], p[:, 2]), C.POS_CLASS[cat])
|
| 170 |
+
return np.where(valid[:, None, None], x, 0.0).astype(np.float32), valid
|
| 171 |
+
|
| 172 |
+
polygon, poly_valid = lines("polygons", "polygon", C.POLYGON_POS_INDEX)
|
| 173 |
+
line_string, ls_valid = lines("line_strings", "line_string", C.LINE_STRING_POS_INDEX)
|
| 174 |
+
|
| 175 |
+
goal = f32("goal_pose").reshape(-1)
|
| 176 |
+
pos["goal"] = _onehot_pos(goal[None], C.POS_CLASS["goal"])
|
| 177 |
+
unit = np.array([[0.0, 0.0, 1.0, 0.0]], np.float32)
|
| 178 |
+
pos["ego_shape"] = _onehot_pos(unit, C.POS_CLASS["ego_shape"])
|
| 179 |
+
pos["turn"] = _onehot_pos(unit, C.POS_CLASS["turn"])
|
| 180 |
+
turn = np.asarray(norm["turn_indicators"], np.float32)[0, :C.TURN_INDICATOR_HISTORY].copy()
|
| 181 |
+
|
| 182 |
+
one = np.ones(1, bool)
|
| 183 |
+
valid = {"ego": one, "neighbor": nb_valid, "static": st_valid, "lane": lane_valid, "route": route_valid,
|
| 184 |
+
"polygon": poly_valid, "line_string": ls_valid, "goal": one, "ego_shape": one.copy(), "turn": one.copy()}
|
| 185 |
+
token_valid = np.concatenate([valid[name] for name, _ in C.TOKEN_LAYOUT])
|
| 186 |
+
key_valid = token_valid.copy()
|
| 187 |
+
key_valid[0] = True
|
| 188 |
+
pos_all = np.concatenate([pos[name] for name, _ in C.TOKEN_LAYOUT], axis=0).astype(np.float32)
|
| 189 |
+
pos_all[~token_valid] = 0.0
|
| 190 |
+
return EncoderFeatures(ego, feat.astype(np.float32), nb_type, static, lane, lane_attr, lane_speed, lane_has,
|
| 191 |
+
route, route_attr, route_speed, route_has, polygon, line_string, goal,
|
| 192 |
+
f32("ego_shape").reshape(-1).copy(), turn, valid, token_valid, key_valid, pos_all)
|
| 193 |
+
|
| 194 |
+
|
| 195 |
+
def decoder_masks(norm: Mapping[str, np.ndarray]) -> DecoderMasks:
|
| 196 |
+
"""Agent key mask of the DiT self-attention (``dit.py:155-156``; the ego is always valid) and the current states
|
| 197 |
+
of the prefix constraint (``multi_step_inference.cpp:300-325``), from the normalised inputs."""
|
| 198 |
+
nb_now = np.asarray(norm["neighbor_agents_past"], np.float32)[0, :, C.INPUT_T, :C.POSE_DIM]
|
| 199 |
+
agent_valid = np.concatenate([np.ones(1, bool), np.any(nb_now != 0, axis=-1)])
|
| 200 |
+
cs = np.zeros((C.MAX_NUM_AGENTS, C.POSE_DIM), np.float32)
|
| 201 |
+
cs[0] = np.asarray(norm["ego_current_state"], np.float32)[0, :C.POSE_DIM]
|
| 202 |
+
cs[1:] = nb_now
|
| 203 |
+
return DecoderMasks(agent_valid, cs)
|
code/tt_diffusion_planner/host/normalize.py
ADDED
|
@@ -0,0 +1,67 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-License-Identifier: Apache-2.0
|
| 2 |
+
"""Input normalization of the Autoware node (``preprocessing_utils.cpp:34-84``) and the speed-limit masks.
|
| 3 |
+
|
| 4 |
+
The rule cannot be folded into the first-layer weights (SPEC 9.1, PLAN 2.12): each row of the last dimension is
|
| 5 |
+
normalised as ``(v - mean) / std`` in float32 unless every value of the row satisfies ``|v| < FLT_EPSILON``, in which
|
| 6 |
+
case the row is left untouched (padding stays zero while ``x`` has mean 10 m). ``ego_shape``,
|
| 7 |
+
``sampled_trajectories``, ``turn_indicators`` and ``delay`` are never normalised.
|
| 8 |
+
"""
|
| 9 |
+
from __future__ import annotations
|
| 10 |
+
|
| 11 |
+
from typing import Dict, Mapping, Tuple
|
| 12 |
+
|
| 13 |
+
import numpy as np
|
| 14 |
+
|
| 15 |
+
from ..reference import config as C
|
| 16 |
+
|
| 17 |
+
__all__ = ["FLT_EPSILON", "normalize_inputs", "normalize_array", "speed_masks"]
|
| 18 |
+
|
| 19 |
+
FLT_EPSILON = np.float32(np.finfo(np.float32).eps) # std::numeric_limits<float>::epsilon()
|
| 20 |
+
|
| 21 |
+
|
| 22 |
+
def normalize_array(value: np.ndarray, mean: np.ndarray, std: np.ndarray) -> np.ndarray:
|
| 23 |
+
"""``normalize_vector`` of ``preprocessing_utils.cpp:36-63`` on one tensor (a new float32 array).
|
| 24 |
+
|
| 25 |
+
Rows have ``std.size`` columns; a single-value mean / std broadcasts over the columns (the C++ ``mean.size() == 1``
|
| 26 |
+
branch); a zero standard deviation raises like the C++."""
|
| 27 |
+
data = np.array(value, dtype=np.float32, copy=True)
|
| 28 |
+
mean = np.asarray(mean, np.float32).reshape(-1)
|
| 29 |
+
std = np.asarray(std, np.float32).reshape(-1)
|
| 30 |
+
if mean.size != std.size:
|
| 31 |
+
raise ValueError("Mean and std must be same size")
|
| 32 |
+
cols = std.size
|
| 33 |
+
if cols == 0 or data.size % cols:
|
| 34 |
+
raise ValueError(f"data size {data.size} is not divisible by the normalizer size {cols}")
|
| 35 |
+
if np.any(np.abs(std) < FLT_EPSILON):
|
| 36 |
+
raise ValueError("Standard deviation is zero, cannot normalize data")
|
| 37 |
+
rows = data.reshape(-1, cols)
|
| 38 |
+
zero_row = np.all(np.abs(rows) < FLT_EPSILON, axis=1)
|
| 39 |
+
m = np.broadcast_to(mean if mean.size > 1 else mean[:1], (cols,))
|
| 40 |
+
s = np.broadcast_to(std if std.size > 1 else std[:1], (cols,))
|
| 41 |
+
normed = ((rows - m) / s).astype(np.float32) # float32 subtract then divide, element by element
|
| 42 |
+
rows[~zero_row] = normed[~zero_row]
|
| 43 |
+
return rows.reshape(data.shape)
|
| 44 |
+
|
| 45 |
+
|
| 46 |
+
def normalize_inputs(raw: Mapping[str, np.ndarray],
|
| 47 |
+
observation: Mapping[str, Tuple[np.ndarray, np.ndarray]]) -> Dict[str, np.ndarray]:
|
| 48 |
+
"""``normalize_input_data(input_data_map, normalization_map)``: every key except the four skipped ones must have
|
| 49 |
+
a normalizer (``Missing key ... from normalization map`` otherwise)."""
|
| 50 |
+
out: Dict[str, np.ndarray] = {}
|
| 51 |
+
for key, value in raw.items():
|
| 52 |
+
if key in C.SKIP_NORMALIZATION:
|
| 53 |
+
out[key] = np.array(value, dtype=np.float32, copy=True)
|
| 54 |
+
continue
|
| 55 |
+
if key not in observation:
|
| 56 |
+
raise KeyError(f"Missing key {key} from normalization map")
|
| 57 |
+
mean, std = observation[key]
|
| 58 |
+
out[key] = normalize_array(value, mean, std)
|
| 59 |
+
return out
|
| 60 |
+
|
| 61 |
+
|
| 62 |
+
def speed_masks(norm_inputs: Mapping[str, np.ndarray]) -> Dict[str, np.ndarray]:
|
| 63 |
+
"""``lanes_has_speed_limit`` / ``route_lanes_has_speed_limit`` of the TensorRT path: normalised speed limit
|
| 64 |
+
``> FLT_EPSILON`` (``inference/utils.hpp:112-123``; the ORT backend uses ``> 0``,
|
| 65 |
+
``onnxruntime_inference.cpp:41-48``; Autoware's default backend is TensorRT)."""
|
| 66 |
+
return {"lanes_has_speed_limit": np.asarray(norm_inputs["lanes_speed_limit"]) > FLT_EPSILON,
|
| 67 |
+
"route_lanes_has_speed_limit": np.asarray(norm_inputs["route_lanes_speed_limit"]) > FLT_EPSILON}
|
code/tt_diffusion_planner/host/pipeline.py
ADDED
|
@@ -0,0 +1,107 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-License-Identifier: Apache-2.0
|
| 2 |
+
"""The host halves of one plan, shared by the Python API, the HTTP server and the CPU reference.
|
| 3 |
+
|
| 4 |
+
``prepare(raw, normalization)`` -> :class:`Prepared`: the raw ONNX-named tensors (already checked against
|
| 5 |
+
``config.INPUT_SCHEMA``) normalised like the node, the speed masks, the encoder's host features, the decoder masks
|
| 6 |
+
and the solver's initial ``x``. The network (device trace or CPU reference) turns it into ``final_x0`` (normalised
|
| 7 |
+
``[321, 81, 4]``, prefix constraint applied) and the turn-indicator logits; ``make_output(...)`` applies the node's
|
| 8 |
+
post-processing and returns the ``ttaw.outputs.Trajectory`` that ``model(...)`` and ``POST /predict`` return.
|
| 9 |
+
"""
|
| 10 |
+
from __future__ import annotations
|
| 11 |
+
|
| 12 |
+
from dataclasses import dataclass
|
| 13 |
+
from typing import Any, Dict, Mapping, Optional, Sequence
|
| 14 |
+
|
| 15 |
+
import numpy as np
|
| 16 |
+
|
| 17 |
+
from ..reference import config as C
|
| 18 |
+
from ..ttaw.io import InputError, encode_array
|
| 19 |
+
from ..ttaw.outputs import Trajectory
|
| 20 |
+
from .features import DecoderMasks, EncoderFeatures, decoder_masks, encoder_features
|
| 21 |
+
from .normalize import normalize_inputs, speed_masks
|
| 22 |
+
from .postprocess import (TurnIndicatorManager, denoising_steps_ego, denormalize, predicted_paths,
|
| 23 |
+
trajectory_from_poses)
|
| 24 |
+
|
| 25 |
+
__all__ = ["Prepared", "prepare", "make_output", "TRAJECTORY_COLUMNS", "PREDICTED_AGENT_COLUMNS", "RUNTIME_PARAMS"]
|
| 26 |
+
|
| 27 |
+
TRAJECTORY_COLUMNS = ("x", "y", "yaw", "cos", "sin", "velocity", "acceleration")
|
| 28 |
+
PREDICTED_AGENT_COLUMNS = ("x", "y", "yaw", "cos", "sin")
|
| 29 |
+
# RT-host parameters (request params / call kwargs): name -> (type, min, max, default); the node's YAML defaults
|
| 30 |
+
RUNTIME_PARAMS = {
|
| 31 |
+
"velocity_smoothing_window": (int, 1, C.OUTPUT_T - 1, C.VELOCITY_SMOOTHING_WINDOW),
|
| 32 |
+
"stopping_threshold": (float, 0.0, None, C.STOPPING_THRESHOLD),
|
| 33 |
+
"turn_indicator_keep_offset": (float, None, None, C.TURN_INDICATOR_KEEP_OFFSET),
|
| 34 |
+
"return_denoising_steps": (bool, None, None, False),
|
| 35 |
+
}
|
| 36 |
+
|
| 37 |
+
|
| 38 |
+
@dataclass
|
| 39 |
+
class Prepared:
|
| 40 |
+
"""Host state of one plan."""
|
| 41 |
+
|
| 42 |
+
raw: Dict[str, np.ndarray] # the validated raw tensors (batch 1)
|
| 43 |
+
norm: Dict[str, np.ndarray] # normalised tensors + the two speed masks
|
| 44 |
+
features: EncoderFeatures
|
| 45 |
+
decoder: DecoderMasks
|
| 46 |
+
x_T: np.ndarray # [321, 81, 4] initial solver state (normalised; correction not yet applied)
|
| 47 |
+
neighbor_rows: np.ndarray # indices of the non-empty neighbour rows (the node's emitted agents)
|
| 48 |
+
enable_force_stop: bool # ego vx > DBL_EPSILON (diffusion_planner_core.cpp:628-629)
|
| 49 |
+
prev_report: int # the latest TurnIndicatorsReport (turn_indicators[30])
|
| 50 |
+
|
| 51 |
+
|
| 52 |
+
def prepare(raw: Mapping[str, np.ndarray], observation: Mapping[str, Any]) -> Prepared:
|
| 53 |
+
"""``raw``: the 15 tensors of ``config.INPUT_SCHEMA`` (float32, batch 1, ego frame, before normalization);
|
| 54 |
+
``observation``: ``Normalization.observation``."""
|
| 55 |
+
missing = sorted(set(C.INPUT_NAMES) - set(raw))
|
| 56 |
+
if missing:
|
| 57 |
+
raise InputError(f"inputs: missing {missing}")
|
| 58 |
+
raw = {k: np.ascontiguousarray(np.asarray(raw[k], np.float32)) for k in C.INPUT_NAMES}
|
| 59 |
+
for k, shape in C.INPUT_SHAPES.items():
|
| 60 |
+
if raw[k].shape != shape:
|
| 61 |
+
raise InputError(f"inputs[{k!r}] has shape {raw[k].shape}, expected {shape}")
|
| 62 |
+
try:
|
| 63 |
+
norm = normalize_inputs(raw, observation)
|
| 64 |
+
except (KeyError, ValueError) as e: # a normalizer problem is a weights-file problem, not a client mistake
|
| 65 |
+
raise RuntimeError(f"normalization failed: {e}") from e
|
| 66 |
+
norm.update(speed_masks(norm))
|
| 67 |
+
feats = encoder_features(norm, norm)
|
| 68 |
+
dec = decoder_masks(norm)
|
| 69 |
+
nb = raw["neighbor_agents_past"][0].reshape(C.MAX_NUM_NEIGHBORS, -1)
|
| 70 |
+
rows = np.flatnonzero(np.any(nb != 0, axis=1))
|
| 71 |
+
vx = float(raw["ego_current_state"][0, 4])
|
| 72 |
+
prev_report = int(round(float(raw["turn_indicators"][0, C.INPUT_T])))
|
| 73 |
+
x_T = norm["sampled_trajectories"][0].copy()
|
| 74 |
+
return Prepared(raw, norm, feats, dec, x_T, rows, vx > np.finfo(np.float64).eps, prev_report)
|
| 75 |
+
|
| 76 |
+
|
| 77 |
+
def make_output(final_x0: np.ndarray, logit: np.ndarray, prepared: Prepared, normalization: Any,
|
| 78 |
+
params: Mapping[str, Any], *, model: str = "", denoising_steps: Optional[Sequence[np.ndarray]] = None,
|
| 79 |
+
timing_ms: Optional[Dict[str, float]] = None, meta: Optional[Dict[str, Any]] = None) -> Trajectory:
|
| 80 |
+
"""The node's post-processing of one plan -> ``Trajectory`` (base_link).
|
| 81 |
+
|
| 82 |
+
``final_x0``: ``[321, 81, 4]`` normalised (t = 0 = current state); ``logit``: ``[5]``; ``params``: the validated
|
| 83 |
+
``RUNTIME_PARAMS``; ``denoising_steps``: the 11 iterates ``[321, 81, 4]`` when ``return_denoising_steps``."""
|
| 84 |
+
x0 = np.asarray(final_x0, np.float32).reshape(C.MAX_NUM_AGENTS, C.OUTPUT_T + 1, C.POSE_DIM)
|
| 85 |
+
mean, std = normalization.state()
|
| 86 |
+
denorm = denormalize(x0, mean, std) # [321, 80, 4]
|
| 87 |
+
traj = trajectory_from_poses(denorm[0], (0.0, 0.0, 0.0),
|
| 88 |
+
velocity_smoothing_window=int(params["velocity_smoothing_window"]),
|
| 89 |
+
enable_force_stop=prepared.enable_force_stop,
|
| 90 |
+
stopping_threshold=float(params["stopping_threshold"]))
|
| 91 |
+
manager = TurnIndicatorManager(keep_offset=float(params["turn_indicator_keep_offset"]))
|
| 92 |
+
decision = manager.evaluate(np.asarray(logit, np.float32).reshape(-1), 0.0, prepared.prev_report)
|
| 93 |
+
rows = prepared.neighbor_rows
|
| 94 |
+
paths = predicted_paths(denorm, C.MAX_NUM_NEIGHBORS)[rows] if rows.size else np.zeros(
|
| 95 |
+
(0, C.OUTPUT_T, len(PREDICTED_AGENT_COLUMNS)), np.float32)
|
| 96 |
+
out_meta: Dict[str, Any] = {"predicted_agent_columns": list(PREDICTED_AGENT_COLUMNS),
|
| 97 |
+
"predicted_agent_rows": [int(r) for r in rows],
|
| 98 |
+
"force_stop": bool(traj.force_stop),
|
| 99 |
+
"time_from_start_s": [round(float(t), 3) for t in traj.time_from_start],
|
| 100 |
+
"valid_counts": prepared.features.counts()}
|
| 101 |
+
if params.get("return_denoising_steps") and denoising_steps is not None:
|
| 102 |
+
ego_steps = denoising_steps_ego(np.stack(denoising_steps), mean, std)
|
| 103 |
+
out_meta["denoising_steps"] = encode_array(ego_steps, key="denoising_steps")
|
| 104 |
+
out_meta.update(meta or {})
|
| 105 |
+
return Trajectory(traj.as_columns(), columns=TRAJECTORY_COLUMNS, turn_indicator=decision.to_dict(),
|
| 106 |
+
predicted_agents=paths, model=model, frame_id="base_link", timing_ms=dict(timing_ms or {}),
|
| 107 |
+
meta=out_meta)
|
code/tt_diffusion_planner/host/postprocess.py
ADDED
|
@@ -0,0 +1,222 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-License-Identifier: Apache-2.0
|
| 2 |
+
"""Post-processing of the Autoware node (``postprocessing_utils.cpp``, ``turn_indicator_manager.cpp``,
|
| 3 |
+
``diffusion_planner_core.cpp:605-694``), in the ego frame of the input tensors.
|
| 4 |
+
|
| 5 |
+
The node transforms every predicted pose to the map frame before building messages; the bundle has no map pose, so
|
| 6 |
+
it applies the same rules with an identity ego-to-map transform (``base_link`` output). Distances, velocities and
|
| 7 |
+
accelerations are invariant to that rigid transform; the published orientation is not exactly (see below).
|
| 8 |
+
|
| 9 |
+
Rules reproduced, with their numeric types:
|
| 10 |
+
|
| 11 |
+
- denormalisation: drop t = 0 (unless kept), ``v * std + mean`` per agent and pose dim, float32;
|
| 12 |
+
- poses: 4x4 matrices in double from the RAW (unnormalised) network cos / sin; the quaternion is Eigen's conversion
|
| 13 |
+
of that 3x3 block, which does not normalise, so for ``|(cos, sin)| != 1`` (0.84-1.09 in the goldens, SPEC 4.6.6)
|
| 14 |
+
the heading a consumer reads with ``tf2::getYaw`` differs from ``atan2(sin, cos)``; ``yaw`` here is that
|
| 15 |
+
``tf2::getYaw`` value of the identity-transform quaternion;
|
| 16 |
+
- ego trajectory (``get_trajectory_from_poses``): ``time_from_start = 0.1 (i + 1)``; velocity = 3-D distance to the
|
| 17 |
+
previous point (the first point to the ego position) / 0.1, stored as float; forward moving average over
|
| 18 |
+
``velocity_smoothing_window`` points summed in double; force stop (only when the ego moves: ``vx > DBL_EPSILON``)
|
| 19 |
+
once the smoothed velocity crosses below ``stopping_threshold``: velocity 0 and the pose frozen to the previous
|
| 20 |
+
pose; the last ``window - 1`` points keep the last smoothed velocity; acceleration = forward difference / 0.1 in
|
| 21 |
+
double stored as float, 0 for the last point;
|
| 22 |
+
- predicted objects: one 80-point path per emitted neighbour (the first ``n`` rows of the neighbour tensor, sorted
|
| 23 |
+
by distance), window 1, no force stop (only the poses are published);
|
| 24 |
+
- turn indicator (``TurnIndicatorManager::evaluate``): hold the last non-KEEP command for ``hold_duration`` (state
|
| 25 |
+
across calls), else add ``keep_offset`` to the KEEP logit, ``p_i = expf(l_i - max)``, ``p_i /= (1e-4f + sum)``,
|
| 26 |
+
argmax (first maximum); KEEP repeats the previous TurnIndicatorsReport, otherwise the command is the index;
|
| 27 |
+
- ``~/debug/denoising_steps``: the ego row of every solver iterate, denormalised with t = 0 kept.
|
| 28 |
+
"""
|
| 29 |
+
from __future__ import annotations
|
| 30 |
+
|
| 31 |
+
import math
|
| 32 |
+
from dataclasses import dataclass, field
|
| 33 |
+
from typing import Dict, Optional, Tuple
|
| 34 |
+
|
| 35 |
+
import numpy as np
|
| 36 |
+
|
| 37 |
+
from ..reference import config as C
|
| 38 |
+
from .solver import SCALAR_MATH
|
| 39 |
+
|
| 40 |
+
__all__ = ["denormalize", "quaternion_from_cos_sin", "tf2_yaw", "EgoTrajectory", "trajectory_from_poses",
|
| 41 |
+
"predicted_paths", "TurnIndicatorManager", "TurnDecision", "denoising_steps_ego"]
|
| 42 |
+
|
| 43 |
+
DBL_EPSILON = float(np.finfo(np.float64).eps)
|
| 44 |
+
|
| 45 |
+
|
| 46 |
+
def denormalize(x: np.ndarray, state_mean: np.ndarray, state_std: np.ndarray, *,
|
| 47 |
+
keep_current_state: bool = False) -> np.ndarray:
|
| 48 |
+
"""``denormalize_prediction``: ``x`` ``[321, 81, 4]`` (normalised, t = 0 = current state) -> ``[321, 80, 4]``
|
| 49 |
+
metres (``[321, 81, 4]`` with ``keep_current_state``). ``state_mean`` / ``state_std`` broadcast as
|
| 50 |
+
``[321 or 1, 1, 4]``."""
|
| 51 |
+
x = np.asarray(x, np.float32)
|
| 52 |
+
if x.shape[-3:] != (C.MAX_NUM_AGENTS, C.OUTPUT_T + 1, C.POSE_DIM):
|
| 53 |
+
raise ValueError(f"unsupported prediction shape {x.shape}")
|
| 54 |
+
p = x if keep_current_state else x[..., 1:, :]
|
| 55 |
+
return (p * np.asarray(state_std, np.float32) + np.asarray(state_mean, np.float32)).astype(np.float32)
|
| 56 |
+
|
| 57 |
+
|
| 58 |
+
def quaternion_from_cos_sin(cos_yaw: np.ndarray, sin_yaw: np.ndarray) -> np.ndarray:
|
| 59 |
+
"""Eigen ``Quaterniond(rotation_matrix)`` of ``[[c, -s, 0], [s, c, 0], [0, 0, 1]]`` (no normalisation), in double:
|
| 60 |
+
``[..., 4]`` = (x, y, z, w). Eigen's branch: trace = 2c + 1 > 0 -> w = sqrt(trace + 1) / 2, z = (m10 - m01) / (4 w);
|
| 61 |
+
else the largest diagonal is m22 -> z = sqrt(m22 - m00 - m11 + 1) / 2, w = (m10 - m01) / (4 z)."""
|
| 62 |
+
c = np.asarray(cos_yaw, np.float64)
|
| 63 |
+
s = np.asarray(sin_yaw, np.float64)
|
| 64 |
+
q = np.zeros(c.shape + (4,), np.float64)
|
| 65 |
+
trace = c + c + 1.0
|
| 66 |
+
pos = trace > 0.0
|
| 67 |
+
with np.errstate(divide="ignore", invalid="ignore"):
|
| 68 |
+
t = np.sqrt(trace + 1.0)
|
| 69 |
+
w_pos, z_pos = 0.5 * t, (s - (-s)) * (0.5 / t)
|
| 70 |
+
# trace <= 0: i = 2 (m22 = 1 > m00 = c, since c <= -0.5), j = 0, k = 1
|
| 71 |
+
t2 = np.sqrt(1.0 - c - c + 1.0)
|
| 72 |
+
z_neg, w_neg = 0.5 * t2, (s - (-s)) * (0.5 / t2)
|
| 73 |
+
q[..., 2] = np.where(pos, z_pos, z_neg)
|
| 74 |
+
q[..., 3] = np.where(pos, w_pos, w_neg)
|
| 75 |
+
return q
|
| 76 |
+
|
| 77 |
+
|
| 78 |
+
def tf2_yaw(q: np.ndarray) -> np.ndarray:
|
| 79 |
+
"""``tf2::getYaw`` (scale-invariant): ``atan2(2 (x y + w z), w^2 + x^2 - y^2 - z^2)``."""
|
| 80 |
+
x, y, z, w = (q[..., i] for i in range(4))
|
| 81 |
+
return np.arctan2(2.0 * (x * y + w * z), w * w + x * x - y * y - z * z)
|
| 82 |
+
|
| 83 |
+
|
| 84 |
+
@dataclass
|
| 85 |
+
class EgoTrajectory:
|
| 86 |
+
"""One trajectory as the node publishes it (in the ego frame here): positions [N, 3] (double), the raw cos / sin
|
| 87 |
+
of each pose [N, 2] (after force-stop freezing), quaternions [N, 4] (x, y, z, w), yaw [N] (``tf2::getYaw``),
|
| 88 |
+
``velocity`` / ``acceleration`` [N] float32, ``time_from_start`` [N] seconds, and whether force stop fired."""
|
| 89 |
+
|
| 90 |
+
position: np.ndarray
|
| 91 |
+
cos_sin: np.ndarray
|
| 92 |
+
quaternion: np.ndarray
|
| 93 |
+
yaw: np.ndarray
|
| 94 |
+
velocity: np.ndarray
|
| 95 |
+
acceleration: np.ndarray
|
| 96 |
+
time_from_start: np.ndarray
|
| 97 |
+
force_stop: bool
|
| 98 |
+
|
| 99 |
+
def as_columns(self) -> np.ndarray:
|
| 100 |
+
"""``[N, 7]`` float32: x, y, yaw, cos, sin, velocity, acceleration (the ``Trajectory.poses`` layout)."""
|
| 101 |
+
return np.stack([self.position[:, 0], self.position[:, 1], self.yaw, self.cos_sin[:, 0], self.cos_sin[:, 1],
|
| 102 |
+
self.velocity, self.acceleration], axis=-1).astype(np.float32)
|
| 103 |
+
|
| 104 |
+
|
| 105 |
+
def trajectory_from_poses(poses_xycs: np.ndarray, base_position: Tuple[float, float, float] = (0.0, 0.0, 0.0), *,
|
| 106 |
+
velocity_smoothing_window: int = C.VELOCITY_SMOOTHING_WINDOW,
|
| 107 |
+
enable_force_stop: bool = True,
|
| 108 |
+
stopping_threshold: float = C.STOPPING_THRESHOLD) -> EgoTrajectory:
|
| 109 |
+
"""``get_trajectory_from_poses`` (``postprocessing_utils.cpp:362-455``) for poses ``[N, 4]`` = denormalised
|
| 110 |
+
(x, y, cos, sin) float32 of one agent, positions at z = ``base_position[2]`` (identity transform)."""
|
| 111 |
+
p = np.asarray(poses_xycs, np.float32)
|
| 112 |
+
n = p.shape[0]
|
| 113 |
+
dt = C.TRAJECTORY_DT
|
| 114 |
+
pos = np.zeros((n, 3), np.float64)
|
| 115 |
+
pos[:, 0] = p[:, 0].astype(np.float64)
|
| 116 |
+
pos[:, 1] = p[:, 1].astype(np.float64)
|
| 117 |
+
pos[:, 2] = float(base_position[2])
|
| 118 |
+
cs = p[:, 2:4].astype(np.float64).copy()
|
| 119 |
+
quat = quaternion_from_cos_sin(cs[:, 0], cs[:, 1])
|
| 120 |
+
vel = np.zeros(n, np.float32)
|
| 121 |
+
prev = np.asarray(base_position, np.float64)
|
| 122 |
+
for i in range(n):
|
| 123 |
+
d = math.hypot(pos[i, 0] - prev[0], pos[i, 1] - prev[1], pos[i, 2] - prev[2])
|
| 124 |
+
vel[i] = np.float32(d / dt)
|
| 125 |
+
prev = pos[i]
|
| 126 |
+
w = int(velocity_smoothing_window)
|
| 127 |
+
if n <= w:
|
| 128 |
+
raise ValueError("velocity_smoothing_window must be smaller than number of points")
|
| 129 |
+
thr = np.float32(stopping_threshold)
|
| 130 |
+
force_stop = False
|
| 131 |
+
for i in range(0, n - w + 1):
|
| 132 |
+
acc = 0.0
|
| 133 |
+
for k in range(w):
|
| 134 |
+
acc += float(vel[i + k])
|
| 135 |
+
vel[i] = np.float32(acc / float(w))
|
| 136 |
+
if enable_force_stop and i > 0 and abs(vel[i - 1]) > thr and abs(vel[i]) < thr:
|
| 137 |
+
force_stop = True
|
| 138 |
+
if i > 0 and force_stop:
|
| 139 |
+
vel[i] = np.float32(0.0)
|
| 140 |
+
pos[i], cs[i], quat[i] = pos[i - 1], cs[i - 1], quat[i - 1]
|
| 141 |
+
last = vel[n - w]
|
| 142 |
+
for i in range(n - w + 1, n):
|
| 143 |
+
vel[i] = last
|
| 144 |
+
if force_stop:
|
| 145 |
+
vel[i] = np.float32(0.0)
|
| 146 |
+
pos[i], cs[i], quat[i] = pos[i - 1], cs[i - 1], quat[i - 1]
|
| 147 |
+
accel = np.zeros(n, np.float32)
|
| 148 |
+
for i in range(n - 1):
|
| 149 |
+
accel[i] = np.float32((float(vel[i + 1]) - float(vel[i])) / dt)
|
| 150 |
+
tfs = dt * (np.arange(n, dtype=np.float64) + 1.0)
|
| 151 |
+
return EgoTrajectory(pos, cs, quat, tf2_yaw(quat), vel, accel, tfs, force_stop)
|
| 152 |
+
|
| 153 |
+
|
| 154 |
+
def predicted_paths(denorm: np.ndarray, n_neighbors: int) -> np.ndarray:
|
| 155 |
+
"""``create_predicted_objects`` poses for the first ``n_neighbors`` neighbours: ``[n, 80, 5]`` float32 =
|
| 156 |
+
x, y, yaw (``tf2::getYaw``), cos, sin. ``denorm`` is ``[321, 80, 4]`` (agent 0 = ego)."""
|
| 157 |
+
nb = np.asarray(denorm, np.float32)[1:1 + int(n_neighbors)]
|
| 158 |
+
q = quaternion_from_cos_sin(nb[..., 2], nb[..., 3])
|
| 159 |
+
return np.stack([nb[..., 0], nb[..., 1], tf2_yaw(q).astype(np.float32), nb[..., 2], nb[..., 3]],
|
| 160 |
+
axis=-1).astype(np.float32)
|
| 161 |
+
|
| 162 |
+
|
| 163 |
+
def denoising_steps_ego(iterates: np.ndarray, state_mean: np.ndarray, state_std: np.ndarray) -> np.ndarray:
|
| 164 |
+
"""``create_denoising_steps_message`` (batch 1): the ego row of each iterate, denormalised with t = 0 kept:
|
| 165 |
+
``[steps, 81, 4]``. ``iterates`` is ``[steps, 321, 81, 4]`` or just the ego rows ``[steps, 1, 81, 4]`` (what the
|
| 166 |
+
device plan reads back); the float32 arithmetic is the same element by element (``x * std + mean``)."""
|
| 167 |
+
it = np.asarray(iterates, np.float32)
|
| 168 |
+
if it.ndim != 4 or it.shape[1] not in (1, C.MAX_NUM_AGENTS) or it.shape[2:] != (C.OUTPUT_T + 1, C.POSE_DIM):
|
| 169 |
+
raise ValueError(f"unsupported iterates shape {it.shape}")
|
| 170 |
+
shape = (C.MAX_NUM_AGENTS, 1, C.POSE_DIM)
|
| 171 |
+
mean0 = np.broadcast_to(np.asarray(state_mean, np.float32), shape)[0]
|
| 172 |
+
std0 = np.broadcast_to(np.asarray(state_std, np.float32), shape)[0]
|
| 173 |
+
return (it[:, 0] * std0 + mean0).astype(np.float32)
|
| 174 |
+
|
| 175 |
+
|
| 176 |
+
@dataclass
|
| 177 |
+
class TurnDecision:
|
| 178 |
+
command: int
|
| 179 |
+
logits: Tuple[float, ...]
|
| 180 |
+
probabilities: Tuple[float, ...]
|
| 181 |
+
keep_selected: bool
|
| 182 |
+
held: bool = False
|
| 183 |
+
|
| 184 |
+
def to_dict(self) -> Dict[str, object]:
|
| 185 |
+
name = C.TURN_INDICATOR_COMMAND_NAMES.get(self.command, str(self.command))
|
| 186 |
+
return {"command": int(self.command), "command_name": name, "keep_selected": bool(self.keep_selected),
|
| 187 |
+
"held": bool(self.held), "logits": [float(v) for v in self.logits],
|
| 188 |
+
"probabilities": [float(v) for v in self.probabilities]}
|
| 189 |
+
|
| 190 |
+
|
| 191 |
+
@dataclass
|
| 192 |
+
class TurnIndicatorManager:
|
| 193 |
+
"""``TurnIndicatorManager`` (stateful: the last non-KEEP command and its stamp). A fresh manager has no held
|
| 194 |
+
command, which is what a single stateless request sees."""
|
| 195 |
+
|
| 196 |
+
hold_duration_s: float = C.TURN_INDICATOR_HOLD_DURATION_S
|
| 197 |
+
keep_offset: float = C.TURN_INDICATOR_KEEP_OFFSET
|
| 198 |
+
_last_command: int = field(default=0, repr=False)
|
| 199 |
+
_last_stamp_s: Optional[float] = field(default=None, repr=False)
|
| 200 |
+
|
| 201 |
+
def evaluate(self, logit: np.ndarray, stamp_s: float = 0.0,
|
| 202 |
+
prev_report: int = C.TURN_INDICATORS_REPORT_DISABLE) -> TurnDecision:
|
| 203 |
+
lg = np.asarray(logit, np.float32).reshape(-1).copy()
|
| 204 |
+
raw = tuple(float(v) for v in lg)
|
| 205 |
+
if lg.size == 0:
|
| 206 |
+
return TurnDecision(1, raw, (), False) # TurnIndicatorsCommand::DISABLE
|
| 207 |
+
if self._last_stamp_s is not None and self._last_stamp_s > 0 and stamp_s <= self._last_stamp_s + float(
|
| 208 |
+
self.hold_duration_s):
|
| 209 |
+
return TurnDecision(self._last_command, raw, (), False, held=True)
|
| 210 |
+
lg[C.TURN_INDICATOR_OUTPUT_KEEP] = np.float32(lg[C.TURN_INDICATOR_OUTPUT_KEEP] + np.float32(self.keep_offset))
|
| 211 |
+
mx = lg.max()
|
| 212 |
+
prob = np.array([SCALAR_MATH.exp(np.float32(v - mx)) for v in lg], np.float32)
|
| 213 |
+
total = np.float32(0.0001)
|
| 214 |
+
for v in prob:
|
| 215 |
+
total = np.float32(total + v)
|
| 216 |
+
prob = (prob / total).astype(np.float32)
|
| 217 |
+
idx = int(np.argmax(prob)) # std::max_element: the first maximum
|
| 218 |
+
keep = idx == C.TURN_INDICATOR_OUTPUT_KEEP
|
| 219 |
+
command = (int(prev_report) & 0xFF) if keep else idx
|
| 220 |
+
if not keep:
|
| 221 |
+
self._last_command, self._last_stamp_s = command, float(stamp_s)
|
| 222 |
+
return TurnDecision(command, raw, tuple(float(v) for v in prob), keep)
|
code/tt_diffusion_planner/host/solver.py
ADDED
|
@@ -0,0 +1,273 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-License-Identifier: Apache-2.0
|
| 2 |
+
"""DPM-Solver++(2M) with denoise-to-zero, as Autoware runs it around the decoder (multi_step mode).
|
| 3 |
+
|
| 4 |
+
Port of ``PKG/src/inference/solver/dpm_solver.cpp`` (guidance disabled, Autoware's default and decision D11) and of
|
| 5 |
+
the prefix constraint of ``PKG/src/inference/multi_step_inference.cpp:300-340``. The scalar schedule is float32 math
|
| 6 |
+
evaluated with the C library's ``expf`` / ``logf`` / ``expm1f`` / ``sqrtf`` (what ``std::exp(float)`` etc. call in
|
| 7 |
+
the node) in the C++ operation order, so the timesteps and update coefficients are the node's bit for bit; numpy's
|
| 8 |
+
own float32 transcendentals differ from glibc in the last bit for many arguments. Without a loadable libm the module
|
| 9 |
+
falls back to numpy (``SCALAR_MATH`` says which).
|
| 10 |
+
|
| 11 |
+
For ``steps = 10``: 11 decoder evaluations at ``timesteps[0..9]`` and ``1/N = 0.001`` (denoise-to-zero); the 11
|
| 12 |
+
published iterates (``denoising_steps``) are the corrected ``x`` after each update and after the final evaluation.
|
| 13 |
+
The update formulas, with ``x`` and the model outputs float32 arrays and every scalar a float32::
|
| 14 |
+
|
| 15 |
+
first (step 1): x = (sigma_t / sigma_s) * x - (alpha_t * phi_1) * m
|
| 16 |
+
second (steps 2..): d = (m0 - m1) / r0
|
| 17 |
+
x = (sigma_t / sigma_0) * x - (alpha_t * phi_1) * m0 - (0.5 * (alpha_t * phi_1)) * d
|
| 18 |
+
|
| 19 |
+
:class:`SolverPlan` lists these scalars per update (the RT-dev coefficient tables of the on-device loop).
|
| 20 |
+
"""
|
| 21 |
+
from __future__ import annotations
|
| 22 |
+
|
| 23 |
+
import ctypes
|
| 24 |
+
import ctypes.util
|
| 25 |
+
from dataclasses import dataclass, field
|
| 26 |
+
from typing import Callable, List, Optional, Tuple
|
| 27 |
+
|
| 28 |
+
import numpy as np
|
| 29 |
+
|
| 30 |
+
from ..reference import config as C
|
| 31 |
+
|
| 32 |
+
__all__ = ["SCALAR_MATH", "marginal_log_mean_coeff", "marginal_alpha", "marginal_std", "marginal_lambda",
|
| 33 |
+
"inverse_lambda", "log_snr_timesteps", "SolverUpdate", "SolverPlan", "solver_plan", "first_update",
|
| 34 |
+
"second_update", "dpm_solver_sample", "apply_prefix_constraint", "SampleResult"]
|
| 35 |
+
|
| 36 |
+
f32 = np.float32
|
| 37 |
+
_T, _N = C.NOISE_SCHEDULE_T, C.NOISE_SCHEDULE_TOTAL_N
|
| 38 |
+
_B0, _B1 = C.NOISE_SCHEDULE_BETA0, C.NOISE_SCHEDULE_BETA1
|
| 39 |
+
|
| 40 |
+
|
| 41 |
+
class _ScalarMath:
|
| 42 |
+
"""float32 ``exp`` / ``log`` / ``expm1`` / ``sqrt`` of the C library, loaded on first use (numpy fallback)."""
|
| 43 |
+
|
| 44 |
+
def __init__(self) -> None:
|
| 45 |
+
self._fns = None
|
| 46 |
+
self.source = "unloaded"
|
| 47 |
+
|
| 48 |
+
def _load(self) -> None:
|
| 49 |
+
fns = None
|
| 50 |
+
try:
|
| 51 |
+
name = ctypes.util.find_library("m") or "libm.so.6"
|
| 52 |
+
libm = ctypes.CDLL(name)
|
| 53 |
+
fns = {}
|
| 54 |
+
for fn in ("expf", "logf", "expm1f", "sqrtf"):
|
| 55 |
+
f = getattr(libm, fn)
|
| 56 |
+
f.restype, f.argtypes = ctypes.c_float, [ctypes.c_float]
|
| 57 |
+
fns[fn] = f
|
| 58 |
+
self.source = f"libm ({name})"
|
| 59 |
+
except (OSError, AttributeError):
|
| 60 |
+
fns = None
|
| 61 |
+
self.source = "numpy (no C math library found)"
|
| 62 |
+
self._fns = fns or {}
|
| 63 |
+
|
| 64 |
+
def _call(self, fn: str, x: float, fallback: Callable) -> np.float32:
|
| 65 |
+
if self._fns is None:
|
| 66 |
+
self._load()
|
| 67 |
+
if fn in self._fns:
|
| 68 |
+
return f32(self._fns[fn](float(f32(x))))
|
| 69 |
+
return f32(fallback(f32(x)))
|
| 70 |
+
|
| 71 |
+
def exp(self, x):
|
| 72 |
+
return self._call("expf", x, np.exp)
|
| 73 |
+
|
| 74 |
+
def log(self, x):
|
| 75 |
+
return self._call("logf", x, np.log)
|
| 76 |
+
|
| 77 |
+
def expm1(self, x):
|
| 78 |
+
return self._call("expm1f", x, np.expm1)
|
| 79 |
+
|
| 80 |
+
def sqrt(self, x):
|
| 81 |
+
return self._call("sqrtf", x, np.sqrt)
|
| 82 |
+
|
| 83 |
+
|
| 84 |
+
SCALAR_MATH = _ScalarMath()
|
| 85 |
+
_m = SCALAR_MATH
|
| 86 |
+
|
| 87 |
+
|
| 88 |
+
# ---- noise schedule (dpm_solver.cpp:49-94), float32 in the C++ operation order ---------------------------------------
|
| 89 |
+
def marginal_log_mean_coeff(t) -> np.float32:
|
| 90 |
+
"""``-0.25f * t * t * (beta1 - beta0) - 0.5f * t * beta0``."""
|
| 91 |
+
t = f32(t)
|
| 92 |
+
return f32(f32(f32(f32(-0.25) * t) * t) * f32(_B1 - _B0)) - f32(f32(f32(0.5) * t) * _B0)
|
| 93 |
+
|
| 94 |
+
|
| 95 |
+
def marginal_alpha(t) -> np.float32:
|
| 96 |
+
return _m.exp(marginal_log_mean_coeff(t))
|
| 97 |
+
|
| 98 |
+
|
| 99 |
+
def marginal_std(t) -> np.float32:
|
| 100 |
+
return _m.sqrt(f32(f32(1.0) - _m.exp(f32(f32(2.0) * marginal_log_mean_coeff(t)))))
|
| 101 |
+
|
| 102 |
+
|
| 103 |
+
def marginal_lambda(t) -> np.float32:
|
| 104 |
+
lmc = marginal_log_mean_coeff(t)
|
| 105 |
+
log_std = f32(f32(0.5) * _m.log(f32(f32(1.0) - _m.exp(f32(f32(2.0) * lmc)))))
|
| 106 |
+
return f32(lmc - log_std)
|
| 107 |
+
|
| 108 |
+
|
| 109 |
+
def _log_add_exp(a, b) -> np.float32:
|
| 110 |
+
m = max(f32(a), f32(b))
|
| 111 |
+
return f32(m + _m.log(f32(_m.exp(f32(f32(a) - m)) + _m.exp(f32(f32(b) - m)))))
|
| 112 |
+
|
| 113 |
+
|
| 114 |
+
def inverse_lambda(lam) -> np.float32:
|
| 115 |
+
beta_delta = f32(_B1 - _B0)
|
| 116 |
+
tmp = f32(f32(f32(2.0) * beta_delta) * _log_add_exp(f32(f32(-2.0) * f32(lam)), f32(0.0)))
|
| 117 |
+
delta = f32(f32(_B0 * _B0) + tmp)
|
| 118 |
+
return f32(f32(tmp / f32(_m.sqrt(delta) + _B0)) / beta_delta)
|
| 119 |
+
|
| 120 |
+
|
| 121 |
+
def log_snr_timesteps(steps: int) -> List[np.float32]:
|
| 122 |
+
"""``steps + 1`` times uniform in log-SNR between t = 1 and t = 1/N (``dpm_solver.cpp:80-94``)."""
|
| 123 |
+
t0 = f32(f32(1.0) / _N)
|
| 124 |
+
lam_t, lam_0 = marginal_lambda(_T), marginal_lambda(t0)
|
| 125 |
+
out = []
|
| 126 |
+
for i in range(steps + 1):
|
| 127 |
+
ratio = f32(f32(i) / f32(steps))
|
| 128 |
+
out.append(inverse_lambda(f32(lam_t + f32(f32(lam_0 - lam_t) * ratio))))
|
| 129 |
+
return out
|
| 130 |
+
|
| 131 |
+
|
| 132 |
+
# ---- updates (dpm_solver.cpp:96-136) ------------------------------------------------------------------------------
|
| 133 |
+
def first_update(x_s: np.ndarray, model_s: np.ndarray, s, t) -> np.ndarray:
|
| 134 |
+
h = f32(marginal_lambda(t) - marginal_lambda(s))
|
| 135 |
+
sigma_s, sigma_t, alpha_t = marginal_std(s), marginal_std(t), marginal_alpha(t)
|
| 136 |
+
phi_1 = _m.expm1(f32(-h))
|
| 137 |
+
a, b = f32(sigma_t / sigma_s), f32(alpha_t * phi_1)
|
| 138 |
+
return (a * x_s.astype(np.float32) - b * model_s.astype(np.float32)).astype(np.float32)
|
| 139 |
+
|
| 140 |
+
|
| 141 |
+
def second_update(x_s: np.ndarray, model_prev: Tuple[np.ndarray, np.ndarray], t_prev: Tuple[float, float],
|
| 142 |
+
t) -> np.ndarray:
|
| 143 |
+
m1, m0 = model_prev # model_prev_list[0] (older), model_prev_list[1] (newer)
|
| 144 |
+
t1, t0 = t_prev
|
| 145 |
+
lam1, lam0, lam_t = marginal_lambda(t1), marginal_lambda(t0), marginal_lambda(t)
|
| 146 |
+
sigma0, sigma_t, alpha_t = marginal_std(t0), marginal_std(t), marginal_alpha(t)
|
| 147 |
+
h0, h = f32(lam0 - lam1), f32(lam_t - lam0)
|
| 148 |
+
r0 = f32(h0 / h)
|
| 149 |
+
phi_1 = _m.expm1(f32(-h))
|
| 150 |
+
a, b = f32(sigma_t / sigma0), f32(alpha_t * phi_1)
|
| 151 |
+
c = f32(f32(0.5) * b)
|
| 152 |
+
d1_0 = ((m0.astype(np.float32) - m1.astype(np.float32)) / r0).astype(np.float32)
|
| 153 |
+
return (a * x_s.astype(np.float32) - b * m0.astype(np.float32) - c * d1_0).astype(np.float32)
|
| 154 |
+
|
| 155 |
+
|
| 156 |
+
@dataclass(frozen=True)
|
| 157 |
+
class SolverUpdate:
|
| 158 |
+
"""One solver update to time ``t``: ``x = a * x - b * m0 - c * (m0 - m1) / r0`` (``c = 0`` for the first-order
|
| 159 |
+
update, which has no ``m1``)."""
|
| 160 |
+
|
| 161 |
+
order: int
|
| 162 |
+
t: float
|
| 163 |
+
a: float
|
| 164 |
+
b: float
|
| 165 |
+
c: float
|
| 166 |
+
r0: float
|
| 167 |
+
|
| 168 |
+
|
| 169 |
+
@dataclass(frozen=True)
|
| 170 |
+
class SolverPlan:
|
| 171 |
+
"""Everything the on-device loop needs for ``steps`` (the RT-dev tables of PLAN.md 2.12): the decoder evaluation
|
| 172 |
+
times (``eval_times[k]`` feeds evaluation ``k``; the t-embedding / adaLN tables are built for these values), the
|
| 173 |
+
update after each of the first ``steps`` evaluations, and the published iterate times."""
|
| 174 |
+
|
| 175 |
+
steps: int
|
| 176 |
+
timesteps: Tuple[float, ...]
|
| 177 |
+
eval_times: Tuple[float, ...]
|
| 178 |
+
updates: Tuple[SolverUpdate, ...]
|
| 179 |
+
denoising_timesteps: Tuple[float, ...]
|
| 180 |
+
scalar_math: str = field(default="")
|
| 181 |
+
|
| 182 |
+
|
| 183 |
+
def solver_plan(steps: int = C.DPM_SOLVER_STEPS) -> SolverPlan:
|
| 184 |
+
if steps < C.DPM_SOLVER_ORDER:
|
| 185 |
+
raise ValueError("DpmSolver steps must be greater than or equal to solver order.")
|
| 186 |
+
ts = log_snr_timesteps(steps)
|
| 187 |
+
updates = []
|
| 188 |
+
for step in range(1, steps + 1):
|
| 189 |
+
t = ts[step]
|
| 190 |
+
if step < C.DPM_SOLVER_ORDER:
|
| 191 |
+
s = _T if step == 1 else ts[step - 1]
|
| 192 |
+
h = f32(marginal_lambda(t) - marginal_lambda(s))
|
| 193 |
+
a = f32(marginal_std(t) / marginal_std(s))
|
| 194 |
+
b = f32(marginal_alpha(t) * _m.expm1(f32(-h)))
|
| 195 |
+
updates.append(SolverUpdate(1, float(t), float(a), float(b), 0.0, 1.0))
|
| 196 |
+
else:
|
| 197 |
+
t1 = _T if step - 2 == 0 else ts[step - 2]
|
| 198 |
+
t0 = ts[step - 1]
|
| 199 |
+
lam1, lam0, lam_t = marginal_lambda(t1), marginal_lambda(t0), marginal_lambda(t)
|
| 200 |
+
h0, h = f32(lam0 - lam1), f32(lam_t - lam0)
|
| 201 |
+
b = f32(marginal_alpha(t) * _m.expm1(f32(-h)))
|
| 202 |
+
updates.append(SolverUpdate(2, float(t), float(f32(marginal_std(t) / marginal_std(t0))), float(b),
|
| 203 |
+
float(f32(f32(0.5) * b)), float(f32(h0 / h))))
|
| 204 |
+
final_t = f32(f32(1.0) / _N)
|
| 205 |
+
return SolverPlan(steps, tuple(float(t) for t in ts), tuple(float(t) for t in ts[:steps]) + (float(final_t),),
|
| 206 |
+
tuple(updates), tuple(float(t) for t in ts[1:]) + (float(final_t),), SCALAR_MATH.source)
|
| 207 |
+
|
| 208 |
+
|
| 209 |
+
# ---- the loop (dpm_solver.cpp:138-234) ---------------------------------------------------------------------------
|
| 210 |
+
@dataclass
|
| 211 |
+
class SampleResult:
|
| 212 |
+
final_x: np.ndarray
|
| 213 |
+
denoising_steps: List[np.ndarray]
|
| 214 |
+
denoising_timesteps: List[float]
|
| 215 |
+
eval_times: List[float]
|
| 216 |
+
nfe: int
|
| 217 |
+
|
| 218 |
+
|
| 219 |
+
def apply_prefix_constraint(x: np.ndarray, current_states: np.ndarray) -> np.ndarray:
|
| 220 |
+
"""``x[agent, 0, :] = current_states[agent]`` in place (``multi_step_inference.cpp:327-340``); ``x`` is
|
| 221 |
+
``[321, 81, 4]`` (or with a leading batch dim)."""
|
| 222 |
+
x[..., 0, :] = current_states
|
| 223 |
+
return x
|
| 224 |
+
|
| 225 |
+
|
| 226 |
+
def dpm_solver_sample(initial_x: np.ndarray, model_fn: Callable[[np.ndarray, np.float32], np.ndarray],
|
| 227 |
+
correcting_fn: Callable[[np.ndarray], None], steps: int = C.DPM_SOLVER_STEPS,
|
| 228 |
+
on_iterate: Optional[Callable[[int, np.ndarray], None]] = None) -> SampleResult:
|
| 229 |
+
"""``DpmSolver::sample`` without guidance. ``model_fn(x, t)`` is one decoder evaluation (x0 prediction),
|
| 230 |
+
``correcting_fn(x)`` the in-place prefix constraint."""
|
| 231 |
+
if steps < C.DPM_SOLVER_ORDER:
|
| 232 |
+
raise ValueError("DpmSolver steps must be greater than or equal to solver order.")
|
| 233 |
+
x = np.array(initial_x, dtype=np.float32, copy=True)
|
| 234 |
+
correcting_fn(x)
|
| 235 |
+
ts = log_snr_timesteps(steps)
|
| 236 |
+
steps_out: List[np.ndarray] = []
|
| 237 |
+
step_ts: List[float] = []
|
| 238 |
+
eval_times: List[float] = []
|
| 239 |
+
|
| 240 |
+
def evaluate(xx: np.ndarray, t) -> np.ndarray:
|
| 241 |
+
eval_times.append(float(t))
|
| 242 |
+
return np.asarray(model_fn(xx, f32(t)), np.float32)
|
| 243 |
+
|
| 244 |
+
def record(xx: np.ndarray, t) -> None:
|
| 245 |
+
steps_out.append(xx.copy())
|
| 246 |
+
step_ts.append(float(t))
|
| 247 |
+
if on_iterate is not None:
|
| 248 |
+
on_iterate(len(steps_out) - 1, xx)
|
| 249 |
+
|
| 250 |
+
t_prev = [_T]
|
| 251 |
+
m_prev = [evaluate(x, ts[0])]
|
| 252 |
+
for step in range(1, C.DPM_SOLVER_ORDER):
|
| 253 |
+
t = ts[step]
|
| 254 |
+
x = first_update(x, m_prev[-1], t_prev[-1], t)
|
| 255 |
+
correcting_fn(x)
|
| 256 |
+
record(x, t)
|
| 257 |
+
t_prev.append(t)
|
| 258 |
+
m_prev.append(evaluate(x, t))
|
| 259 |
+
for step in range(C.DPM_SOLVER_ORDER, steps + 1):
|
| 260 |
+
t = ts[step]
|
| 261 |
+
x = second_update(x, (m_prev[0], m_prev[1]), (t_prev[0], t_prev[1]), t)
|
| 262 |
+
correcting_fn(x)
|
| 263 |
+
record(x, t)
|
| 264 |
+
t_prev = [t_prev[1], t]
|
| 265 |
+
m_prev = [m_prev[1], None]
|
| 266 |
+
if step < steps:
|
| 267 |
+
m_prev[1] = evaluate(x, t)
|
| 268 |
+
final_t = f32(f32(1.0) / _N)
|
| 269 |
+
x = evaluate(x, final_t) # denoise to zero: the model output itself, no guidance on the last step
|
| 270 |
+
x = np.array(x, dtype=np.float32, copy=True)
|
| 271 |
+
correcting_fn(x)
|
| 272 |
+
record(x, final_t)
|
| 273 |
+
return SampleResult(x, steps_out, step_ts, eval_times, len(eval_times))
|
code/tt_diffusion_planner/io.py
ADDED
|
@@ -0,0 +1,32 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-License-Identifier: Apache-2.0
|
| 2 |
+
"""Input decoding and output encoding shared by the Python API and the HTTP server of diffusion-planner-p150.
|
| 3 |
+
|
| 4 |
+
The implementation is the vendored ``ttaw.io`` (C08); numpy only, data parsing only, and every client mistake raises
|
| 5 |
+
:class:`InputError`, which the server maps to HTTP 400. This model's only input is ``inputs``: the 15 ONNX-named
|
| 6 |
+
planner tensors of the node's ``create_input_data()`` (``reference.config.INPUT_SCHEMA``), as a ``{name: array}``
|
| 7 |
+
mapping, an ``.npz`` path or its bytes, or the JSON envelope ``{"format": "npz", "data": <base64>}`` /
|
| 8 |
+
``{"format": "json", "arrays": {...}}``. :func:`load_inputs` / :func:`decode_inputs` bind the schema, so
|
| 9 |
+
``model(inputs=...)`` and ``POST /predict`` accept and refuse the same inputs (names, shapes, finite values).
|
| 10 |
+
"""
|
| 11 |
+
from __future__ import annotations
|
| 12 |
+
|
| 13 |
+
from typing import Any, Mapping, Optional
|
| 14 |
+
|
| 15 |
+
from .reference.config import INPUT_SCHEMA
|
| 16 |
+
from .ttaw import io as _io
|
| 17 |
+
from .ttaw.io import * # noqa: F401,F403 (the decoders and encoders: ttaw/API.md section 9)
|
| 18 |
+
|
| 19 |
+
__all__ = list(_io.__all__) + ["DEFAULT_POINT_FIELDS", "INPUT_SCHEMA", "load_inputs", "decode_inputs"]
|
| 20 |
+
|
| 21 |
+
DEFAULT_POINT_FIELDS = () # no point cloud input
|
| 22 |
+
|
| 23 |
+
|
| 24 |
+
def load_inputs(source: Any) -> dict:
|
| 25 |
+
"""Python-API planner input (mapping, ``.npz`` path or bytes, JSON envelope) -> ``{name: float32 array}``
|
| 26 |
+
checked against ``INPUT_SCHEMA``."""
|
| 27 |
+
return _io.load_named_arrays(source, INPUT_SCHEMA)
|
| 28 |
+
|
| 29 |
+
|
| 30 |
+
def decode_inputs(spec: Mapping[str, Any], *, max_bytes: Optional[int] = None) -> dict:
|
| 31 |
+
"""The ``inputs`` envelope of ``/predict`` -> ``{name: float32 array}`` checked against ``INPUT_SCHEMA``."""
|
| 32 |
+
return _io.decode_named_arrays(spec, INPUT_SCHEMA, max_bytes=max_bytes)
|
code/tt_diffusion_planner/reference/__init__.py
ADDED
|
@@ -0,0 +1,34 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-License-Identifier: Apache-2.0
|
| 2 |
+
"""CPU reference of Diffusion Planner v5.0 (Autoware diffusion_planner) -- the ground truth every PCC /
|
| 3 |
+
output-agreement gate of the TT port compares against. Importable without ttnn (torch, numpy, onnx only).
|
| 4 |
+
|
| 5 |
+
- ``config.py`` dimensions, token layout, constants and node parameters, each with its Autoware / ONNX source.
|
| 6 |
+
- ``weights.py`` the three ONNX files and ``diffusion_planner.param.json`` read as DATA with the vendored
|
| 7 |
+
``ttaw.weights.OnnxWeights``, every tensor addressed through its consuming node, into one canonical
|
| 8 |
+
``{name: float32}`` dict that the reference and the ttnn graph share (no BatchNorm to fold).
|
| 9 |
+
- ``model.py`` pure-PyTorch fp32 encoder / DiT decoder / turn head with per-module taps.
|
| 10 |
+
- ``rewrites.py`` the exact graph rewrites the TT port applies (per-step adaLN tables folded into the LayerNorm
|
| 11 |
+
affine, hoisted cross-attention K/V, the pad-relative fp32 pre-projection island) as float64-built
|
| 12 |
+
constants plus CPU forwards that use them, tested against ``model.py``.
|
| 13 |
+
- ``pipeline.py`` ``ReferencePlanner``: host pre-processing (``..host``) -> encoder -> DPM-Solver++(2M) loop over 11
|
| 14 |
+
decoder evaluations -> turn head -> host post-processing; ``run()`` records taps for goldens.
|
| 15 |
+
- ``ort.py`` ONNX Runtime on the shipped ONNX (the reference's oracle; research venv and tests only).
|
| 16 |
+
- ``goldens.py`` golden generation (taps + outputs per scene) and the small goldens kept in ``tests/goldens``.
|
| 17 |
+
|
| 18 |
+
The ``to_dict()`` of ``ReferencePlanner()(inputs=sample)`` is stored as ``samples/<stem>.reference.json``:
|
| 19 |
+
``server/smoke_test.py`` compares the served output with it.
|
| 20 |
+
"""
|
| 21 |
+
|
| 22 |
+
__all__ = ["ReferencePlanner", "load_weights", "find_weights_dir"]
|
| 23 |
+
|
| 24 |
+
|
| 25 |
+
def __getattr__(name):
|
| 26 |
+
if name == "ReferencePlanner":
|
| 27 |
+
from .pipeline import ReferencePlanner
|
| 28 |
+
|
| 29 |
+
return ReferencePlanner
|
| 30 |
+
if name in ("load_weights", "find_weights_dir"):
|
| 31 |
+
from . import weights
|
| 32 |
+
|
| 33 |
+
return getattr(weights, name)
|
| 34 |
+
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
code/tt_diffusion_planner/reference/config.py
ADDED
|
@@ -0,0 +1,175 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-License-Identifier: Apache-2.0
|
| 2 |
+
"""Constants of the Autoware Diffusion Planner v5.0 contract, each with its source.
|
| 3 |
+
|
| 4 |
+
Shared by the host pre/post-processing (``tt_diffusion_planner.host``), the CPU reference (``reference``) and the
|
| 5 |
+
ttnn graph (``tt``). Every value that changes a device shape is a COMPILE parameter (PLAN.md section 0.4): a change
|
| 6 |
+
means a new trace and a new image.
|
| 7 |
+
|
| 8 |
+
Source abbreviations: ``PKG`` = autoware_universe @ 9ceaccf ``planning/autoware_diffusion_planner``; ``HFD`` = the
|
| 9 |
+
weights ``AutowareFoundation/diffusion_planner@v5.0`` (423efde6); ``SPEC`` = ``research/diffusion-planner/SPEC.md`` of
|
| 10 |
+
the porting workspace; ``T4M`` = tier4/Diffusion-Planner @ 40114a8 ``diffusion_planner/diffusion_planner`` (read
|
| 11 |
+
for semantics only, never imported).
|
| 12 |
+
"""
|
| 13 |
+
from __future__ import annotations
|
| 14 |
+
|
| 15 |
+
from typing import Dict, Tuple
|
| 16 |
+
|
| 17 |
+
import numpy as np
|
| 18 |
+
|
| 19 |
+
# ---- dimensions (PKG/include/autoware/diffusion_planner/dimensions.hpp:26-105) --------------------------------------
|
| 20 |
+
NUM_SEGMENTS_IN_LANE = 140
|
| 21 |
+
NUM_SEGMENTS_IN_ROUTE = 25
|
| 22 |
+
NUM_POLYGONS = 10
|
| 23 |
+
NUM_LINE_STRINGS = 60
|
| 24 |
+
NUM_STATIC_OBJECTS = 5
|
| 25 |
+
MAX_NUM_NEIGHBORS = 320
|
| 26 |
+
MAX_NUM_AGENTS = MAX_NUM_NEIGHBORS + 1 # ego + neighbours
|
| 27 |
+
HIDDEN_DIM = 256
|
| 28 |
+
POINTS_PER_SEGMENT = 20
|
| 29 |
+
POINTS_PER_POLYGON = 40
|
| 30 |
+
POINTS_PER_LINE_STRING = 20
|
| 31 |
+
LINE_TYPE_NUM = 10
|
| 32 |
+
POLYGON_TYPE_NUM = 1 # intersection_area
|
| 33 |
+
LINE_STRING_TYPE_NUM = 2 # stop_line, road_border
|
| 34 |
+
SEGMENT_POINT_DIM = 13 + 2 * LINE_TYPE_NUM # 33
|
| 35 |
+
INPUT_T = 30 # history steps before the current one (31 samples with the current)
|
| 36 |
+
OUTPUT_T = 80 # future steps (8 s at 0.1 s)
|
| 37 |
+
POSE_DIM = 4 # x, y, cos(yaw), sin(yaw)
|
| 38 |
+
AGENT_STATE_DIM = 11 # x, y, cos, sin, vx, vy, width, length, is_vehicle, is_pedestrian, is_bicycle
|
| 39 |
+
EGO_CURRENT_STATE_DIM = 10
|
| 40 |
+
STATIC_OBJECT_DIM = 10
|
| 41 |
+
EGO_SHAPE_DIM = 3 # wheel_base, length, width
|
| 42 |
+
TURN_INDICATOR_OUTPUT_DIM = 5
|
| 43 |
+
# logit order (dimensions.hpp:74-79); the published command is the index for 0..3 (TurnIndicatorsCommand)
|
| 44 |
+
TURN_INDICATOR_LABELS = ("NONE", "DISABLE", "ENABLE_LEFT", "ENABLE_RIGHT", "KEEP")
|
| 45 |
+
TURN_INDICATOR_OUTPUT_KEEP = 4
|
| 46 |
+
TURN_INDICATOR_COMMAND_NAMES = {0: "NO_COMMAND", 1: "DISABLE", 2: "ENABLE_LEFT", 3: "ENABLE_RIGHT"}
|
| 47 |
+
TURN_INDICATORS_REPORT_DISABLE = 1 # TurnIndicatorsReport::DISABLE, the prev_report default (core.cpp:635-637)
|
| 48 |
+
|
| 49 |
+
# ---- the encoder's 564 scene tokens, in concatenation order (T4M/model/module/encoder.py:251-265; SPEC 2.5) -------
|
| 50 |
+
TOKEN_LAYOUT: Tuple[Tuple[str, int], ...] = (
|
| 51 |
+
("ego", 1),
|
| 52 |
+
("neighbor", MAX_NUM_NEIGHBORS),
|
| 53 |
+
("static", NUM_STATIC_OBJECTS),
|
| 54 |
+
("lane", NUM_SEGMENTS_IN_LANE),
|
| 55 |
+
("route", NUM_SEGMENTS_IN_ROUTE),
|
| 56 |
+
("polygon", NUM_POLYGONS),
|
| 57 |
+
("line_string", NUM_LINE_STRINGS),
|
| 58 |
+
("goal", 1),
|
| 59 |
+
("ego_shape", 1),
|
| 60 |
+
("turn", 1),
|
| 61 |
+
)
|
| 62 |
+
ENCODING_TOKEN_NUM = sum(n for _, n in TOKEN_LAYOUT) # 564 (dimensions.hpp:33-35)
|
| 63 |
+
|
| 64 |
+
|
| 65 |
+
def token_slices() -> Dict[str, slice]:
|
| 66 |
+
"""``{category: slice of the 564 tokens}``."""
|
| 67 |
+
out, start = {}, 0
|
| 68 |
+
for name, n in TOKEN_LAYOUT:
|
| 69 |
+
out[name] = slice(start, start + n)
|
| 70 |
+
start += n
|
| 71 |
+
return out
|
| 72 |
+
|
| 73 |
+
|
| 74 |
+
TOKEN_SLICES = token_slices()
|
| 75 |
+
|
| 76 |
+
# class ids of the 14-dim positional feature (T4M encoder.py:10-21). The turn-indicator token reuses the ego-shape
|
| 77 |
+
# id 8 (FloatsEncoder hard-codes CLASS_TYPE_EGO_SHAPE, encoder.py:808; confirmed by the ONNX ConstantOfShape values).
|
| 78 |
+
POS_CLASS = {"ego": 0, "neighbor": 1, "static": 2, "lane": 3, "route": 4, "polygon": 5, "line_string": 6,
|
| 79 |
+
"goal": 7, "ego_shape": 8, "turn": 8}
|
| 80 |
+
POS_CLASS_NUM = 10
|
| 81 |
+
POS_FEATURE_DIM = 4 + POS_CLASS_NUM # 14
|
| 82 |
+
|
| 83 |
+
# ---- in-graph pre-processing of the encoder (SPEC 3.8; T4M encoder.py) ------------------------------------------
|
| 84 |
+
EGO_HISTORY_KEEP = slice(0, 6) # the 6 OLDEST ego samples are kept, rows 6..30 zeroed (encoder.py:170-175)
|
| 85 |
+
NEIGHBOR_HISTORY_KEEP = slice(INPUT_T + 1 - 6, INPUT_T + 1) # the 6 newest neighbour samples (rows 25..30)
|
| 86 |
+
TURN_INDICATOR_HISTORY = INPUT_T # turn_indicators[:, :-1]: the 30 values before the current report (encoder.py:209)
|
| 87 |
+
LANE_POS_INDEX = POINTS_PER_SEGMENT // 2 # 10: point used for the lane / route position feature (encoder.py:592)
|
| 88 |
+
POLYGON_POS_INDEX = POINTS_PER_POLYGON // 2 # 20 (LineEncoder, encoder.py:685)
|
| 89 |
+
LINE_STRING_POS_INDEX = POINTS_PER_LINE_STRING // 2 # 10
|
| 90 |
+
LANE_FEATURE_DIM = 8 # x, y, dx, dy, left - centre (x, y), right - centre (x, y)
|
| 91 |
+
LANE_ATTRIBUTE_DIM = SEGMENT_POINT_DIM - LANE_FEATURE_DIM # 25: traffic light (5) + line types (2 x 10) of point 0
|
| 92 |
+
NEIGHBOR_FEATURE_DIM = 9 # x, y, cos, sin, 0, 0 (velocities zeroed), width, length, valid-step flag
|
| 93 |
+
POLYGON_FEATURE_DIM = 2 + POLYGON_TYPE_NUM + 2 # x, y, type, dx, dy
|
| 94 |
+
LINE_STRING_FEATURE_DIM = 2 + LINE_STRING_TYPE_NUM + 2 # x, y, stop_line, road_border, dx, dy
|
| 95 |
+
|
| 96 |
+
# ---- network constants (ONNX, SPEC 4.1-4.4) ------------------------------------------------------------------------
|
| 97 |
+
LN_EPS = 1e-5 # every LayerNormalization of the three graphs (109 nodes; ttnn's default is 1e-12: pass it explicitly)
|
| 98 |
+
NUM_HEADS = 8
|
| 99 |
+
HEAD_DIM = HIDDEN_DIM // NUM_HEADS # 32
|
| 100 |
+
ATTN_SCALE = np.float32(0.17677669) # ONNX constant Mul_3 / Mul_2 = 1/sqrt(32) (applied to Q before Q.K^T)
|
| 101 |
+
MIXER_TOKENS = 64 # token_pre_project out, tokens_mlp width
|
| 102 |
+
MIXER_CHANNELS = 128 # channel_pre_project out, channels_mlp width
|
| 103 |
+
MIXER_DEPTH = 6 # encoder_mixer_depth (HFD param.json)
|
| 104 |
+
FUSION_DEPTH = 6 # encoder_fusion_depth
|
| 105 |
+
FUSION_MLP_DIM = 4 * HIDDEN_DIM
|
| 106 |
+
DIT_DEPTH = 3 # decoder_depth
|
| 107 |
+
DIT_MLP_DIM = 4 * HIDDEN_DIM
|
| 108 |
+
DIT_INPUT_DIM = (OUTPUT_T + 1) * POSE_DIM # 324 = preproj input / final projection output
|
| 109 |
+
DIT_TIME_DIM = OUTPUT_T + 1 # 81 = t_embedder input (one diffusion time per trajectory point)
|
| 110 |
+
TURN_HEAD_STEPS = tuple(range(1, OUTPUT_T, 10)) # final_x0[0, 1::10, :2] -> 8 points, 16 values (SPEC 4.4)
|
| 111 |
+
|
| 112 |
+
# ---- DPM-Solver++(2M) (PKG/src/inference/solver/dpm_solver.cpp:29-33; PKG/config/diffusion_planner.param.yaml) -----
|
| 113 |
+
# yaml l.17 model.multi_step_model.dpm_solver_steps: NFE = steps + 1 (COMPILE: the loop length of the trace)
|
| 114 |
+
DPM_SOLVER_STEPS = 10
|
| 115 |
+
DPM_SOLVER_ORDER = 2
|
| 116 |
+
NOISE_SCHEDULE_T = np.float32(1.0)
|
| 117 |
+
NOISE_SCHEDULE_TOTAL_N = np.float32(1000.0)
|
| 118 |
+
NOISE_SCHEDULE_BETA0 = np.float32(0.1)
|
| 119 |
+
NOISE_SCHEDULE_BETA1 = np.float32(20.0)
|
| 120 |
+
# SPEC 5.1 (recomputed from dpm_solver.cpp:80-94): the 11 solver timesteps for steps = 10, float32
|
| 121 |
+
DPM_TIMESTEPS_STEPS10 = (1.0, 0.89912426, 0.78557116, 0.65344393, 0.49344844, 0.30464348, 0.14064588, 0.05360911,
|
| 122 |
+
0.01809745, 0.00499264, 0.00100062)
|
| 123 |
+
|
| 124 |
+
# ---- node parameters this port reproduces (PKG/config/diffusion_planner.param.yaml, effective YAML defaults) -------
|
| 125 |
+
VELOCITY_SMOOTHING_WINDOW = 8 # yaml l.32
|
| 126 |
+
STOPPING_THRESHOLD = 0.3 # yaml l.33, m/s
|
| 127 |
+
TURN_INDICATOR_KEEP_OFFSET = -1.25 # yaml l.34
|
| 128 |
+
TURN_INDICATOR_HOLD_DURATION_S = 1.0 # yaml l.35 (stateful; see host.postprocess.TurnIndicatorManager)
|
| 129 |
+
TRAJECTORY_DT = 0.1 # postprocessing_utils.cpp:370 (constexpr double dt)
|
| 130 |
+
DELAY_STEP_MAX = OUTPUT_T // 2 # core.cpp:437 clamps delay_step to [0, 40]
|
| 131 |
+
|
| 132 |
+
# ---- raw input tensors: the InputDataMap of DiffusionPlannerCore::create_input_data (core.cpp:414-595), batch 1 ----
|
| 133 |
+
# All float32 and in the ego (base_link) frame, BEFORE normalization. ``sampled_trajectories`` is already in the
|
| 134 |
+
# normalised state space (x_T: zeros with the default temperature 0), ``delay`` is only read by the single-step graph.
|
| 135 |
+
# The speed-limit masks (lanes_has_speed_limit, route_lanes_has_speed_limit) are not inputs: the inference backend
|
| 136 |
+
# derives them from the normalised speed limits (PKG/include/.../inference/utils.hpp:112-123).
|
| 137 |
+
INPUT_SHAPES: Dict[str, Tuple[int, ...]] = {
|
| 138 |
+
"sampled_trajectories": (1, MAX_NUM_AGENTS, OUTPUT_T + 1, POSE_DIM),
|
| 139 |
+
"ego_agent_past": (1, INPUT_T + 1, POSE_DIM),
|
| 140 |
+
"ego_current_state": (1, EGO_CURRENT_STATE_DIM),
|
| 141 |
+
"neighbor_agents_past": (1, MAX_NUM_NEIGHBORS, INPUT_T + 1, AGENT_STATE_DIM),
|
| 142 |
+
"static_objects": (1, NUM_STATIC_OBJECTS, STATIC_OBJECT_DIM),
|
| 143 |
+
"lanes": (1, NUM_SEGMENTS_IN_LANE, POINTS_PER_SEGMENT, SEGMENT_POINT_DIM),
|
| 144 |
+
"lanes_speed_limit": (1, NUM_SEGMENTS_IN_LANE, 1),
|
| 145 |
+
"route_lanes": (1, NUM_SEGMENTS_IN_ROUTE, POINTS_PER_SEGMENT, SEGMENT_POINT_DIM),
|
| 146 |
+
"route_lanes_speed_limit": (1, NUM_SEGMENTS_IN_ROUTE, 1),
|
| 147 |
+
"polygons": (1, NUM_POLYGONS, POINTS_PER_POLYGON, 2 + POLYGON_TYPE_NUM),
|
| 148 |
+
"line_strings": (1, NUM_LINE_STRINGS, POINTS_PER_LINE_STRING, 2 + LINE_STRING_TYPE_NUM),
|
| 149 |
+
"goal_pose": (1, POSE_DIM),
|
| 150 |
+
"ego_shape": (1, EGO_SHAPE_DIM),
|
| 151 |
+
"turn_indicators": (1, INPUT_T + 1),
|
| 152 |
+
"delay": (1, 1),
|
| 153 |
+
}
|
| 154 |
+
INPUT_NAMES = tuple(INPUT_SHAPES)
|
| 155 |
+
# name -> (shape, dtype): the ModelBase.INPUT_SCHEMA of the API and the server (ttaw API.md section 10)
|
| 156 |
+
INPUT_SCHEMA = {name: (shape, np.float32) for name, shape in INPUT_SHAPES.items()}
|
| 157 |
+
# preprocessing_utils.cpp:34-84 normalises every tensor except these (core.cpp / node.cpp:633-634)
|
| 158 |
+
SKIP_NORMALIZATION = ("ego_shape", "sampled_trajectories", "turn_indicators", "delay")
|
| 159 |
+
# the inputs of the encoder graph, in its input order (HFD diffusion_planner_encoder.onnx)
|
| 160 |
+
ENCODER_INPUTS = ("ego_agent_past", "neighbor_agents_past", "static_objects", "lanes", "lanes_speed_limit",
|
| 161 |
+
"lanes_has_speed_limit", "route_lanes", "route_lanes_speed_limit", "route_lanes_has_speed_limit",
|
| 162 |
+
"polygons", "line_strings", "goal_pose", "ego_shape", "turn_indicators")
|
| 163 |
+
|
| 164 |
+
# ---- weight files (HFD; BUNDLE_CONVENTIONS.md section 12) -----------------------------------------------------------
|
| 165 |
+
ENCODER_ONNX = "diffusion_planner_encoder.onnx"
|
| 166 |
+
DECODER_ONNX = "diffusion_planner_decoder.onnx"
|
| 167 |
+
TURN_INDICATOR_ONNX = "diffusion_planner_turn_indicator.onnx"
|
| 168 |
+
PARAM_JSON = "diffusion_planner.param.json"
|
| 169 |
+
WEIGHT_MAJOR_VERSION = 5 # PKG/include/autoware/diffusion_planner/constants.hpp:22 (arg_reader.hpp:54-78)
|
| 170 |
+
FILE_SHA256 = { # SPEC 1.4
|
| 171 |
+
ENCODER_ONNX: "2856886a3ed63b963cb18b457876d0bbd8916d345549ee5f199ceb49e3fcca49",
|
| 172 |
+
DECODER_ONNX: "eb30c0c0c8e80b8460d293d600c7b8c9e21f615ff6ce8de5ed01670c3d1057ca",
|
| 173 |
+
TURN_INDICATOR_ONNX: "07acfb58a5de1f25587fa86deb39a6de5605f8d93518e97f960ee27145f4a732",
|
| 174 |
+
PARAM_JSON: "ee3145b68fd1e1e44e532933dfe66cfee4384fbd637382c87ab5190c66a8e268",
|
| 175 |
+
}
|
code/tt_diffusion_planner/reference/goldens.py
ADDED
|
@@ -0,0 +1,162 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-License-Identifier: Apache-2.0
|
| 2 |
+
"""Golden tensors of the CPU reference: per-module taps and final outputs of one planning scene.
|
| 3 |
+
|
| 4 |
+
Two products per scene (``code/scripts/ref_golden.py`` writes both):
|
| 5 |
+
|
| 6 |
+
- the **full goldens** (``<dir>/<scene>.npz``, 10-25 MB, kept OUT of the bundle, by default under
|
| 7 |
+
``research/diffusion-planner/goldens``): the raw inputs, the host features, every encoder tap (mixer taps compacted
|
| 8 |
+
to the valid entities, with their row indices), the encoding, for each of the 11 decoder evaluations its input
|
| 9 |
+
``x`` (all 321 agents, for teacher forcing), its time, its block taps and output (valid agents), the solver
|
| 10 |
+
iterates, ``final_x0``, the turn pool and logits, the post-processed outputs, and the port's rewrite constants
|
| 11 |
+
(adaLN tables, pad-relative island constants);
|
| 12 |
+
- the **small goldens** (``tests/goldens/<stem>.outputs.npz``, ~0.1 MB, shipped): ``final_x0`` of the valid agents,
|
| 13 |
+
the logits, the ego trajectory columns, the turn command and the ego row of every solver iterate.
|
| 14 |
+
|
| 15 |
+
Tap names follow ``reference.model.TAP_NAMES``; arrays of valid entities carry ``<name>.rows`` (indices into the full
|
| 16 |
+
tensor). Device tests compare TT replay outputs against these, never TT against TT.
|
| 17 |
+
"""
|
| 18 |
+
from __future__ import annotations
|
| 19 |
+
|
| 20 |
+
import datetime as _dt
|
| 21 |
+
from pathlib import Path
|
| 22 |
+
from typing import Any, Dict, Optional
|
| 23 |
+
|
| 24 |
+
import numpy as np
|
| 25 |
+
|
| 26 |
+
from . import config as C
|
| 27 |
+
from ..host import pipeline as hp
|
| 28 |
+
from ..host.postprocess import denoising_steps_ego
|
| 29 |
+
from ..host.solver import solver_plan
|
| 30 |
+
from ..ttaw.golden import TapRegistry, save_goldens
|
| 31 |
+
from . import rewrites as R
|
| 32 |
+
from .pipeline import ReferencePlanner
|
| 33 |
+
|
| 34 |
+
__all__ = ["scene_goldens", "small_goldens", "lite_goldens", "write_scene", "SMALL_KEYS"]
|
| 35 |
+
|
| 36 |
+
SMALL_KEYS = ("final_x0", "final_x0.rows", "logit", "trajectory", "turn_command", "denoising_ego")
|
| 37 |
+
MIXER_CATS = ("ego", "neighbor", "lane", "route", "polygon", "line_string")
|
| 38 |
+
|
| 39 |
+
|
| 40 |
+
def _rows(valid: np.ndarray) -> np.ndarray:
|
| 41 |
+
return np.flatnonzero(np.asarray(valid, bool)).astype(np.int32)
|
| 42 |
+
|
| 43 |
+
|
| 44 |
+
def scene_goldens(ref: ReferencePlanner, raw: Any, *, params: Optional[Dict[str, Any]] = None) -> Dict[str, np.ndarray]:
|
| 45 |
+
"""Run the reference on one scene and return the full golden dict (numpy)."""
|
| 46 |
+
taps = TapRegistry()
|
| 47 |
+
res = ref.run(raw, taps=taps, keep_eval_io=True)
|
| 48 |
+
prep = res.prepared
|
| 49 |
+
f = prep.features
|
| 50 |
+
g: Dict[str, np.ndarray] = {}
|
| 51 |
+
for k, v in prep.raw.items():
|
| 52 |
+
g[f"in.{k}"] = v
|
| 53 |
+
# host features (what the device consumes)
|
| 54 |
+
for name in ("ego", "neighbor", "neighbor_type", "static", "lane", "lane_attr", "lane_speed", "lane_has_speed",
|
| 55 |
+
"route", "route_attr", "route_speed", "route_has_speed", "polygon", "line_string", "goal",
|
| 56 |
+
"ego_shape", "turn", "token_valid", "key_valid", "pos"):
|
| 57 |
+
g[f"host.{name}"] = np.asarray(getattr(f, name))
|
| 58 |
+
for cat, v in f.valid.items():
|
| 59 |
+
g[f"host.valid.{cat}"] = np.asarray(v, bool)
|
| 60 |
+
g["host.agent_valid"] = prep.decoder.agent_valid
|
| 61 |
+
g["host.current_states"] = prep.decoder.current_states
|
| 62 |
+
t = taps.to_dict()
|
| 63 |
+
# encoder taps: mixer internals compacted to the valid entities
|
| 64 |
+
for cat in MIXER_CATS:
|
| 65 |
+
rows = _rows(f.valid[cat])
|
| 66 |
+
for suffix in ("pre", "mixer"):
|
| 67 |
+
g[f"enc.{cat}.{suffix}"] = t[f"enc.{cat}.{suffix}"][rows]
|
| 68 |
+
g[f"enc.{cat}.{suffix}.rows"] = rows
|
| 69 |
+
for name, _ in C.TOKEN_LAYOUT:
|
| 70 |
+
g[f"enc.{name}"] = t[f"enc.{name}"]
|
| 71 |
+
g["enc.tokens"] = t["enc.tokens"]
|
| 72 |
+
for i in range(C.FUSION_DEPTH):
|
| 73 |
+
g[f"enc.fusion.{i}"] = t[f"enc.fusion.{i}"]
|
| 74 |
+
g["enc.encoding"] = res.encoding
|
| 75 |
+
# decoder evaluations: inputs for teacher forcing, outputs and block taps on the valid agents
|
| 76 |
+
arows = _rows(prep.decoder.agent_valid)
|
| 77 |
+
g["dec.rows"] = arows
|
| 78 |
+
g["dec.t"] = np.asarray(res.eval_times, np.float32)
|
| 79 |
+
g["dec.x_in"] = np.stack(res.eval_inputs).astype(np.float32) # [11, 321, 81, 4]
|
| 80 |
+
g["dec.out"] = np.stack(res.eval_outputs)[:, arows].astype(np.float32)
|
| 81 |
+
for k in range(len(res.eval_times)):
|
| 82 |
+
g[f"dec.{k}.temb"] = t[f"dec.{k}.temb"][0]
|
| 83 |
+
g[f"dec.{k}.x"] = t[f"dec.{k}.x"][arows]
|
| 84 |
+
for i in range(C.DIT_DEPTH):
|
| 85 |
+
g[f"dec.{k}.block{i}"] = t[f"dec.{k}.block{i}"][arows]
|
| 86 |
+
g["solver.iterates"] = np.stack(res.denoising_steps).astype(np.float32)[:, arows]
|
| 87 |
+
g["solver.timesteps"] = np.asarray(res.denoising_timesteps, np.float32)
|
| 88 |
+
g["final_x0"] = res.final_x0
|
| 89 |
+
g["turn.pool"] = t["turn.pool"]
|
| 90 |
+
g["turn.logit"] = res.logit
|
| 91 |
+
# post-processed outputs (the API's Trajectory)
|
| 92 |
+
p = {k: spec[3] for k, spec in hp.RUNTIME_PARAMS.items()}
|
| 93 |
+
p.update(params or {})
|
| 94 |
+
out = hp.make_output(res.final_x0, res.logit, prep, ref.normalization, p, denoising_steps=res.denoising_steps)
|
| 95 |
+
g["out.trajectory"] = out.poses
|
| 96 |
+
g["out.predicted_agents"] = out.predicted_agents
|
| 97 |
+
g["out.turn_command"] = np.asarray(out.turn_indicator["command"], np.int32)
|
| 98 |
+
g["out.denoising_ego"] = denoising_steps_ego(np.stack(res.denoising_steps), *ref.normalization.state())
|
| 99 |
+
# the port's rewrite constants (built in float64, rounded once)
|
| 100 |
+
plan = solver_plan(C.DPM_SOLVER_STEPS)
|
| 101 |
+
tab = R.adaln_tables(ref.weights.params, plan.eval_times)
|
| 102 |
+
g["port.adaln.temb"] = tab.temb
|
| 103 |
+
for i, blk in enumerate(tab.blocks):
|
| 104 |
+
for k, v in blk.items():
|
| 105 |
+
g[f"port.adaln.block{i}.{k}"] = v
|
| 106 |
+
g["port.adaln.final_gamma"], g["port.adaln.final_beta"] = tab.final_gamma, tab.final_beta
|
| 107 |
+
for cat in ("ego", "neighbor"):
|
| 108 |
+
cst = R.island_constants(ref.weights.params, cat)
|
| 109 |
+
for k in ("gelu_b1", "c0", "t1_pad", "g_pad", "t2_pad", "w_t1_valid"):
|
| 110 |
+
g[f"port.island.{cat}.{k}"] = getattr(cst, k)
|
| 111 |
+
g["port.solver"] = np.asarray([[u.order, u.t, u.a, u.b, u.c, u.r0] for u in plan.updates], np.float64)
|
| 112 |
+
return g
|
| 113 |
+
|
| 114 |
+
|
| 115 |
+
LITE_PREFIXES = ("host.valid.", "host.token_valid", "host.agent_valid", "enc.", "dec.rows", "dec.t", "final_x0",
|
| 116 |
+
"turn.", "out.", "solver.timesteps")
|
| 117 |
+
|
| 118 |
+
|
| 119 |
+
def lite_goldens(g: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]:
|
| 120 |
+
"""The compact per-scene goldens of the public-data instants (~1 MB): validity, the encoder category outputs and
|
| 121 |
+
the encoding on valid rows, final x0 of the valid agents, logits and the post-processed outputs (no inputs: they
|
| 122 |
+
stay in ``public_data/inputs``; no mixer or per-evaluation decoder taps)."""
|
| 123 |
+
out: Dict[str, np.ndarray] = {}
|
| 124 |
+
tok = np.flatnonzero(g["host.token_valid"])
|
| 125 |
+
for k, v in g.items():
|
| 126 |
+
if not k.startswith(LITE_PREFIXES) or k.startswith(("enc.fusion.", "enc.tokens")):
|
| 127 |
+
continue
|
| 128 |
+
if k.endswith((".pre", ".mixer", ".pre.rows", ".mixer.rows")):
|
| 129 |
+
continue
|
| 130 |
+
out[k] = v
|
| 131 |
+
for name, _ in C.TOKEN_LAYOUT:
|
| 132 |
+
rows = np.flatnonzero(g[f"host.valid.{name}"])
|
| 133 |
+
out[f"enc.{name}"], out[f"enc.{name}.rows"] = g[f"enc.{name}"][rows], rows.astype(np.int32)
|
| 134 |
+
out["enc.encoding"], out["enc.encoding.rows"] = g["enc.encoding"][tok], tok.astype(np.int32)
|
| 135 |
+
out["final_x0"] = g["final_x0"][g["dec.rows"]]
|
| 136 |
+
return out
|
| 137 |
+
|
| 138 |
+
|
| 139 |
+
def small_goldens(g: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]:
|
| 140 |
+
"""The shipped subset of a full golden dict (see the module docstring)."""
|
| 141 |
+
rows = g["dec.rows"]
|
| 142 |
+
return {"final_x0": g["final_x0"][rows], "final_x0.rows": rows, "logit": g["turn.logit"],
|
| 143 |
+
"trajectory": g["out.trajectory"], "turn_command": g["out.turn_command"],
|
| 144 |
+
"denoising_ego": g["out.denoising_ego"]}
|
| 145 |
+
|
| 146 |
+
|
| 147 |
+
def write_scene(ref: ReferencePlanner, raw: Any, scene: str, full_dir: Optional[Path], small_dir: Optional[Path],
|
| 148 |
+
meta: Optional[Dict[str, Any]] = None, *, lite: bool = False) -> Dict[str, Any]:
|
| 149 |
+
g = scene_goldens(ref, raw)
|
| 150 |
+
info = {"scene": scene, "created": _dt.datetime.now(_dt.timezone.utc).isoformat(timespec="seconds"),
|
| 151 |
+
"weights_sha256": ref.weights.sha256, "solver_steps": C.DPM_SOLVER_STEPS,
|
| 152 |
+
"valid_counts": {k: int(np.asarray(v).sum()) for k, v in g.items() if k.startswith("host.valid.")},
|
| 153 |
+
"producer": "tt_diffusion_planner.reference (fp32 CPU, torch)", **(meta or {})}
|
| 154 |
+
paths = {}
|
| 155 |
+
if full_dir is not None:
|
| 156 |
+
tensors = lite_goldens(g) if lite else g
|
| 157 |
+
info["lite"] = bool(lite)
|
| 158 |
+
paths["full"] = str(save_goldens(Path(full_dir) / f"{scene}.npz", tensors, info, compress=True))
|
| 159 |
+
if small_dir is not None:
|
| 160 |
+
paths["small"] = str(save_goldens(Path(small_dir) / f"{scene}.outputs.npz", small_goldens(g), info,
|
| 161 |
+
compress=True))
|
| 162 |
+
return {"info": info, "paths": paths, "goldens": g}
|
code/tt_diffusion_planner/reference/model.py
ADDED
|
@@ -0,0 +1,226 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-License-Identifier: Apache-2.0
|
| 2 |
+
"""Pure-PyTorch fp32 re-implementation of the deployed v5.0 graphs (encoder, DiT decoder, turn-indicator head).
|
| 3 |
+
|
| 4 |
+
Semantics follow the exported ONNX (SPEC 4.2-4.4; tier4/Diffusion-Planner ``encoder.py`` / ``mixer.py`` / ``dit.py``
|
| 5 |
+
/ ``decoder.py`` read for reference, never imported), operation by operation where the order matters for float32:
|
| 6 |
+
attention scales Q before ``Q.K^T`` and adds the ``-inf`` key bias before the softmax; the fusion takes Q from
|
| 7 |
+
``LN1(x)`` but K and V from the un-normalised ``x``; SiLU is ``x * sigmoid(x)``; GELU is exact (erf) everywhere in
|
| 8 |
+
the encoder and in the decoder's ``preproj`` / ``t_embedder``, tanh-approximate in the six DiT MLPs and the final
|
| 9 |
+
projection; every LayerNorm uses epsilon 1e-5.
|
| 10 |
+
|
| 11 |
+
Inputs are the host features of :mod:`tt_diffusion_planner.host.features` (computed from normalised inputs), so
|
| 12 |
+
this module holds only the network. Weights are the canonical dict of :mod:`.weights` (``w`` is ``[in, out]``).
|
| 13 |
+
Every module output can be recorded with a ``ttaw.golden.TapRegistry`` (names in :data:`TAP_NAMES`); valid-row
|
| 14 |
+
compaction of the taps is left to the caller.
|
| 15 |
+
"""
|
| 16 |
+
from __future__ import annotations
|
| 17 |
+
|
| 18 |
+
from typing import Dict, Mapping, Optional, Tuple
|
| 19 |
+
|
| 20 |
+
import numpy as np
|
| 21 |
+
import torch
|
| 22 |
+
import torch.nn.functional as F
|
| 23 |
+
|
| 24 |
+
from . import config as C
|
| 25 |
+
from ..ttaw.golden import NULL_TAPS, TapRegistry
|
| 26 |
+
|
| 27 |
+
__all__ = ["Encoder", "Decoder", "TurnHead", "torch_params", "TAP_NAMES"]
|
| 28 |
+
|
| 29 |
+
Tensor = torch.Tensor
|
| 30 |
+
|
| 31 |
+
TAP_NAMES = {
|
| 32 |
+
"enc.<cat>.pre": "mixer categories: token_pre_project output transposed back to [E, 64, 128]",
|
| 33 |
+
"enc.<cat>.mixer": "mixer categories: output of the 6 MixerBlocks [E, 64, 128]",
|
| 34 |
+
"enc.<cat>": "category output [E, 256] (masked; route with its position embedding), the fusion input rows",
|
| 35 |
+
"enc.tokens": "fusion input [564, 256] (category outputs + masked positional embedding)",
|
| 36 |
+
"enc.fusion.<i>": "fusion block i output [564, 256]",
|
| 37 |
+
"enc.encoding": "final LayerNorm [564, 256] (the ``encoding`` graph output)",
|
| 38 |
+
"dec.<k>.temb": "t_embedder output [321, 256] of evaluation k",
|
| 39 |
+
"dec.<k>.x": "preproj + agent embedding [321, 256]",
|
| 40 |
+
"dec.<k>.block<i>": "DiT block i output [321, 256]",
|
| 41 |
+
"dec.<k>.out": "model_output [321, 81, 4]",
|
| 42 |
+
"turn.pool": "mean of the encoding over the 564 tokens [256]",
|
| 43 |
+
"turn.logit": "turn_indicator_logit [5]",
|
| 44 |
+
}
|
| 45 |
+
|
| 46 |
+
|
| 47 |
+
def torch_params(params: Mapping[str, np.ndarray], dtype: torch.dtype = torch.float32) -> Dict[str, Tensor]:
|
| 48 |
+
"""``{name: tensor}`` copies (the ONNX-backed arrays are read-only)."""
|
| 49 |
+
return {k: torch.from_numpy(np.array(v, copy=True)).to(dtype) for k, v in params.items()}
|
| 50 |
+
|
| 51 |
+
|
| 52 |
+
def _t(a, dtype=torch.float32) -> Tensor:
|
| 53 |
+
if isinstance(a, torch.Tensor):
|
| 54 |
+
return a.to(dtype)
|
| 55 |
+
return torch.from_numpy(np.ascontiguousarray(a)).to(dtype)
|
| 56 |
+
|
| 57 |
+
|
| 58 |
+
class _Ops:
|
| 59 |
+
"""Shared building blocks over the canonical parameter dict."""
|
| 60 |
+
|
| 61 |
+
def __init__(self, p: Mapping[str, Tensor]):
|
| 62 |
+
self.p = p
|
| 63 |
+
|
| 64 |
+
def linear(self, x: Tensor, name: str) -> Tensor:
|
| 65 |
+
return torch.matmul(x, self.p[f"{name}.w"]) + self.p[f"{name}.b"]
|
| 66 |
+
|
| 67 |
+
def ln(self, x: Tensor, name: str) -> Tensor:
|
| 68 |
+
return F.layer_norm(x, x.shape[-1:], self.p[f"{name}.gamma"], self.p[f"{name}.beta"], eps=C.LN_EPS)
|
| 69 |
+
|
| 70 |
+
def mlp(self, x: Tensor, name: str, approximate: str = "none") -> Tensor:
|
| 71 |
+
return self.linear(F.gelu(self.linear(x, f"{name}.fc1"), approximate=approximate), f"{name}.fc2")
|
| 72 |
+
|
| 73 |
+
def attention(self, q_in: Tensor, kv: Tensor, name_q: str, key_bias: Optional[Tensor]) -> Tensor:
|
| 74 |
+
"""Multi-head attention with 8 heads of 32: ``q_in`` [Sq, 256] -> Q (``name_q`` linear); ``kv`` [Sk, 512] is
|
| 75 |
+
the already projected K|V; ``key_bias`` [Sk] is 0 / -inf. Returns the heads concatenated [Sq, 256]
|
| 76 |
+
(before the output projection)."""
|
| 77 |
+
q = self.linear(q_in, name_q)
|
| 78 |
+
return self.attend(q, kv[:, :C.HIDDEN_DIM], kv[:, C.HIDDEN_DIM:], key_bias)
|
| 79 |
+
|
| 80 |
+
@staticmethod
|
| 81 |
+
def attend(q: Tensor, k: Tensor, v: Tensor, key_bias: Optional[Tensor]) -> Tensor:
|
| 82 |
+
sq, sk = q.shape[0], k.shape[0]
|
| 83 |
+
qh = q.reshape(sq, C.NUM_HEADS, C.HEAD_DIM).transpose(0, 1) * float(C.ATTN_SCALE) # [H, Sq, 32]
|
| 84 |
+
kh = k.reshape(sk, C.NUM_HEADS, C.HEAD_DIM).permute(1, 2, 0) # [H, 32, Sk]
|
| 85 |
+
vh = v.reshape(sk, C.NUM_HEADS, C.HEAD_DIM).transpose(0, 1) # [H, Sk, 32]
|
| 86 |
+
scores = torch.matmul(qh, kh)
|
| 87 |
+
if key_bias is not None:
|
| 88 |
+
scores = scores + key_bias.reshape(1, 1, sk)
|
| 89 |
+
o = torch.matmul(torch.softmax(scores, dim=-1), vh) # [H, Sq, 32]
|
| 90 |
+
return o.transpose(0, 1).reshape(sq, C.HIDDEN_DIM)
|
| 91 |
+
|
| 92 |
+
|
| 93 |
+
def _key_bias(valid: np.ndarray) -> Tensor:
|
| 94 |
+
b = np.where(np.asarray(valid, bool), np.float32(0.0), np.float32(-np.inf)).astype(np.float32)
|
| 95 |
+
return torch.from_numpy(b)
|
| 96 |
+
|
| 97 |
+
|
| 98 |
+
class Encoder(_Ops):
|
| 99 |
+
"""``diffusion_planner_encoder.onnx`` on host features -> ``encoding`` [564, 256]."""
|
| 100 |
+
|
| 101 |
+
def mixer_trunk(self, x: Tensor, mod: str, taps: TapRegistry, cat: str) -> Tensor:
|
| 102 |
+
"""channel_pre (C_in -> 128 -> 128), token_pre over the time / point axis (T -> 64 -> 64), 6 MixerBlocks,
|
| 103 |
+
mean over the 64 tokens -> [E, 128]."""
|
| 104 |
+
N = f"encoder.{mod}"
|
| 105 |
+
x = self.mlp(x, f"{N}.channel_pre_project") # [E, T, 128]
|
| 106 |
+
x = self.mlp(x.transpose(1, 2), f"{N}.token_pre_project").transpose(1, 2) # [E, 64, 128]
|
| 107 |
+
taps.tap(f"enc.{cat}.pre", x)
|
| 108 |
+
for i in range(C.MIXER_DEPTH):
|
| 109 |
+
B = f"{N}.blocks.{i}"
|
| 110 |
+
y = self.mlp(self.ln(x, f"{B}.norm1").transpose(1, 2), f"{B}.tokens_mlp").transpose(1, 2)
|
| 111 |
+
x = x + y
|
| 112 |
+
x = x + self.mlp(self.ln(x, f"{B}.norm2"), f"{B}.channels_mlp")
|
| 113 |
+
taps.tap(f"enc.{cat}.mixer", x)
|
| 114 |
+
return x.mean(dim=1)
|
| 115 |
+
|
| 116 |
+
def head(self, x: Tensor, mod: str) -> Tensor:
|
| 117 |
+
"""LayerNorm(128) + emb_project (128 -> 256 -> 256)."""
|
| 118 |
+
return self.mlp(self.ln(x, f"encoder.{mod}.norm"), f"encoder.{mod}.emb_project")
|
| 119 |
+
|
| 120 |
+
def lanes(self, x: Tensor, attr: Tensor, speed: Tensor, has_speed: Tensor, valid: Tensor, mod: str,
|
| 121 |
+
taps: TapRegistry, cat: str) -> Tensor:
|
| 122 |
+
N = f"encoder.{mod}"
|
| 123 |
+
h = self.mixer_trunk(x, mod, taps, cat)
|
| 124 |
+
speed_emb = torch.where(has_speed, self.linear(speed, f"{N}.speed_limit_emb"),
|
| 125 |
+
self.p[f"{N}.unknown_speed_emb"].reshape(1, -1))
|
| 126 |
+
h = h + speed_emb + self.linear(attr, f"{N}.attribute_emb")
|
| 127 |
+
return self.head(h, mod) * valid
|
| 128 |
+
|
| 129 |
+
def small(self, x: Tensor, mod: str) -> Tensor:
|
| 130 |
+
"""goal / ego-shape / turn-indicator encoders: channel MLP -> LayerNorm -> projection."""
|
| 131 |
+
return self.head(self.mlp(x, f"encoder.{mod}.channel_pre_project"), mod)
|
| 132 |
+
|
| 133 |
+
def forward(self, f, taps: TapRegistry = NULL_TAPS) -> Tensor:
|
| 134 |
+
valid = {k: _t(v.astype(np.float32)).reshape(-1, 1) for k, v in f.valid.items()}
|
| 135 |
+
out: Dict[str, Tensor] = {}
|
| 136 |
+
e = self.mixer_trunk(_t(f.ego)[None], "ego_encoder", taps, "ego")
|
| 137 |
+
out["ego"] = self.head(e, "ego_encoder")
|
| 138 |
+
n = self.mixer_trunk(_t(f.neighbor), "neighbor_encoder", taps, "neighbor")
|
| 139 |
+
n = n + self.linear(_t(f.neighbor_type), "encoder.neighbor_encoder.type_emb")
|
| 140 |
+
out["neighbor"] = self.head(n, "neighbor_encoder") * valid["neighbor"]
|
| 141 |
+
out["static"] = self.mlp(_t(f.static), "encoder.static_encoder.projection") * valid["static"]
|
| 142 |
+
out["lane"] = self.lanes(_t(f.lane), _t(f.lane_attr), _t(f.lane_speed), torch.from_numpy(f.lane_has_speed),
|
| 143 |
+
valid["lane"], "lane_encoder", taps, "lane")
|
| 144 |
+
route = self.lanes(_t(f.route), _t(f.route_attr), _t(f.route_speed), torch.from_numpy(f.route_has_speed),
|
| 145 |
+
valid["route"], "route_encoder", taps, "route")
|
| 146 |
+
out["route"] = route + self.p["encoder.route_position_embedding"] * valid["route"]
|
| 147 |
+
poly = self.mixer_trunk(_t(f.polygon), "polygon_encoder", taps, "polygon")
|
| 148 |
+
out["polygon"] = self.head(poly, "polygon_encoder") * valid["polygon"]
|
| 149 |
+
ls = self.mixer_trunk(_t(f.line_string), "line_string_encoder", taps, "line_string")
|
| 150 |
+
out["line_string"] = self.head(ls, "line_string_encoder") * valid["line_string"]
|
| 151 |
+
out["goal"] = self.small(_t(f.goal)[None], "goal_pose_encoder")
|
| 152 |
+
out["ego_shape"] = self.small(_t(f.ego_shape)[None], "ego_shape_encoder")
|
| 153 |
+
out["turn"] = self.small(_t(f.turn)[None], "turn_indicator_encoder")
|
| 154 |
+
for name, _ in C.TOKEN_LAYOUT:
|
| 155 |
+
taps.tap(f"enc.{name}", out[name])
|
| 156 |
+
x = torch.cat([out[name] for name, _ in C.TOKEN_LAYOUT], dim=0) # [564, 256]
|
| 157 |
+
pos = self.linear(_t(f.pos), "encoder.pos_emb") * _t(f.token_valid.astype(np.float32)).reshape(-1, 1)
|
| 158 |
+
x = taps.tap("enc.tokens", x + pos)
|
| 159 |
+
key_bias = _key_bias(f.key_valid)
|
| 160 |
+
for i in range(C.FUSION_DEPTH):
|
| 161 |
+
B = f"encoder.fusion.blocks.{i}"
|
| 162 |
+
kv = self.linear(x, f"{B}.attn.kv") # K and V from the un-normalised x
|
| 163 |
+
heads = self.attention(self.ln(x, f"{B}.norm1"), kv, f"{B}.attn.q", key_bias)
|
| 164 |
+
x = x + self.linear(heads, f"{B}.attn.out")
|
| 165 |
+
x = x + self.mlp(self.ln(x, f"{B}.norm2"), f"{B}.mlp")
|
| 166 |
+
taps.tap(f"enc.fusion.{i}", x)
|
| 167 |
+
return taps.tap("enc.encoding", self.ln(x, "encoder.fusion.norm"))
|
| 168 |
+
|
| 169 |
+
|
| 170 |
+
class Decoder(_Ops):
|
| 171 |
+
"""``diffusion_planner_decoder.onnx``: one DiT evaluation (x0 prediction) per call.
|
| 172 |
+
|
| 173 |
+
``cross_kv(encoding)`` computes the K|V of the three cross-attention blocks once per plan (the decoder graph
|
| 174 |
+
recomputes the same product at every call; hoisting it is exact, SPEC 4.6.1)."""
|
| 175 |
+
|
| 176 |
+
def cross_kv(self, encoding: Tensor) -> Tuple[Tensor, ...]:
|
| 177 |
+
return tuple(self.linear(encoding, f"decoder.dit.blocks.{i}.cross_attn.kv") for i in range(C.DIT_DEPTH))
|
| 178 |
+
|
| 179 |
+
def time_embedding(self, t, agents: int = C.MAX_NUM_AGENTS) -> Tensor:
|
| 180 |
+
"""t_embedder of the 81 per-point diffusion times [P, 81] (uniform in the multi-step mode) -> [P, 256]."""
|
| 181 |
+
tt = torch.full((agents, C.DIT_TIME_DIM), float(np.float32(t))) if np.ndim(t) == 0 else _t(t)
|
| 182 |
+
return self.mlp(tt, "decoder.dit.t_embedder")
|
| 183 |
+
|
| 184 |
+
def forward(self, x_t: np.ndarray, t, kv: Tuple[Tensor, ...], agent_valid: np.ndarray,
|
| 185 |
+
taps: TapRegistry = NULL_TAPS, prefix: str = "dec") -> Tensor:
|
| 186 |
+
"""``x_t`` [P, 81, 4] with P = 321 as exported, or any agent bucket ``P >= 1 + valid neighbours`` (ego first;
|
| 187 |
+
exact for the agents present, SPEC 4.6.4); ``agent_valid`` [P]; ``t`` a scalar time (or [P, 81])."""
|
| 188 |
+
P = int(np.shape(x_t)[0])
|
| 189 |
+
x = self.mlp(_t(x_t).reshape(P, C.DIT_INPUT_DIM), "decoder.dit.preproj")
|
| 190 |
+
c = taps.tap(f"{prefix}.temb", self.time_embedding(t, P))
|
| 191 |
+
emb = self.p["decoder.dit.agent_embedding"]
|
| 192 |
+
x = x + torch.cat([emb[0:1], emb[1:2].expand(P - 1, -1)], dim=0)
|
| 193 |
+
x = taps.tap(f"{prefix}.x", x)
|
| 194 |
+
silu_c = c * torch.sigmoid(c)
|
| 195 |
+
key_bias = _key_bias(agent_valid)
|
| 196 |
+
for i in range(C.DIT_DEPTH):
|
| 197 |
+
B = f"decoder.dit.blocks.{i}"
|
| 198 |
+
mod = self.linear(silu_c, f"{B}.adaLN_modulation")
|
| 199 |
+
shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = torch.split(mod, C.HIDDEN_DIM, dim=-1)
|
| 200 |
+
h = self.ln(x, f"{B}.norm1") * (scale_msa + 1.0) + shift_msa
|
| 201 |
+
qkv = self.linear(h, f"{B}.attn.qkv")
|
| 202 |
+
q, k, v = torch.split(qkv, C.HIDDEN_DIM, dim=-1)
|
| 203 |
+
x = x + gate_msa * self.linear(self.attend(q, k, v, key_bias), f"{B}.attn.out")
|
| 204 |
+
h = self.ln(x, f"{B}.norm2") * (scale_mlp + 1.0) + shift_mlp
|
| 205 |
+
x = x + gate_mlp * self.mlp(h, f"{B}.mlp1", "tanh")
|
| 206 |
+
heads = self.attention(self.ln(x, f"{B}.norm3"), kv[i], f"{B}.cross_attn.q", None) # no key mask
|
| 207 |
+
x = x + self.linear(heads, f"{B}.cross_attn.out")
|
| 208 |
+
x = x + self.mlp(self.ln(x, f"{B}.norm4"), f"{B}.mlp2", "tanh")
|
| 209 |
+
taps.tap(f"{prefix}.block{i}", x)
|
| 210 |
+
Fn = "decoder.dit.final_layer"
|
| 211 |
+
shift, scale = torch.split(self.linear(silu_c, f"{Fn}.adaLN_modulation"), C.HIDDEN_DIM, dim=-1)
|
| 212 |
+
h = self.ln(x, f"{Fn}.norm_final") * (scale + 1.0) + shift
|
| 213 |
+
h = F.gelu(self.linear(self.ln(h, f"{Fn}.proj.0"), f"{Fn}.proj.1"), approximate="tanh")
|
| 214 |
+
out = self.linear(self.ln(h, f"{Fn}.proj.3"), f"{Fn}.proj.4").reshape(P, C.OUTPUT_T + 1, C.POSE_DIM)
|
| 215 |
+
return taps.tap(f"{prefix}.out", out)
|
| 216 |
+
|
| 217 |
+
|
| 218 |
+
class TurnHead(_Ops):
|
| 219 |
+
"""``diffusion_planner_turn_indicator.onnx``: ``W [272 -> 5]`` over ``final_x0[0, 1::10, :2]`` (16 values) and the
|
| 220 |
+
token mean of the encoding (256)."""
|
| 221 |
+
|
| 222 |
+
def forward(self, encoding: Tensor, final_x0: np.ndarray, taps: TapRegistry = NULL_TAPS) -> Tensor:
|
| 223 |
+
pool = taps.tap("turn.pool", encoding.mean(dim=0))
|
| 224 |
+
ego = _t(final_x0)[0, 1::10, :2].reshape(-1)
|
| 225 |
+
feat = torch.cat([ego, pool], dim=0)[None]
|
| 226 |
+
return taps.tap("turn.logit", self.linear(feat, "decoder.turn_indicator_predictor")[0])
|
code/tt_diffusion_planner/reference/ort.py
ADDED
|
@@ -0,0 +1,137 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-License-Identifier: Apache-2.0
|
| 2 |
+
"""ONNX Runtime on the shipped v5.0 ONNX files: the oracle of the CPU reference (research and tests only;
|
| 3 |
+
onnxruntime is not a runtime dependency of the bundle or the image).
|
| 4 |
+
|
| 5 |
+
``OrtPlanner.run(raw)`` is the Autoware ``multi_step`` pipeline with ORT as the backend (a port of the verified
|
| 6 |
+
research script ``research/diffusion-planner/scripts/dp_reference.py``): the node's normalization and speed masks,
|
| 7 |
+
encoder.onnx, the DPM-Solver++(2M) loop of :mod:`tt_diffusion_planner.host.solver` calling decoder.onnx 11 times,
|
| 8 |
+
then turn_indicator.onnx. ``taps=[...]`` exposes intermediate tensors of the encoder / decoder graphs by ONNX tensor
|
| 9 |
+
name (the decoder taps are recorded per evaluation), for the per-module PCC test of the reference.
|
| 10 |
+
"""
|
| 11 |
+
from __future__ import annotations
|
| 12 |
+
|
| 13 |
+
from dataclasses import dataclass, field
|
| 14 |
+
from pathlib import Path
|
| 15 |
+
from typing import Any, Dict, List, Sequence
|
| 16 |
+
|
| 17 |
+
import numpy as np
|
| 18 |
+
|
| 19 |
+
from . import config as C
|
| 20 |
+
from ..host.normalize import normalize_inputs, speed_masks
|
| 21 |
+
from ..host.solver import apply_prefix_constraint, dpm_solver_sample
|
| 22 |
+
from ..host.features import decoder_masks
|
| 23 |
+
from ..ttaw import io as tio
|
| 24 |
+
from .weights import load_param_json
|
| 25 |
+
|
| 26 |
+
__all__ = ["OrtPlanner", "OrtResult", "available"]
|
| 27 |
+
|
| 28 |
+
|
| 29 |
+
def available() -> bool:
|
| 30 |
+
import importlib.util
|
| 31 |
+
|
| 32 |
+
return importlib.util.find_spec("onnxruntime") is not None
|
| 33 |
+
|
| 34 |
+
|
| 35 |
+
def _session(model: Any, threads: int, optimize: bool = True):
|
| 36 |
+
import onnxruntime as ort
|
| 37 |
+
|
| 38 |
+
so = ort.SessionOptions()
|
| 39 |
+
so.graph_optimization_level = (ort.GraphOptimizationLevel.ORT_ENABLE_ALL if optimize
|
| 40 |
+
else ort.GraphOptimizationLevel.ORT_DISABLE_ALL)
|
| 41 |
+
so.intra_op_num_threads = threads
|
| 42 |
+
so.inter_op_num_threads = 1
|
| 43 |
+
so.log_severity_level = 3
|
| 44 |
+
src = model.SerializeToString() if hasattr(model, "SerializeToString") else str(model)
|
| 45 |
+
return ort.InferenceSession(src, so, providers=["CPUExecutionProvider"])
|
| 46 |
+
|
| 47 |
+
|
| 48 |
+
def _with_outputs(path: Path, names: Sequence[str]):
|
| 49 |
+
"""The model with extra graph outputs (intermediate tensors by name)."""
|
| 50 |
+
import onnx
|
| 51 |
+
|
| 52 |
+
m = onnx.load(str(path))
|
| 53 |
+
have = {o.name for o in m.graph.output}
|
| 54 |
+
m.graph.output.extend([onnx.ValueInfoProto(name=n) for n in names if n not in have])
|
| 55 |
+
return m
|
| 56 |
+
|
| 57 |
+
|
| 58 |
+
@dataclass
|
| 59 |
+
class OrtResult:
|
| 60 |
+
norm: Dict[str, np.ndarray]
|
| 61 |
+
encoding: np.ndarray # [1, 564, 256]
|
| 62 |
+
final_x0: np.ndarray # [1, 321, 81, 4]
|
| 63 |
+
logit: np.ndarray # [1, 5]
|
| 64 |
+
denoising_steps: List[np.ndarray]
|
| 65 |
+
denoising_timesteps: List[float]
|
| 66 |
+
eval_times: List[float]
|
| 67 |
+
eval_inputs: List[np.ndarray] = field(default_factory=list)
|
| 68 |
+
eval_outputs: List[np.ndarray] = field(default_factory=list)
|
| 69 |
+
taps: Dict[str, np.ndarray] = field(default_factory=dict)
|
| 70 |
+
|
| 71 |
+
|
| 72 |
+
class OrtPlanner:
|
| 73 |
+
"""``OrtPlanner(weights_dir, threads=4, encoder_taps=(...), decoder_taps=(...))``. With taps the sessions run
|
| 74 |
+
without graph optimizations (every intermediate kept as computed); without, ``ORT_ENABLE_ALL`` like
|
| 75 |
+
``onnxruntime_inference.cpp:172``."""
|
| 76 |
+
|
| 77 |
+
def __init__(self, weights_dir: Path, *, threads: int = 4, encoder_taps: Sequence[str] = (),
|
| 78 |
+
decoder_taps: Sequence[str] = (), turn_taps: Sequence[str] = ()):
|
| 79 |
+
wd = Path(weights_dir)
|
| 80 |
+
self.normalization = load_param_json(wd / C.PARAM_JSON)
|
| 81 |
+
self.encoder_taps, self.decoder_taps, self.turn_taps = list(encoder_taps), list(decoder_taps), list(turn_taps)
|
| 82 |
+
opt = not (encoder_taps or decoder_taps or turn_taps)
|
| 83 |
+
self.enc = _session(_with_outputs(wd / C.ENCODER_ONNX, self.encoder_taps), threads, opt)
|
| 84 |
+
self.dec = _session(_with_outputs(wd / C.DECODER_ONNX, self.decoder_taps), threads, opt)
|
| 85 |
+
self.turn = _session(_with_outputs(wd / C.TURN_INDICATOR_ONNX, self.turn_taps), threads, opt)
|
| 86 |
+
|
| 87 |
+
def normalize(self, raw: Any) -> Dict[str, np.ndarray]:
|
| 88 |
+
arrays = tio.load_named_arrays(raw, C.INPUT_SCHEMA)
|
| 89 |
+
norm = normalize_inputs(arrays, self.normalization.observation)
|
| 90 |
+
norm.update(speed_masks(norm))
|
| 91 |
+
return norm
|
| 92 |
+
|
| 93 |
+
def encode(self, norm: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]:
|
| 94 |
+
feed = {k: norm[k] for k in C.ENCODER_INPUTS}
|
| 95 |
+
outs = self.enc.run(["encoding"] + self.encoder_taps, feed)
|
| 96 |
+
return dict(zip(["encoding"] + self.encoder_taps, outs))
|
| 97 |
+
|
| 98 |
+
def decode(self, encoding: np.ndarray, x: np.ndarray, t, neighbor_agents_past: np.ndarray) -> Dict[str, Any]:
|
| 99 |
+
dt = np.full((1, C.MAX_NUM_AGENTS, C.OUTPUT_T + 1, 1), np.float32(t), np.float32)
|
| 100 |
+
feed = {"encoding": encoding, "sampled_trajectories": np.asarray(x, np.float32).reshape(
|
| 101 |
+
1, C.MAX_NUM_AGENTS, C.OUTPUT_T + 1, C.POSE_DIM), "diffusion_time": dt,
|
| 102 |
+
"neighbor_agents_past": neighbor_agents_past}
|
| 103 |
+
outs = self.dec.run(["model_output"] + self.decoder_taps, feed)
|
| 104 |
+
return dict(zip(["model_output"] + self.decoder_taps, outs))
|
| 105 |
+
|
| 106 |
+
def turn_logit(self, encoding: np.ndarray, final_x0: np.ndarray) -> Dict[str, np.ndarray]:
|
| 107 |
+
outs = self.turn.run(["turn_indicator_logit"] + self.turn_taps,
|
| 108 |
+
{"encoding": encoding, "final_x0": np.asarray(final_x0, np.float32)})
|
| 109 |
+
return dict(zip(["turn_indicator_logit"] + self.turn_taps, outs))
|
| 110 |
+
|
| 111 |
+
def run(self, raw: Any, *, steps: int = C.DPM_SOLVER_STEPS, keep_eval_io: bool = False) -> OrtResult:
|
| 112 |
+
norm = self.normalize(raw)
|
| 113 |
+
enc = self.encode(norm)
|
| 114 |
+
encoding = enc["encoding"]
|
| 115 |
+
taps: Dict[str, np.ndarray] = {f"enc:{k}": v for k, v in enc.items() if k != "encoding"}
|
| 116 |
+
cs = decoder_masks(norm).current_states
|
| 117 |
+
nb = norm["neighbor_agents_past"]
|
| 118 |
+
eval_in: List[np.ndarray] = []
|
| 119 |
+
eval_out: List[np.ndarray] = []
|
| 120 |
+
|
| 121 |
+
def model_fn(x: np.ndarray, t) -> np.ndarray:
|
| 122 |
+
k = len(eval_out)
|
| 123 |
+
out = self.decode(encoding, x, t, nb)
|
| 124 |
+
for name in self.decoder_taps:
|
| 125 |
+
taps[f"dec{k}:{name}"] = out[name]
|
| 126 |
+
y = out["model_output"][0]
|
| 127 |
+
eval_out.append(y if keep_eval_io else np.empty(0))
|
| 128 |
+
if keep_eval_io:
|
| 129 |
+
eval_in.append(np.array(x, np.float32, copy=True))
|
| 130 |
+
return y
|
| 131 |
+
|
| 132 |
+
res = dpm_solver_sample(norm["sampled_trajectories"][0], model_fn,
|
| 133 |
+
lambda x: apply_prefix_constraint(x, cs), steps)
|
| 134 |
+
tl = self.turn_logit(encoding, res.final_x[None])
|
| 135 |
+
taps.update({f"turn:{k}": v for k, v in tl.items() if k != "turn_indicator_logit"})
|
| 136 |
+
return OrtResult(norm, encoding, res.final_x[None], tl["turn_indicator_logit"], res.denoising_steps,
|
| 137 |
+
res.denoising_timesteps, res.eval_times, eval_in, eval_out if keep_eval_io else [], taps)
|
code/tt_diffusion_planner/reference/pipeline.py
ADDED
|
@@ -0,0 +1,123 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-License-Identifier: Apache-2.0
|
| 2 |
+
"""The fp32 CPU reference planner: host pre-processing -> encoder -> DPM-Solver++(2M) over the DiT decoder (11
|
| 3 |
+
evaluations) -> turn head -> host post-processing, i.e. the Autoware ``multi_step`` mode with guidance off.
|
| 4 |
+
|
| 5 |
+
from tt_diffusion_planner.reference import ReferencePlanner
|
| 6 |
+
ref = ReferencePlanner() # weights: find_weights_dir() or weights_dir=...
|
| 7 |
+
out = ref(inputs="code/tt_diffusion_planner/samples/kashiwanoha_dense.npz") # ttaw.outputs.Trajectory
|
| 8 |
+
res = ref.run(raw_inputs, taps=TapRegistry()) # every intermediate (golden generation)
|
| 9 |
+
|
| 10 |
+
``__call__`` returns exactly what the device model returns (same ``host.prepare`` / ``host.make_output``), so the
|
| 11 |
+
stored ``samples/<stem>.reference.json`` is this class's ``to_dict()``.
|
| 12 |
+
"""
|
| 13 |
+
from __future__ import annotations
|
| 14 |
+
|
| 15 |
+
import time
|
| 16 |
+
from dataclasses import dataclass, field
|
| 17 |
+
from pathlib import Path
|
| 18 |
+
from typing import Any, Dict, List, Optional
|
| 19 |
+
|
| 20 |
+
import numpy as np
|
| 21 |
+
import torch
|
| 22 |
+
|
| 23 |
+
from . import config as C
|
| 24 |
+
from ..host import pipeline as hp
|
| 25 |
+
from ..host.solver import apply_prefix_constraint, dpm_solver_sample
|
| 26 |
+
from ..ttaw import io as tio
|
| 27 |
+
from ..ttaw.golden import NULL_TAPS, TapRegistry
|
| 28 |
+
from .model import Decoder, Encoder, TurnHead, torch_params
|
| 29 |
+
from .weights import PlannerWeights, find_weights_dir, load_weights
|
| 30 |
+
|
| 31 |
+
__all__ = ["ReferencePlanner", "ReferenceResult"]
|
| 32 |
+
|
| 33 |
+
MODEL_NAME = "diffusion-planner-p150"
|
| 34 |
+
|
| 35 |
+
|
| 36 |
+
@dataclass
|
| 37 |
+
class ReferenceResult:
|
| 38 |
+
prepared: hp.Prepared
|
| 39 |
+
encoding: np.ndarray # [564, 256]
|
| 40 |
+
final_x0: np.ndarray # [321, 81, 4] normalised
|
| 41 |
+
logit: np.ndarray # [5]
|
| 42 |
+
denoising_steps: List[np.ndarray]
|
| 43 |
+
denoising_timesteps: List[float]
|
| 44 |
+
eval_times: List[float]
|
| 45 |
+
eval_inputs: List[np.ndarray] = field(default_factory=list) # x fed to each decoder evaluation
|
| 46 |
+
eval_outputs: List[np.ndarray] = field(default_factory=list) # model_output of each evaluation
|
| 47 |
+
timing_ms: Dict[str, float] = field(default_factory=dict)
|
| 48 |
+
|
| 49 |
+
|
| 50 |
+
class ReferencePlanner:
|
| 51 |
+
"""CPU fp32 reference of the deployed network plus the shared host code. ``threads`` bounds torch's intra-op
|
| 52 |
+
threads (the workspace rule is <= 4)."""
|
| 53 |
+
|
| 54 |
+
def __init__(self, weights_dir: Optional[str] = None, *, weights: Optional[PlannerWeights] = None,
|
| 55 |
+
threads: Optional[int] = 4, steps: int = C.DPM_SOLVER_STEPS):
|
| 56 |
+
if weights is None:
|
| 57 |
+
wd = find_weights_dir(weights_dir)
|
| 58 |
+
if wd is None:
|
| 59 |
+
raise FileNotFoundError("Diffusion Planner v5.0 weights not found: pass weights_dir= or set "
|
| 60 |
+
"DIFFUSION_PLANNER_WEIGHTS_DIR")
|
| 61 |
+
weights = load_weights(wd)
|
| 62 |
+
self.weights = weights
|
| 63 |
+
self.normalization = weights.normalization
|
| 64 |
+
if threads:
|
| 65 |
+
torch.set_num_threads(int(threads))
|
| 66 |
+
p = torch_params(weights.params)
|
| 67 |
+
self.encoder, self.decoder, self.turn = Encoder(p), Decoder(p), TurnHead(p)
|
| 68 |
+
self.steps = steps
|
| 69 |
+
|
| 70 |
+
# ---- network ---------------------------------------------------------------------------------------------
|
| 71 |
+
@torch.no_grad()
|
| 72 |
+
def run(self, raw: Any, *, taps: TapRegistry = NULL_TAPS, keep_eval_io: bool = False) -> ReferenceResult:
|
| 73 |
+
"""One plan on raw ONNX-named inputs (mapping, ``.npz`` path / bytes or the JSON envelope)."""
|
| 74 |
+
arrays = tio.load_named_arrays(raw, C.INPUT_SCHEMA)
|
| 75 |
+
t0 = time.perf_counter()
|
| 76 |
+
prep = hp.prepare(arrays, self.normalization.observation)
|
| 77 |
+
t1 = time.perf_counter()
|
| 78 |
+
encoding = self.encoder.forward(prep.features, taps)
|
| 79 |
+
t2 = time.perf_counter()
|
| 80 |
+
kv = self.decoder.cross_kv(encoding)
|
| 81 |
+
cs = prep.decoder.current_states
|
| 82 |
+
eval_in: List[np.ndarray] = []
|
| 83 |
+
eval_out: List[np.ndarray] = []
|
| 84 |
+
|
| 85 |
+
def model_fn(x: np.ndarray, t) -> np.ndarray:
|
| 86 |
+
k = len(eval_out)
|
| 87 |
+
y = self.decoder.forward(x, t, kv, prep.decoder.agent_valid, taps, prefix=f"dec.{k}").numpy()
|
| 88 |
+
if keep_eval_io:
|
| 89 |
+
eval_in.append(x.copy())
|
| 90 |
+
eval_out.append(y if keep_eval_io else np.empty(0))
|
| 91 |
+
return y
|
| 92 |
+
|
| 93 |
+
res = dpm_solver_sample(prep.x_T, model_fn, lambda x: apply_prefix_constraint(x, cs), self.steps)
|
| 94 |
+
t3 = time.perf_counter()
|
| 95 |
+
logit = self.turn.forward(encoding, res.final_x, taps).numpy()
|
| 96 |
+
t4 = time.perf_counter()
|
| 97 |
+
for k, x in enumerate(res.denoising_steps):
|
| 98 |
+
taps.tap(f"solver.x{k}", x)
|
| 99 |
+
taps.tap("final_x0", res.final_x)
|
| 100 |
+
return ReferenceResult(prep, encoding.numpy(), res.final_x, logit, res.denoising_steps,
|
| 101 |
+
res.denoising_timesteps, res.eval_times, eval_in, eval_out if keep_eval_io else [],
|
| 102 |
+
{"preprocess": (t1 - t0) * 1e3, "encoder": (t2 - t1) * 1e3, "solver": (t3 - t2) * 1e3,
|
| 103 |
+
"turn": (t4 - t3) * 1e3})
|
| 104 |
+
|
| 105 |
+
# ---- the API-shaped call ----------------------------------------------------------------------------------
|
| 106 |
+
def __call__(self, inputs: Any, **params: Any):
|
| 107 |
+
"""``model(inputs=...)`` of the CPU reference: the same host pre / post-processing as the device model."""
|
| 108 |
+
p = {k: spec[3] for k, spec in hp.RUNTIME_PARAMS.items()}
|
| 109 |
+
unknown = sorted(set(params) - set(p))
|
| 110 |
+
if unknown:
|
| 111 |
+
raise tio.InputError(f"unknown parameter(s) {unknown}; allowed: {sorted(p)}")
|
| 112 |
+
p.update(params)
|
| 113 |
+
t0 = time.perf_counter()
|
| 114 |
+
res = self.run(inputs)
|
| 115 |
+
total = (time.perf_counter() - t0) * 1e3
|
| 116 |
+
steps = res.denoising_steps if p.get("return_denoising_steps") else None
|
| 117 |
+
return hp.make_output(res.final_x0, res.logit, res.prepared, self.normalization, p, model=MODEL_NAME,
|
| 118 |
+
denoising_steps=steps, timing_ms={**res.timing_ms, "total": total},
|
| 119 |
+
meta={"reference": "fp32 CPU (torch)"})
|
| 120 |
+
|
| 121 |
+
@staticmethod
|
| 122 |
+
def sample_path(name: str = "kashiwanoha_dense.npz") -> Path:
|
| 123 |
+
return Path(__file__).resolve().parents[1] / "samples" / name
|
code/tt_diffusion_planner/reference/rewrites.py
ADDED
|
@@ -0,0 +1,183 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-License-Identifier: Apache-2.0
|
| 2 |
+
"""Exact rewrites of the deployed graph that the TT port applies, as constants and CPU forwards (SPEC 4.6, PLAN 2.12).
|
| 3 |
+
|
| 4 |
+
Each rewrite is exact in real arithmetic; its constants are computed in float64 from the float32 weights and rounded
|
| 5 |
+
once to float32 (the "fold in fp64, round once" rule of ``ttaw.weights``). The CPU forwards here use them and are
|
| 6 |
+
tested against the as-exported :mod:`.model` (``tests/test_reference_host.py``), so the device port has a CPU oracle in
|
| 7 |
+
its own parameterisation.
|
| 8 |
+
|
| 9 |
+
1. **Per-step adaLN tables folded into the LayerNorm affine** (SPEC 4.6.2, 8.1). With a uniform diffusion time (the
|
| 10 |
+
multi-step mode) the t-embedding and every adaLN output are the same for all 321 agents and depend only on t, so
|
| 11 |
+
for the 11 evaluation times ``modulate(LN(x; g, b), shift, scale) = LN(x; g (1 + scale), b (1 + scale) + shift)``
|
| 12 |
+
and the gates are per-step ``[256]`` rows. :func:`adaln_tables` builds them for ``SolverPlan.eval_times``.
|
| 13 |
+
2. **Cross-attention K/V hoisted** (SPEC 4.6.1): ``encoding @ W_kv + b_kv`` of the three blocks once per plan
|
| 14 |
+
(the decoder graph recomputes it at each of the 11 calls).
|
| 15 |
+
3. **Pad-relative fp32 pre-projection island** (SPEC 4.6.7, probe P12): after the in-graph history truncation the
|
| 16 |
+
ego (rows 6..30) and neighbour (rows 0..24) inputs of ``channel_pre_project`` are zero rows, which all map to
|
| 17 |
+
``c0 = fc2(gelu(b1)) + b2``. With the all-zero ("pad") agent's outputs ``t1_pad = b_t1 + c0 (x) sum_t W_t1[t]``,
|
| 18 |
+
``g_pad = gelu(t1_pad)``, ``t2_pad = g_pad @ W_t2 + b_t2`` the island becomes
|
| 19 |
+
``t1 = t1_pad + (z_valid - c0)^T @ W_t1[valid rows]`` and ``t2 = t2_pad + (gelu(t1) - g_pad) @ W_t2``, so the
|
| 20 |
+
device's TF32-like fp32 matmuls only see deviations from the pad agent (P12: deviation PCC 0.995706 -> 0.999991
|
| 21 |
+
on ``straight``). ``gelu_b1 = gelu(b_c1)`` also lets ``channel_pre`` itself run pad-relative:
|
| 22 |
+
``z - c0 = (gelu(x @ W_c1 + b_c1) - gelu_b1) @ W_c2``.
|
| 23 |
+
"""
|
| 24 |
+
from __future__ import annotations
|
| 25 |
+
|
| 26 |
+
from dataclasses import dataclass
|
| 27 |
+
from typing import Dict, Mapping, Sequence, Tuple
|
| 28 |
+
|
| 29 |
+
import numpy as np
|
| 30 |
+
import torch
|
| 31 |
+
import torch.nn.functional as F
|
| 32 |
+
|
| 33 |
+
from . import config as C
|
| 34 |
+
from .model import Decoder, _t
|
| 35 |
+
|
| 36 |
+
__all__ = ["AdaLNTables", "adaln_tables", "decoder_forward_folded", "IslandConstants", "island_constants",
|
| 37 |
+
"island_forward_pad_relative", "island_forward_exported", "cross_kv"]
|
| 38 |
+
|
| 39 |
+
|
| 40 |
+
# ---- 1. adaLN tables ---------------------------------------------------------------------------------------------
|
| 41 |
+
@dataclass
|
| 42 |
+
class AdaLNTables:
|
| 43 |
+
"""float32 tables for K evaluation times: ``temb [K, 256]``; per DiT block ``i``: ``norm1_gamma``,
|
| 44 |
+
``norm1_beta``, ``gate_msa``, ``norm2_gamma``, ``norm2_beta``, ``gate_mlp`` ``[K, 256]``; the final layer's
|
| 45 |
+
``final_gamma`` / ``final_beta`` ``[K, 256]``."""
|
| 46 |
+
|
| 47 |
+
eval_times: Tuple[float, ...]
|
| 48 |
+
temb: np.ndarray
|
| 49 |
+
blocks: Tuple[Dict[str, np.ndarray], ...]
|
| 50 |
+
final_gamma: np.ndarray
|
| 51 |
+
final_beta: np.ndarray
|
| 52 |
+
|
| 53 |
+
|
| 54 |
+
def adaln_tables(params: Mapping[str, np.ndarray], eval_times: Sequence[float]) -> AdaLNTables:
|
| 55 |
+
p = {k: torch.from_numpy(np.array(v, np.float64)) for k, v in params.items() if k.startswith("decoder.dit")}
|
| 56 |
+
|
| 57 |
+
def lin(x, n):
|
| 58 |
+
return x @ p[f"{n}.w"] + p[f"{n}.b"]
|
| 59 |
+
|
| 60 |
+
t = torch.tensor([[float(np.float32(v))] * C.DIT_TIME_DIM for v in eval_times], dtype=torch.float64)
|
| 61 |
+
c = lin(F.gelu(lin(t, "decoder.dit.t_embedder.fc1")), "decoder.dit.t_embedder.fc2") # [K, 256]
|
| 62 |
+
silu = c * torch.sigmoid(c)
|
| 63 |
+
blocks = []
|
| 64 |
+
for i in range(C.DIT_DEPTH):
|
| 65 |
+
B = f"decoder.dit.blocks.{i}"
|
| 66 |
+
sh_msa, sc_msa, g_msa, sh_mlp, sc_mlp, g_mlp = torch.split(lin(silu, f"{B}.adaLN_modulation"), C.HIDDEN_DIM, -1)
|
| 67 |
+
g1, b1 = p[f"{B}.norm1.gamma"], p[f"{B}.norm1.beta"]
|
| 68 |
+
g2, b2 = p[f"{B}.norm2.gamma"], p[f"{B}.norm2.beta"]
|
| 69 |
+
blocks.append({k: v.to(torch.float32).numpy() for k, v in {
|
| 70 |
+
"norm1_gamma": g1 * (1 + sc_msa), "norm1_beta": b1 * (1 + sc_msa) + sh_msa, "gate_msa": g_msa,
|
| 71 |
+
"norm2_gamma": g2 * (1 + sc_mlp), "norm2_beta": b2 * (1 + sc_mlp) + sh_mlp, "gate_mlp": g_mlp}.items()})
|
| 72 |
+
Fn = "decoder.dit.final_layer"
|
| 73 |
+
sh, sc = torch.split(lin(silu, f"{Fn}.adaLN_modulation"), C.HIDDEN_DIM, -1)
|
| 74 |
+
gf, bf = p[f"{Fn}.norm_final.gamma"], p[f"{Fn}.norm_final.beta"]
|
| 75 |
+
return AdaLNTables(tuple(float(v) for v in eval_times), c.to(torch.float32).numpy(), tuple(blocks),
|
| 76 |
+
(gf * (1 + sc)).to(torch.float32).numpy(), (bf * (1 + sc) + sh).to(torch.float32).numpy())
|
| 77 |
+
|
| 78 |
+
|
| 79 |
+
def cross_kv(params: Mapping[str, np.ndarray], encoding: np.ndarray) -> Tuple[np.ndarray, ...]:
|
| 80 |
+
"""Hoisted cross-attention K|V ``[564, 512]`` of each DiT block (float32 matmul, as the exported graph)."""
|
| 81 |
+
enc = np.asarray(encoding, np.float32)
|
| 82 |
+
return tuple((enc @ params[f"decoder.dit.blocks.{i}.cross_attn.kv.w"]
|
| 83 |
+
+ params[f"decoder.dit.blocks.{i}.cross_attn.kv.b"]).astype(np.float32) for i in range(C.DIT_DEPTH))
|
| 84 |
+
|
| 85 |
+
|
| 86 |
+
@torch.no_grad()
|
| 87 |
+
def decoder_forward_folded(dec: Decoder, tables: AdaLNTables, k: int, x_t: np.ndarray, kv: Sequence[torch.Tensor],
|
| 88 |
+
agent_valid: np.ndarray) -> torch.Tensor:
|
| 89 |
+
"""One decoder evaluation with the step-``k`` tables: LayerNorms with folded affine (no SiLU / adaLN matmuls),
|
| 90 |
+
gates as constant rows. Same result as ``Decoder.forward`` up to float32 rounding."""
|
| 91 |
+
from .model import _key_bias
|
| 92 |
+
|
| 93 |
+
P = int(np.shape(x_t)[0])
|
| 94 |
+
x = dec.mlp(_t(x_t).reshape(P, C.DIT_INPUT_DIM), "decoder.dit.preproj")
|
| 95 |
+
emb = dec.p["decoder.dit.agent_embedding"]
|
| 96 |
+
x = x + torch.cat([emb[0:1], emb[1:2].expand(P - 1, -1)], dim=0)
|
| 97 |
+
key_bias = _key_bias(agent_valid)
|
| 98 |
+
|
| 99 |
+
def ln_folded(h, gamma, beta):
|
| 100 |
+
return F.layer_norm(h, h.shape[-1:], _t(gamma), _t(beta), eps=C.LN_EPS)
|
| 101 |
+
|
| 102 |
+
for i in range(C.DIT_DEPTH):
|
| 103 |
+
B, T = f"decoder.dit.blocks.{i}", tables.blocks[i]
|
| 104 |
+
h = ln_folded(x, T["norm1_gamma"][k], T["norm1_beta"][k])
|
| 105 |
+
q, kk, v = torch.split(dec.linear(h, f"{B}.attn.qkv"), C.HIDDEN_DIM, dim=-1)
|
| 106 |
+
x = x + _t(T["gate_msa"][k]) * dec.linear(dec.attend(q, kk, v, key_bias), f"{B}.attn.out")
|
| 107 |
+
h = ln_folded(x, T["norm2_gamma"][k], T["norm2_beta"][k])
|
| 108 |
+
x = x + _t(T["gate_mlp"][k]) * dec.mlp(h, f"{B}.mlp1", "tanh")
|
| 109 |
+
heads = dec.attention(dec.ln(x, f"{B}.norm3"), kv[i], f"{B}.cross_attn.q", None)
|
| 110 |
+
x = x + dec.linear(heads, f"{B}.cross_attn.out")
|
| 111 |
+
x = x + dec.mlp(dec.ln(x, f"{B}.norm4"), f"{B}.mlp2", "tanh")
|
| 112 |
+
Fn = "decoder.dit.final_layer"
|
| 113 |
+
h = ln_folded(x, tables.final_gamma[k], tables.final_beta[k])
|
| 114 |
+
h = F.gelu(dec.linear(dec.ln(h, f"{Fn}.proj.0"), f"{Fn}.proj.1"), approximate="tanh")
|
| 115 |
+
return dec.linear(dec.ln(h, f"{Fn}.proj.3"), f"{Fn}.proj.4").reshape(P, C.OUTPUT_T + 1, C.POSE_DIM)
|
| 116 |
+
|
| 117 |
+
|
| 118 |
+
# ---- 3. pad-relative island ---------------------------------------------------------------------------------------
|
| 119 |
+
@dataclass
|
| 120 |
+
class IslandConstants:
|
| 121 |
+
"""Pad-agent constants of one pre-projection island (float32, built in float64). ``valid_rows`` are the time
|
| 122 |
+
rows that can be non-zero after the truncation (ego 0..5, neighbours 25..30); ``w_t1_valid`` is
|
| 123 |
+
``W_t1[valid_rows]`` ``[6, 64]``."""
|
| 124 |
+
|
| 125 |
+
category: str
|
| 126 |
+
valid_rows: Tuple[int, ...]
|
| 127 |
+
gelu_b1: np.ndarray # [128] gelu(b_c1): channel_pre's first activation of a zero row
|
| 128 |
+
c0: np.ndarray # [128] channel_pre of a zero row
|
| 129 |
+
t1_pad: np.ndarray # [128, 64]
|
| 130 |
+
g_pad: np.ndarray # [128, 64]
|
| 131 |
+
t2_pad: np.ndarray # [128, 64]
|
| 132 |
+
w_t1_valid: np.ndarray # [6, 64]
|
| 133 |
+
|
| 134 |
+
|
| 135 |
+
ISLAND_ROWS = {"neighbor": tuple(range(C.NEIGHBOR_HISTORY_KEEP.start, C.NEIGHBOR_HISTORY_KEEP.stop)),
|
| 136 |
+
"ego": tuple(range(C.EGO_HISTORY_KEEP.start, C.EGO_HISTORY_KEEP.stop))}
|
| 137 |
+
ISLAND_MODULE = {"neighbor": "encoder.neighbor_encoder", "ego": "encoder.ego_encoder"}
|
| 138 |
+
|
| 139 |
+
|
| 140 |
+
def island_constants(params: Mapping[str, np.ndarray], category: str) -> IslandConstants:
|
| 141 |
+
N = ISLAND_MODULE[category]
|
| 142 |
+
w = {k: torch.from_numpy(np.array(v, np.float64)) for k, v in params.items() if k.startswith(N + ".")}
|
| 143 |
+
gelu_b1 = F.gelu(w[f"{N}.channel_pre_project.fc1.b"])
|
| 144 |
+
c0 = gelu_b1 @ w[f"{N}.channel_pre_project.fc2.w"] + w[f"{N}.channel_pre_project.fc2.b"]
|
| 145 |
+
w_t1 = w[f"{N}.token_pre_project.fc1.w"] # [31, 64]
|
| 146 |
+
t1_pad = w[f"{N}.token_pre_project.fc1.b"][None, :] + c0[:, None] * w_t1.sum(0)[None, :]
|
| 147 |
+
g_pad = F.gelu(t1_pad)
|
| 148 |
+
t2_pad = g_pad @ w[f"{N}.token_pre_project.fc2.w"] + w[f"{N}.token_pre_project.fc2.b"]
|
| 149 |
+
rows = ISLAND_ROWS[category]
|
| 150 |
+
f32 = lambda t: t.to(torch.float32).numpy() # noqa: E731
|
| 151 |
+
return IslandConstants(category, rows, f32(gelu_b1), f32(c0), f32(t1_pad), f32(g_pad), f32(t2_pad),
|
| 152 |
+
f32(w_t1[list(rows)]))
|
| 153 |
+
|
| 154 |
+
|
| 155 |
+
def island_forward_exported(params: Mapping[str, np.ndarray], category: str, x: np.ndarray,
|
| 156 |
+
dtype: torch.dtype = torch.float64) -> Tuple[torch.Tensor, torch.Tensor]:
|
| 157 |
+
"""The island as exported: ``x`` ``[E, 31, C_in]`` -> ``t1``, ``t2`` ``[E, 128, 64]`` (``dtype`` math)."""
|
| 158 |
+
N = ISLAND_MODULE[category]
|
| 159 |
+
w = {k: torch.from_numpy(np.array(v, np.float64)).to(dtype) for k, v in params.items() if k.startswith(N + ".")}
|
| 160 |
+
|
| 161 |
+
def lin(h, n):
|
| 162 |
+
return h @ w[f"{N}.{n}.w"] + w[f"{N}.{n}.b"]
|
| 163 |
+
|
| 164 |
+
z = lin(F.gelu(lin(torch.from_numpy(np.asarray(x, np.float64)).to(dtype), "channel_pre_project.fc1")),
|
| 165 |
+
"channel_pre_project.fc2")
|
| 166 |
+
t1 = lin(z.transpose(1, 2), "token_pre_project.fc1")
|
| 167 |
+
return t1, lin(F.gelu(t1), "token_pre_project.fc2")
|
| 168 |
+
|
| 169 |
+
|
| 170 |
+
def island_forward_pad_relative(params: Mapping[str, np.ndarray], consts: IslandConstants, x: np.ndarray,
|
| 171 |
+
dtype: torch.dtype = torch.float64) -> Tuple[torch.Tensor, torch.Tensor]:
|
| 172 |
+
"""The pad-relative island on the valid rows only (``x`` ``[E, 31, C_in]``; rows outside ``valid_rows`` must be
|
| 173 |
+
zero, which the truncation guarantees). Matmuls see deviations from the pad agent only."""
|
| 174 |
+
N = ISLAND_MODULE[consts.category]
|
| 175 |
+
w = {k: torch.from_numpy(np.array(v, np.float64)).to(dtype) for k, v in params.items() if k.startswith(N + ".")}
|
| 176 |
+
cst = {k: torch.from_numpy(np.array(getattr(consts, k), np.float64)).to(dtype)
|
| 177 |
+
for k in ("gelu_b1", "c0", "t1_pad", "g_pad", "t2_pad", "w_t1_valid")}
|
| 178 |
+
xv = torch.from_numpy(np.asarray(x, np.float64)[:, list(consts.valid_rows)]).to(dtype) # [E, 6, C_in]
|
| 179 |
+
h = F.gelu(xv @ w[f"{N}.channel_pre_project.fc1.w"] + w[f"{N}.channel_pre_project.fc1.b"]) - cst["gelu_b1"]
|
| 180 |
+
dz = h @ w[f"{N}.channel_pre_project.fc2.w"] # z - c0, [E, 6, 128]
|
| 181 |
+
t1 = cst["t1_pad"] + dz.transpose(1, 2) @ cst["w_t1_valid"] # [E, 128, 64]
|
| 182 |
+
t2 = cst["t2_pad"] + (F.gelu(t1) - cst["g_pad"]) @ w[f"{N}.token_pre_project.fc2.w"]
|
| 183 |
+
return t1, t2
|
code/tt_diffusion_planner/reference/weights.py
ADDED
|
@@ -0,0 +1,442 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-License-Identifier: Apache-2.0
|
| 2 |
+
"""Weights of the Diffusion Planner v5.0 export, read as DATA from the three ONNX files and the param JSON.
|
| 3 |
+
|
| 4 |
+
Every tensor is addressed through the graph node that consumes it (``ttaw.weights.OnnxWeights``), never by
|
| 5 |
+
initializer name: most MatMul weights are anonymous (``onnx::MatMul_4039``), several biases are deduplicated across
|
| 6 |
+
modules (``static_encoder/projection/fc2`` reads ``...fc1.bias``; ``route_encoder/attribute_emb`` reads
|
| 7 |
+
``...speed_limit_emb.bias``), the decoder cross-attention K/V of all three blocks is one unnamed ``[256, 1536]``
|
| 8 |
+
MatMul followed by a Split, and the fusion / cross-attention Q-K-V biases and weights are constant-folded tensors
|
| 9 |
+
(SPEC 6.2). The result is one flat ``{canonical name: float32 array}`` dict that the CPU reference and the ttnn
|
| 10 |
+
graph both consume:
|
| 11 |
+
|
| 12 |
+
- linear layers: ``<module>.w`` ``[in, out]`` (``y = x @ w + b``) and ``<module>.b`` ``[out]``;
|
| 13 |
+
- LayerNorms: ``<module>.gamma`` / ``<module>.beta``;
|
| 14 |
+
- attention: ``...attn.q`` / ``.kv`` (fusion: Q from LN(x), K|V from x), ``...attn.qkv`` (DiT self-attention),
|
| 15 |
+
``...cross_attn.q`` / ``.kv`` (K|V of one block, a column block of the fused cross K/V MatMul), ``...out``;
|
| 16 |
+
- embeddings: ``decoder.dit.agent_embedding`` ``[2, 256]`` (ego, neighbour), ``encoder.route_position_embedding``
|
| 17 |
+
``[25, 256]``, ``encoder.<lane|route>_encoder.unknown_speed_emb`` ``[128]``.
|
| 18 |
+
|
| 19 |
+
There is no BatchNorm in this network, so nothing is folded here. The exact rewrites the TT port applies on top of
|
| 20 |
+
these tensors (per-step adaLN tables folded into the LayerNorm affine, hoisted cross K/V, the pad-relative fp32
|
| 21 |
+
pre-projection island) are in :mod:`.rewrites`, computed from this dict in float64 with one final rounding.
|
| 22 |
+
|
| 23 |
+
The loader also checks the invariants the reference and the port rely on (LayerNorm epsilon 1e-5 everywhere, exact
|
| 24 |
+
GELU in the encoder and the decoder pre-projection / t-embedder, tanh GELU in the DiT MLPs and the final projection,
|
| 25 |
+
attention scale 1/sqrt(32), the ``-inf`` key mask), so a different export fails loudly instead of silently.
|
| 26 |
+
"""
|
| 27 |
+
from __future__ import annotations
|
| 28 |
+
|
| 29 |
+
import json
|
| 30 |
+
import math
|
| 31 |
+
from dataclasses import dataclass, field
|
| 32 |
+
from pathlib import Path
|
| 33 |
+
from typing import Any, Dict, List, Mapping, Optional, Tuple
|
| 34 |
+
|
| 35 |
+
import numpy as np
|
| 36 |
+
|
| 37 |
+
from . import config as C
|
| 38 |
+
from ..ttaw.weights import OnnxWeights, file_sha256
|
| 39 |
+
|
| 40 |
+
__all__ = ["PlannerWeights", "load_weights", "load_param_json", "Normalization", "find_weights_dir",
|
| 41 |
+
"MIXER_ENCODERS", "SMALL_ENCODERS"]
|
| 42 |
+
|
| 43 |
+
# categories with an MLP-Mixer trunk -> ONNX module prefix
|
| 44 |
+
MIXER_ENCODERS = {"ego": "ego_encoder", "neighbor": "neighbor_encoder", "lane": "lane_encoder",
|
| 45 |
+
"route": "route_encoder", "polygon": "polygon_encoder", "line_string": "line_string_encoder"}
|
| 46 |
+
# categories encoded by a channel MLP + LayerNorm + projection (goal pose, ego shape, turn indicators)
|
| 47 |
+
SMALL_ENCODERS = {"goal": "goal_pose_encoder", "ego_shape": "ego_shape_encoder", "turn": "turn_indicator_encoder"}
|
| 48 |
+
|
| 49 |
+
|
| 50 |
+
# ------------------------------------------------------------------------------------------------ param JSON
|
| 51 |
+
|
| 52 |
+
@dataclass(frozen=True)
|
| 53 |
+
class Normalization:
|
| 54 |
+
"""``observation_normalizer`` (per input tensor: mean / std over the last dim) and ``state_normalizer``
|
| 55 |
+
(per agent and pose dim) of ``diffusion_planner.param.json`` (PKG/include/.../utils/arg_reader.hpp:80-140)."""
|
| 56 |
+
|
| 57 |
+
observation: Dict[str, Tuple[np.ndarray, np.ndarray]]
|
| 58 |
+
state_mean: np.ndarray # [321, 4] (or [4])
|
| 59 |
+
state_std: np.ndarray
|
| 60 |
+
major_version: int
|
| 61 |
+
args: Dict[str, Any] = field(default_factory=dict)
|
| 62 |
+
|
| 63 |
+
def state(self) -> Tuple[np.ndarray, np.ndarray]:
|
| 64 |
+
"""``(mean, std)`` broadcastable to ``[321, T, 4]``."""
|
| 65 |
+
return _per_agent(self.state_mean), _per_agent(self.state_std)
|
| 66 |
+
|
| 67 |
+
|
| 68 |
+
def _per_agent(v: np.ndarray) -> np.ndarray:
|
| 69 |
+
v = np.asarray(v, np.float32).reshape(-1)
|
| 70 |
+
if v.size == C.POSE_DIM:
|
| 71 |
+
return v.reshape(1, 1, C.POSE_DIM)
|
| 72 |
+
if v.size == C.MAX_NUM_AGENTS * C.POSE_DIM:
|
| 73 |
+
return v.reshape(C.MAX_NUM_AGENTS, 1, C.POSE_DIM)
|
| 74 |
+
raise ValueError(f"unsupported state normalizer size {v.size}")
|
| 75 |
+
|
| 76 |
+
|
| 77 |
+
def load_param_json(path: Path) -> Normalization:
|
| 78 |
+
"""Parse ``diffusion_planner.param.json`` like ``arg_reader.hpp``; refuses a major version other than 5."""
|
| 79 |
+
with open(path) as f:
|
| 80 |
+
j = json.load(f)
|
| 81 |
+
major = int(j.get("major_version", -1))
|
| 82 |
+
if major != C.WEIGHT_MAJOR_VERSION:
|
| 83 |
+
raise ValueError(f"{path}: major_version {major}, this port needs {C.WEIGHT_MAJOR_VERSION} (constants.hpp:22)")
|
| 84 |
+
obs = {}
|
| 85 |
+
for key, v in j["observation_normalizer"].items():
|
| 86 |
+
mean = np.asarray(v.get("mean", []), np.float32).reshape(-1)
|
| 87 |
+
std = np.asarray(v.get("std", []), np.float32).reshape(-1)
|
| 88 |
+
if mean.shape != std.shape:
|
| 89 |
+
raise ValueError(f"{path}: normalizer {key!r} mean / std sizes differ")
|
| 90 |
+
obs[key] = (mean, std)
|
| 91 |
+
sn = j["state_normalizer"]
|
| 92 |
+
args = {k: v for k, v in j.items() if k not in ("observation_normalizer", "state_normalizer")}
|
| 93 |
+
return Normalization(obs, np.asarray(sn["mean"], np.float32), np.asarray(sn["std"], np.float32), major, args)
|
| 94 |
+
|
| 95 |
+
|
| 96 |
+
# ------------------------------------------------------------------------------------------------ ONNX weights
|
| 97 |
+
|
| 98 |
+
@dataclass
|
| 99 |
+
class PlannerWeights:
|
| 100 |
+
"""The flat canonical parameter dict plus provenance (file sha256) and the export facts checked at load."""
|
| 101 |
+
|
| 102 |
+
params: Dict[str, np.ndarray]
|
| 103 |
+
normalization: Normalization
|
| 104 |
+
sha256: Dict[str, str]
|
| 105 |
+
facts: Dict[str, Any]
|
| 106 |
+
path: Path
|
| 107 |
+
|
| 108 |
+
def __getitem__(self, name: str) -> np.ndarray:
|
| 109 |
+
return self.params[name]
|
| 110 |
+
|
| 111 |
+
def linear(self, name: str) -> Tuple[np.ndarray, np.ndarray]:
|
| 112 |
+
return self.params[f"{name}.w"], self.params[f"{name}.b"]
|
| 113 |
+
|
| 114 |
+
def num_parameters(self) -> int:
|
| 115 |
+
return int(sum(v.size for v in self.params.values()))
|
| 116 |
+
|
| 117 |
+
def to_torch(self, dtype: Any = None) -> Dict[str, Any]:
|
| 118 |
+
"""``{name: torch.Tensor}`` (float32 by default; float64 for the fp64 rewrites)."""
|
| 119 |
+
import torch
|
| 120 |
+
|
| 121 |
+
dt = dtype or torch.float32
|
| 122 |
+
return {k: torch.from_numpy(np.ascontiguousarray(v)).to(dt) for k, v in self.params.items()}
|
| 123 |
+
|
| 124 |
+
|
| 125 |
+
class _Reader:
|
| 126 |
+
"""Node-addressed access to one ONNX file with the conventions of this export."""
|
| 127 |
+
|
| 128 |
+
def __init__(self, path: Path):
|
| 129 |
+
self.w = OnnxWeights(path)
|
| 130 |
+
self.facts: Dict[str, Any] = {"gelu": {}, "ln_eps": set()}
|
| 131 |
+
|
| 132 |
+
def has_node(self, name: str) -> bool:
|
| 133 |
+
try:
|
| 134 |
+
self.w.node(name)
|
| 135 |
+
return True
|
| 136 |
+
except KeyError:
|
| 137 |
+
return False
|
| 138 |
+
|
| 139 |
+
def const_input(self, node_name: str) -> np.ndarray:
|
| 140 |
+
"""The single constant input of a binary node (the bias of a MatMul + Add pair)."""
|
| 141 |
+
node = self.w.node(node_name)
|
| 142 |
+
consts = [t for t in node.inputs if t and self.w.has(t)]
|
| 143 |
+
if len(consts) != 1:
|
| 144 |
+
raise ValueError(f"{node_name}: expected one constant input, found {len(consts)}")
|
| 145 |
+
return np.asarray(self.w.array(consts[0]), np.float32)
|
| 146 |
+
|
| 147 |
+
def bias_after(self, matmul_name: str) -> np.ndarray:
|
| 148 |
+
add = self.w.consumer_of(self.w.node(matmul_name).outputs[0], "Add")
|
| 149 |
+
return self.const_input(add.name)
|
| 150 |
+
|
| 151 |
+
def linear(self, path: str) -> Tuple[np.ndarray, np.ndarray]:
|
| 152 |
+
"""``(w [in, out], b [out])`` of the torch ``nn.Linear`` exported under ``<path>``: either ``<path>/MatMul``
|
| 153 |
+
followed by an ``Add`` (3-D inputs) or ``<path>/Gemm`` (2-D inputs); exactly one of the two must exist."""
|
| 154 |
+
has_mm, has_gemm = self.has_node(f"{path}/MatMul"), self.has_node(f"{path}/Gemm")
|
| 155 |
+
if has_mm == has_gemm:
|
| 156 |
+
raise KeyError(f"{path}: expected exactly one of MatMul / Gemm, found {has_mm=} {has_gemm=}")
|
| 157 |
+
if has_mm:
|
| 158 |
+
w = np.asarray(self.w.matmul_weight(f"{path}/MatMul"), np.float32)
|
| 159 |
+
return w, self.bias_after(f"{path}/MatMul")
|
| 160 |
+
return self.gemm(f"{path}/Gemm")
|
| 161 |
+
|
| 162 |
+
def gemm(self, node_name: str) -> Tuple[np.ndarray, np.ndarray]:
|
| 163 |
+
"""``(w [in, out], b [out])`` of a ``Gemm`` node (``y = x @ W^T + b`` with ``transB = 1``)."""
|
| 164 |
+
g = self.w.gemm(node_name)
|
| 165 |
+
if g.trans_a or g.alpha != 1.0 or g.beta != 1.0 or g.bias is None:
|
| 166 |
+
raise ValueError(f"{node_name}: unexpected attributes {g.trans_a=} {g.alpha=} {g.beta=}")
|
| 167 |
+
w = np.asarray(g.weight, np.float32)
|
| 168 |
+
return (w.T if g.trans_b else w).copy(), np.asarray(g.bias, np.float32).reshape(-1)
|
| 169 |
+
|
| 170 |
+
def layer_norm(self, path: str) -> Tuple[np.ndarray, np.ndarray]:
|
| 171 |
+
node = self.w.node(f"{path}/LayerNormalization")
|
| 172 |
+
self.facts["ln_eps"].add(float(node.attrs.get("epsilon", 1e-5)))
|
| 173 |
+
if int(node.attrs.get("axis", -1)) != -1:
|
| 174 |
+
raise ValueError(f"{path}: LayerNormalization over axis {node.attrs.get('axis')}")
|
| 175 |
+
return (np.asarray(self.w.param(node.name, 1), np.float32),
|
| 176 |
+
np.asarray(self.w.param(node.name, 2), np.float32))
|
| 177 |
+
|
| 178 |
+
def gelu(self, path: str) -> str:
|
| 179 |
+
node = self.w.node(f"{path}/Gelu")
|
| 180 |
+
mode = str(node.attrs.get("approximate", "none"))
|
| 181 |
+
self.facts["gelu"][path] = mode
|
| 182 |
+
return mode
|
| 183 |
+
|
| 184 |
+
def scalar(self, node_name: str) -> float:
|
| 185 |
+
return float(np.asarray(self.const_input(node_name)).reshape(-1)[0])
|
| 186 |
+
|
| 187 |
+
|
| 188 |
+
def _put_linear(params: Dict[str, np.ndarray], name: str, wb: Tuple[np.ndarray, np.ndarray]) -> None:
|
| 189 |
+
params[f"{name}.w"], params[f"{name}.b"] = wb
|
| 190 |
+
|
| 191 |
+
|
| 192 |
+
def _put_ln(params: Dict[str, np.ndarray], name: str, gb: Tuple[np.ndarray, np.ndarray]) -> None:
|
| 193 |
+
params[f"{name}.gamma"], params[f"{name}.beta"] = gb
|
| 194 |
+
|
| 195 |
+
|
| 196 |
+
def _mlp(r: _Reader, params: Dict[str, np.ndarray], onnx_path: str, name: str, gelu: str) -> None:
|
| 197 |
+
_put_linear(params, f"{name}.fc1", r.linear(f"{onnx_path}/fc1"))
|
| 198 |
+
_put_linear(params, f"{name}.fc2", r.linear(f"{onnx_path}/fc2"))
|
| 199 |
+
mode = r.gelu(f"{onnx_path}/act")
|
| 200 |
+
if mode != gelu:
|
| 201 |
+
raise ValueError(f"{onnx_path}: GELU approximate={mode!r}, expected {gelu!r}")
|
| 202 |
+
|
| 203 |
+
|
| 204 |
+
def _read_encoder(path: Path, params: Dict[str, np.ndarray]) -> Dict[str, Any]:
|
| 205 |
+
r = _Reader(path)
|
| 206 |
+
for cat, mod in {**MIXER_ENCODERS, **SMALL_ENCODERS}.items():
|
| 207 |
+
P, N = f"/encoder/{mod}", f"encoder.{mod}"
|
| 208 |
+
_mlp(r, params, f"{P}/channel_pre_project", f"{N}.channel_pre_project", "none")
|
| 209 |
+
if cat in MIXER_ENCODERS:
|
| 210 |
+
_mlp(r, params, f"{P}/token_pre_project", f"{N}.token_pre_project", "none")
|
| 211 |
+
for i in range(C.MIXER_DEPTH):
|
| 212 |
+
B, BN = f"{P}/blocks.{i}", f"{N}.blocks.{i}"
|
| 213 |
+
_put_ln(params, f"{BN}.norm1", r.layer_norm(f"{B}/norm1"))
|
| 214 |
+
_mlp(r, params, f"{B}/tokens_mlp", f"{BN}.tokens_mlp", "none")
|
| 215 |
+
_put_ln(params, f"{BN}.norm2", r.layer_norm(f"{B}/norm2"))
|
| 216 |
+
_mlp(r, params, f"{B}/channels_mlp", f"{BN}.channels_mlp", "none")
|
| 217 |
+
_put_ln(params, f"{N}.norm", r.layer_norm(f"{P}/norm"))
|
| 218 |
+
_mlp(r, params, f"{P}/emb_project", f"{N}.emb_project", "none")
|
| 219 |
+
_put_linear(params, "encoder.neighbor_encoder.type_emb", r.linear("/encoder/neighbor_encoder/type_emb"))
|
| 220 |
+
for mod in ("lane_encoder", "route_encoder"):
|
| 221 |
+
_put_linear(params, f"encoder.{mod}.speed_limit_emb", r.linear(f"/encoder/{mod}/speed_limit_emb"))
|
| 222 |
+
_put_linear(params, f"encoder.{mod}.attribute_emb", r.linear(f"/encoder/{mod}/attribute_emb"))
|
| 223 |
+
unk = r.w.param(f"/encoder/{mod}/unknown_speed_emb/Gather", 0)
|
| 224 |
+
params[f"encoder.{mod}.unknown_speed_emb"] = np.asarray(unk, np.float32).reshape(-1)
|
| 225 |
+
_mlp(r, params, "/encoder/static_encoder/projection", "encoder.static_encoder.projection", "none")
|
| 226 |
+
_put_linear(params, "encoder.pos_emb", r.linear("/encoder/pos_emb"))
|
| 227 |
+
rpe = np.asarray(r.w.param("/encoder/Slice_5", 0), np.float32)
|
| 228 |
+
params["encoder.route_position_embedding"] = rpe.reshape(C.NUM_SEGMENTS_IN_ROUTE, C.HIDDEN_DIM)
|
| 229 |
+
scales, fills = set(), set()
|
| 230 |
+
for i in range(C.FUSION_DEPTH):
|
| 231 |
+
B, BN = f"/encoder/fusion/blocks.{i}", f"encoder.fusion.blocks.{i}"
|
| 232 |
+
_put_ln(params, f"{BN}.norm1", r.layer_norm(f"{B}/norm1"))
|
| 233 |
+
q_w = np.asarray(r.w.matmul_weight(f"{B}/attn/MatMul"), np.float32)
|
| 234 |
+
_put_linear(params, f"{BN}.attn.q", (q_w, r.bias_after(f"{B}/attn/MatMul")))
|
| 235 |
+
kv_w = np.asarray(r.w.matmul_weight(f"{B}/attn/MatMul_1"), np.float32)
|
| 236 |
+
_put_linear(params, f"{BN}.attn.kv", (kv_w, r.bias_after(f"{B}/attn/MatMul_1")))
|
| 237 |
+
_put_linear(params, f"{BN}.attn.out", r.gemm(f"{B}/attn/Gemm"))
|
| 238 |
+
_put_ln(params, f"{BN}.norm2", r.layer_norm(f"{B}/norm2"))
|
| 239 |
+
_mlp(r, params, f"{B}/mlp", f"{BN}.mlp", "none")
|
| 240 |
+
scales.add(r.scalar(f"{B}/attn/Mul_3"))
|
| 241 |
+
# the key-padding bias Where(mask, -inf, 0) is built once in block 0 and shared by the six blocks
|
| 242 |
+
fills.add(float(np.asarray(r.w.param("/encoder/fusion/blocks.0/attn/Where", 1)).reshape(-1)[0]))
|
| 243 |
+
_put_ln(params, "encoder.fusion.norm", r.layer_norm("/encoder/fusion/norm"))
|
| 244 |
+
return {"gelu": r.facts["gelu"], "ln_eps": r.facts["ln_eps"], "attn_scale": scales, "mask_fill": fills,
|
| 245 |
+
"sha256": r.w.sha256, "nodes": len(r.w.nodes())}
|
| 246 |
+
|
| 247 |
+
|
| 248 |
+
def _read_decoder(path: Path, params: Dict[str, np.ndarray]) -> Dict[str, Any]:
|
| 249 |
+
r = _Reader(path)
|
| 250 |
+
_mlp(r, params, "/dit/preproj", "decoder.dit.preproj", "none")
|
| 251 |
+
_mlp(r, params, "/dit/t_embedder", "decoder.dit.t_embedder", "none")
|
| 252 |
+
ego_row = np.asarray(r.w.param("/dit/Concat_3", 0), np.float32).reshape(-1)
|
| 253 |
+
expand = r.w.producer(r.w.node("/dit/Concat_3").inputs[1])
|
| 254 |
+
nb_row = np.asarray(r.w.param(expand.name, 0), np.float32).reshape(-1)
|
| 255 |
+
params["decoder.dit.agent_embedding"] = np.stack([ego_row, nb_row])
|
| 256 |
+
scales, fills = set(), set()
|
| 257 |
+
for i in range(C.DIT_DEPTH):
|
| 258 |
+
B, BN = f"/dit/blocks.{i}", f"decoder.dit.blocks.{i}"
|
| 259 |
+
_put_linear(params, f"{BN}.adaLN_modulation", r.linear(f"{B}/adaLN_modulation/adaLN_modulation.1"))
|
| 260 |
+
for n in ("norm1", "norm2", "norm3", "norm4"):
|
| 261 |
+
_put_ln(params, f"{BN}.{n}", r.layer_norm(f"{B}/{n}"))
|
| 262 |
+
qkv_w = np.asarray(r.w.matmul_weight(f"{B}/attn/MatMul"), np.float32)
|
| 263 |
+
_put_linear(params, f"{BN}.attn.qkv", (qkv_w, r.bias_after(f"{B}/attn/MatMul")))
|
| 264 |
+
_put_linear(params, f"{BN}.attn.out", r.gemm(f"{B}/attn/Gemm"))
|
| 265 |
+
_mlp(r, params, f"{B}/mlp1", f"{BN}.mlp1", "tanh")
|
| 266 |
+
q_w = np.asarray(r.w.matmul_weight(f"{B}/cross_attn/MatMul"), np.float32)
|
| 267 |
+
_put_linear(params, f"{BN}.cross_attn.q", (q_w, r.bias_after(f"{B}/cross_attn/MatMul")))
|
| 268 |
+
# K|V of this block: the bias is the constant input of cross_attn/Add_5, the weight a column block of the
|
| 269 |
+
# unnamed [256, 1536] MatMul whose output is split three ways (one 512-wide K|V slice per block)
|
| 270 |
+
add5 = r.w.node(f"{B}/cross_attn/Add_5")
|
| 271 |
+
kv_b = r.const_input(add5.name)
|
| 272 |
+
(kv_in,) = [t for t in add5.inputs if t and not r.w.has(t)]
|
| 273 |
+
split = r.w.producer(kv_in)
|
| 274 |
+
if split is None or split.op_type != "Split":
|
| 275 |
+
raise ValueError(f"{B}/cross_attn/Add_5: K|V does not come from a Split")
|
| 276 |
+
part = list(split.outputs).index(kv_in)
|
| 277 |
+
sizes = [int(s) for s in np.asarray(r.w.array(split.inputs[1])).reshape(-1)]
|
| 278 |
+
fused = r.w.producer(split.inputs[0])
|
| 279 |
+
fused_w = np.asarray(r.w.matmul_weight(fused.name), np.float32)
|
| 280 |
+
start = int(sum(sizes[:part]))
|
| 281 |
+
_put_linear(params, f"{BN}.cross_attn.kv", (fused_w[:, start:start + sizes[part]].copy(), kv_b))
|
| 282 |
+
_put_linear(params, f"{BN}.cross_attn.out", r.gemm(f"{B}/cross_attn/Gemm"))
|
| 283 |
+
_mlp(r, params, f"{B}/mlp2", f"{BN}.mlp2", "tanh")
|
| 284 |
+
scales.add(r.scalar(f"{B}/attn/Mul_2"))
|
| 285 |
+
scales.add(r.scalar(f"{B}/cross_attn/Mul_2"))
|
| 286 |
+
fills.add(float(np.asarray(r.w.param("/dit/blocks.0/attn/Where", 1)).reshape(-1)[0])) # shared by the blocks
|
| 287 |
+
F = "/dit/final_layer"
|
| 288 |
+
_put_linear(params, "decoder.dit.final_layer.adaLN_modulation",
|
| 289 |
+
r.linear(f"{F}/adaLN_modulation/adaLN_modulation.1"))
|
| 290 |
+
_put_ln(params, "decoder.dit.final_layer.norm_final", r.layer_norm(f"{F}/norm_final"))
|
| 291 |
+
_put_ln(params, "decoder.dit.final_layer.proj.0", r.layer_norm(f"{F}/proj/proj.0"))
|
| 292 |
+
_put_linear(params, "decoder.dit.final_layer.proj.1", r.linear(f"{F}/proj/proj.1"))
|
| 293 |
+
gelu = r.gelu(f"{F}/proj/proj.2")
|
| 294 |
+
if gelu != "tanh":
|
| 295 |
+
raise ValueError(f"{F}/proj/proj.2: GELU approximate={gelu!r}, expected 'tanh'")
|
| 296 |
+
_put_ln(params, "decoder.dit.final_layer.proj.3", r.layer_norm(f"{F}/proj/proj.3"))
|
| 297 |
+
_put_linear(params, "decoder.dit.final_layer.proj.4", r.linear(f"{F}/proj/proj.4"))
|
| 298 |
+
return {"gelu": r.facts["gelu"], "ln_eps": r.facts["ln_eps"], "attn_scale": scales, "mask_fill": fills,
|
| 299 |
+
"sha256": r.w.sha256, "nodes": len(r.w.nodes())}
|
| 300 |
+
|
| 301 |
+
|
| 302 |
+
def _read_turn(path: Path, params: Dict[str, np.ndarray]) -> Dict[str, Any]:
|
| 303 |
+
r = _Reader(path)
|
| 304 |
+
_put_linear(params, "decoder.turn_indicator_predictor", r.linear("/turn_indicator_predictor"))
|
| 305 |
+
return {"sha256": r.w.sha256, "nodes": len(r.w.nodes())}
|
| 306 |
+
|
| 307 |
+
|
| 308 |
+
EXPECTED_SHAPES = { # spot checks of the canonical layout (SPEC 4.2-4.4)
|
| 309 |
+
"encoder.neighbor_encoder.channel_pre_project.fc1.w": (C.NEIGHBOR_FEATURE_DIM, C.MIXER_CHANNELS),
|
| 310 |
+
"encoder.neighbor_encoder.token_pre_project.fc1.w": (C.INPUT_T + 1, C.MIXER_TOKENS),
|
| 311 |
+
"encoder.lane_encoder.token_pre_project.fc1.w": (C.POINTS_PER_SEGMENT, C.MIXER_TOKENS),
|
| 312 |
+
"encoder.polygon_encoder.channel_pre_project.fc1.w": (C.POLYGON_FEATURE_DIM, C.MIXER_CHANNELS),
|
| 313 |
+
"encoder.polygon_encoder.token_pre_project.fc1.w": (C.POINTS_PER_POLYGON, C.MIXER_TOKENS),
|
| 314 |
+
"encoder.line_string_encoder.channel_pre_project.fc1.w": (C.LINE_STRING_FEATURE_DIM, C.MIXER_CHANNELS),
|
| 315 |
+
"encoder.lane_encoder.attribute_emb.w": (C.LANE_ATTRIBUTE_DIM, C.MIXER_CHANNELS),
|
| 316 |
+
"encoder.turn_indicator_encoder.channel_pre_project.fc1.w": (C.TURN_INDICATOR_HISTORY, C.MIXER_CHANNELS),
|
| 317 |
+
"encoder.pos_emb.w": (C.POS_FEATURE_DIM, C.HIDDEN_DIM),
|
| 318 |
+
"encoder.fusion.blocks.0.attn.kv.w": (C.HIDDEN_DIM, 2 * C.HIDDEN_DIM),
|
| 319 |
+
"decoder.dit.preproj.fc1.w": (C.DIT_INPUT_DIM, 512),
|
| 320 |
+
"decoder.dit.t_embedder.fc1.w": (C.DIT_TIME_DIM, 512),
|
| 321 |
+
"decoder.dit.blocks.0.adaLN_modulation.w": (C.HIDDEN_DIM, 6 * C.HIDDEN_DIM),
|
| 322 |
+
"decoder.dit.blocks.0.attn.qkv.w": (C.HIDDEN_DIM, 3 * C.HIDDEN_DIM),
|
| 323 |
+
"decoder.dit.blocks.2.cross_attn.kv.w": (C.HIDDEN_DIM, 2 * C.HIDDEN_DIM),
|
| 324 |
+
"decoder.dit.final_layer.proj.4.w": (C.DIT_MLP_DIM, C.DIT_INPUT_DIM),
|
| 325 |
+
"decoder.turn_indicator_predictor.w": (2 * len(C.TURN_HEAD_STEPS) + C.HIDDEN_DIM, C.TURN_INDICATOR_OUTPUT_DIM),
|
| 326 |
+
}
|
| 327 |
+
|
| 328 |
+
|
| 329 |
+
def _check_facts(facts: Dict[str, Any]) -> None:
|
| 330 |
+
eps = facts["encoder"]["ln_eps"] | facts["decoder"]["ln_eps"]
|
| 331 |
+
if eps != {C.LN_EPS} and not all(math.isclose(e, C.LN_EPS, rel_tol=1e-6) for e in eps):
|
| 332 |
+
raise ValueError(f"LayerNorm epsilons {sorted(eps)}, expected {C.LN_EPS}")
|
| 333 |
+
scales = facts["encoder"]["attn_scale"] | facts["decoder"]["attn_scale"]
|
| 334 |
+
if len(scales) != 1 or not math.isclose(scales.pop(), 1.0 / math.sqrt(C.HEAD_DIM), rel_tol=1e-6):
|
| 335 |
+
raise ValueError(f"attention scales {facts['encoder']['attn_scale'] | facts['decoder']['attn_scale']}")
|
| 336 |
+
fills = facts["encoder"]["mask_fill"] | facts["decoder"]["mask_fill"]
|
| 337 |
+
if fills != {float("-inf")}:
|
| 338 |
+
raise ValueError(f"attention mask fill values {fills}, expected -inf")
|
| 339 |
+
|
| 340 |
+
|
| 341 |
+
def find_weights_dir(explicit: Optional[str] = None) -> Optional[Path]:
|
| 342 |
+
"""A local directory holding the v5.0 files, or None: ``explicit`` > ``$DIFFUSION_PLANNER_WEIGHTS_DIR`` > the
|
| 343 |
+
workspace download (``assets/diffusion-planner/hf_diffusion_planner``) > the HF cache snapshot of the pinned
|
| 344 |
+
revision (``local_files_only``; never a network access)."""
|
| 345 |
+
import os
|
| 346 |
+
|
| 347 |
+
names = (C.ENCODER_ONNX, C.DECODER_ONNX, C.TURN_INDICATOR_ONNX, C.PARAM_JSON)
|
| 348 |
+
cands: List[Path] = []
|
| 349 |
+
for c in (explicit, os.environ.get("DIFFUSION_PLANNER_WEIGHTS_DIR")):
|
| 350 |
+
if c:
|
| 351 |
+
cands.append(Path(c).expanduser())
|
| 352 |
+
here = Path(__file__).resolve()
|
| 353 |
+
for parent in here.parents:
|
| 354 |
+
cands.append(parent / "assets" / "diffusion-planner" / "hf_diffusion_planner")
|
| 355 |
+
for c in cands:
|
| 356 |
+
if all((c / n).is_file() for n in names):
|
| 357 |
+
return c
|
| 358 |
+
try: # the HF cache of `from_pretrained` (offline lookup only)
|
| 359 |
+
from huggingface_hub import snapshot_download
|
| 360 |
+
|
| 361 |
+
p = Path(snapshot_download("AutowareFoundation/diffusion_planner",
|
| 362 |
+
revision="423efde67f5414734da43a7ad856c17ceb8b51aa",
|
| 363 |
+
allow_patterns=list(names), local_files_only=True))
|
| 364 |
+
if all((p / n).is_file() for n in names):
|
| 365 |
+
return p
|
| 366 |
+
except Exception: # noqa: BLE001 -- not cached / no huggingface_hub: no weights
|
| 367 |
+
pass
|
| 368 |
+
return None
|
| 369 |
+
|
| 370 |
+
|
| 371 |
+
def load_weights(weights_dir: Path, *, verify_sha256: bool = True) -> PlannerWeights:
|
| 372 |
+
"""Read the encoder / decoder / turn-indicator ONNX files and the param JSON of ``weights_dir``."""
|
| 373 |
+
weights_dir = Path(weights_dir)
|
| 374 |
+
sha = {n: file_sha256(weights_dir / n) for n in C.FILE_SHA256}
|
| 375 |
+
if verify_sha256:
|
| 376 |
+
bad = {n: s for n, s in sha.items() if s != C.FILE_SHA256[n]}
|
| 377 |
+
if bad:
|
| 378 |
+
raise ValueError(f"{weights_dir}: files differ from AutowareFoundation/diffusion_planner@v5.0: "
|
| 379 |
+
f"{sorted(bad)} (pass verify_sha256=False to load another export)")
|
| 380 |
+
params: Dict[str, np.ndarray] = {}
|
| 381 |
+
facts = {"encoder": _read_encoder(weights_dir / C.ENCODER_ONNX, params),
|
| 382 |
+
"decoder": _read_decoder(weights_dir / C.DECODER_ONNX, params),
|
| 383 |
+
"turn": _read_turn(weights_dir / C.TURN_INDICATOR_ONNX, params)}
|
| 384 |
+
_check_facts(facts)
|
| 385 |
+
for name, shape in EXPECTED_SHAPES.items():
|
| 386 |
+
if tuple(params[name].shape) != shape:
|
| 387 |
+
raise ValueError(f"{name}: shape {params[name].shape}, expected {shape}")
|
| 388 |
+
for name, arr in params.items():
|
| 389 |
+
if arr.dtype != np.float32 or not np.isfinite(arr).all():
|
| 390 |
+
raise ValueError(f"{name}: not finite float32")
|
| 391 |
+
norm = load_param_json(weights_dir / C.PARAM_JSON)
|
| 392 |
+
return PlannerWeights(params, norm, sha, facts, weights_dir)
|
| 393 |
+
|
| 394 |
+
|
| 395 |
+
def param_count(params: Mapping[str, np.ndarray]) -> int:
|
| 396 |
+
return int(sum(v.size for v in params.values()))
|
| 397 |
+
|
| 398 |
+
|
| 399 |
+
def coverage(weights: PlannerWeights) -> Dict[str, Any]:
|
| 400 |
+
"""Proof that the canonical dict is a re-labelling of the export: every float initializer (size > 1) of the three
|
| 401 |
+
files equals one canonical tensor (or the concatenation it was split from: the fused cross K/V MatMul, the two
|
| 402 |
+
agent-embedding rows), and no canonical tensor holds the same data twice except the biases the export itself
|
| 403 |
+
deduplicated. Returns ``{"unused": [...], "duplicates": [...], "initializers": n}`` (tests require both lists
|
| 404 |
+
to be empty / the known pair)."""
|
| 405 |
+
import hashlib
|
| 406 |
+
|
| 407 |
+
import onnx
|
| 408 |
+
from onnx import numpy_helper
|
| 409 |
+
|
| 410 |
+
def h(a: np.ndarray) -> str:
|
| 411 |
+
a = np.ascontiguousarray(np.asarray(a, np.float32))
|
| 412 |
+
return hashlib.sha1(a.tobytes() + str(a.shape).encode()).hexdigest()
|
| 413 |
+
|
| 414 |
+
p = weights.params
|
| 415 |
+
known = {h(v): k for k, v in p.items()}
|
| 416 |
+
kv = np.concatenate([p[f"decoder.dit.blocks.{i}.cross_attn.kv.w"] for i in range(C.DIT_DEPTH)], axis=1)
|
| 417 |
+
known[h(kv)] = "decoder.dit.blocks.*.cross_attn.kv.w (fused)"
|
| 418 |
+
emb = p["decoder.dit.agent_embedding"]
|
| 419 |
+
known[h(emb[0:1])] = "decoder.dit.agent_embedding[0]"
|
| 420 |
+
known[h(emb[1:2])] = "decoder.dit.agent_embedding[1]"
|
| 421 |
+
rpe = p["encoder.route_position_embedding"]
|
| 422 |
+
known[h(rpe.reshape(1, *rpe.shape))] = "encoder.route_position_embedding"
|
| 423 |
+
for mod in ("lane_encoder", "route_encoder"):
|
| 424 |
+
unk = p[f"encoder.{mod}.unknown_speed_emb"]
|
| 425 |
+
known[h(unk.reshape(1, -1))] = f"encoder.{mod}.unknown_speed_emb"
|
| 426 |
+
for k, v in p.items(): # Gemm weights are stored [out, in]
|
| 427 |
+
if k.endswith(".w"):
|
| 428 |
+
known.setdefault(h(v.T), k)
|
| 429 |
+
unused, total = [], 0
|
| 430 |
+
for f in (C.ENCODER_ONNX, C.DECODER_ONNX, C.TURN_INDICATOR_ONNX):
|
| 431 |
+
for init in onnx.load(str(weights.path / f)).graph.initializer:
|
| 432 |
+
a = numpy_helper.to_array(init)
|
| 433 |
+
if a.dtype != np.float32 or a.size <= 1:
|
| 434 |
+
continue
|
| 435 |
+
total += 1
|
| 436 |
+
if h(a) not in known:
|
| 437 |
+
unused.append(f"{f}:{init.name}{list(a.shape)}")
|
| 438 |
+
seen: Dict[str, List[str]] = {}
|
| 439 |
+
for k, v in p.items():
|
| 440 |
+
seen.setdefault(h(v), []).append(k)
|
| 441 |
+
dups = sorted(tuple(sorted(ks)) for ks in seen.values() if len(ks) > 1)
|
| 442 |
+
return {"unused": unused, "duplicates": dups, "initializers": total}
|
code/tt_diffusion_planner/samples/README.md
ADDED
|
@@ -0,0 +1,23 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Sample inputs shipped with diffusion-planner-p150
|
| 2 |
+
|
| 3 |
+
Only REDISTRIBUTABLE data goes here (Apache-2.0 / MIT / CC-BY with attribution). Data whose license is unstated or
|
| 4 |
+
non-commercial (the Autoware demo rosbag, nuScenes, Argoverse 2, ...) is NOT shipped; the nuScenes-derived planner
|
| 5 |
+
instants used for local agreement tests live in the porting workspace only (`research/diffusion-planner/public_data`).
|
| 6 |
+
|
| 7 |
+
Never use the `.bin` suffix here (or `.pt`, `.pth`, `.ckpt`, `.safetensors`): tt-model's staging silently drops
|
| 8 |
+
those suffixes from `code/` (`CODE_IGNORE`, tt-model-manager `src/tt_kernel/build.py:328-332`).
|
| 9 |
+
|
| 10 |
+
Each sample is one `.npz` holding exactly the 15 raw (pre-normalization, ego-frame, batch-1, float32) tensors of the
|
| 11 |
+
node's `DiffusionPlannerCore::create_input_data()` (`reference/config.py` `INPUT_SCHEMA`): pass it as
|
| 12 |
+
`model(inputs="<file>.npz")` or as the `/predict` field `inputs` (`server/client.py --inputs <file>.npz`).
|
| 13 |
+
|
| 14 |
+
Next to each sample, `<stem>.reference.json` is the `/predict` body of the fp32 CPU reference
|
| 15 |
+
(`tt_diffusion_planner.reference.ReferencePlanner`, `Output.to_dict()`, timing removed) on it;
|
| 16 |
+
`server/smoke_test.py` compares the served trajectory with it (ADE / FDE gates). Regenerate both the reference bodies
|
| 17 |
+
and `tests/goldens/` with `code/scripts/ref_golden.py` whenever the reference, the weights or the post-processing
|
| 18 |
+
changes.
|
| 19 |
+
|
| 20 |
+
| file | content | source | license |
|
| 21 |
+
|---|---|---|---|
|
| 22 |
+
| `kashiwanoha_dense.npz` (default sample) | one planning instant on the kashiwanoha test map: ego at 6 m/s on a 17-lanelet route, 88 neighbours (constant-speed vehicles), 123 lanes (24 lanelets with traffic lights in the map), 60 line strings (stop lines and road borders), no intersection polygon; 113,893 B, sha256 `d8c2aaef...99f7` | Lanelet2 map [AutowareFoundation/map-carla-kashiwanoha](https://huggingface.co/datasets/AutowareFoundation/map-carla-kashiwanoha) 0.2.0 (`lanelet2_map.osm`, sha256 `4fe358f2...4e6d`, byte-identical to autoware_universe `planning/autoware_diffusion_planner/test_map`); scene scripted by `research/diffusion-planner/scripts/dp_scene.py` (`dp_reference.py --scene osm --osm .../lanelet2_map.osm --n-agents 120 --seed 7`), a Python port of the node's tensor construction (simplifications: SPEC 7) | Apache-2.0 (map: Apache-2.0 per its dataset card) |
|
| 23 |
+
| `straight_road.npz` | procedural 3-lane straight road: ego at 8 m/s, 12 neighbours (vehicles and pedestrians), 33 lanes with a green traffic light ahead, 6 route lanes, 1 intersection polygon, 9 line strings (borders + stop line); the precision-sensitive scene of SPEC 0.5; 10,832 B, sha256 `b9979b90...d381` | generated by `research/diffusion-planner/scripts/dp_scene.py` (`dp_reference.py --scene straight`), no external data | Apache-2.0 (this repo) |
|
code/tt_diffusion_planner/samples/kashiwanoha_dense.npz
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:d8c2aaefae61e155f5281c8b31920820d442416c653b9354be33d1fe089399f7
|
| 3 |
+
size 113893
|
code/tt_diffusion_planner/samples/kashiwanoha_dense.reference.json
ADDED
|
The diff for this file is too large to render.
See raw diff
|
|
|
code/tt_diffusion_planner/samples/straight_road.npz
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:b9979b9007df6ca1190b4b9ef61e24534f862c41b2635d0219dd22c2dfd9d381
|
| 3 |
+
size 10832
|
code/tt_diffusion_planner/samples/straight_road.reference.json
ADDED
|
@@ -0,0 +1,887 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"model": "diffusion-planner-p150",
|
| 3 |
+
"frame_id": "base_link",
|
| 4 |
+
"meta": {
|
| 5 |
+
"predicted_agent_columns": [
|
| 6 |
+
"x",
|
| 7 |
+
"y",
|
| 8 |
+
"yaw",
|
| 9 |
+
"cos",
|
| 10 |
+
"sin"
|
| 11 |
+
],
|
| 12 |
+
"predicted_agent_rows": [
|
| 13 |
+
0,
|
| 14 |
+
1,
|
| 15 |
+
2,
|
| 16 |
+
3,
|
| 17 |
+
4,
|
| 18 |
+
5,
|
| 19 |
+
6,
|
| 20 |
+
7,
|
| 21 |
+
8,
|
| 22 |
+
9,
|
| 23 |
+
10,
|
| 24 |
+
11
|
| 25 |
+
],
|
| 26 |
+
"force_stop": false,
|
| 27 |
+
"time_from_start_s": [
|
| 28 |
+
0.1,
|
| 29 |
+
0.2,
|
| 30 |
+
0.3,
|
| 31 |
+
0.4,
|
| 32 |
+
0.5,
|
| 33 |
+
0.6,
|
| 34 |
+
0.7,
|
| 35 |
+
0.8,
|
| 36 |
+
0.9,
|
| 37 |
+
1.0,
|
| 38 |
+
1.1,
|
| 39 |
+
1.2,
|
| 40 |
+
1.3,
|
| 41 |
+
1.4,
|
| 42 |
+
1.5,
|
| 43 |
+
1.6,
|
| 44 |
+
1.7,
|
| 45 |
+
1.8,
|
| 46 |
+
1.9,
|
| 47 |
+
2.0,
|
| 48 |
+
2.1,
|
| 49 |
+
2.2,
|
| 50 |
+
2.3,
|
| 51 |
+
2.4,
|
| 52 |
+
2.5,
|
| 53 |
+
2.6,
|
| 54 |
+
2.7,
|
| 55 |
+
2.8,
|
| 56 |
+
2.9,
|
| 57 |
+
3.0,
|
| 58 |
+
3.1,
|
| 59 |
+
3.2,
|
| 60 |
+
3.3,
|
| 61 |
+
3.4,
|
| 62 |
+
3.5,
|
| 63 |
+
3.6,
|
| 64 |
+
3.7,
|
| 65 |
+
3.8,
|
| 66 |
+
3.9,
|
| 67 |
+
4.0,
|
| 68 |
+
4.1,
|
| 69 |
+
4.2,
|
| 70 |
+
4.3,
|
| 71 |
+
4.4,
|
| 72 |
+
4.5,
|
| 73 |
+
4.6,
|
| 74 |
+
4.7,
|
| 75 |
+
4.8,
|
| 76 |
+
4.9,
|
| 77 |
+
5.0,
|
| 78 |
+
5.1,
|
| 79 |
+
5.2,
|
| 80 |
+
5.3,
|
| 81 |
+
5.4,
|
| 82 |
+
5.5,
|
| 83 |
+
5.6,
|
| 84 |
+
5.7,
|
| 85 |
+
5.8,
|
| 86 |
+
5.9,
|
| 87 |
+
6.0,
|
| 88 |
+
6.1,
|
| 89 |
+
6.2,
|
| 90 |
+
6.3,
|
| 91 |
+
6.4,
|
| 92 |
+
6.5,
|
| 93 |
+
6.6,
|
| 94 |
+
6.7,
|
| 95 |
+
6.8,
|
| 96 |
+
6.9,
|
| 97 |
+
7.0,
|
| 98 |
+
7.1,
|
| 99 |
+
7.2,
|
| 100 |
+
7.3,
|
| 101 |
+
7.4,
|
| 102 |
+
7.5,
|
| 103 |
+
7.6,
|
| 104 |
+
7.7,
|
| 105 |
+
7.8,
|
| 106 |
+
7.9,
|
| 107 |
+
8.0
|
| 108 |
+
],
|
| 109 |
+
"valid_counts": {
|
| 110 |
+
"ego": 1,
|
| 111 |
+
"neighbor": 12,
|
| 112 |
+
"static": 0,
|
| 113 |
+
"lane": 33,
|
| 114 |
+
"route": 6,
|
| 115 |
+
"polygon": 1,
|
| 116 |
+
"line_string": 9,
|
| 117 |
+
"goal": 1,
|
| 118 |
+
"ego_shape": 1,
|
| 119 |
+
"turn": 1
|
| 120 |
+
},
|
| 121 |
+
"reference": "fp32 CPU (torch)"
|
| 122 |
+
},
|
| 123 |
+
"timing_ms": {},
|
| 124 |
+
"num_poses": 80,
|
| 125 |
+
"columns": [
|
| 126 |
+
"x",
|
| 127 |
+
"y",
|
| 128 |
+
"yaw",
|
| 129 |
+
"cos",
|
| 130 |
+
"sin",
|
| 131 |
+
"velocity",
|
| 132 |
+
"acceleration"
|
| 133 |
+
],
|
| 134 |
+
"trajectory": [
|
| 135 |
+
[
|
| 136 |
+
0.8299,
|
| 137 |
+
0.0007,
|
| 138 |
+
-0.0007,
|
| 139 |
+
1.0033,
|
| 140 |
+
-0.0007,
|
| 141 |
+
8.0591,
|
| 142 |
+
-0.3468
|
| 143 |
+
],
|
| 144 |
+
[
|
| 145 |
+
1.6377,
|
| 146 |
+
0.0009,
|
| 147 |
+
-0.0006,
|
| 148 |
+
1.0026,
|
| 149 |
+
-0.0006,
|
| 150 |
+
8.0244,
|
| 151 |
+
-0.0349
|
| 152 |
+
],
|
| 153 |
+
[
|
| 154 |
+
2.4372,
|
| 155 |
+
0.0009,
|
| 156 |
+
-0.0012,
|
| 157 |
+
1.0021,
|
| 158 |
+
-0.0012,
|
| 159 |
+
8.0209,
|
| 160 |
+
0.0592
|
| 161 |
+
],
|
| 162 |
+
[
|
| 163 |
+
3.2377,
|
| 164 |
+
0.0016,
|
| 165 |
+
-0.0016,
|
| 166 |
+
1.0021,
|
| 167 |
+
-0.0016,
|
| 168 |
+
8.0268,
|
| 169 |
+
-0.0312
|
| 170 |
+
],
|
| 171 |
+
[
|
| 172 |
+
4.0421,
|
| 173 |
+
0.0016,
|
| 174 |
+
-0.0023,
|
| 175 |
+
1.0018,
|
| 176 |
+
-0.0023,
|
| 177 |
+
8.0237,
|
| 178 |
+
-0.0017
|
| 179 |
+
],
|
| 180 |
+
[
|
| 181 |
+
4.8417,
|
| 182 |
+
0.0019,
|
| 183 |
+
-0.0021,
|
| 184 |
+
1.0017,
|
| 185 |
+
-0.0021,
|
| 186 |
+
8.0236,
|
| 187 |
+
0.0849
|
| 188 |
+
],
|
| 189 |
+
[
|
| 190 |
+
5.643,
|
| 191 |
+
0.0022,
|
| 192 |
+
-0.0018,
|
| 193 |
+
1.0014,
|
| 194 |
+
-0.0018,
|
| 195 |
+
8.032,
|
| 196 |
+
0.0807
|
| 197 |
+
],
|
| 198 |
+
[
|
| 199 |
+
6.4473,
|
| 200 |
+
0.0021,
|
| 201 |
+
-0.0013,
|
| 202 |
+
1.0011,
|
| 203 |
+
-0.0013,
|
| 204 |
+
8.0401,
|
| 205 |
+
-0.0193
|
| 206 |
+
],
|
| 207 |
+
[
|
| 208 |
+
7.2494,
|
| 209 |
+
0.0013,
|
| 210 |
+
-0.0014,
|
| 211 |
+
1.0009,
|
| 212 |
+
-0.0014,
|
| 213 |
+
8.0382,
|
| 214 |
+
0.0377
|
| 215 |
+
],
|
| 216 |
+
[
|
| 217 |
+
8.0545,
|
| 218 |
+
0.0012,
|
| 219 |
+
-0.0013,
|
| 220 |
+
1.0007,
|
| 221 |
+
-0.0013,
|
| 222 |
+
8.0419,
|
| 223 |
+
-0.0051
|
| 224 |
+
],
|
| 225 |
+
[
|
| 226 |
+
8.8587,
|
| 227 |
+
-0.0027,
|
| 228 |
+
-0.0003,
|
| 229 |
+
1.0005,
|
| 230 |
+
-0.0003,
|
| 231 |
+
8.0414,
|
| 232 |
+
0.0536
|
| 233 |
+
],
|
| 234 |
+
[
|
| 235 |
+
9.6567,
|
| 236 |
+
-0.0027,
|
| 237 |
+
-0.0001,
|
| 238 |
+
1.0003,
|
| 239 |
+
-0.0001,
|
| 240 |
+
8.0468,
|
| 241 |
+
0.1724
|
| 242 |
+
],
|
| 243 |
+
[
|
| 244 |
+
10.4609,
|
| 245 |
+
-0.0016,
|
| 246 |
+
0.0004,
|
| 247 |
+
1.0002,
|
| 248 |
+
0.0004,
|
| 249 |
+
8.064,
|
| 250 |
+
0.0914
|
| 251 |
+
],
|
| 252 |
+
[
|
| 253 |
+
11.2674,
|
| 254 |
+
-0.002,
|
| 255 |
+
0.0002,
|
| 256 |
+
1.0002,
|
| 257 |
+
0.0002,
|
| 258 |
+
8.0732,
|
| 259 |
+
-0.0142
|
| 260 |
+
],
|
| 261 |
+
[
|
| 262 |
+
12.0751,
|
| 263 |
+
-0.0002,
|
| 264 |
+
0.0009,
|
| 265 |
+
1.0003,
|
| 266 |
+
0.0009,
|
| 267 |
+
8.0718,
|
| 268 |
+
0.2182
|
| 269 |
+
],
|
| 270 |
+
[
|
| 271 |
+
12.8778,
|
| 272 |
+
0.0009,
|
| 273 |
+
0.0008,
|
| 274 |
+
1.0005,
|
| 275 |
+
0.0008,
|
| 276 |
+
8.0936,
|
| 277 |
+
0.1081
|
| 278 |
+
],
|
| 279 |
+
[
|
| 280 |
+
13.6829,
|
| 281 |
+
0.0038,
|
| 282 |
+
0.0015,
|
| 283 |
+
1.0007,
|
| 284 |
+
0.0015,
|
| 285 |
+
8.1044,
|
| 286 |
+
0.1823
|
| 287 |
+
],
|
| 288 |
+
[
|
| 289 |
+
14.4876,
|
| 290 |
+
0.0044,
|
| 291 |
+
0.0018,
|
| 292 |
+
1.0007,
|
| 293 |
+
0.0018,
|
| 294 |
+
8.1226,
|
| 295 |
+
0.3151
|
| 296 |
+
],
|
| 297 |
+
[
|
| 298 |
+
15.2961,
|
| 299 |
+
0.0073,
|
| 300 |
+
0.0023,
|
| 301 |
+
1.0007,
|
| 302 |
+
0.0023,
|
| 303 |
+
8.1541,
|
| 304 |
+
0.2434
|
| 305 |
+
],
|
| 306 |
+
[
|
| 307 |
+
16.1079,
|
| 308 |
+
0.0092,
|
| 309 |
+
0.0022,
|
| 310 |
+
1.0006,
|
| 311 |
+
0.0022,
|
| 312 |
+
8.1785,
|
| 313 |
+
0.2183
|
| 314 |
+
],
|
| 315 |
+
[
|
| 316 |
+
16.9194,
|
| 317 |
+
0.0097,
|
| 318 |
+
0.0021,
|
| 319 |
+
1.0007,
|
| 320 |
+
0.0021,
|
| 321 |
+
8.2003,
|
| 322 |
+
0.2977
|
| 323 |
+
],
|
| 324 |
+
[
|
| 325 |
+
17.7247,
|
| 326 |
+
0.012,
|
| 327 |
+
0.0019,
|
| 328 |
+
1.0003,
|
| 329 |
+
0.0019,
|
| 330 |
+
8.2301,
|
| 331 |
+
0.5019
|
| 332 |
+
],
|
| 333 |
+
[
|
| 334 |
+
18.5499,
|
| 335 |
+
0.0152,
|
| 336 |
+
0.0019,
|
| 337 |
+
1.0005,
|
| 338 |
+
0.0019,
|
| 339 |
+
8.2803,
|
| 340 |
+
0.294
|
| 341 |
+
],
|
| 342 |
+
[
|
| 343 |
+
19.3613,
|
| 344 |
+
0.018,
|
| 345 |
+
0.0018,
|
| 346 |
+
1.0002,
|
| 347 |
+
0.0018,
|
| 348 |
+
8.3097,
|
| 349 |
+
0.5381
|
| 350 |
+
],
|
| 351 |
+
[
|
| 352 |
+
20.181,
|
| 353 |
+
0.0202,
|
| 354 |
+
0.0019,
|
| 355 |
+
1.0002,
|
| 356 |
+
0.0019,
|
| 357 |
+
8.3635,
|
| 358 |
+
0.4827
|
| 359 |
+
],
|
| 360 |
+
[
|
| 361 |
+
21.0109,
|
| 362 |
+
0.0217,
|
| 363 |
+
0.0018,
|
| 364 |
+
1.0002,
|
| 365 |
+
0.0018,
|
| 366 |
+
8.4117,
|
| 367 |
+
0.4352
|
| 368 |
+
],
|
| 369 |
+
[
|
| 370 |
+
21.8389,
|
| 371 |
+
0.0219,
|
| 372 |
+
0.0013,
|
| 373 |
+
1.0003,
|
| 374 |
+
0.0013,
|
| 375 |
+
8.4553,
|
| 376 |
+
0.5616
|
| 377 |
+
],
|
| 378 |
+
[
|
| 379 |
+
22.6681,
|
| 380 |
+
0.0248,
|
| 381 |
+
0.0012,
|
| 382 |
+
1.0002,
|
| 383 |
+
0.0012,
|
| 384 |
+
8.5114,
|
| 385 |
+
0.5495
|
| 386 |
+
],
|
| 387 |
+
[
|
| 388 |
+
23.5035,
|
| 389 |
+
0.0257,
|
| 390 |
+
0.0008,
|
| 391 |
+
1.0001,
|
| 392 |
+
0.0008,
|
| 393 |
+
8.5664,
|
| 394 |
+
0.6455
|
| 395 |
+
],
|
| 396 |
+
[
|
| 397 |
+
24.3489,
|
| 398 |
+
0.0303,
|
| 399 |
+
0.0013,
|
| 400 |
+
1.0002,
|
| 401 |
+
0.0013,
|
| 402 |
+
8.6309,
|
| 403 |
+
0.562
|
| 404 |
+
],
|
| 405 |
+
[
|
| 406 |
+
25.1976,
|
| 407 |
+
0.0307,
|
| 408 |
+
0.0015,
|
| 409 |
+
1.0003,
|
| 410 |
+
0.0015,
|
| 411 |
+
8.6871,
|
| 412 |
+
0.5735
|
| 413 |
+
],
|
| 414 |
+
[
|
| 415 |
+
26.052,
|
| 416 |
+
0.0323,
|
| 417 |
+
0.0015,
|
| 418 |
+
1.0002,
|
| 419 |
+
0.0015,
|
| 420 |
+
8.7445,
|
| 421 |
+
0.6293
|
| 422 |
+
],
|
| 423 |
+
[
|
| 424 |
+
26.9104,
|
| 425 |
+
0.0347,
|
| 426 |
+
0.0017,
|
| 427 |
+
1.0005,
|
| 428 |
+
0.0017,
|
| 429 |
+
8.8074,
|
| 430 |
+
0.6244
|
| 431 |
+
],
|
| 432 |
+
[
|
| 433 |
+
27.7751,
|
| 434 |
+
0.037,
|
| 435 |
+
0.0021,
|
| 436 |
+
1.0007,
|
| 437 |
+
0.0021,
|
| 438 |
+
8.8698,
|
| 439 |
+
0.6727
|
| 440 |
+
],
|
| 441 |
+
[
|
| 442 |
+
28.648,
|
| 443 |
+
0.0421,
|
| 444 |
+
0.0022,
|
| 445 |
+
1.0006,
|
| 446 |
+
0.0022,
|
| 447 |
+
8.9371,
|
| 448 |
+
0.5461
|
| 449 |
+
],
|
| 450 |
+
[
|
| 451 |
+
29.5212,
|
| 452 |
+
0.0447,
|
| 453 |
+
0.0023,
|
| 454 |
+
1.0008,
|
| 455 |
+
0.0023,
|
| 456 |
+
8.9917,
|
| 457 |
+
0.8209
|
| 458 |
+
],
|
| 459 |
+
[
|
| 460 |
+
30.4082,
|
| 461 |
+
0.0479,
|
| 462 |
+
0.003,
|
| 463 |
+
1.001,
|
| 464 |
+
0.003,
|
| 465 |
+
9.0738,
|
| 466 |
+
0.585
|
| 467 |
+
],
|
| 468 |
+
[
|
| 469 |
+
31.2986,
|
| 470 |
+
0.047,
|
| 471 |
+
0.0028,
|
| 472 |
+
1.001,
|
| 473 |
+
0.0028,
|
| 474 |
+
9.1323,
|
| 475 |
+
0.6027
|
| 476 |
+
],
|
| 477 |
+
[
|
| 478 |
+
32.1931,
|
| 479 |
+
0.0503,
|
| 480 |
+
0.0031,
|
| 481 |
+
1.0009,
|
| 482 |
+
0.0031,
|
| 483 |
+
9.1926,
|
| 484 |
+
0.8186
|
| 485 |
+
],
|
| 486 |
+
[
|
| 487 |
+
33.0979,
|
| 488 |
+
0.0518,
|
| 489 |
+
0.0032,
|
| 490 |
+
1.0013,
|
| 491 |
+
0.0032,
|
| 492 |
+
9.2744,
|
| 493 |
+
0.7072
|
| 494 |
+
],
|
| 495 |
+
[
|
| 496 |
+
34.0062,
|
| 497 |
+
0.0534,
|
| 498 |
+
0.0028,
|
| 499 |
+
1.0012,
|
| 500 |
+
0.0028,
|
| 501 |
+
9.3451,
|
| 502 |
+
0.7215
|
| 503 |
+
],
|
| 504 |
+
[
|
| 505 |
+
34.9247,
|
| 506 |
+
0.0575,
|
| 507 |
+
0.0033,
|
| 508 |
+
1.0012,
|
| 509 |
+
0.0033,
|
| 510 |
+
9.4173,
|
| 511 |
+
0.7555
|
| 512 |
+
],
|
| 513 |
+
[
|
| 514 |
+
35.8413,
|
| 515 |
+
0.0606,
|
| 516 |
+
0.0032,
|
| 517 |
+
1.0011,
|
| 518 |
+
0.0032,
|
| 519 |
+
9.4928,
|
| 520 |
+
0.7423
|
| 521 |
+
],
|
| 522 |
+
[
|
| 523 |
+
36.7802,
|
| 524 |
+
0.061,
|
| 525 |
+
0.0034,
|
| 526 |
+
1.0015,
|
| 527 |
+
0.0034,
|
| 528 |
+
9.5671,
|
| 529 |
+
0.4912
|
| 530 |
+
],
|
| 531 |
+
[
|
| 532 |
+
37.714,
|
| 533 |
+
0.0635,
|
| 534 |
+
0.0032,
|
| 535 |
+
1.0015,
|
| 536 |
+
0.0032,
|
| 537 |
+
9.6162,
|
| 538 |
+
0.774
|
| 539 |
+
],
|
| 540 |
+
[
|
| 541 |
+
38.6526,
|
| 542 |
+
0.0643,
|
| 543 |
+
0.0027,
|
| 544 |
+
1.0016,
|
| 545 |
+
0.0027,
|
| 546 |
+
9.6936,
|
| 547 |
+
0.8083
|
| 548 |
+
],
|
| 549 |
+
[
|
| 550 |
+
39.6127,
|
| 551 |
+
0.0679,
|
| 552 |
+
0.003,
|
| 553 |
+
1.002,
|
| 554 |
+
0.003,
|
| 555 |
+
9.7744,
|
| 556 |
+
0.5879
|
| 557 |
+
],
|
| 558 |
+
[
|
| 559 |
+
40.574,
|
| 560 |
+
0.069,
|
| 561 |
+
0.0024,
|
| 562 |
+
1.002,
|
| 563 |
+
0.0024,
|
| 564 |
+
9.8332,
|
| 565 |
+
0.6322
|
| 566 |
+
],
|
| 567 |
+
[
|
| 568 |
+
41.54,
|
| 569 |
+
0.0739,
|
| 570 |
+
0.0025,
|
| 571 |
+
1.0023,
|
| 572 |
+
0.0025,
|
| 573 |
+
9.8964,
|
| 574 |
+
0.7491
|
| 575 |
+
],
|
| 576 |
+
[
|
| 577 |
+
42.5189,
|
| 578 |
+
0.0753,
|
| 579 |
+
0.0022,
|
| 580 |
+
1.0026,
|
| 581 |
+
0.0022,
|
| 582 |
+
9.9714,
|
| 583 |
+
0.4674
|
| 584 |
+
],
|
| 585 |
+
[
|
| 586 |
+
43.4949,
|
| 587 |
+
0.0774,
|
| 588 |
+
0.0017,
|
| 589 |
+
1.0025,
|
| 590 |
+
0.0017,
|
| 591 |
+
10.0181,
|
| 592 |
+
0.7418
|
| 593 |
+
],
|
| 594 |
+
[
|
| 595 |
+
44.4731,
|
| 596 |
+
0.0796,
|
| 597 |
+
0.002,
|
| 598 |
+
1.0028,
|
| 599 |
+
0.002,
|
| 600 |
+
10.0923,
|
| 601 |
+
0.7677
|
| 602 |
+
],
|
| 603 |
+
[
|
| 604 |
+
45.4688,
|
| 605 |
+
0.0837,
|
| 606 |
+
0.0017,
|
| 607 |
+
1.0026,
|
| 608 |
+
0.0017,
|
| 609 |
+
10.169,
|
| 610 |
+
0.583
|
| 611 |
+
],
|
| 612 |
+
[
|
| 613 |
+
46.4721,
|
| 614 |
+
0.0856,
|
| 615 |
+
0.0019,
|
| 616 |
+
1.0025,
|
| 617 |
+
0.0019,
|
| 618 |
+
10.2274,
|
| 619 |
+
0.5995
|
| 620 |
+
],
|
| 621 |
+
[
|
| 622 |
+
47.4792,
|
| 623 |
+
0.0883,
|
| 624 |
+
0.0013,
|
| 625 |
+
1.0028,
|
| 626 |
+
0.0013,
|
| 627 |
+
10.2873,
|
| 628 |
+
0.56
|
| 629 |
+
],
|
| 630 |
+
[
|
| 631 |
+
48.4911,
|
| 632 |
+
0.0938,
|
| 633 |
+
0.0013,
|
| 634 |
+
1.0026,
|
| 635 |
+
0.0013,
|
| 636 |
+
10.3433,
|
| 637 |
+
0.5217
|
| 638 |
+
],
|
| 639 |
+
[
|
| 640 |
+
49.517,
|
| 641 |
+
0.0969,
|
| 642 |
+
0.0013,
|
| 643 |
+
1.0025,
|
| 644 |
+
0.0013,
|
| 645 |
+
10.3955,
|
| 646 |
+
0.5525
|
| 647 |
+
],
|
| 648 |
+
[
|
| 649 |
+
50.5334,
|
| 650 |
+
0.0991,
|
| 651 |
+
0.0008,
|
| 652 |
+
1.0021,
|
| 653 |
+
0.0008,
|
| 654 |
+
10.4507,
|
| 655 |
+
0.5985
|
| 656 |
+
],
|
| 657 |
+
[
|
| 658 |
+
51.5687,
|
| 659 |
+
0.1035,
|
| 660 |
+
0.0005,
|
| 661 |
+
1.002,
|
| 662 |
+
0.0005,
|
| 663 |
+
10.5106,
|
| 664 |
+
0.424
|
| 665 |
+
],
|
| 666 |
+
[
|
| 667 |
+
52.6083,
|
| 668 |
+
0.1078,
|
| 669 |
+
0.0006,
|
| 670 |
+
1.002,
|
| 671 |
+
0.0006,
|
| 672 |
+
10.553,
|
| 673 |
+
0.5447
|
| 674 |
+
],
|
| 675 |
+
[
|
| 676 |
+
53.6507,
|
| 677 |
+
0.109,
|
| 678 |
+
-0.0002,
|
| 679 |
+
1.0017,
|
| 680 |
+
-0.0002,
|
| 681 |
+
10.6074,
|
| 682 |
+
0.3741
|
| 683 |
+
],
|
| 684 |
+
[
|
| 685 |
+
54.7019,
|
| 686 |
+
0.113,
|
| 687 |
+
-0.0004,
|
| 688 |
+
1.0014,
|
| 689 |
+
-0.0004,
|
| 690 |
+
10.6448,
|
| 691 |
+
0.3706
|
| 692 |
+
],
|
| 693 |
+
[
|
| 694 |
+
55.7538,
|
| 695 |
+
0.1194,
|
| 696 |
+
-0.0005,
|
| 697 |
+
1.0012,
|
| 698 |
+
-0.0005,
|
| 699 |
+
10.6819,
|
| 700 |
+
0.4656
|
| 701 |
+
],
|
| 702 |
+
[
|
| 703 |
+
56.8074,
|
| 704 |
+
0.1212,
|
| 705 |
+
-0.0005,
|
| 706 |
+
1.0012,
|
| 707 |
+
-0.0005,
|
| 708 |
+
10.7285,
|
| 709 |
+
0.3772
|
| 710 |
+
],
|
| 711 |
+
[
|
| 712 |
+
57.8776,
|
| 713 |
+
0.1224,
|
| 714 |
+
-0.0012,
|
| 715 |
+
1.0006,
|
| 716 |
+
-0.0012,
|
| 717 |
+
10.7662,
|
| 718 |
+
0.2135
|
| 719 |
+
],
|
| 720 |
+
[
|
| 721 |
+
58.9418,
|
| 722 |
+
0.126,
|
| 723 |
+
-0.0008,
|
| 724 |
+
1.0007,
|
| 725 |
+
-0.0008,
|
| 726 |
+
10.7875,
|
| 727 |
+
0.4392
|
| 728 |
+
],
|
| 729 |
+
[
|
| 730 |
+
60.011,
|
| 731 |
+
0.1313,
|
| 732 |
+
-0.0007,
|
| 733 |
+
1.0008,
|
| 734 |
+
-0.0007,
|
| 735 |
+
10.8314,
|
| 736 |
+
0.2563
|
| 737 |
+
],
|
| 738 |
+
[
|
| 739 |
+
61.0942,
|
| 740 |
+
0.1372,
|
| 741 |
+
-0.0007,
|
| 742 |
+
1.0008,
|
| 743 |
+
-0.0007,
|
| 744 |
+
10.8571,
|
| 745 |
+
0.1101
|
| 746 |
+
],
|
| 747 |
+
[
|
| 748 |
+
62.1665,
|
| 749 |
+
0.14,
|
| 750 |
+
-0.0007,
|
| 751 |
+
1.0008,
|
| 752 |
+
-0.0007,
|
| 753 |
+
10.8681,
|
| 754 |
+
0.4464
|
| 755 |
+
],
|
| 756 |
+
[
|
| 757 |
+
63.2474,
|
| 758 |
+
0.1409,
|
| 759 |
+
-0.0002,
|
| 760 |
+
1.0007,
|
| 761 |
+
-0.0002,
|
| 762 |
+
10.9127,
|
| 763 |
+
0.3852
|
| 764 |
+
],
|
| 765 |
+
[
|
| 766 |
+
64.3365,
|
| 767 |
+
0.144,
|
| 768 |
+
-0.0007,
|
| 769 |
+
1.0007,
|
| 770 |
+
-0.0007,
|
| 771 |
+
10.9513,
|
| 772 |
+
0.1574
|
| 773 |
+
],
|
| 774 |
+
[
|
| 775 |
+
65.4203,
|
| 776 |
+
0.1503,
|
| 777 |
+
-0.0001,
|
| 778 |
+
1.0009,
|
| 779 |
+
-0.0001,
|
| 780 |
+
10.967,
|
| 781 |
+
0.2386
|
| 782 |
+
],
|
| 783 |
+
[
|
| 784 |
+
66.5075,
|
| 785 |
+
0.149,
|
| 786 |
+
-0.0005,
|
| 787 |
+
1.0001,
|
| 788 |
+
-0.0005,
|
| 789 |
+
10.9909,
|
| 790 |
+
0.0
|
| 791 |
+
],
|
| 792 |
+
[
|
| 793 |
+
67.6069,
|
| 794 |
+
0.154,
|
| 795 |
+
-0.0003,
|
| 796 |
+
1.0002,
|
| 797 |
+
-0.0003,
|
| 798 |
+
10.9909,
|
| 799 |
+
0.0
|
| 800 |
+
],
|
| 801 |
+
[
|
| 802 |
+
68.6966,
|
| 803 |
+
0.1604,
|
| 804 |
+
0.0001,
|
| 805 |
+
0.9999,
|
| 806 |
+
0.0001,
|
| 807 |
+
10.9909,
|
| 808 |
+
0.0
|
| 809 |
+
],
|
| 810 |
+
[
|
| 811 |
+
69.7886,
|
| 812 |
+
0.1624,
|
| 813 |
+
0.0008,
|
| 814 |
+
0.9997,
|
| 815 |
+
0.0008,
|
| 816 |
+
10.9909,
|
| 817 |
+
0.0
|
| 818 |
+
],
|
| 819 |
+
[
|
| 820 |
+
70.8966,
|
| 821 |
+
0.1674,
|
| 822 |
+
0.001,
|
| 823 |
+
0.9997,
|
| 824 |
+
0.001,
|
| 825 |
+
10.9909,
|
| 826 |
+
0.0
|
| 827 |
+
],
|
| 828 |
+
[
|
| 829 |
+
72.0083,
|
| 830 |
+
0.1712,
|
| 831 |
+
0.0015,
|
| 832 |
+
0.9997,
|
| 833 |
+
0.0015,
|
| 834 |
+
10.9909,
|
| 835 |
+
0.0
|
| 836 |
+
],
|
| 837 |
+
[
|
| 838 |
+
73.11,
|
| 839 |
+
0.1757,
|
| 840 |
+
0.0015,
|
| 841 |
+
0.999,
|
| 842 |
+
0.0015,
|
| 843 |
+
10.9909,
|
| 844 |
+
0.0
|
| 845 |
+
],
|
| 846 |
+
[
|
| 847 |
+
74.2129,
|
| 848 |
+
0.179,
|
| 849 |
+
0.0022,
|
| 850 |
+
0.9989,
|
| 851 |
+
0.0022,
|
| 852 |
+
10.9909,
|
| 853 |
+
0.0
|
| 854 |
+
]
|
| 855 |
+
],
|
| 856 |
+
"turn_indicator": {
|
| 857 |
+
"command": 1,
|
| 858 |
+
"command_name": "DISABLE",
|
| 859 |
+
"keep_selected": true,
|
| 860 |
+
"held": false,
|
| 861 |
+
"logits": [
|
| 862 |
+
-16.221363067626953,
|
| 863 |
+
-6.708434581756592,
|
| 864 |
+
-2.558558940887451,
|
| 865 |
+
-3.9311578273773193,
|
| 866 |
+
6.994775295257568
|
| 867 |
+
],
|
| 868 |
+
"probabilities": [
|
| 869 |
+
2.8843449850768366e-10,
|
| 870 |
+
3.903549895767355e-06,
|
| 871 |
+
0.00024758701329119503,
|
| 872 |
+
6.275029591051862e-05,
|
| 873 |
+
0.9995858073234558
|
| 874 |
+
]
|
| 875 |
+
},
|
| 876 |
+
"predicted_agents": {
|
| 877 |
+
"format": "npz",
|
| 878 |
+
"key": "predicted_agents",
|
| 879 |
+
"dtype": "float32",
|
| 880 |
+
"shape": [
|
| 881 |
+
12,
|
| 882 |
+
80,
|
| 883 |
+
5
|
| 884 |
+
],
|
| 885 |
+
"data": "UEsDBC0AAAAAAAAAIQB33n+g//////////8UABQAcHJlZGljdGVkX2FnZW50cy5ucHkBABAAgEsAAAAAAACASwAAAAAAAJNOVU1QWQEAdgB7J2Rlc2NyJzogJzxmNCcsICdmb3J0cmFuX29yZGVyJzogRmFsc2UsICdzaGFwZSc6ICgxMiwgODAsIDUpLCB9ICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAKBGidwERxyUB6Gom8cXt6Pwihh7zc6JfACwzKQC5TszyHTHs/fa+xPMLCkcBCq8pAYj9YPfBefD/Cw1Y9IKGLwBJqy0BTyJQ9bhF9P+H+kz18oYXAGyzMQCB9sD1vwX0/7tKvPagkf8D12MxAJzXCPWYXfj/4oME9YHRzwBJyzUAoQ9I9s0x+P5G/0T1YhmfARBjOQCdp4z02ZH4/9e3iPaTTW8BRqM5Ah5bxPWBffj+IGfE9+P1PwP4Xz0D4m/49FCt+P5AG/j3sokTAt6HPQG4GBz646H0/O6sGPsQOOcDGGtBAdxUQPguEfT8unw8+OMUtwC6B0EBFzBY+MnR9P95RFj7EUiLA9gzRQNxtHT4//Hw/qc8cPmyhF8DdeNFA2MgiPsDXfD9PHyI+kDENwO+10UDkKCQ+2mp8P29cIz7g6QHALknSQA0fJT4JEXw/bjUkPhAz77/4tNJA4e8jPooVfD8VCCM+4FHav1wF00BK3yE+rLh7P8jaID5wAca/klfTQBZJHj47wns/oEkdPrClsb9v2tNALzUaPqSsez/OMRk+uJmcv1co1EBXdRM+Deh7P1+IEj54FYm/6nDUQBHdCz5YCHw/9P4KPiCqaL+MqtRAw+cDPuhqfD/wKQM+8CFBv7Xy1ECOPfc9oex8P9wN9j3whxe/kzLVQEr85j19PH0/hvvlPcAg4L6DYdVAj/PWPZ30fT8yStY94OyPvuKE1UCKzsQ9Nwl+P+8zxD0ATwa+GJ3VQFNPtD14YX4/FtuzPQAQlTx719VAAMGnPVKgfj+9Zac9AIkzPgjj1UA1AJo95N9+Pxi8mT0AxqQ+GA/WQDmziz029H4/AHiLPWBe+D7iLtZAN8N+PQHlfj9kS349wBYgP/k01kB4EGw93M1+PwOUaz2AvUc/SDDWQGq1UT150H4/zERRPaDZcT8+S9ZAzd1APU/Nfj9gc0A9YIOKP5tf1kCuHS49E25+P7CbLT1o7J0/WFvWQAgzFz3EmX4/oM0WPcDNsT9ibdZAcwoFPbArfj+/kwQ9gADEP5081kB3XvA8yiV+PzyE7zw4rtY/TGPWQC5BzzwP830/gW/OPPAg6j+Nc9ZAbC+vPMjgfT9Cd648SHz9P09i1kCbaIY8sv19P1zihTxUDghAT07WQNvkXjzRvn0/dupdPArNEUD4MtZA+XYrPB/jfT88wio8FNUaQEwt1kD/p7s7kZ99PyDJujtMcCRAWP/VQBUYRju7130/ckJFOyZhLUAq99VAqmONuP6VfT8AuYy4HnI3QFK51UBz47i72HV9P8z4t7uGlUBAzKvVQC0HvruPKX0/svm8u9CsSUC3rtVAL+AivHHJfD/b2iG8IJlTQCuA1UBL4GG8bqp8P6NoYLz0VFxAn23VQKMLiLxiL3w/7giHvNL+ZEAYRtVAARyivIn+ez+q2KC8cFduQAwI1UAbt7+89M57P4gnvrw61XdAXOPUQJ9A3rzJk3s/mljcvF5TgEBWxNRAGzD4vN3Xez8YMfa8ZAmFQFtq1ECjNwS9e2h7P/0KA70uk4lAylPUQBhpDr3ni3s/ny8NvTNJjkBl+dNA4awavQa+ez9EaBm9WTGTQLMU1EBhIR69eAt8P6btHL2Uo5dAI+XTQNTEKL07+3s/yXcnvfBmnEAattNApAIzvXuJfD/v0zG959+gQD3a00DiITK9RJZ8Pwr5ML0Bo6VAzb3TQLZ3Qr1SD30/IWNBvVKCqkATedNAckNDvUhffT9PTEK9i7quQDRG00Dx8Ue9ytt9P/slR73pr7NAeDbTQNXbTr2CCn4/bxxOveBGuEBMKdNAynhUvUAxfj/qxFO9Q+W8QB4o00AtzVu90oZ+P7g4W73kX8FAhibTQJsGXb0vzX4/249cvei6xUDEMdNAXGFhvUj2fj/s+mC9cpLKQDxG00BX8GW9LBx/P3aZZb05585AmEXTQNPPar3nNn8/CoRqvQu400C0adNA5LBwvWg0fz/kYnC9Zn/YQOJN00BmhG69il1/P/RJbr1WxNxAykvTQJ12b70kjX8/WFJvve4h4UCzKNNAdEp1vWE7fz8A/3S9WsjlQK1Q00ADOnG9GAl/P4HXcL1EcepAgmLTQFUWb72Wun4/t49uvfR8JkEZ72JAmkcmPWuSfj+01iU98LI6QeTpYUCrWSI9+Lt/P4lJIj3T5E9BXR9hQLHxGT2EO4A/IBoaPToSZUGE82BASiAYPXuWgD8zfhg9CPV5QbrPYEB0QhM95seAP4S5Ez1gRIdBEpRgQN8SCD3LtoA/P3cIPZmXkUHpmGBAcIEBPWKSgD9CzgE9m8SbQXejYED3he08AWOAPxfm7TzGCKZB23FgQP/t2jxcJIA/bhDbPEAwsEFesmBArELKPPQMgD+JT8o8NjK6QeGEX0C6zbQ8ReZ/P4XGtDx7LcRBYrVfQBYdqDww338/0hOoPAwOzkHOnF9An2mGPDsBgD8Ka4Y80MvXQcWwX0DO4XY8lRqAP6P8djxWa+FBN1hfQJwmOTw8WoA/YGg5PEIF60HO315AoeUePN51gD8cLx889Gb0QTPBXkBc77Q7lbSAPxxvtTtQvP1BW3ReQDhQGzszAIE/rOsbO/p9A0Id611AP35cOtM5gT+MjF06HB8IQto9XUCFI3+63WSBP5hDgLrnpwxCVFVcQFYVl7qxmoE/uAeYuk00EUKbeVtAaZm8upW9gT+w4b2616MVQj/aWkBTai67sbGBP9SRL7tbDhpCsOZZQEBfQbu3xYE/AbZCu0RhHkIt81hAVexsu7ixgT/MfW67dqoiQsIDWEB5LX+7nauBP+ZrgLsk+yZCHJJWQEa4k7tJjoE/KZ6Uu4QrK0KedFVALZimu3dpgT9/g6e75EkvQr/+U0C+cMu70UiBPzh2zLsCbDNCW3dSQBXSALw0MIE/VGsBvLd1N0JXzVBAg/MhvKoagT+spiK8LHM7QoQxT0ApwkG84QaBP7WJQrwIZz9CYm5NQOQ0dbx62YA/YAZ2vE9TQ0I6uUtADoCRvBDCgD9W75G81ChHQt3RSUA6cKa8AtSAP4n7przsDUtC8CVIQKcWxbwxu4A/NanFvDDRTkLHIUZAUD3UvOfDgD/H4tS8Zo9SQi80REBO8Pe8GNSAP5XC+LzsOFZCHBJCQIhgBb2ozYA/tc4FvdfPWUJQA0BA8KUVvTDzgD9gOBa9u3RdQjb5PUCcniK9ROuAP4w5I73GBWFCv+g7QGSULb1ZBoE/9EwuvbOAZEJOhzlAkT0+vZQhgT+OHT+9gv5nQhRFN0ALoU29XjyBP0aqTr3nYmtCa2I1QC8pVr3oWoE/9FdXvfzQbkJrgTNAoplcvZlpgT/13l29YCVyQhhOMUBodma95HuBPw/cZ73AinVC5qgvQGyXbr0ddoE/eAVwvZO/eEIzJy5AYE50vZxRgT81o3W9RvF7QvCNLEB/UXO9CFyBP7WudL06NH9CppsqQMlYd72XN4E/L5l4vd8tgULjXyhA9Vd4vWkhgT9IhHm9Mc+CQjRUJ0Dy1nO9fCeBP+sCdb3xWIRCIrklQInac730/4A/3OB0vYjohUKAxyNA1KtvvYr1gD9Jo3C9THKHQpExIkCptmq9K86APy+Ea70O/4hC5EkhQPx7ab3q6YA/k2FqvdR4ikKgyh9A+CxnvbXJgD/h8me90gWMQvjhHkCH7mK9mceAP1+uY70yg41ClKsdQKFMXr2R6YA/eCVfvWz+jkLu1RxAQ99avbfogD+cs1u9yHGQQkUtHED/TlS9JveAPy4oVb1j6pFCcvgaQOhFUb3pFYE/zDRSvVxik0KngxpA6ItNvdg2gT+WkE69FNqUQllTGUAmGUi9NVeBP6cvSb0mO5ZCz+cYQHWdQb1KeYE/F8RCvbmvl0Jq+RdAt8Y7vYmhgT90AT29KRuZQstNF0CmPzi9xcOBP86MOb3khZpC6gUWQJq5Nb1qxoE/3AM3vT7om0JniBVAI3Auvbb7gT/l0C+9YFKdQs4cFEDaazC9QhCCP/DeMb3SvJ5C1ccTQO/aKr24IYI/i00svSwUoEKC4xJANw8ovc4agj8Edym9iHChQnLkEUAzeiG9xSmCP+fcIr340aJCvJYRQLUFHr2lMII/12QfvVorpEKKIBBAJe0bvbklgj/WQB29JYqlQjx5D0Cm9Ra99BOCP78zGL0c4KZCeywOQM4FFL1YE4I/MT0VvVMgqELMPA1A+Q0PvWrbgT9jGxC9wXipQiklDEAp4gy9rr6BP4/bDb3SV3NBi6VTwLDtAD1fdXw/GAwAPSgXgUHjl1PAaEABPVQnfT8wiwA9C8yIQSrUU8Bwq/w8pth9P3Sg+zzwiJBBcadTwBT6Az0qsH4/baYDPVQxmEFnklPAGp3/POdvfz93Wv88kMufQT+SU8Dvq+U8ktd/P6ad5Txyb6dBIWZTwEXr1jzMAIA/GO/WPG4Hr0GTP1PAZza9PHbifz+kLb08gLm2QZZAU8A/Wa88zpp/P004rzwcZL5BU/dSwHUTmzyzVn8/XeGaPO4FxkHlRFTAgiCDPFv2fj8x3YI8Zq7NQSr7U8AIxW08Cdh+P6c8bTyQSNVBjfhTwIJrRjzAzX4/cvVFPPvh3EEy/lPAaAc9PN7Nfj/rljw8yXXkQbBFVMCD0ww89jF/PxCbDDw4C+xB6pBUwIpA4zstQH8/oeviO8GG80F0bFTACKSSOy50fz8MfJI7vP36QRbAVMCSv986AvB/P5i43zrhQAFCRglVwM+qCDouE4A/DLUIOjL8BELzlFXA6PBQul4tgD/wFVG6KK4IQtvzVcDmzSe6pz2AP1D2J7qYbQxCE1FWwKGrD7o4YoA/wOIPut8gEEJByFbAPX3juiZrgD943OO6dM0TQrAWV8DDXcC6YJCAP0LKwLrHdBdCP6ZXwCBo37pFlIA/iOnfut4gG0JfCljA6mCZus+wgD/aypm6L8YeQiWXWMB+Dve5e7eAP5C/97lDYyJCju9YwHe0NLmFqYA/ICw1ucgDJkKcg1nAQ4sGOvK0gD9c6gY6bpwpQsrtWcDCTjA63caAP7jXMDr0Oy1CmmlawDwaqDqH1IA/zKWoOi/LMEL21lrAjijwOiTWgD928fA6L1k0QoRqW8BR0+Y6atCAP0CP5zrC8jdCE8tbwIfTAjvqzYA/xTwDO3F7O0J3Q1zABLAjO2XcgD/1PCQ7CAw/Qv+DXMB4Veo6XtiAP4ob6zpljUJCEMNcwEabVTvwyIA//EJWO+YaRkIQVl3AdxQMO02/gD8ofQw7E5BJQtzjXcAA5hs7M7mAP81WHDtbD01CJidewODrBztov4A/g1EIO92MUEIDvl7AmnbGOuKygD9KAcc6ZwZUQsS8XsD9nBY7CKaAP7D+Fjuwe1dC/ylfwI5eYDo1v4A/JAZhOg7qWkJy5F/AmEFDuNXagD+A6EO4IlleQsYoYMBj6TO6UOSAP9iJNLriyGFCOK9gwNc2hrou84A/VraGur4zZUKa8GDA2sicuuYTgT/UcZ26nK1oQiZJYcBovhu73vuAP6hXHLsvAWxCBn5hwON1WLsKC4E/u1dZu3xlb0Ik6GHASfxVuwgqgT909Va7GNFyQixvYsA4t4C7RSOBP7ZJgbv4K3ZCaTFjwJ0Jj7vjJoE/cK6Pu5yWeUJBXWPAvZemu28vgT9KXae76vN8QuJmZMBp1cS71x6BPxyyxbskKIBCUsJkwASCy7vXJIE/+mrMu9TXgUIjw2XA4/r8u3AZgT9UEf67boeDQk8CZsA/+gO89BKBPy6IBLypLoVCXvNmwErVF7yG8IA/OWQYvNrehkKbXWfAD3wnvDnZgD+MCii8voaIQo9QaMBvLTa8I8iAP1e8Nrx/NIpCYQtpwN3JWbxxpoA/SVhavEbci0J0ZmnATmhnvB6NgD/a6Ge86YSNQr1MasDyZIS8kZKAP3yxhLwgMo9CIvtqwMe6j7zqgIA/GgSQvKbikELt+GvAfyCavFtwgD9OZZq8T4OSQiYLbcB3y568fGqAP8sOn7wKLZRC1JRtwJKbpby4YoA/4dylvNXVlUIt2W7AK+ewvJFCgD/uFrG8IXyXQgcDcMCU1rK8fiWAP5fysryyJZlC8gRxwMxdvbzbH4A/hXe9vA7PmkLIXXLA33jGvP8SgD8Wisa8aXOcQmp9c8ApQsO8neB/P444w7yQEp5CftR0wMubzLwpZX8/o2DMvJe4n0LxBHbA14rHvDgdfz/6NMe8tmOhQsj9dsBOIci81NB+P1itx7zWAaNCb4x4wK0OxLw6i34/UoLDvOqopEK+w3nA6uPCvMPyfT9UHsK8qU6mQl45e8AZbri8usZ9Pwajt7y936dCUJB8wIDvrryKLn0/sfqtvGyBqUJJzX3AOUymvGO4fD/4PKW8zygJwq/VSb3z+QO7W397P9nQArvqngfCNDccvanFYDq+VHw/UilfOvELBsIJWue8t2GCO9myfD+KioE7j24Ewi3fg7zd89s7x1t9P5LR2jtmywLC4P6iu1FiQjxdF34/YqlBPKQZAcIa2kc7c8+FPPiufj8geIU8ks/+wWzbSTzD/Lk8hxZ/P/6puTzkTvvBlYygPLws5zxrUn8/SOLmPAzN98GuusE8o9gPPctFfz8aqA89+ET0wX89+zzRCyM9ryl/PxLNIj3Ur/DBygvcPBAZOz1y634/Ubw6PbYT7cGK4AY92TtPPROSfj8Hs049vGzpwYHzMT2iAmY9GK9+P7h6ZT3CsuXBx0FTPZ+ocT21e34/QANxPXzu4cG6lIA9izx+PUzHfj8ftn093ibewcxNlz2wv4E9bhV/P1ePgT02OdrBqpW0PYL6gj1gQH8/59SCPRJG1sEyl9E9nR2GPfJ9fz/PB4Y9DEDSwWnI8D2cUYY9QJZ/PzBChj0IMM7B/DcJPpM1hj1L3X8/xjiGPQDpycEQ7hk+wniDPebJfz9tdoM9/rbFwV0pLD5f+YA9TeN/Pw79gD2oTsHBhcE+Pj4Zej3i738/RCV6PWYBvcHKLlA+wcFtPXYOgD9J4G09yHW4wfTSYT4GHmQ9dTCAP1BYZD2c57PB/nl1Pp2YWj2qPIA/ttlaPXJVr8GEDoQ+mxtOPc17gD9wik49brqqwQL+jT6gNkA9kJOAP3muQD1k86XBud2UPlkOOD3bnIA/Eoc4PZQ7ocEmnaA+oYwtPUS4gD83EC49dk6cwXJ+qT7tbyY9WuWAP+wKJz14UJfBhKixPm7DGz1kBIE/sWYcPQljksFNPLs+hkgQPXQOgT/K5BA9gWmNwdldxz4hkgg9qgaBP4IhCT0cTIjBNYDNPhAc/DwvCYE/VSb9PBJGg8G0N9c+Qjf0PHQRgT/GQPU8agd8wc1l3D6bGuU8wQyBP/YO5jxciXHBcnjlPhJb0DywAIE/3y7RPNr0ZsEqD+4+wg3IPADfgD+Uvsg8BFZcwbah8j48DLs8ctaAPwKruzyKoVHBqrb7PgWvqzwb2IA/kEGsPPCrRsHKXwA/oqeXPNC0gD/cE5g8yLI7wSa7BT/zA5A8wdyAPxeBkDywtTDBB7oFP/bDgzzQ44A/8zmEPMSvJcGcSAc/S1tvPBbMgD80G3A8rlUawcK/BT+e6E48SdaAP4SWTzzAWw/By8YFP3KnQTwC6IA/iFdCPIb/A8HyygQ/Yf0lPLDcgD/WjCY8QPjwwPK+Az8OFwU8vMaAP5B+BTwQ69nAs90DP+LH5juY2IA/YovnOxDkwsC2GgM/B1nCO5jPgD/G9sI7FOCrwHWcAz9FH347t7GAP8PPfjs4UJXAD80DP8/LKDt3nIA/ADMpO4hse8CHEQI/Upn7OveZgD+oMPw6jMRMwDQ6AT9Pswq1yH6APwD4CrVElx3ApQwBPwkKUrrFcIA/kGZSurAK3b+ikP0+8nh3ujNjgD/Y2He6YAt7v/Yz/D4qFFe7KlqAP/dfV7ug+YG+4h76Pn96jLvHW4A/6qyMuwCP/T6K+fk+1qeluyFUgD9e3qW7wEmePyrO8T5OIMy7Z1KAPy5izLvgpv0/GnDwPvxEALzpYYA/NnYAvNQRL0B0k+8+RVwZvOJsgD/JnRm8Ch5fQHjP6z71Rja8ln6AP5KhNrz8yIdAvOPsPv3XXrzqh4A/Lk9fvAtOoEA5s+Q+rfFkvAuwgD8SkGW8ok64QDCP4D4E5ni8FbyAPxyeebwq9NBAPQDcPjZEjrxawYA/lLCOvFjY6EDuvtk+hNGSvH7SgD8/S5O8E8QAQVZS0j4EHZ6888aAPyeZnrzbeQxBy6jSPpUXrrxkuoA/BZiuvMqyGEFiQc0+b5O3vLS2gD9sGLi8QxclQWDSzD6rssS8td+AP/tgxbwnFTFBgvDJPtJyx7zN5YA/YyjIvAQ+PUGcT8g+gbLLvBTpgD+pbsy8OTRJQQaZxT6KOdG8UxSBP0we0rwoc1VBytG/PhdE2rzLLYE/uEjbvAsuYUEygLo+RlvUvHMygT+KXNW8Gm1tQbQduz51LNm84kqBP25I2rw0dnlBFsu4PrxC2Lzeh4E/ApHZvKjSD8K6V2HAeHQxvAUXez9AwS+8swoOwqzCYMBPSS68tAN8P2ruLLz0PgzCHVZgwC2EGbx3gHw/8ncYvB5pCsKZ0V/A/c77u881fT8OcPq7yY8IwhZwX8C02FG7y/V9P7gCUbsUowbC4U1fwELBPrlKlX4/IDo+uV+4BMKlBl/Al/qpO7cNfz9Aqqk7qb8CwhLjXsBJeR88p1p/PxtGHzzhxgDCcOtewHh9gjzOgX8/A16CPLKO/cGQtF7AZBqlPPeWfz/z+aQ84I/5wefjX8Aw1dQ83Jt/P5+u1DweevXBfrFfwN/K/Tyxdn8/AIz9PG5b8cEDJ1/AAhEZPZS6fz/QABk96iXtwRrXXsBa+yc95Zl/P+DfJz2o6+jBYDRewEnbOD0J8H8/jt04PXai5MEOjV3AG3pCPb8egD/SmkI9jjPgwbuzXMDMbEs9NyuAP9qZSz1QxtvBT9RbwHymVT1CP4A/sudVPbQ/18EDvlrAlA1cPcNIgD+xWVw9tLLSwVyYWcAcrWI9HlmAP9oKYz0w983BBVpYwETsYz1ZQ4A/TjdkPR5IycEgF1fAZJBhPaw8gD921GE99G3EwS7BVcCsh149uTKAP8rBXj2epr/BG4xUwAM0Vz0hP4A/xXVXPeCyusF+UVPA2iBPPfRFgD/EZE89Fqm1wdzfUcCWrEk9gkWAP8rtST3MrrDBxJRQwBvkPT19dYA/+kM+PSulq8EcPE/A/mgwPX9+gD8oxzA9OXSmwUhBTsCTqyk9FIWAPwAKKj1cUaHBUsRMwC5VHT0XmYA/O7gdPWAAnMFReUvAFOMVPXW3gD/JUhY94pqWwSmJSsASLAk9BM2APzedCT29S5HBBjZJwCy1+TyY34A/PZT6PEvni8FK1kfAtGTlPHrZgD9tK+Y8EHaGweISR8AsP8w8xuCAPzn1zDw2B4HBOwJGwGXjvjy19YA/1Jy/PBbXdsFSZkXAzhitPHfygD9ovq08iJlrwcx8RMCJ7pI8hPaAPwl9kzygVWDBmqFDwC2HhDxb4oA/Gv2EPIDwVMEAD0PAs0dwPJvpgD8RJHE8RppJwRIlQsB+1008xfmAPwWhTjyM9z3BMudBwFyXIzzy3oA/LSYkPJJCMsEWpkDALzUiPKcWgT8W5iI86KQmwTb+QMBFjAU8Bh2BPyYhBjyo9BrBs/lAwPcb6Dv9CoE/SQ7pOwD/DsETy0HA/8y9O9kdgT8Qob47llIDwUr5QcDJf6s7pyuBP6FIrDu8ve7AvIlCwM2qkDu8G4E/NEuROyB61sAcA0PAYqA5OwQCgT+BWzo7wFK+wOs8Q8CU3is7bheBPzSaLDuQFKbASoNDwJp6KDubB4E/HCgpO6TgjcDbgkPAMOvEOqXogD8mnsU6ELRswMBtQ8Dt7Ew6ib+AP0CGTTqQKDvAnxpEwG8bYDp8s4A/kLhgOuB1CsCqIETAOAjatyKfgD/Aj9q3wKCyv9NeRMBfcik5gIKAP8DIKTnQeCK/z/JEwD3DPzpbbYA/KBVAOoCXGD4g4kTAhiUMuqBkgD+cXAy6YBJpP0H9RMALNCm6q2eAP5B4KbrAetc/SvxEwCRRZbrfZ4A/MK5lutxKHUADxEXAI9GwukZlgD8YF7G6UOtOQISgRcACsA67zHiAP1rzDrtlb4BAdpRFwJOlU7tCioA/7RdUu2I8mUDh2UXAYOCTu0eRgD9cNJS79GKyQCNdRcAGE9q7lKOAP5Oe2ruqmstA5xZGwJzryLuY2oA/VZfJuxKt5ED5PkbAoarhu47ogD/bd+K7ITb+QEZRRsASDQi8wvOAP9GOCLxHawtBmh9GwJ+MBLzCDIE/9hcFvEdBGEE/lkbAgX0RvBMJgT9mFBK8am8kQbTeRcBKqCe88waBP+BUKLzE9jBB2F1GwNTsNrx0CIE/SKo3vGzDPUH3u0XAUpFDvCgvgT+CeUS86DtKQbqLRcBlXUC8XSqBPyk+QbwytFZBvStFwFkSPryWKoE/lvA+vBclY0FqxkTAiDE6vI5KgT93Iju81MJvQarCRMAy50G8LFaBP/PqQrzt7HtBYJtEwFPiLbxmT4E/j8YuvA8/hEEfyEPAIIsuvBZigT/4fC+8t3KKQURQQ8DinyW8OI2BPzyhJryC3RjCD6/EQFD0zLpoA38/No/MuobrFMKvs8RAXydxO1MZgD9MP3E7MNoQwv+yxEAbpww88JCAP/b2DDx0uwzCCtjEQOoPTzzSzYA/GbdPPCqUCMLaD8VAUXKHPObPgD8b4Yc8umUEwjgpxUCEqZU8NKqAPxYOljy3JgDCLWjFQH5krTzfgIA/cL2tPNTY98FRwMVAhzC5PDpXgD+mcbk8PEbvwb/dxUAqlcI8ZC6AP8W6wjwKwObBAkfGQMJExzwyNYA/r3DHPMQx3sEmP8ZAq+LHPPMUgD+R9cc8vqPVwfmYxkAtUsk8bvd/P2ZRyTxMD83Bcq7GQJXrvDzi5X8/F+S8PMCKxMFk/sZA4mK7PDzHfz8zULs8pgm8wZQNx0BdU6Y8K9l/PzZIpjyEn7PBMwnHQHR2mTxqtn8/jGGZPA0bq8FWNMdAv7WIPD7Bfz/NpYg8EriiweU4x0AlyHU8ogmAP5PSdTwXTprBSEPHQCpgZjzwGoA/YXlmPPjekcEcQcdARjZTPIg0gD9dYlM8koKJwU4Jx0AL4jg8NEiAP7EWOTwfN4HByAfHQJo7Nzx9aIA/4YY3PNjSccFhFMdAZPQiPExvgD+VOyM8KlZhwZbQxkDtHfk7z4uAP0im+TuA2lDBqdrGQN1BzzuckIA/HrfPO2JfQMGEr8ZA86+2O2qGgD/+D7c7kPEvwZiDxkB3K5478W6APxdwnjuEuB/BsWXGQIA2hzujYIA/l2mHO0hbD8FSHsZAmUQ7OxlNgD8HfTs7oEH+wL/zxUC1O8Y6hhSAP5xLxjqY6d3ALbPFQHCmk7e2AoA/AKiTtyRSvcDdNMVA2P2vuozOfz/a7K+6KKidwAwRxUB5ih+7KYF/P/hiH7swb3vASdjEQBObnrupOn8/Bl6eu7RfOsA6bMRAB+C9u+kqfz8kkb27cNb1v35CxEATgvC7Pgl/P3AO8LsA12y/M6rDQNIQHLwxB38/SMUbvICCgT1cbMNA4fpFvB4Tfz/mn0W8cPuGP0z0wkBsRlq8XRp/P1nlWbysvgJAVXXCQJiZfLwKan8/41B8vOo3QkDo6MFARWCQvC10fz/MOZC8UcqAQEmRwUBOVJu8XeF/PzRMm7wGS6BAVDTBQB+KqbzwO4A/XrOpvDMGwEAVssBAdPu4vNBegD/7Qbm8yHffQFFIwEALVMG8mWuAP5qnwbyYBP9A662/QLKj07ygmoA/jCbUvC0VD0EKKr9A05jSvGTFgD8vPtO8WfUeQXOvvkCMsdq8OrCAP21L27z8oi5BhB++QOse5rzaooA/MbXmvF5kPkEk3L1AuyLYvMKugD9+udi8mANOQQtovUBgM928F3uAPy6h3bwgrF1BDfq8QAi54bxhWIA/ngrivOZEbUG4j7xAMrnWvFozgD9r59a8Xyh9QQg6vECaP9G8QzeAP7Bv0bwfeIZBcIW7QASGzbwlLYA/Ba3NvNZPjkFLe7tA92jIvIoEgD8Ub8i89gyWQTfwukALcLq8F+h/P2VpurxY4J1Bnau6QE7NubyT6X8/Nce5vHSspUEkXbpAeYq0vD8DgD+ijrS8mnKtQUYUukD8P6y8SQKAPyZDrLx+OrVBNI25QE/xqrzEIYA/cgmrvN7ovEHeh7lAA3ygvJRFgD/yqKC8LKbEQY9XuUCcBp+85nSAP4JQn7zGdMxBWxS5QDhSn7yYnoA/NbafvGwv1EGXCrlA9C6fvEi2gD+VoZ+80MzbQah+uEDIW5W8y+WAP+rilbx9cONB2GG4QLjjkrzY94A/8XKTvOpJ60GA+LdAEaKWvPI0gT/yWJe8osbyQVyst0C7loW8NTCBPz02hrxyPfpBfia3QADlhrwbE4E/v3aHvOHuAEK9DrdAW96HvBkUgT+xcYi8nLYEQta2tkD+kIa8BwiBP44ch7yMhAhCXPq2QNoteLzVFYE/ajx5vDREDEJftrZAT1JtvEHwgD8ZMm68Ze0PQnDmtkADME688K+AP2q+TryvrBNCeK+2QJAlSbyKjIA/o5RJvBJ0F0JsdbZAts86vJ9vgD+vITu8fhobQlUntkBnnTW8+E6AP+fVNbwWvh5C8wu2QJmtLbxpFIA/3LstvDSOIkKM2rVANQEUvEcNgD8kCRS8RUgzQhJB2EAZ1Ri7gQ99P3/0F7v0cDZCpg7XQLaLlbthGn4/8P2Uuye2OUKb2dVAZ8XXu0nKfj8UQ9e7+vY8QmHU1EBUDe27pCx/P7yr7LuOOUBCfNbTQESX/7tqRn8/9Dr/uzdyQ0KEztJAGC0/vD7Vfj8Yvj68RbhGQqvh0UBZgGm8kJR+P5zbaLxl8klCkOTQQG9ZkbyOTn4/XN+QvBg8TUK51c9Ab1WpvCYofj/ruqi8TXxQQpv6zkBSS828bVR+P5+izLwUuVNCNvDNQAAQ4rwZd34/L2bhvGjyVkIiBc1AipX0vP3jfj+CEvS8ZS1aQsobzEDSvAK9thJ/PxGDAr3CXl1CUiDLQPdGA73ELH8/rRMDvaKSYEK8JspAYxoLva9wfz/e9gq9Vb1jQgoKyUBWtRC9DHF/P8iQEL2z3GZCyhzIQG8AG70db38/TdkavXYFakIU8MZAbcUmvSDAfz+Etia9iCFtQqrexUBhey+9Rs9/P4xxL70ZTHBCsrPEQL1JNb1n/H8/C1A1vbFec0JMgMNAxjo9vQACgD/fRD29hHZ2QpQ5wkBN7US9OyqAP4EXRb1kiXlC2BXBQHJvTb2GMoA/BqNNvbWZfELfxL9AepVUveh0gD/MAlW9dqJ/QrydvkC5AWC9XoaAP6CFYL0EV4FCKEC9QFGdZL1Wt4A/SFBlvS7agkIj37tAnyhnvWLEgD+16We9HVmEQseuukBcMmm9LtKAPwACar3p3oVCuFe5QNX1bL1b4YA/addtvelfh0JmBrhAkxJxvfDmgD/u/XG9PNyIQk5/tkB05nK9weWAP77Sc71+WopCdiW1QPPPeL2O6oA/lcd5veTXi0Kbu7NAuah6vW/bgD+uk3u9+FqNQjZbskDaYHy9HN2AP1pPfb2z045CT/CwQFzIfL162IA/u7J9vdhWkEIYoa9AxeOAvfDSgD/nWIG9QtSRQggXrkCTToC9xdCAP/7BgL0qUZNCQKSsQP8+hL1Kw4A/r6+EvSjGlEI1UatA04aEvUXBgD/B9oS9/UaWQhXiqUC9u4e9a9uAP9g8iL0qwZdCpD6oQATKiL330oA/z0eJvXg/mUJSIKdAKEWJvUTXgD/KxYm97bmaQsihpUCOM4y9MtaAP+62jL3MMZxC3wmkQGS1kb0q5YA/n0eSvcWvnUJbrKJArQCVvZvcgD/4kZW9xSufQrkToUCn1Za9i+uAP/Zxl73FoaBCp3afQE4CnL0Q9YA/C6ucvT0jokKq3p1AG/GevdXjgD8Rk5+995ajQlxlnED4VqS9L72APxLnpL1VEaVCmdaaQNKFpb1dyIA/hR6mvf6JpkKmdplAv0qqvfWqgD+r1aq95wCoQsrul0ARLa69WqKAP322rr2jgqlCMHuWQGWSr70CkYA/dBGwvbD8qkI9G5VAj4KwvXKIgD+w/LC92XmsQsKIk0DQi7K9hXyAP6//sr387a1CnSaSQFs8s738XYA/gpuzvUBpr0K9mJBAp8G0vc1mgD9ZKLW9/dqwQkknj0Dxfbe9g0+AP2vWt71YVrJCLqyNQLLGt72eZIA/ki64vafHs0I+LoxA3zC3vWFIgD/+g7e9aES1QtHZikDs2rm9ZGaAP/1Fur0ttrZCNmGJQM0Jur0SXoA//W66vSguuEKu24dAjp66vVJ4gD9pF7u9k625QrGPhkD0G7u9w6CAP+Kyu72+JrtCJhqFQNivu70iwIA/dl68vVSSvEJMsINA7d63vUXbgD8qnLi9KAy+Qn2HgkCalLa9Tt2AP4ZRt70kh79CATOBQE/ltb3i/YA/dLi2vd70wEIbO39A3kO0vR//gD9yFbW9YmTCQpd/fED2G7K9TyyBP8oJs71o3cNCqMp5QNrlsL0NGIE/qsOxvUBQxUIcbndAIByuvRYxgT+cBq+92sjGQtSXdEDSL6y9tkKBP/0irb3DM8hCpeFxQIvEqL3qOoE/vKypvQKtyUJY4G9ATEGjvS5DgT+mJaS9OhnLQoL3bEDt0qK94T6BP9Kzo722ksxCAMxqQHAFn72EGIE/PsifvbcAzkKXimhAQJ+bvQMsgT/laJy95mvPQttuZkBH9Ze9KwKBP3OgmL0B4tBCsvhjQDMak72P24A/naiTvSSnXEJReXQ9VoelPCzSez/kLqQ8xX5eQol0jD36row8IGR9P1j4izxKXGBCi5SaPeg3YTx0R34/BndgPFg6YkKoLrM9A5dcPF7tfj+KIVw8PSRkQqtHyz0ammo8mFR/P5ZMajx/DWZCmPzhPeVqeDyiYH8/yh54PFEDaEKd8Pk9U8OPPClEfz+Gj488KPlpQtUmBT5Iy6Q8ORl/P2yCpDxo/mtC0+oLPu3vsjwE134/84myPFEFbkK7dBY+Jtu6POyQfj9AV7o8cBJwQoL2DT5Dx8A8YVx+P4krwDytJHJCbjkVPmlUyDwbY34/aLXHPPg9dEKz8hs+WS3JPBhJfj94g8g8Glt2QprSGz7IRsk8jVZ+Px2iyDxAiHhCv9odPiUoxjyXu34/D63FPDqyekKmBR0+5Vy1POwUfz+EC7U84d98Qm5jIT5O6aE8GWt/P5G7oTxyG39CcWkcPonqizzu5X8/SOSLPG6sgELRbhk+UQ9pPE0igD+NL2k8eM6BQnQhEz7XZ0c8pV6APzGyRzxw8YJCblgQPlqKHTxPfYA/x9cdPKUWhEJomAg+kn72O+ObgD/4FPc7HUWFQlc5AD5Ddq073MKAP2b6rTvscYZCOpbyPVrDXDtP5YA/J4ldO0Glh0IJIuM9ZWmbOr7qgD/o95s6UNqIQg2wyT2qcCe5y/GAP9AOKLkIEopChqi0PXQP1rq28YA/lNnWusZJi0JCPqE9iFkhu6vcgD+i5CG76Y2MQg5Wjj0vjUe7cNWAP5kzSLsvzo1CReGIPeFpULtA14A/KhlRu8EVj0Jeplc96AuCu0vFgD8scIK7K12QQruAST3BTYa7Zr2APyqxhrsLqZFCJyEkPeqhU7sCrYA//DBUuxT6kkLtchk9b/L5uiKxgD9mn/q6e0+UQvUPAT0wRPy6SZuAPzrd/LqcnpVCMpLxPKD+hbrGi4A/ykeGulf4lkK8xsU8XHBbOimHgD845Fs6u1OYQubwzzzbPfs6lmmAP4Cl+zrAq5lC9WjAPI4PZzvIaoA//29nOz0Om0IDItA8ZPCpO5lugD/oOao76G+cQlovxTxtsd07L3qAP3Qb3jv5051Cv+AAPY6iETyJcoA/9eMRPLs3n0Jjvh09KPk7PJV5gD/2Ujw8yJqgQsb0CD1TtUo8k4uAP4IkSzxKB6JCRy0kPUIhYTyJl4A/bqdhPD94o0JqeTQ95faAPMW5gD8qVYE8eOKkQqmiPT09sI08wL+AP0YbjjwgWaZC9HtUPXyLlzyl14A/QAyYPG7Hp0Ij83c9ewylPL7cgD87nKU8oT6pQrZUez3PJbM87/2AP1jZszxJt6pCsoeMPUEuujwcJYE/fAW7PPUqrEKgpYw9ACfBPIgzgT9WEcI8AKatQucxqT16KMo8YkCBPx4oyzxWH69CUeOxPVs82jyMQYE/ylHbPBycsEJZS8894hfhPJlPgT+cQuI8GxiyQmJw3T3Pe+g8pT6BPzKh6Tx2mLNCe3b/PS4q8jxNVoE/hXLzPL8StUKKvQM+4zHvPDJFgT8bZvA8UJG2Qn8wET5iq+88qF2BPyL38DwlErhCRfcbPtF+/DzbVIE/JtT9PCKduUJy1SU+vl/1PDdegT8ktPY8ox+7QvIkMT63rPM86EyBPzfu9DzZqbxCLNw8PjbC8Tw+VoE/7gnzPDQzvkJuqks+0FPvPCdegT+Ln/A8OL+/QhA5WT6Hz+s8/FaBP6cP7TzwRMFChJVlPq4r8jxMUYE/S2/zPA7VwkK4FH4+kN7uPFRYgT81JPA8r2PEQgdwhD4QTeo8AjWBP/1r6zwi8sVCk2+KPpk65TzUGYE/zTrmPIiDx0JxbpQ+gtrgPEj6gD/3ueE8YhTJQh3xmz7mE9c8O9mAP5TN1zzOpMpCFV2mPsLJ2jwjvoA/mG/bPPI3zEJ+96w+9yLRPLySgD/BndE8RM3NQgOjtz5wW848Hm6APwC3zjyKb89CKeTAPvgF0DybTYA/5UfQPEz90EKYMs8+j23KPC4qgD+Nkco8pJrSQtDU1z7knsU8ksl/P1aMxTyjMdRC/mPjPn5uyTw3rH8/IFDJPNbL1ULfX/E+7ibCPNRafz+f6sE8KmbXQpIL/z7pIMU8jux+P0q5xDxgel1C6FFpQDGCyTyME34/+MLIPOrFYELcZGhA15GyPMEffz9wRbI8jDNkQuOTZ0AUGpg8saJ/P3r/lzyMmWdCElBnQHvfljxh638/f9qWPJb9akJlDmdA3XibPIX/fz/peZs8sFVuQonAZkCke408/41/PwldjTwZt3FCyYFmQNMwgTxmKX8/XPuAPDQJdUIcP2ZAD0ZWPNatfj9RuVU8UmR4QtLRZUAptSg8TTN+P70dKDxWuntCLM5lQEGZ6jtkAH4/F6/pO8v+fkLKJWRAvDeKO8/YfT/9ook7wiSBQsLxY0BBGjM7mwp+P+RqMjv4w4JC4qFjQEbckzguLn4/wFWTOANghEKQHWNASZRAuqxnfj+0+j+66viFQnmGYkCy7y27Yc1+P46HLbsNjYdCQMdhQCEomLvSGH8/f+OXu38XiUKJT2FA4tz/uzmAfz9cnf+7m6SKQq5PYEDYXTi8zv5/P+ldOLyYKoxCz5VfQNMEYLwaK4A/bytgvASyjUINb15AbEJ/vOtkgD9fqH+8Yi6PQnxkXUD3Co689ICAP25TjrysqZBCYRhcQLwnnLwdpYA/qo2cvPsfkkKx9FpAzySrvIWkgD9klKu81pOTQjWCWUBFrLG8fM+APw8+srxn/5RCDktYQBSBxLyOwYA/EhjFvO9plkJj1FZABsHJvKzVgD8KbMq8DNSXQl8vVUBcU8+8ssKAP+Hzz7ySMplCw+RTQAc00Ly+uYA/+M3QvLmSmkJwTlJAkRPSvKe+gD/5stK8i+ubQnMGUUCu09C8cbyAP01w0bxMP51CjitPQJ0q1rwIuYA/isjWvOWSnkJGvk1AVH/cvPykgD/ZEN281OGfQrH+S0A5HOC8PZeAPzWk4LyJL6FCqMZKQPeL4bxXn4A/AhzivF9zokLTNElA6AbgvBecgD8Wk+C8nLyjQoYFSEAtaem8i5mAPzn56byu+qRC749GQH0R37y/nYA/e57fvD80pkLXTkVACE/ovMuagD9/3+i8MWmnQtzzQ0Czu+K8uoyAPw0847xdnKhCl/pCQNjg5bzHp4A/YHvmvG7OqULnjkFATaHjvNqigD/dNeS8GfmqQrjeQEAPKOK8pZqAP1y04rxaHaxCzu8/QAil37wzl4A/ryzgvCtArUL6sD5A0WPqvCmtgD93Buu8ml2uQqswPkCG9Oa8NqGAP+SJ57yjfK9C0kU9QLxp47yvroA/qQjkvM2NsEKDaTxAs/7jvN+5gD8EqOS8u6WxQvKaO0DCr928C6+AP9BK3rzSsrJCQBM7QOE95bw+lYA/XMflvIO5s0L0lTpAeyHXvMKYgD8Gpde82MW0QsfsOUBoadm8wbaAP+QH2rzmx7VC5rQ5QJ1T37x2r4A/O/DfvE/LtkJ7ajlAnBnVvNOigD87pNW84cW3QiFCOUC6gM68AKeAPz4Kz7z4w7hCLMg4QAuQxbwbr4A/oxnGvJS1uULcqjhA3Qq/vAWegD8Cg7+83aS6QraFOEBTprm838WAP9w3urwukrtCtv83QILZw7wgv4A/H27EvGl+vEKIYzhAGTjDvGTWgD/23cO8e2G9QhgkOEAIE768pd2AP8q5vrxMSb5C8S44QCodxLyWD4E/oO/EvBAiv0IEQzhAmKfCvNsKgT/bdMO8MwPAQggIOEButMG8BC2BP4Sawrx25sBCZkw4QLPhxbxeSIE//+HGvCi+wULX3zdAOabFvIp7gT+5zca8UIfCQj04OEDrCru8upaBPy42vLxuX8NCL7k4QBphuLz4roE/gZm5vMozxELYBDlAgd20vPrPgT8zJ7a84vzEQgBeOEDYm7i8Ht6BP6T2ubyAwsVCaLA4QK0es7wxCII/fYy0vE+PxkKnwjhAniy0vIIIgj/UnLW8nVrHQlI7OUAnaq68rDGCP4bqr7x8G8hCVtQ4QGxMrbxbOYI/hM+uvDPXyELtBDlAVr+gvHJFgj/FLaK8TZ3JQkGQOUBSD5G8hFWCP+BikrxhV8pClaY5QGuokLy0UYI/3/iRvLQXy0JbKjpA8X+AvKYugj8JmYG8+sjLQpWIOkBmUXG8wDmCP5lrc7zgc8xCgbM6QHJiUbx3CoI/gg5TvGkxzUIvIztAPYUxvBDhgT9G0zK8vCeFQv+wcD1z+qk8Tgx6PyQCqDzgLIZC7S2LPQyCjjx66ns/9F+NPDc1h0I3uZo9xdFpPBkjfT8UhGg8Dj2IQtoftT05bXM85gJ+P1F8cjzcS4lCkkvPPWTSiDx4in4/Ym+IPMpZikLbQ+k9pT+UPKikfj8b3JM8Cm+LQl7vAT5rPKs8K5p+P1bGqjz0goxCfbsLPmcjwjy2cn4/FI/BPGqejUIZJBQ+fYHQPFpAfj8Qzs88+7mOQhJgID7PW9c8Ww9+PxKO1jz21o9Cj0sYPoTE2Ty1730/F+fYPKz3kELePyE+aw/fPGAEfj/HNd48IBuSQqnGKD6R59s8PvV9P2kK2zzLPpNCdhspPvqv3DxP/n0/99XbPP1qlELUbCs+nwHXPKdUfj9QUdY8SJSVQrWtKj4CasQ8NJ5+P7DkwzxRvpZCUI4uPiu5rjwl9n4/I2CuPL/vl0Ly1Cg+KKaVPEpsfz8MfJU8lCCZQo4qJT7/gnM8gst/Py1rczzCVZpCL64ePl36UDyPIoA/TRdRPKGIm0K30Rs+OcslPFNHgD/H+SU8/b+cQtXIET5Bmfs7XmmAPyAB/Dve/Z1CpekIPjA+tTtZjoA/FqO1O/A7n0Jy+f89gSx3O+y5gD8X4Hc7In2gQv6H7z104NY6Gr6APwiA1zrvv6FC7h7UPfWdtzkUxIA/mCq4OeYFo0Kpf7w9p/RnumK+gD8ooWi6zkqkQvVKpz1vA8a6S6mAP2SGxrq/maVCKZuUPYUQ2rpmpIA/kpzauh/lpkIXqY494bqtuoingD+ULK66xzaoQmWjXT2eh+K6V5KAPyAJ47pliKlCTn1QPdqn4Lp0hIA/GhzhuoPdqkJXuyw99h6nuXp0gD8Aa6e5PjesQr6ZJj2TesI6JHiAP9rVwjrkkq1CD0INPX1qADtfX4A/V5oAO4HrrkK/Bwg9Zl9jO3RRgD/Np2M7pkmwQvb7+jztcr47fkuAPzqrvjvtqrFC/RUIPfIJ6jvaK4A/SjLqO78Js0J1Mgc9Fh4XPNkqgD+nNxc8fnC0QorVGj3KVDw8li+AP1R4PDzD17VCa4ogPawFXDxeNYA/YzRcPJQ/t0KW3Uk9cm+DPCIygD/oiYM8Zaa4QsS1ej1/qZ48uDiAP+vNnjzTDrpCJ515PTHhqTxoR4A/IxKqPNZ8u0J43JI9z6K6PFpZgD8F5ro8AfG8Qk2QpD1fts48nHuAP/8czzy4Xr5CYLm2Pa/14DwMiYA/vnHhPIfXv0JgdNA9xJzxPE6kgD9VPPI81EnBQrl+6z1Lkv48EK2AP6ZD/zwMwMJC/yX8PSVqCj0dyoA/ztoKPcw8xEK+LQs+uLkPPdH4gD8zSRA9MbPFQsqTEz7UuRM90QSBP3NUFD1IL8dC/QkpPu5dGj0yEIE/wgYbPXiryEJI2DQ+pnwjPdMZgT81NiQ9XyrKQjJNSj6N2Cc9yCmBP9WhKD2kqctCajBaPqBNLT14HoE/NBYuPdspzULtKnE+0GAzPUE6gT9iRDQ9uaXOQiacez7sjTI9HyiBP7xjMz14JtBCHv+IPuunMj36QIE/NY8zPY+n0ULWKpE+Frk4PbY8gT+spTk9ojTTQr5tmj6abDU9iU2BP5tgNj0Lt9RCn9SkPu/WND0kOYE/srs1PT1C1kKafK0+atAzPcpIgT/JvjQ9pMzXQoXTuT6XPTI9JlSBP6gxMz3HWtlCXU3DPjJBLz2dU4E/kjAwPVzg2kJ2Pc0+RRYzPZVPgT9fCDQ9ZHDcQqs73T5o0zA9/luBP9jKMT0o/91CuNrmPqEhLj2UOYE/rP0uPbKN30LwS/A+fIgqPZ0lgT9pUis91RzhQui//T6Yjyc9VAaBP05BKD0Fr+JCF3YEPwjPIj3G6YA/NmkjPXpA5ELenAs/v4kiPbzPgD8eEyM9mNPlQqVEED9aGh09OaaAP06FHT1uaOdC3eMXPxfJGz3fhYA/YB8cPQ4J6ULk6B0/vZkbPQtngD8r3Rs9B5XqQk7yJj92bBc9aj6AP8yVFz0lM+xCa6EsP8FJFT3g538/80YVPcbG7UJJvDM/IisVPdu6fz81GxU9C17vQpJIPD+wYRI9ZGl/P506Ej2s+fBCBK9EP9SWEz3M6X4/t0oTPWA9kkJLgmDA5zV+vFxbez8I6Xu8oMuTQuquYMDTdYi8TXl9P0TKh7yqXZVCAdNgwKTPdrwV+X4/GFJ2vCztlkISdmDAl0wUvCfVfz9wQBS8M3+YQgIbYMDMy0a7ixmAP6zfRrssEZpC/blfwI56eTqRCIA/6IJ5Ok+nm0KoS1/AZmDLOz/Nfz9oTMs7QDedQmrQXsB2Jhw8+F5/P6f1Gzzzy55Co3RewK5ZTDyD8X4/Zu5LPJRfoEI5u13Aq2hxPNymfj8Mx3A8i+2hQpa9XsD0PoQ8snF+P8/YgzztgaNCGQNewJoMmDygqH4/vqeXPFcRpULed13ALr2iPH60fj8tVaI8mp+mQt0wXcCeErw8Yeh+PwSuuzyeMahCEsxcwCpqxDwjPX8/0SHEPGm+qULCZ1zAAFjJPCtnfz9+Hsk8kEarQuDCW8ApLco8o6l/P68NyjzM0axC1plbwFiswzxE8n8/eqnDPBJbrkLyMFvAE9/BPKUJgD+y6ME8AuWvQqIAW8CElso8oxyAP9KvyjwwZbFCtFpawFjm0jwFHYA/PAHTPIbtskIQJFrA3sbWPPgqgD8S7tY8nm+0QrrUWcAyC9w85SeAP+Aw3DxE8rVC0aNZwFaF6DwSP4A/n8LoPDpvt0IiNFnArh7qPFcwgD/4Tuo8me64Qvf8WMCIJ/E8XjOAP2Jc8TzgbbpCtcdYwAWU+jyxIYA/Abr6PH3ku0JQgVjAA8kBPewGgD9NzwE9cV69QrQ0WMBCDgg95QeAP6gVCD2u1L5C95NXwAEMDT33DYA/RBcNPT9LwEIUlFfA5T8QPaoDgD/HRRA9csDBQiRDV8DPQRI9u8Z/P201Ej26MMNCgvFWwIMrFj3hsH8/nBgWPSumxEJlbVbAs8IbPfqrfz/xrRs98xPGQu8lVsA7LB49Ppl/P4QRHj01hsdC0cZVwAAwHj1bqH8/9BkePc/uyELIIVXA9+AkPRiSfz9EwyQ95VnKQo/IVMB/8SM9zmZ/PwrGIz20wstCC4lUwIbdKD0RXX8/560oPU4vzUJps1PAnZMrPRR7fz99bSs9+ZbOQvJ2U8C+xyw9v3J/P6GeLD3r/M9CQJ9SwKPdMT11Vn8/46kxPfBg0UK31lHA4Nw2PQt9fz/ftTY9EMPSQmOmUcBGujU9mpJ/PxKbNT3kJ9RC5P5QwPNeNz2Ps38/aUs3PfCK1UIYjVDAgdI5PaDtfz//0zk9SO7WQrzJT8DQ0z09JguAP8jkPT08UdhC8vxOwIPLQT0OIIA/Cu1BPQ6v2UIjX07AEpM9PUwogD+VuT093gvbQhi2TcDGt0Q9eD2AP7PwRD3IcdxC/0VNwEo4RD0UaoA/OJNEPabQ3UJJTUzAobo/PaV9gD+2IUA9sjDfQq6HS8Bel0A9X4qAP5IIQT3mjeBCFPVKwHauQT2Km4A/Zi1CPZzu4UL3KErARHQ/PcW4gD9mB0A91E7jQsmKScDu3T49DcyAP/B+Pz2yo+RCyNJIwGArPj1f+IA/qOw+PTj95UJCjUjAtIY3PUz5gD9RQTg9CFvnQgKWR8BW6zA9DRGBPxyvMT0ArehCPj9HwI14LT2NIYE/bUMuPZcK6kLdYEbAU2kkPYVMgT+PRCU9RVvrQqd/RcBbvR89RkOBP0iMID2ktuxCdwlFwMCyGD0nZIE/vYsZPZwQ7kJaIUTAeucRPWh9gT/UxBI9PGvvQk3jQ8CntAo9gpqBP32WCz2Os/BCGUZDwFsDDD19oYE/NOsMPaQL8kImMkLAHMIGPZKygT8Aqgc9smLzQtqPQcDaIgQ91KGBP3j9BD2EtPRClmRBwPG0+jzInIE/O078PLwA9kLrfEDASQH1PG+RgT8vhvY8Lln3QjXnP8AjH/E8eoqBPyyX8jxKqfhCrhI/wJYq6zyWe4E/cIvsPPT0+UJapD7AnGDmPAxNgT87kOc8q0X7QvSePcDsAus8/DSBP7gi7DzpmvxCj/o8wI997DzOI4E/WY/tPDnh/UKd3TvAO4LpPELrgD/iXOo8BDb/QmMxO8B26fQ8CJWAP7h89TwWQQBDUak6wHvt+zwKcoA/zGL8PEPgAEP8YjnAwToEPSgzgD8fWAQ9nIgBQ7I8OMBUTw49Vq5/P0s8Dj0BHp9CxihgwCc9YLwNxHo/MPNdvCCwoEJ9X2DAkaV7vOkefT+MPHq80UWiQnKHYMBdX2W8t8l+P1DVZLwV2KNCVzBgwMEYAbwUvH8/zQcBvGJtpUI63l/AzcfxuhwQgD8I1/G6dQKnQsyAX8C2DBI7Jv1/P+oLEjsDnKhCMBxfwDzq9Tu8tX8/3Mb1O0UvqkLdqF7AJwY0PPs/fz8ZwzM81MarQn5YXsD7X2Q8ns9+PyjZYzxUXa1CpqVdwHG8hTxJgH4/+ViFPADurkIomV7ALgeRPPtJfj8TjJA8S4WwQubjXcAqSaU8GHh+PxTMpDzzF7JCY2BdwPyRsjwwgH4/7Q2yPIKos0JYIl3AgSbMPESnfj+/n8s8NT21QkC7XMDNltY8m/d+PyAr1jx3zLZCd1pcwECY2zyeG38/qDnbPPVWuEK8u1vAKQHePLlZfz+JvN08sOS5QhSQW8Dw2dg8aJV/Pwiw2Dz+b7tCwSpbwIDH1zyirn8/ZqjXPP77vEKY8VrAeAXiPA7Jfz/i8OE83n2+QpI8WsDbVOw8Ocl/P8Q/7DztB8BCbgVawDkL8Tzo238/rv7wPFCLwUIqqFnADvf4POzNfz+c4/g8FA/DQrB3WcAurwM9kvl/P26wAz06jcRC1vNYwMI9BT132H8/ejYFPb4MxkKBq1jARxAKPcfYfz8MCQo9Lo3HQmdrWMBjpw89+7F/P0SVDz0VBMlCnBRYwFFvFT04en8/g0wVPW19ykJMqVfAlmscPb51fz82Rhw99vLLQqLsVsARPSM9RH5/PzsZIz1yaM1C2OZWwHh9Jz2sZn8/R1EnPQfdzkJKfFbAyPQpPcYffz+VsCk9CkzQQiUJVsAkBy89bgl/P6i5Lj29v9FCDWJVwHEuNj3hBH8/w9w1PTcr00IgB1XA9rM4PYXqfj/cVzg91ZvUQmyTVMDNKTo9P/x+P4vTOT3gAdZChMlTwOaPQT046n4/FjBBPb5q10JnUlPAEbNBPfTAfj+WQ0E9adHYQvj9UsCrfEY9pbV+P4UGRj3ROtpCyQpSwDViSj2O1H4/XPZJPRCg20Iyt1HAG29LPQ/Dfj/b+0o9jQPdQsHHUMDr7VA9tKl+P89tUD3DY95CDtVPwGumVj3XyX4/8DBWPSbE30Lfg0/ArFhVPf/Xfj+n6VQ92SXhQizETsA6q1c9sf5+P5NLVz2/heJCAkROwA+wWT3hMH8/GWVZPWTn40InVU3A+SBePQpgfz9/6V09AEblQshuTMC0tGE97Yx/P5aQYT1GouZC3sVLwJ9uXD0xnH8/Q1FcPSX750KDAkvA2bFjPS6/fz8HpGM9YV/pQip3SsDZNGI9XBCAPwhSYj3UvOpCRWRJwC8lXD3YIYA/3U9cPZMa7EKpjEjAm1BcPRoygD9WiVw9l3XtQjLUR8Cdllw9I0eAP5ThXD1B1O5CLAVHwJ6dWD0hZYA/IwBZPWgz8EJCSEbArrBXPRaBgD86Klg9WYbxQsOfRcAGS1U9NLOAP7XsVT0F3/JCWV9FwCxjTT3AtoA/1wBOPXk79EI/X0TA0DtFPYjQgD9D5kU9h4v1QgYFRMC+2z89a+qAP3GUQD3s5/ZCPhlDwJr1NT1HHoE/x8g2PeI3+EJNLELAACQwPcUYgT8q7DA9oZL5QqTEQcCfBig9HT6BP3rdKD2s6/pCK9JAwHJaID2fYoE/2D0hPfpF/ELVoEDAzWwXPQyFgT9eVxg9eY79QrwTQMCqbxg9N5aBPxRmGT3j5f5Cqgo/wIdGEj3irIE/mD8TPYweAEO9bj7AIxIPPa6jgT9uABA9FscAQ8RNPsD11AY98qaBP9u2Bz0ybQFD0nE9wIS7Az0Wn4E/CZQEPRIZAkNy2DzAtIMBPQObgT9sVgI9iMECQyAEPMAcWfs8gI+BP27m/DwQaAND+Zo7wJZw9jwRZ4E/B8/3PNAPBENEljrAl9z6PEBSgT8YLfw8zroEQy4JOsBn2/o81EOBP8Ud/DzqXQVDoec4wLZJ9zyRC4E/AFH4PBYJBkPRPTjAcswBPSWzgD8RKgI98q4GQ9K6N8AO1gQ9bI6AP/IiBT0CTgdD0G82wBODCz0hUIA/M7ILPSj3B0PAUzXAoe4VPXvefz8Z6RU9UEsBAi0DLQAAAAAAAAAhAHfef6CASwAAgEsAABQAAAAAAAAAAAAAAIABAAAAAHByZWRpY3RlZF9hZ2VudHMubnB5UEsFBgAAAAABAAEAQgAAAMZLAAAAAA=="
|
| 886 |
+
}
|
| 887 |
+
}
|
code/tt_diffusion_planner/server/__init__.py
ADDED
|
@@ -0,0 +1,5 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-License-Identifier: Apache-2.0
|
| 2 |
+
"""HTTP serving of diffusion-planner-p150, thin bindings of the vendored ``ttaw.server``: ``app.py`` (the ASGI app
|
| 3 |
+
tt-model
|
| 4 |
+
runs with uvicorn), ``client.py`` (stdlib request builder, runnable as a script), ``smoke_test.py`` (stdlib PASS/FAIL
|
| 5 |
+
check of a running server: ETH dispatch and the 12x10 grid asserted, agreement with the stored CPU reference)."""
|
code/tt_diffusion_planner/server/app.py
ADDED
|
@@ -0,0 +1,51 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-License-Identifier: Apache-2.0
|
| 2 |
+
"""ASGI serving app of diffusion-planner-p150 on one Tenstorrent Blackhole p150: the HTTP contract of the Autoware
|
| 3 |
+
collection (vendored ``ttaw.server.app``, BUNDLE_CONVENTIONS.md section 7) bound to :class:`DiffusionPlanner`.
|
| 4 |
+
|
| 5 |
+
Served by tt-model-manager as ``kind: tt-dit-server``::
|
| 6 |
+
|
| 7 |
+
python -m uvicorn --host 0.0.0.0 --port <p> --lifespan on tt_diffusion_planner.server.app:app
|
| 8 |
+
|
| 9 |
+
Routes: ``GET /``, ``/health`` and ``/v1/health`` (always 200: ``ok`` / ``starting`` / ``error``), ``/info``,
|
| 10 |
+
``/v1/models`` (stub), ``POST /predict``; errors 400 / 422 / 503 / 500 (SERVING.md section 3). Everything that
|
| 11 |
+
touches the device happens in the lifespan: weights -> device (ETH dispatch, 12x10) -> graph -> trace capture of
|
| 12 |
+
every warm-up variant, so uvicorn's ``Application startup complete`` (the line ``tt-model serve`` waits for) means
|
| 13 |
+
warm; SIGTERM (``tt-model stop``, 120 s) closes the model under the lock. ``/predict`` calls the Python API, so it
|
| 14 |
+
returns exactly what ``model(...)`` returns. Importing this module has no side effects beyond importing fastapi and
|
| 15 |
+
pydantic (the image's ``verify:`` imports it without a device); the environment is read in the lifespan only.
|
| 16 |
+
|
| 17 |
+
Model-specific request fields (e.g. PointPainting ``rois``): subclass ``PredictRequest``, pass ``request_model=`` and
|
| 18 |
+
``decode_extra=`` (which adds the decoded field to the call kwargs) to :class:`ServerSpec`, and list the keyword in
|
| 19 |
+
``EXTRA_INPUTS`` of the model class. Host tests swap the model before starting the app:
|
| 20 |
+
``app.state.ttaw.model_factory = Stub``; ``app.state.ttaw.predict(request)`` is the route handler itself.
|
| 21 |
+
"""
|
| 22 |
+
from pathlib import Path
|
| 23 |
+
|
| 24 |
+
from .. import __version__
|
| 25 |
+
from ..api import DiffusionPlanner
|
| 26 |
+
from ..ttaw.server.app import PredictRequest, ServerSpec, create_app, parse_mesh_shape
|
| 27 |
+
|
| 28 |
+
__all__ = ["SPEC", "app", "PredictRequest", "parse_mesh_shape"]
|
| 29 |
+
|
| 30 |
+
SPEC = ServerSpec(
|
| 31 |
+
model_name=DiffusionPlanner.MODEL_NAME,
|
| 32 |
+
env_prefix=DiffusionPlanner.ENV_PREFIX,
|
| 33 |
+
model_cls=DiffusionPlanner,
|
| 34 |
+
task="ego trajectory planning with a diffusion model (DPM-Solver++ 10 steps), neighbour prediction and a "
|
| 35 |
+
"turn-indicator command",
|
| 36 |
+
default_weights=DiffusionPlanner.DEFAULT_REPO,
|
| 37 |
+
owner="changh95",
|
| 38 |
+
io="the Autoware planner tensors (ego and neighbour histories, lanes, route, polygons, line strings, goal, ego "
|
| 39 |
+
"shape, turn-indicator history) in, an 8 s ego trajectory, predicted paths of the valid neighbours and a "
|
| 40 |
+
"turn-indicator command out",
|
| 41 |
+
autoware={"package": "autoware_diffusion_planner",
|
| 42 |
+
"path": "planning/autoware_diffusion_planner",
|
| 43 |
+
"autoware_universe": "9ceaccf026c31ffc5319bc9eeb4bd7bede0af3fd"},
|
| 44 |
+
source={"repo": "https://huggingface.co/changh95/diffusion-planner-p150", "license": "Apache-2.0"},
|
| 45 |
+
calib_dir=Path(__file__).resolve().parents[1] / "calib",
|
| 46 |
+
version=__version__,
|
| 47 |
+
description="Diffusion Planner v5.0 (Autoware diffusion_planner) on one Tenstorrent Blackhole p150. "
|
| 48 |
+
"Not an OpenAI-compatible API.",
|
| 49 |
+
)
|
| 50 |
+
|
| 51 |
+
app = create_app(SPEC)
|
code/tt_diffusion_planner/server/client.py
ADDED
|
@@ -0,0 +1,22 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#!/usr/bin/env python3
|
| 2 |
+
# SPDX-License-Identifier: Apache-2.0
|
| 3 |
+
"""Build (and optionally send) a ``POST /predict`` request for diffusion-planner-p150. Standard library only.
|
| 4 |
+
|
| 5 |
+
# the planner tensors (.npz holding the 15 raw tensors of INPUT_SCHEMA), with an optional runtime param
|
| 6 |
+
python3 code/tt_diffusion_planner/server/client.py \
|
| 7 |
+
--inputs code/tt_diffusion_planner/samples/kashiwanoha_dense.npz --param stopping_threshold=0.3 --out req.json
|
| 8 |
+
curl -s localhost:20000/predict -H 'Content-Type: application/json' -d @req.json
|
| 9 |
+
|
| 10 |
+
# send it directly and print the response
|
| 11 |
+
python3 code/tt_diffusion_planner/server/client.py \
|
| 12 |
+
--inputs code/tt_diffusion_planner/samples/kashiwanoha_dense.npz --url http://127.0.0.1:20000
|
| 13 |
+
|
| 14 |
+
The implementation is the vendored ``ttaw/server/client.py`` (C08); this file runs it from the model repository
|
| 15 |
+
with any Python 3.9+, without numpy and without installing the package. In Python, use
|
| 16 |
+
``tt_diffusion_planner.ttaw.server.client`` (``build_request``, ``post``, ``wait_ready``, ...).
|
| 17 |
+
"""
|
| 18 |
+
import runpy
|
| 19 |
+
from pathlib import Path
|
| 20 |
+
|
| 21 |
+
if __name__ == "__main__":
|
| 22 |
+
runpy.run_path(str(Path(__file__).resolve().parents[1] / "ttaw" / "server" / "client.py"), run_name="__main__")
|
code/tt_diffusion_planner/server/smoke_test.py
ADDED
|
@@ -0,0 +1,162 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#!/usr/bin/env python3
|
| 2 |
+
# SPDX-License-Identifier: Apache-2.0
|
| 3 |
+
"""Smoke test of a running diffusion-planner-p150 server (standard library only; any Python 3.9+, no numpy).
|
| 4 |
+
|
| 5 |
+
python3 code/tt_diffusion_planner/server/smoke_test.py --url http://127.0.0.1:20000 --wait 1800
|
| 6 |
+
# a served container package, one serve profile (what code/scripts/container_smoke.sh runs):
|
| 7 |
+
python3 code/tt_diffusion_planner/server/smoke_test.py --url http://127.0.0.1:20000 --wait 600 \
|
| 8 |
+
--manifest <out>/diffusion-planner-p150/tt_kernel_manifest.json [--profile <name>] \
|
| 9 |
+
--out /tmp/diffusion-planner.json
|
| 10 |
+
|
| 11 |
+
Checks (all failures are collected; prints ONE line ``PASS ...`` / ``FAIL ...``, exit code 0 / 1):
|
| 12 |
+
* ``/health`` reports ``ok`` (waiting up to ``--wait`` seconds), ``/info`` and ``/v1/models`` answer;
|
| 13 |
+
* ``/info`` reports ETH dispatch and the 12x10 grid, the p150 target of every published number (PLAN.md 0.3
|
| 14 |
+
item 6, D14): WORKER dispatch (e.g. an ETH open that fell back) or any other grid FAILS;
|
| 15 |
+
* with ``--manifest``: ``/info`` runs what the staged package pins for the serve profile
|
| 16 |
+
(``DIFFUSION_PLANNER_DISPATCH``,
|
| 17 |
+
``DIFFUSION_PLANNER_NUM_CQS``, ``DIFFUSION_PLANNER_VARIANT``, the weights revision);
|
| 18 |
+
* ``POST /predict`` of the shipped sample (the planner tensors as ``inputs``) answers 200 with the documented fields
|
| 19 |
+
and passes this model's output gates (:func:`output_gates`: 80 finite trajectory rows of 7 columns, a valid
|
| 20 |
+
turn-indicator command, the predicted neighbour paths);
|
| 21 |
+
* the served output agrees with the stored CPU-reference output of the sample within :data:`REFERENCE_GATES`:
|
| 22 |
+
``--reference``, else the first of ``<sample stem>.<profile>.reference.json``,
|
| 23 |
+
``<stem>.<variant>.reference.json``, ``<stem>.reference.json`` next to the sample (skipped when none exists);
|
| 24 |
+
* a malformed request answers 400 (not 500).
|
| 25 |
+
|
| 26 |
+
The checks themselves are the vendored ``ttaw/server/client.py`` and ``ttaw/server/smoke.py`` (stdlib only, loaded
|
| 27 |
+
by path); this file holds the model-specific parts: the sample request, the output gates and their thresholds.
|
| 28 |
+
"""
|
| 29 |
+
from __future__ import annotations
|
| 30 |
+
|
| 31 |
+
import argparse
|
| 32 |
+
import json
|
| 33 |
+
import math
|
| 34 |
+
import sys
|
| 35 |
+
import time
|
| 36 |
+
from pathlib import Path
|
| 37 |
+
from typing import Any, Dict, List, Optional
|
| 38 |
+
|
| 39 |
+
HERE = Path(__file__).resolve().parent
|
| 40 |
+
sys.path.insert(0, str(HERE.parent / "ttaw" / "server"))
|
| 41 |
+
import client as ttaw_client # noqa: E402 (stdlib-only modules of the vendored ttaw)
|
| 42 |
+
import smoke as ttaw_smoke # noqa: E402
|
| 43 |
+
|
| 44 |
+
MODEL = "diffusion-planner-p150"
|
| 45 |
+
ENV_PREFIX = "DIFFUSION_PLANNER"
|
| 46 |
+
EXPECT_DISPATCH, EXPECT_GRID = "eth", "12x10" # the p150 target; deliberately not a command-line option
|
| 47 |
+
DEFAULT_INPUT = HERE.parent / "samples" / "kashiwanoha_dense.npz"
|
| 48 |
+
DEFAULT_EXPECT = "" # unused for a planner (no detections); kept for the shared command line
|
| 49 |
+
# Agreement with the stored CPU reference (names of ttaw_smoke.DEFAULT_GATES): the ego trajectory's average / final
|
| 50 |
+
# displacement (the predicted_agents array is reported, not gated). Keep them in line with tests/test_e2e_device.py
|
| 51 |
+
# (ego mean error <= 0.3 m; the max error <= 1.0 m gate needs the whole array and lives in the device test).
|
| 52 |
+
REFERENCE_GATES: Dict[str, Optional[float]] = {"max_ade": 0.3, "max_fde": 1.0}
|
| 53 |
+
REQUIRED_KEYS = ("model", "frame_id", "timing_ms", "num_poses", "columns", "trajectory", "turn_indicator",
|
| 54 |
+
"predicted_agents")
|
| 55 |
+
TRAJECTORY_COLUMNS = ["x", "y", "yaw", "cos", "sin", "velocity", "acceleration"]
|
| 56 |
+
|
| 57 |
+
|
| 58 |
+
def build_sample_request(path: Path) -> Dict[str, Any]:
|
| 59 |
+
"""The ``/predict`` body for the sample: the planner tensors ``.npz`` as ``inputs``."""
|
| 60 |
+
if not Path(path).is_file():
|
| 61 |
+
raise FileNotFoundError(str(path))
|
| 62 |
+
return ttaw_client.build_request(inputs=str(path))
|
| 63 |
+
|
| 64 |
+
|
| 65 |
+
def output_gates(body: Dict[str, Any], expect: str) -> List[str]:
|
| 66 |
+
"""This model's plausibility gates on a 200 response: 80 trajectory rows of the documented 7 finite columns,
|
| 67 |
+
a turn-indicator command in 0..3 with 5 logits, and an encoded predicted-agents array of 80 x 5 per agent."""
|
| 68 |
+
fails: List[str] = []
|
| 69 |
+
traj = body.get("trajectory") or []
|
| 70 |
+
if body.get("columns") != TRAJECTORY_COLUMNS:
|
| 71 |
+
fails.append(f"columns {body.get('columns')} != {TRAJECTORY_COLUMNS}")
|
| 72 |
+
if len(traj) != 80 or body.get("num_poses") != 80:
|
| 73 |
+
fails.append(f"{len(traj)} trajectory rows, expected 80")
|
| 74 |
+
if not all(len(row) == 7 and all(math.isfinite(v) for v in row) for row in traj):
|
| 75 |
+
fails.append("non-finite or malformed trajectory row")
|
| 76 |
+
turn = body.get("turn_indicator") or {}
|
| 77 |
+
if turn.get("command") not in (0, 1, 2, 3) or len(turn.get("logits") or []) != 5:
|
| 78 |
+
fails.append(f"turn_indicator {turn}")
|
| 79 |
+
agents = body.get("predicted_agents") or {}
|
| 80 |
+
shape = agents.get("shape") or []
|
| 81 |
+
if len(shape) != 3 or shape[1:] != [80, 5]:
|
| 82 |
+
fails.append(f"predicted_agents shape {shape}")
|
| 83 |
+
return fails
|
| 84 |
+
|
| 85 |
+
|
| 86 |
+
def main(argv: Optional[List[str]] = None) -> int:
|
| 87 |
+
ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
|
| 88 |
+
ap.add_argument("--url", default="http://127.0.0.1:20000")
|
| 89 |
+
ap.add_argument("--input", type=Path, default=DEFAULT_INPUT, help="sample to POST (default: the shipped one)")
|
| 90 |
+
ap.add_argument("--expect", default=DEFAULT_EXPECT, help="unused by this planner (kept for the shared CLI)")
|
| 91 |
+
ap.add_argument("--reference", type=Path, help="stored CPU-reference /predict body of --input "
|
| 92 |
+
"(default: looked up next to it)")
|
| 93 |
+
ap.add_argument("--manifest", type=Path, help="staged tt_kernel_manifest.json: /info must run what it pins")
|
| 94 |
+
ap.add_argument("--profile", help="serve profile being smoke-tested (default: the package's default)")
|
| 95 |
+
ap.add_argument("--wait", type=float, default=0.0, help="seconds to wait for /health == ok")
|
| 96 |
+
ap.add_argument("--out", type=Path, help="write the /predict response here")
|
| 97 |
+
a = ap.parse_args(argv)
|
| 98 |
+
base = a.url.rstrip("/")
|
| 99 |
+
|
| 100 |
+
health = ttaw_client.wait_ready(base, wait_s=a.wait)
|
| 101 |
+
if health.get("status") != "ok":
|
| 102 |
+
print(ttaw_client.smoke_line(MODEL, f"/health status={health.get('status')!r}",
|
| 103 |
+
[f"not ready after {a.wait:.0f} s (error: {health.get('error')})"]))
|
| 104 |
+
return 1
|
| 105 |
+
info, fails = ttaw_client.check_service(base, expect_dispatch=EXPECT_DISPATCH, expect_grid=EXPECT_GRID)
|
| 106 |
+
if not info:
|
| 107 |
+
print(ttaw_client.smoke_line(MODEL, "/info unreachable", fails))
|
| 108 |
+
return 1
|
| 109 |
+
profile = a.profile
|
| 110 |
+
if a.manifest:
|
| 111 |
+
try:
|
| 112 |
+
pinned = ttaw_smoke.pinned_config(a.manifest, a.profile)
|
| 113 |
+
except (OSError, ValueError) as e:
|
| 114 |
+
fails.append(f"manifest: {e}")
|
| 115 |
+
else:
|
| 116 |
+
profile = pinned["profile"]
|
| 117 |
+
fails += ttaw_smoke.check_pinned(info, pinned, ENV_PREFIX)
|
| 118 |
+
|
| 119 |
+
code, body = 0, {}
|
| 120 |
+
t0 = time.perf_counter()
|
| 121 |
+
try:
|
| 122 |
+
code, body = ttaw_client.post(base, build_sample_request(a.input))
|
| 123 |
+
except (OSError, ValueError) as e: # unreadable sample, unknown suffix, server gone
|
| 124 |
+
fails.append(f"/predict: {type(e).__name__}: {e}")
|
| 125 |
+
rtt_ms = (time.perf_counter() - t0) * 1e3
|
| 126 |
+
if a.out:
|
| 127 |
+
a.out.write_text(json.dumps(body, indent=1))
|
| 128 |
+
ref_summary = "none"
|
| 129 |
+
if code and code != 200:
|
| 130 |
+
fails.append(f"/predict HTTP {code}: {str(body)[:300]}")
|
| 131 |
+
elif code == 200:
|
| 132 |
+
fails += [f"missing key {k!r}" for k in REQUIRED_KEYS if k not in body]
|
| 133 |
+
fails += output_gates(body, a.expect)
|
| 134 |
+
ref_path = a.reference or ttaw_smoke.find_reference(a.input, profile, info.get("variant"))
|
| 135 |
+
if ref_path:
|
| 136 |
+
try:
|
| 137 |
+
metrics, more = ttaw_smoke.compare_with_reference(body, ttaw_smoke.load_json(ref_path),
|
| 138 |
+
gates=REFERENCE_GATES)
|
| 139 |
+
except (OSError, ValueError) as e:
|
| 140 |
+
metrics, more = {}, [f"reference {ref_path}: {e}"]
|
| 141 |
+
fails += more
|
| 142 |
+
ref_summary = f"{ref_path.name} ({ttaw_smoke.describe_metrics(metrics)})"
|
| 143 |
+
try:
|
| 144 |
+
bad = ttaw_client.check_bad_request(base)
|
| 145 |
+
except OSError as e:
|
| 146 |
+
bad = f"malformed-request check: {type(e).__name__}: {e}"
|
| 147 |
+
if bad:
|
| 148 |
+
fails.append(bad)
|
| 149 |
+
|
| 150 |
+
device, timing = info.get("device") or {}, body.get("timing_ms") or {}
|
| 151 |
+
turn = (body.get("turn_indicator") or {}).get("command_name")
|
| 152 |
+
summary = (f"profile={profile or '-'} variant={info.get('variant')} dispatch={device.get('dispatch')} "
|
| 153 |
+
f"grid={device.get('grid')} cqs={device.get('num_command_queues')} n={body.get('num_poses')} "
|
| 154 |
+
f"turn={turn} "
|
| 155 |
+
f"reference={ref_summary} device_ms={timing.get('device')} total_ms={timing.get('total')} "
|
| 156 |
+
f"rtt_ms={rtt_ms:.1f}")
|
| 157 |
+
print(ttaw_client.smoke_line(MODEL, summary, fails))
|
| 158 |
+
return 1 if fails else 0
|
| 159 |
+
|
| 160 |
+
|
| 161 |
+
if __name__ == "__main__":
|
| 162 |
+
sys.exit(main())
|
code/tt_diffusion_planner/tests/__init__.py
ADDED
|
@@ -0,0 +1 @@
|
|
|
|
|
|
|
| 1 |
+
# SPDX-License-Identifier: Apache-2.0
|
code/tt_diffusion_planner/tests/_research.py
ADDED
|
@@ -0,0 +1,57 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-License-Identifier: Apache-2.0
|
| 2 |
+
"""Access to the verified research scripts of the porting workspace (``research/diffusion-planner/scripts``), the
|
| 3 |
+
independent implementations the host code is unit-tested against. Absent in an installed package or the image:
|
| 4 |
+
tests that need them skip."""
|
| 5 |
+
from __future__ import annotations
|
| 6 |
+
|
| 7 |
+
import importlib.util
|
| 8 |
+
import sys
|
| 9 |
+
from pathlib import Path
|
| 10 |
+
from typing import Optional
|
| 11 |
+
|
| 12 |
+
import numpy as np
|
| 13 |
+
|
| 14 |
+
PKG = Path(__file__).resolve().parents[1]
|
| 15 |
+
RESEARCH = PKG.parents[3] / "research" / "diffusion-planner"
|
| 16 |
+
SCRIPTS = RESEARCH / "scripts"
|
| 17 |
+
ORT_GOLDENS = RESEARCH / "ort"
|
| 18 |
+
FULL_GOLDENS = RESEARCH / "goldens"
|
| 19 |
+
SAMPLES = PKG / "samples"
|
| 20 |
+
SMALL_GOLDENS = PKG / "tests" / "goldens"
|
| 21 |
+
|
| 22 |
+
|
| 23 |
+
def load_script(name: str):
|
| 24 |
+
"""Import ``research/diffusion-planner/scripts/<name>.py`` by path (None when absent). Its directory is put on
|
| 25 |
+
``sys.path`` only while it imports (``dp_scene`` imports ``dp_common`` by name)."""
|
| 26 |
+
path = SCRIPTS / f"{name}.py"
|
| 27 |
+
if not path.is_file():
|
| 28 |
+
return None
|
| 29 |
+
key = f"_dp_research_{name}"
|
| 30 |
+
if key in sys.modules:
|
| 31 |
+
return sys.modules[key]
|
| 32 |
+
spec = importlib.util.spec_from_file_location(key, path)
|
| 33 |
+
mod = importlib.util.module_from_spec(spec)
|
| 34 |
+
sys.path.insert(0, str(SCRIPTS))
|
| 35 |
+
try:
|
| 36 |
+
spec.loader.exec_module(mod)
|
| 37 |
+
finally:
|
| 38 |
+
sys.path.remove(str(SCRIPTS))
|
| 39 |
+
sys.modules[key] = mod
|
| 40 |
+
return mod
|
| 41 |
+
|
| 42 |
+
|
| 43 |
+
def sample_raw(stem: str) -> dict:
|
| 44 |
+
with np.load(SAMPLES / f"{stem}.npz") as z:
|
| 45 |
+
return {k: np.array(z[k]) for k in z.files}
|
| 46 |
+
|
| 47 |
+
|
| 48 |
+
def research_scene_raw(scene: str) -> Optional[dict]:
|
| 49 |
+
path = ORT_GOLDENS / f"golden_{scene}.npz"
|
| 50 |
+
if not path.is_file():
|
| 51 |
+
return None
|
| 52 |
+
with np.load(path) as z:
|
| 53 |
+
return {k[len("raw/"):]: np.array(z[k]) for k in z.files if k.startswith("raw/")}
|
| 54 |
+
|
| 55 |
+
|
| 56 |
+
def research_scenes() -> list:
|
| 57 |
+
return sorted(p.stem[len("golden_"):] for p in ORT_GOLDENS.glob("golden_*.npz"))
|