Repository navigation
onnx: KV cache, AMD-protocol bench mode, and our own ORT build - #683
Conversation
The exported graph re-ran the whole sequence per step, so `--cached` now traces the cache as an input/output: one stacked `[2*layers,1,kv_heads,past,head_dim]` tensor instead of 56 named ones (dynamo cannot place 56 axes against its argument tree), and no attention_mask/position_ids, which the graph derives from the cache length -- with past=0 the traced mask path expands a zero-length axis and ONNX Runtime rejects the shape. `ort_run --cached` drives it, and `--bench` emits AMD's `model_benchmark` block: prompt processing, token generation and token sampling as avg/p50/stddev, E2E and peak RSS, over their prompt-length sweep and their -g/-w/-r defaults, so a row from this box and a row from theirs mean the same thing. build_ort.sh pins v1.23.2 and builds it from source (one `#include <cstdint>` for GCC 15, oneDNN's CMake-4 floor, libdnnl into the prefix or the DNNL EP silently never registers). BENCHMARK-PARITY.md carries the protocol, the flag mapping and the measured rows; bench_npu_serve.py measures a server on the same three stages, because the NPU row comes from our engine and not from someone else's.
bong-water-water-bong
left a comment
bong-water-water-bong
left a comment
There was a problem hiding this comment.
VERDICT: CHANGES (posted as a comment: GitHub refuses --request-changes on a PR owned by the same account; no tprs-ready). One blocker, and it is CI's first gate: tools/copyright.py --check fails on this head. Everything else I was asked to check passes.
Blocker — three files without the notice, so build cannot go green.
$ python3 tools/copyright.py --check
missing copyright notice: onnx/BENCHMARK-PARITY.md
missing copyright notice: onnx/bench_npu_serve.py
missing copyright notice: onnx/build_ort.sh
ci.yml's build job runs exactly this as its first step, and has_notice() (tools/copyright.py:89) wants all three of the copyright line, SPDX-License-Identifier: Apache-2.0, and "Licensed under the Apache License, Version 2.0" inside the first 20 lines. build_ort.sh has the first two but not the third (its header is the two-line form followed by prose); BENCHMARK-PARITY.md needs the <!-- … --> block that onnx/README.md opens with; bench_npu_serve.py needs the same as the other Python files. python3 tools/copyright.py --fix will do all three.
1. The bench reproduces AMD's protocol as BENCHMARK-PARITY.md documents it — checked both sides. The doc quotes their command line (model_benchmark.exe -i <model> -p <model>\prompts.txt -l "2048,1024,512,256,128"), maps every flag with their defaults to ours (-l/-g/-w/-r → --lens/--gen/--warmup/--repeat, same 2048,1024,512,256,128 / 128 / 1 / 5), copies their stage definitions, and states plainly that -t (peak-CPU sampling) is not implemented and that peak working set is VmHWM. The code matches the doc: --bench prints their block in their names and units —
Prompt processing (time to first token): avg (us) … avg (tokens/s) … p50 (us) … stddev (us)
Token generation: avg (us) … avg (tokens/s) … p50 … stddev
Token sampling: avg (us) … avg (tokens/s) … p50 … stddev
E2E generation (entire generation loop):
Peak working set size (bytes): <VmHWM from /proc/self/status>
and it refuses to bench an uncached graph, because "token generation" there would be the prefill wearing another name.
2. No vendor binary is linked or required. build_ort.sh clones https://github.com/microsoft/onnxruntime at v1.23.2 and builds it with --use_xnnpack --use_dnnl (the same three builtin EPs the distro package serves), with one guarded fix — grep -q before, then it asserts the #include <cstdint> is in core/common/semver.h afterwards. There is no providers_ryzenai, no VitisAI wheel, no .dll or prebuilt blob in the diff; those strings appear only in prose explaining why AMD's tool cannot be run here.
3. NOTICE is right. It now covers both paths — the distro libonnxruntime1.23 or the one we build from the published tag with build_ort.sh, naming the single #include <cstdint> source change and that the sources are the published tag (MIT covers it) — and it keeps the standing statement that AMD's VitisAI EP wheel is never linked or redistributed.
4. The numbers carry box, device, dtype and engine — the measured tables are literally those columns (ryzen / CPU 8 threads / fp32 / our ORT 1.23.2 + ort_run --cached; strixhalo / CPU 32 threads / fp32 / same graph, same runner), prefill on its own row set, and the doc explicitly separates definition-parity from speed-parity (theirs is int4 NPU Windows, ours fp32 CPU Linux). The one number that could mislead is labelled: "not comparable to AMD's published one as a speed claim".
Nothing else looked wrong: the cache is one stacked past/present [56,1,8,n,128] tensor carried as a single OrtValue, --bench requires --cached, and the doc names the cost it creates (224 KB/token copy in the graph, 459 MB at n=2048) rather than hiding it.
Fix the three headers and this is READY on the next head — I re-check only the copyright step and the diff, since everything above is already verified.
Decode and prefill are unchanged within noise (16.7->16.7 tok/s at 128, 6.7->6.6 at 2048), so the per-step allocation was not the O(n) cost — the graph is. Kept because the cache is now exactly two stacks instead of whatever the arena grew to; the doc says plainly that it is not a speed fix. Adds the NPU row from our own dx route (decode 72.6 tok/s at 162 tokens, 47.7 at 1335; prefill 63-78, the weak stage), bench_npu.sh to make the round trip reproducible, and --mode json|stream because the route streams every token in reasoning_content after the generation completes.
…ivate state file (#685) The paragraph sent public readers to ~/briefs/TORTEP-STATE.md, a lane file that only exists on our boxes. The numbers are in onnx/BENCHMARK-PARITY.md (#683) and docs/onnx.md (#647); link those. Co-authored-by: bong-water-water-bong <bong-water-water-bong@1bit.gg> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
From the owner's pi session on ryzen (commit 72e67de, 8 files, +774/-49): the ORT runner gets a KV cache (the decode step no longer re-runs the whole sequence), an AMD-protocol benchmark mode that reproduces model_benchmark's stages and units (onnx/BENCHMARK-PARITY.md), onnx/build_ort.sh for our own ORT build, onnx/bench_npu_serve.py, NOTICE and docs. Opened by the orchestrator for review; the squash merge replaces the local merge commits the session made.
🤖 Generated with Claude Code