anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

Package this fork as vllm-embedding: README, docs, roadmap, harness

Browse Source 
Turns a branch of accumulated edge/CPU work into something someone else can
read, reproduce and continue.

README.md states what the fork is for (multimodal models on edge CPUs rather
than servers), the headline memory and speed results with their costs, and why
the numbers are worth reading -- five retracted conclusions, all measurement
errors rather than coding errors. Upstream's README is preserved at
docs/README.vllm-omni.md and NOTICE records the Apache-2.0 derivation and the
files this fork changes.

docs/edge/:
  methodology.md  the measurement discipline, which is the most transferable
                  thing here: what each of the five retractions was, and the
                  rule it left behind. Includes the two traps specific to this
                  stack -- vLLM's AOT compile cache keying on config so two
                  arms sharing an artifact load the wrong graph, and allocated
                  not being resident until mincore(2) says so.
  results.md      every measured figure with its spread and baseline, and an
                  explicit note that the four memory levers overlap and must
                  not be added.
  cpu-4bit.md     the two 4-bit paths, the packed layout as measured rather
                  than assumed, why the path choice is a config field, and the
                  GPTQ output-head trap that invalidated a whole round of W4A8
                  numbers.
  harness.md      how to run a comparison, and the two arguments that matter
                  more than they look (--context-len, --prompt-len), each of
                  which has produced a wrong conclusion by being left at its
                  default.

ROADMAP.md orders by whether the verification method is correct rather than by
expected size, lists what is done with results, what is next with expected
gains and how each would be verified, and -- at equal length -- what was
considered and rejected, including the CPU megakernel with the measurements
that rule it out for now.

benchmarks/edge_harness/ brings the tooling into the repo so the numbers are
reproducible: the grid runner, the single-run benchmarks, the memory ledger,
the fidelity metric, the Android agent suite, the quantizers, and the tests
that pin the harness's judgement. 82 tests pass from the new location.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01L1ax5wYLA6XbqEL1i8L9gZ
This commit is contained in:
tzhouamandClaude Opus 5 committed 2026-09-14 13:25:32 +08:00
1 parent 289c13df68
commit b36a3431bd
32 files changed
+5641 -86
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

NOTICE
+23
View File
@@ -0,0 +1,23 @@
vllm-embedding
Copyright 2026 tzhouam
This product is a derivative work of vLLM-Omni
(https://github.com/vllm-project/vllm-omni), licensed under the Apache License,
Version 2.0. Upstream code is Copyright the vLLM-Omni contributors, and vLLM-Omni
is itself a derivative of vLLM (https://github.com/vllm-project/vllm).
Modifications in this fork are confined to edge/CPU inference and are listed in
README.md; the substantive ones are:
vllm_omni/model_executor/layers/quantization/cpu_int4.py
4-bit weight-only quantization for the CPU backend
vllm_omni/model_executor/models/spark2_5/spark2_5.py
RoPE table sized to max_model_len rather than the advertised context
vllm_omni/edge/
device-class weight-path selection, KV budget sizing, memory release
vllm_omni/platforms/cpu/, vllm_omni/patch.py
CPU platform and CPU-backend patches
benchmarks/edge_harness/
measurement tooling (original to this fork)
Upstream's README is preserved at docs/README.vllm-omni.md.
README.md
+64 -86
View File
@@ -1,109 +1,87 @@
<p align="center">
<picture>
<source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/logos/vllm-omni-logo.png">
<img alt="vllm-omni" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/logos/vllm-omni-logo.png" width=55%>
</picture>
</p>
<h3 align="center">
Easy, fast, and cheap omni-modality model serving for everyone
</h3>
# vllm-embedding
<p align="center">
| <a href="https://vllm-omni.readthedocs.io/en/latest/"><b>Documentation</b></a> | <a href="https://deepwiki.com/vllm-project/vllm-omni"><b>DeepWiki</b></a> | <a href="https://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> | <a href="docs/assets/WeChat.jpg"><b>WeChat</b></a> | <a href="https://arxiv.org/abs/2602.02204"><b>Paper</b></a> | <a href="https://docs.google.com/presentation/d/1aPj0OGl_-ZVoib-Qne5dGDAlrRFB-PdHl6E-EE99g8E/edit?usp=sharing"><b>Slides</b></a> |
</p>
A fork of [vLLM-Omni](https://github.com/vllm-project/vllm-omni) for running
multimodal models on **edge CPUs** — machines with tens of gigabytes of RAM and
no GPU, rather than servers.
---
The upstream project targets datacentre inference. Getting the same models onto
a mini-PC or a gateway box turns out to be less about kernels than about three
things nobody profiles: how much memory is resident while the model *loads*,
how much is resident that nothing ever reads, and whether the number you just
measured came from the code path you thought it did.
*Latest News* 🔥
Upstream's own README is preserved at [docs/README.vllm-omni.md](docs/README.vllm-omni.md).
- [2026/08] We released [0.28.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.28.0), featuring production-ready [MiniMax H3](recipes/MiniMaxAI/MiniMax-H3.md) serving on GPU and NPU, a unified AR/DiT paged KV cache runtime, and enhanced realtime full-duplex serving for the [MiniCPM-o series](recipes/OpenBMB/MiniCPM-o-4_5.md).
- [2026/08] [VeRL-Omni](https://github.com/verl-project/verl-omni) `v0.2.0` is released: faster diffusion RL powered by vLLM-Omni (request-level/step-wise batching with FA3), rebuilt Qwen3-Omni multimodal training (DPO & GSPO), plus LTX-2.3, Qwen-Image-Edit support and more. See the [release notes](https://github.com/verl-project/verl-omni/releases/tag/v0.2.0).
- [2026/08] We released [0.26.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.26.0) - aligned with the vLLM 0.26 release line, featuring [MiniMax H3](recipes/MiniMaxAI/MiniMax-H3.md) joint video/audio generation, an experimental full-duplex realtime runtime for [MiniCPM-o 4.5](recipes/OpenBMB/MiniCPM-o-4_5.md), distributed layerwise diffusion offload, and broader model, hardware, streaming, TTS, and quantization support.
- [2026/07] We released [0.24.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.24.0) - aligned with the vLLM 0.24 release line, expanding production-ready coverage across TTS, speech, diffusion, image/video generation, and robot-policy serving, with major Omni stage runtime refactoring, diffusion request-level batching, async output materialization, quantization/cache/memory improvements, and broad CUDA/ROCm/XPU/NPU support.
- [2026/06] Starting with [0.14.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.14.0), vLLM-Omni publishes a stable release aligned with every even-numbered upstream vLLM minor version. [0.16.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.16.0), [0.18.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.18.0), [0.20.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.20.0), and [0.22.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.22.0) continued this cadence, expanding omni and world-model support with [NVIDIA Cosmos3](recipes/cosmos3/Cosmos3-Nano.md) and DreamZero, adding models such as MiniCPM-o 4.5, MOSS-TTS, and Lance, and advancing TTS, diffusion, distributed execution, quantization, RL integration through [VeRL-Omni](https://github.com/verl-project/verl-omni), and CUDA/ROCm/MUSA/NPU/XPU coverage.
- [2026/03] Check out our first public [project deepdive](https://youtu.be/sgwNfsNnR9I) at the vLLM Hong Kong Meetup!
- [2025/11] vLLM community officially released [vllm-project/vllm-omni](https://github.com/vllm-project/vllm-omni) in order to support omni-modality models serving.
## What this fork adds
---
| area | change |
|---|---|
| 4-bit CPU quantization | `cpu_int4` — a W4A16 path built on PyTorch's tinygemm kernel, with the prefill path reading the kernel's own weight layout instead of a duplicate |
| Memory | RoPE tables sized to the servable context; weights repacked in blocks; KV budget that is actually honoured on CPU |
| Device routing | `vllm_omni/edge/weight_path.py` picks the weight format per device class, from measurements rather than assumption |
| Measurement | `benchmarks/edge_harness/` — the benchmark, profiling, memory-attribution and fidelity tooling every number here was produced with |
## About
## Headline results
[vLLM](https://github.com/vllm-project/vllm) was originally designed to support large language models for text-based autoregressive generation tasks. vLLM-Omni is a framework that extends its support for omni-modality model inference and serving:
Spark-X2.5-1.7B, Xeon Platinum 8480C, 24 pinned node-local cores, W4A16 g128.
Every figure is a median of at least three interleaved passes against a control
arm, on a shared host where end-to-end throughput swings ~30% run to run.
- **Omni-modality**: Text, image, audio, video, and action data processing
- **Non-autoregressive Architectures**: extend the AR support of vLLM to Diffusion Transformers (DiT) and other parallel generation models
- **Heterogeneous outputs**: from traditional text generation to multimodal and action outputs
**Memory** — each measured as its own A/B with the alternative reachable by a flag:
<p align="center">
<picture>
<img alt="vllm-omni" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/omni-modality-model-architecture.png" width=55%>
</picture>
</p>
| lever | saving | cost |
|---|---|---|
| KV budget 1 GiB → 192 MiB | **843 MB** | fewer cached tokens |
| RoPE table capped to `max_model_len` | **618 MB** | none |
| One packed weight layout instead of two | **615 MB** | ~1.6× prefill |
| Blocked weight repack | **489 MB** of load-time peak | none |
vLLM-Omni is fast with:
These overlap and **must not be added** — see [docs/edge/results.md](docs/edge/results.md).
- State-of-the-art AR support by leveraging efficient KV cache management from vLLM
- Pipelined stage execution overlapping for high throughput performance
- Fully disaggregation based on OmniConnector and dynamic resource allocation across stages
**Speed**, same machine:
vLLM-Omni is flexible and easy to use with:
| path | decode | prefill (pp512) |
|---|---|---|
| W4A8 (AMX `int4_scaled_mm_cpu`) | 74.4 tok/s | — |
| W4A16 (tinygemm) | 71.6 tok/s | 2190 tok/s |
| `cpu_gemm_wna16` fallback | 32.3 tok/s | — |
| llama.cpp Q4_K_M (reference) | 90.0 tok/s | 1121 tok/s |
- Heterogeneous pipeline abstraction to manage complex model workflows
- Seamless integration with popular Hugging Face models
- Tensor, pipeline, data and expert parallelism support for distributed inference
- Streaming outputs
- OpenAI-compatible API server
- [Full-duplex realtime serving](docs/serving/realtime_duplex_api.md) with streaming audio input and output
## Why the numbers here are worth reading
vLLM-Omni seamlessly supports most popular open-source models on HuggingFace, including:
This project has retracted five conclusions, every one of them a measurement
error rather than a coding error. The harness exists because of them:
- **Omni-modality models** (e.g. Qwen3-Omni, MiniCPM-o 4.5, Cosmos3, HunyuanImage, BAGEL)
- **TTS models** (e.g. Qwen3-TTS, IndexTTS 2.5, dots.tts, CosyVoice3)
- **Diffusion models** — image, video, and audio generation (e.g. MiniMax H3, LTX-2.5, SANA-Video, Wan2.2)
- **Robot-policy and action models** (e.g. π0, GR00T-N1.7, DreamZero-DROID, InternVLA-A1)
- a "3.1 GB saving" that was `VmRSS` counting the mmap'd checkpoint;
- W4A8 timings taken on a checkpoint whose output head was silently left in bf16;
- a memory benchmark reading `/proc/self/status` while the model ran in a *worker process*;
- "KV sizing saves only 188 MB" — measured through a code path that discarded the setting, so both arms ran the same cache;
- a 1.9× prefill cost quoted from the one pass of a grid whose other passes had crashed.
## Getting Started
[docs/edge/methodology.md](docs/edge/methodology.md) is the resulting discipline,
and it is the most transferable thing in this repository.
Visit our [documentation](https://vllm-omni.readthedocs.io/en/latest/) to learn more.
## Quick start
- [Installation](https://vllm-omni.readthedocs.io/en/latest/getting_started/installation/)
- [Quickstart](https://vllm-omni.readthedocs.io/en/latest/getting_started/quickstart/)
- [Nightly CUDA Docker images](https://hub.docker.com/r/vllm/vllm-omni/tags?name=nightly) built from the latest `main`
- [List of Supported Models](https://vllm-omni.readthedocs.io/en/latest/models/supported_models/)
- [Deployment Recipes](https://recipes.vllm.ai) for vLLM-Omni model serving
```bash
# CPU build of vLLM, then this fork
pip install vllm==0.28.0+cpu
VLLM_TARGET_DEVICE=cpu pip install -e .
## Contributing
We welcome and value any contributions and collaborations.
Please check out [Contributing to vLLM-Omni](https://vllm-omni.readthedocs.io/en/latest/contributing/) for how to get involved.
## Citation
If you use vLLM-Omni for your research, please cite our [paper](https://arxiv.org/abs/2602.02204):
```bibtex
@article{yin2026vllmomni,
title={vLLM-Omni: Fully Disaggregated Serving for Any-to-Any Multimodal Models},
author={Peiqi Yin, Jiangyun Zhu, Han Gao, Chenguang Zheng, Yongxiang Huang, Taichang Zhou, Ruirui Yang, Weizhi Liu, Weiqing Chen, Canlin Guo, Didan Deng, Zifeng Mo, Cong Wang, James Cheng, Roger Wang, Hongsheng Liu},
journal={arXiv preprint arXiv:2602.02204},
year={2026}
}
# a right-sized edge deployment
VLLM_CPU_KVCACHE_SPACE=1 python -c "
from vllm import LLM, SamplingParams
llm = LLM(model='<your-cpu_int4-checkpoint>', quantization='cpu_int4',
dtype='bfloat16', max_model_len=2048, max_num_seqs=1,
kv_cache_memory_bytes=192 << 20) # the single largest memory lever
print(llm.generate(['Hello'], SamplingParams(max_tokens=32))[0].outputs[0].text)
"
```
## Join the Community
See [docs/edge/cpu-4bit.md](docs/edge/cpu-4bit.md) for building a checkpoint and
choosing between the 4-bit paths, and [ROADMAP.md](ROADMAP.md) for what is
measured, what is next, and what was tried and abandoned.
Feel free to ask questions, provide feedbacks and discuss with fellow users of vLLM-Omni in `#sig-omni` slack channel at [slack.vllm.ai](https://slack.vllm.ai) or vLLM user forum at [discuss.vllm.ai](https://discuss.vllm.ai).
## Licence
## Star History
<a href="https://www.star-history.com/?repos=vllm-project%2Fvllm-omni&type=date&legend=top-left">
<picture>
<source media="(prefers-color-scheme: dark)" srcset="https://api.star-history.com/chart?repos=vllm-project/vllm-omni&type=date&theme=dark&legend=top-left&sealed_token=ExgLDZJoQEg27Zfhhut2LqN0GYO6Fw2PWLwPE6JYBUp2BgM3hmsYlwaIVopnUEfbRXidQ4nisumrTdKYydiKhy1SZXipw47qY2_tiUDhCpsPXeXtPuEVKVzBwKs3pw0tiHsJgtSfwXx5yjHXck0Y2SblzFWeJYCkTe1WLGTbUAOIETjXXQJjyCGZvKz5" />
<source media="(prefers-color-scheme: light)" srcset="https://api.star-history.com/chart?repos=vllm-project/vllm-omni&type=date&legend=top-left&sealed_token=ExgLDZJoQEg27Zfhhut2LqN0GYO6Fw2PWLwPE6JYBUp2BgM3hmsYlwaIVopnUEfbRXidQ4nisumrTdKYydiKhy1SZXipw47qY2_tiUDhCpsPXeXtPuEVKVzBwKs3pw0tiHsJgtSfwXx5yjHXck0Y2SblzFWeJYCkTe1WLGTbUAOIETjXXQJjyCGZvKz5" />
<img alt="Star History Chart" src="https://api.star-history.com/chart?repos=vllm-project/vllm-omni&type=date&legend=top-left&sealed_token=ExgLDZJoQEg27Zfhhut2LqN0GYO6Fw2PWLwPE6JYBUp2BgM3hmsYlwaIVopnUEfbRXidQ4nisumrTdKYydiKhy1SZXipw47qY2_tiUDhCpsPXeXtPuEVKVzBwKs3pw0tiHsJgtSfwXx5yjHXck0Y2SblzFWeJYCkTe1WLGTbUAOIETjXXQJjyCGZvKz5" />
</picture>
</a>
## License
Apache License 2.0, as found in the [LICENSE](./LICENSE) file.
Apache-2.0, inherited from vLLM-Omni. See [LICENSE](LICENSE) and [NOTICE](NOTICE).
ROADMAP.md
+107
View File
@@ -0,0 +1,107 @@
# Roadmap
Ordered by whether the verification method is *correct*, not by expected size.
That ordering is deliberate: this project retracted five conclusions, and in
every case the fix was to the measurement rather than the code.
Each item states the measurement it starts from, what it should gain, and how
that gain would be verified. An item without a verification method is not on
this list.
## Done
| | item | result |
|---|---|---|
| H1 | Task oracles judged by device state, with variants | alarm oracle needed a *saved* alarm, not a typed `7`; a Wi-Fi round trip now fails the already-satisfied task |
| H2 | Harness semantics | arm order really rotates; `best` respects metric direction; provenance stamped; failures recorded |
| H3 | The three missing measurements | `ready_s`, peak `RssAnon`, memory at a filled context — all previously unmeasurable |
| H4 | Contention guard | aggregate stolen capacity + straggler rule, on the target cores rather than host load average |
| H5 | Regression gates | bands from observed spread; throughput explicitly ungateable at ~30% swing |
| H6 | Fidelity metric that ranks | top-1 agreement and KL, replacing a saturating 12-prompt greedy check |
| H8 | Retracted numbers purged from source | two modules were still shipping withdrawn figures as documentation |
| I0 | Blocked weight repack | **489 MB** of load-time peak, no cost |
| I1 | One packed weight layout | **615 MB** resident, at ~1.63× prefill |
| I2 | Measure `cpu_gemm_wna16` | 2.2× slower than tinygemm; routing was right, now for a measured reason |
| I7 | Memory ledger, by attribution | found a **618 MB** RoPE table and an **843 MB** KV cache, both inside an "unexplained 1.6 GB" |
## Next
**I6 — import cost.** `import vllm_omni` is ~30 s against `import vllm`'s 5.2 s,
with transformers' auto image-processing (4.6 s) and the `openai` client (1.3 s)
identifiable inside it, for a text model that needs neither. Every stage
subprocess pays it again.
*Expect:* an upper bound of ~6 s per process; deferred imports often move cost
rather than remove it. *Verify:* `ready_s`, which now exists — the old timer
started after imports, so the benchmark meant to prove this could not see it.
**Drop the row-major embedding copy.** `embed_tokens.weight_packed` is 128 MB,
kept solely because the tinygemm layout was believed un-addressable for a
single-row lookup. That premise is now known false. *Verify:* `RssAnon` plus an
embedding-output SNR assertion.
**I10 first step — why is the GEMM at 60%?** The decode GEMM sustains 132 GB/s
against a measured 222 GB/s ceiling. Part is arithmetic (tinygemm dequantizes
per activation row) and part may be that each of 113 calls streams only a few MB,
too little for the prefetcher to reach steady state. Attribute before optimising.
*Expect:* unknown — a diagnosis, whose output is whether the 2.76 ms implied by
the gap is reachable at all.
**I3 — asymmetric W4A8.** The symmetric grid is our exporter's choice, not the
kernel's limit; `models/spark2_5-gptq-asym-g128` already exists. The cheapest
untried fidelity experiment, and it tests the *stated* cause of W4A8's error.
*Verify:* `fidelity_kl.py` and the scored agent suite.
**I5 — attention wrapper consolidation.** Four ops per layer, 1.31 ms/step, of
which 0.83 ms is wrapper self time; at batch 1 the cache-update loop runs two
iterations while the wrapper rebuilds context and allocates a dummy tensor for
compile ordering. *Expect:* 0.3–0.5 ms, revised down from an earlier 0.6–0.9 ms
that assumed a workload never profiled. *Verify:* per-op profile **and**
unprofiled step latency — one outer call that still launches two threaded
kernels removes neither thread launch.
**I4 — a better quantizer.** Ours reproduces 2/12 greedy prompts against its own
bf16; Q4_K_M manages 4/12 — but at 4.5 bits against our 4.25, so close the
budget gap first (g64, or fp16 scale+zero, both config changes) before
implementing error-compensated rounding. Timebox it; a null result is worth
recording.
## Considered and rejected
**A CPU megakernel.** On GPU this fuses a forward pass into one launch to remove
launch overhead and the cross-kernel prefetch bubble. On CPU the analogues are
OpenMP fork/join and ATen dispatch — and llama.cpp already is one:
`ggml_threadpool` starts threads once and they walk the whole graph,
synchronising at `ggml_barrier` through an atomic spin, where vLLM enters a
separate `omp parallel` region per kernel.
Measured before deciding: a parallel region costs ~1.2 µs at 24 threads, and
spin-waiting does not help (`OMP_WAIT_POLICY=ACTIVE` moved it 1.23 → 1.18 µs), so
it is barrier cost rather than thread wake-up. At a few hundred regions per
decode step that is 0.3–0.4 ms, 2–3% — the same money as I5, far more cheaply.
The larger prize is the streaming bubble, but capturing it means replacing the
execution model inside vLLM's CPU backend, which reimplements ggml. Not next.
**More kernel fusion.** Four attempts, four losses: vLLM's fused CPU RMSNorm is
3.6–4.5× faster in isolation and slower in situ (inductor already fuses the
native form with the residual add); fused `gelu_and_mul` 2.7× slower at Spark's
shape; `cpp.min_chunk_size` −6% decode; an inline gate reduction −44%.
**MiniCPM-o CPU overlay.** The 8 B thinker at bf16 is ~19 GB, so it buys no edge
deployability. The argument for it — that three-stage parallel init should beat
the two-stage 2.1× — is also not how overlap works: the bound is ΣT/max(T).
**Speculative decoding for tg128.** It would inflate a fixed-prompt `ignore_eos`
benchmark without representing anything. Measure it on the agent suite or not at
all.
**Thread count and warm-up batch size as memory levers.** 64 MB and 2 MB.
## Harness work still open
- **A real device in the agent loop.** The toolkit has 38 tests and has never
touched a phone — this host has no `adb`, no emulator, no `/dev/kvm`. The
real-device path and per-task resets exist; what is untested is whether a real
Settings app matches the fixtures.
- **Success rates, not single runs.** The build comparison that produced
"navigation survives quantization, completion does not" is one run per cell.
Variants now exist so repeats are trials; they need running.
benchmarks/edge_harness/android_agent/README.md
+92
View File
@@ -0,0 +1,92 @@
# Android agent toolkit for Spark-X2.5
Lets a text-only model operate a phone: observe the accessibility tree, call
one tool, execute it on the device, observe again.
ui_tree.py uiautomator XML -> indexed element list -> back to a tap point
device.py AdbDevice (a real phone) | SimulatedDevice (scripted, in-process)
executor.py one tool call -> one device action, errors as sentences
tool_parse.py recover a call from whatever shape the model emitted
agent.py the loop: budget, no-op detection, undo detection, fallbacks
run_episode.py CLI
## Running it
Serve the model (Spark emits GLM-4.5 tool format, so the stock parsers apply):
vllm serve <model> --port 8077 --enable-auto-tool-choice \
--tool-call-parser glm45 --reasoning-parser glm45
Then:
python run_episode.py --task "Turn on Wi-Fi." --device sim # no phone needed
python run_episode.py --task "Turn on Wi-Fi." --device adb # a real phone
`pytest test_android_agent.py` runs 21 tests with no phone and no server.
## What the loop does that a plain ReAct loop does not
Each one is a response to a measured failure of *this* model, not generic
scaffolding:
- **Retry without thinking on an empty reply.** The model can reason past its
token budget without closing `<think>`; vLLM's reasoning parser only emits
content once the block closes, so the reply arrives empty with no tool call.
Raising the cap buys a longer silence. A non-thinking retry costs ~0.6 s.
- **Report actions that changed nothing.** "the screen did not change" is
information the model cannot otherwise have.
- **Report and stop on self-undo.** Tapping a switch that is already on undoes
the task. Each tap really does change the screen, so a no-change check
cannot see it; the loop watches for the *same* action returning the screen to
its previous state. Deliberately narrow -- opening a submenu and pressing
back revisits a screen and must not trip it.
- **Recover malformed tool calls.** See below.
- **A step budget**, so a confused episode ends instead of running forever.
The loop never decides a task is finished. It reports what happened and lets
the model call `task_complete`; a harness that completes tasks on the model's
behalf measures itself.
## Measured: which build can actually finish a task
Simulated phone, one run per cell, `max_steps` 10, Xeon 8480C 24 threads.
| build | decode | "Turn on Wi-Fi." | "Set an alarm for 7:30 AM." |
|---|---|---|---|
| bf16 | 41.0 tok/s | done in 4 steps | wrong app (Settings, not Clock), toggles Airplane mode, stuck |
| **W4A16 g128** | 80.0 tok/s | done in 4 steps | right app, full navigation (Clock -> Alarm -> Add -> Save), then no tool call |
| W4A8 g128 | 86.6 tok/s | turns Wi-Fi on, then toggles it off again (2 runs, same failure) | taps "Phone" twice, reaches Alarm, then loops |
Two things fall out, and both matter more than the throughput numbers:
**Navigation survives quantization; completion does not.** Every build walks
home -> Settings -> Network -> Wi-Fi correctly. Only some of them notice that
the switch now reads `checked=true` and stop. W4A8 -- the fastest build, and
the one the speed-first selector recommends -- turns Wi-Fi on and then taps it
off again, ending with the task undone.
**The fidelity number had a task-level meaning all along.** W4A8 reproduced
1/12 greedy prompts against its own bf16 and W4A16 2/12
(analysis/spark_edge_support.md 2c). That looked like a small difference on a
proxy metric. It is the difference between finishing and not.
So for agentic use the recommendation is **W4A16**, not W4A8: 8% slower, and
it is the only 4-bit build here that completed anything. `weight_path.py`'s
`priority="fidelity"` selects it.
## Honest limits
- **No phone was involved.** This host has no `adb`, no emulator and no
`/dev/kvm`. The simulator emits real uiautomator XML, so the parser, the
index-to-coordinate mapping and the stopping rules are exercised on the path
hardware drives -- but whether a real Settings app is laid out the way these
fixtures assume is untested.
- **Two tasks, one run each.** Enough to show a failure mode reproduces, not
enough to rank builds by success rate.
- **Stock `glm45` does not always parse Spark.** On 4 of 5 screens it does. On
the Wi-Fi toggle screen the model emitted
`[tap]\n{"tool": "tap", "parameters": {"index": 1}` -- unterminated, and
`tool_calls` came back empty. `tool_parse.py` recovers that and reports every
recovery, because an agent that silently repairs its model's output hides
regressions in that model. The earlier claim that the stock parsers work
"as-is" is true most of the time, which is not the same thing.
benchmarks/edge_harness/android_agent/action_space.py
+152
View File
@@ -0,0 +1,152 @@
"""Android control action space for a text-only LLM agent.
The tool surface mirrors what the mainstream ADB-based harnesses expose
(droidrun/mobilerun's Portal tools, ghost-in-the-droid's touch/see tools,
AndroidWorld's `JSONAction`), so an agent written against these schemas can be
pointed at any of them. Every action is expressed over the accessibility tree
rather than pixel coordinates, because Spark-X2.5 is a text-only model: it
never sees a screenshot, it sees the uiautomator dump.
`ADB_TEMPLATES` gives the literal `adb` invocation each tool maps to, which is
what a driver executes against a real device.
"""
TOOLS = [
{
"type": "function",
"function": {
"name": "tap",
"description": (
"Tap a UI element. Use the element's index exactly as shown in "
"the accessibility tree."
),
"parameters": {
"type": "object",
"properties": {
"index": {
"type": "integer",
"description": "Index of the element in the tree",
}
},
"required": ["index"],
},
},
},
{
"type": "function",
"function": {
"name": "type_text",
"description": (
"Type text into a focused input field. Tap the field first."
),
"parameters": {
"type": "object",
"properties": {
"index": {"type": "integer", "description": "Index of the input field"},
"text": {"type": "string", "description": "Text to enter"},
},
"required": ["index", "text"],
},
},
},
{
"type": "function",
"function": {
"name": "swipe",
"description": (
"Scroll the screen. The direction is where the content should "
"move: use 'down' to see what is further down the page."
),
"parameters": {
"type": "object",
"properties": {
"direction": {
"type": "string",
"enum": ["up", "down", "left", "right"],
}
},
"required": ["direction"],
},
},
},
{
"type": "function",
"function": {
"name": "press_key",
"description": "Press a hardware or system key.",
"parameters": {
"type": "object",
"properties": {
"key": {
"type": "string",
"enum": ["back", "home", "enter", "delete", "recent"],
}
},
"required": ["key"],
},
},
},
{
"type": "function",
"function": {
"name": "open_app",
"description": "Launch an app by its visible name.",
"parameters": {
"type": "object",
"properties": {"name": {"type": "string"}},
"required": ["name"],
},
},
},
{
"type": "function",
"function": {
"name": "task_complete",
"description": (
"Call when the user's request has been fully carried out."
),
"parameters": {
"type": "object",
"properties": {
"summary": {"type": "string", "description": "What was done"}
},
"required": ["summary"],
},
},
},
]
# How each tool becomes a real device command. `{bounds}` is filled from the
# element's centre point in the accessibility tree.
ADB_TEMPLATES = {
"tap": "adb shell input tap {x} {y}",
"type_text": "adb shell input text {text_escaped}",
"swipe": "adb shell input swipe {x1} {y1} {x2} {y2} 300",
"press_key": "adb shell input keyevent {keycode}",
"open_app": "adb shell monkey -p {package} -c android.intent.category.LAUNCHER 1",
}
KEYCODES = {
"back": "KEYCODE_BACK",
"home": "KEYCODE_HOME",
"enter": "KEYCODE_ENTER",
"delete": "KEYCODE_DEL",
"recent": "KEYCODE_APP_SWITCH",
}
SYSTEM_PROMPT = (
"You control an Android phone. At each step you are given the current "
"screen as an accessibility tree: one line per element, with its index, "
"type, text and whether it is clickable. Call exactly one tool per step to "
"make progress on the user's task. Refer to elements only by the index "
"shown.\n"
# Both rules below exist because the model failed them in measurement: it
# tapped a Wi-Fi switch that already read checked=true, turning it back
# off, and never called task_complete. A switch's state is on the screen,
# so the rule can be stated in terms of what it can see.
"A switch or checkbox shows its state as checked=true or checked=false. "
"Tapping it flips that state.\n"
"Before acting, check whether the screen already shows what the task "
"asked for. If it does, call task_complete immediately and do not tap "
"anything else -- acting again would undo it."
)
benchmarks/edge_harness/android_agent/agent.py
+271
View File
@@ -0,0 +1,271 @@
"""The observe-decide-act loop.
Three things in here are not generic agent scaffolding; they are direct
responses to how this model failed when it was measured one step at a time
(analysis/spark_edge_support.md section 5):
* **A step budget, not a bigger token cap.** On the one genuinely ambiguous
screen the model reasoned past 1200 tokens without closing `<think>`, and
vLLM's reasoning parser only emits content once the block closes -- so the
reply came back empty, with no tool call at all. Raising the cap buys a
longer silence. The loop instead retries that step with thinking disabled,
which measured 0.61 s against 8.9 s, and takes the weaker-but-present answer.
* **Detecting the no-op.** The clearest wrong action was tapping a Wi-Fi
switch that was already on -- undoing the task it had just completed. An
action that leaves the screen signature unchanged is reported back as such,
because "nothing happened" is information the model cannot otherwise have.
* **Repeat suppression.** Two identical actions on an identical screen means
stuck, and a third will not help. The loop says so in the transcript and
stops, rather than burning the budget.
"""
from __future__ import annotations
import json
import time
from dataclasses import dataclass, field
from urllib import error as urlerror
from urllib import request as urlrequest
from action_space import SYSTEM_PROMPT, TOOLS
from device import Device
from executor import ActionResult, execute
from tool_parse import extract_tool_call
from ui_tree import Element, parse_ui_dump, render_tree, screen_signature
@dataclass
class AgentConfig:
base_url: str = "http://127.0.0.1:8077/v1"
model: str = "XHToken/Spark-X2.5-1.7B"
max_steps: int = 12
think: bool = True
think_fallback: bool = True
"""On an empty reply, retry the step without thinking rather than fail."""
max_tokens: int = 1024
temperature: float = 0.0
timeout: float = 900.0
history: int = 6
retries: int = 3
"""Transient HTTP failures retried before a step is abandoned."""
retry_backoff_s: float = 2.0
"""Recent actions shown back to the model."""
@dataclass
class Step:
n: int
screen: str
tool: str | None
args: dict
result: ActionResult | None
latency_s: float
thinking: bool
note: str = ""
parsed_via: str = "tool_calls"
"""Which parser recovered the call; anything but tool_calls is worth seeing."""
@dataclass
class Episode:
task: str
steps: list[Step] = field(default_factory=list)
done: bool = False
summary: str = ""
stop_reason: str = ""
@property
def n_steps(self) -> int:
return len(self.steps)
def transcript(self) -> str:
lines = [f"task: {self.task}"]
for s in self.steps:
call = f"{s.tool}({json.dumps(s.args, ensure_ascii=False)})" if s.tool else "(no tool call)"
outcome = s.result.message if s.result else s.note
if s.parsed_via not in ("tool_calls", "none"):
outcome += f" [recovered via {s.parsed_via}]"
lines.append(f" {s.n:2d}. {call:48s} -> {outcome} [{s.latency_s:.2f}s]")
lines.append(f"stopped: {self.stop_reason}")
return "\n".join(lines)
class ChatError(RuntimeError):
pass
def chat(cfg: AgentConfig, messages: list[dict], *, think: bool) -> tuple[dict, float]:
payload = {
"model": cfg.model,
"messages": messages,
"tools": TOOLS,
"tool_choice": "auto",
"temperature": cfg.temperature,
"max_tokens": cfg.max_tokens,
}
if not think:
# Spark's chat template opens <think> by default; closing it up front
# is what makes the model answer directly.
payload["chat_template_kwargs"] = {"enable_thinking": False}
req = urlrequest.Request(
f"{cfg.base_url}/chat/completions",
data=json.dumps(payload).encode(),
headers={"Content-Type": "application/json"},
)
t0 = time.perf_counter()
last: Exception | None = None
for attempt in range(max(1, cfg.retries)):
try:
with urlrequest.urlopen(req, timeout=cfg.timeout) as r:
return json.loads(r.read()), time.perf_counter() - t0
except (urlerror.URLError, TimeoutError, json.JSONDecodeError) as exc:
# A phone task is a dozen round trips; one dropped connection
# should cost a retry, not the episode.
last = exc
if attempt + 1 < max(1, cfg.retries):
time.sleep(cfg.retry_backoff_s * (attempt + 1))
raise ChatError(
f"{cfg.base_url} failed after {cfg.retries} attempts: {last}") from last
def wait_for_server(cfg: AgentConfig, timeout_s: float = 300.0) -> bool:
"""Block until the endpoint serves a model list, or give up.
Starting an episode against a server that is still loading produces a
confusing transcript full of connection errors; this turns that into one
clear wait.
"""
deadline = time.monotonic() + timeout_s
while time.monotonic() < deadline:
try:
with urlrequest.urlopen(f"{cfg.base_url}/models", timeout=5) as r:
if b'"id"' in r.read():
return True
except Exception:
time.sleep(2.0)
return False
def parse_response(response: dict):
"""The one call to act on, from whichever shape the model used.
Spark does not always stay in GLM-4.5 tool format -- see tool_parse -- and
on a phone a dropped step is a stalled task, so the fallbacks matter.
"""
try:
message = response["choices"][0]["message"]
except (KeyError, IndexError):
message = {}
return extract_tool_call(message)
def first_tool_call(response: dict) -> tuple[str | None, dict]:
"""Name and arguments only, for callers that do not care how it parsed."""
parsed = parse_response(response)
return parsed.name, parsed.args
def _user_message(task: str, screen: str, log: list[str], history: int) -> str:
parts = [f"Task: {task}", "", "Current screen:", screen]
if log:
parts += ["", "Actions so far:"] + [f" {line}" for line in log[-history:]]
parts += ["", "Choose the next action. Call exactly one tool."]
return "\n".join(parts)
def observe(device: Device) -> tuple[list[Element], str, str]:
elements = parse_ui_dump(device.dump())
return elements, render_tree(elements), screen_signature(elements)
def run_episode(device: Device, task: str, cfg: AgentConfig | None = None) -> Episode:
cfg = cfg or AgentConfig()
episode = Episode(task=task)
log: list[str] = []
recent: list[tuple[str, str]] = [] # (action, screen signature before it)
signature_before_last: str | None = None
last_action: str | None = None
oscillations = 0
for n in range(1, cfg.max_steps + 1):
elements, screen, signature = observe(device)
messages = [
{"role": "system", "content": SYSTEM_PROMPT},
{"role": "user",
"content": _user_message(task, screen, log, cfg.history)},
]
thinking = cfg.think
response, latency = chat(cfg, messages, think=thinking)
parsed = parse_response(response)
tool, args = parsed.name, parsed.args
if tool is None and cfg.think_fallback and thinking:
# The measured failure: reasoning never closed, so the parser
# emitted nothing. A second, non-thinking pass is worth ~0.6 s.
thinking = False
retry, retry_latency = chat(cfg, messages, think=False)
latency += retry_latency
parsed = parse_response(retry)
tool, args = parsed.name, parsed.args
if tool is None:
episode.steps.append(Step(n, screen, None, {}, None, latency, thinking,
note="model returned no tool call",
parsed_via=parsed.source))
episode.stop_reason = "no tool call, even without thinking"
return episode
if (tool_key := f"{tool}:{json.dumps(args, sort_keys=True)}") and \
recent.count((tool_key, signature)) >= 2:
episode.steps.append(Step(n, screen, tool, args, None, latency, thinking,
note="repeated a no-op action",
parsed_via=parsed.source))
episode.stop_reason = (
f"stuck: {tool} repeated on an unchanged screen"
)
return episode
recent.append((tool_key, signature))
result = execute(tool, args, elements, device)
if result.ok and not result.done:
_, _, after = observe(device)
if after == signature:
# Say so out loud: the model has no other way to learn that
# its action changed nothing.
result = ActionResult(True, result.message + " -- the screen did not change",
action=result.action)
elif tool_key == last_action and after == signature_before_last:
# Oscillation: the *same* action, undoing what it just did --
# the measured failure, tapping a switch off and on forever.
# It hides from the no-change check because each individual tap
# really does change the screen. Deliberately narrow: revisiting
# a screen by other means (open a submenu, press back) is normal
# navigation and must not trip this.
result = ActionResult(
True,
result.message + " -- that undid your previous action; the "
"setting is now back where it started",
action=result.action,
)
oscillations += 1
if oscillations >= 2:
episode.steps.append(Step(n, screen, tool, args, result, latency,
thinking, parsed_via=parsed.source))
episode.stop_reason = "stuck: repeatedly undoing its own action"
return episode
signature_before_last, last_action = signature, tool_key
episode.steps.append(Step(n, screen, tool, args, result, latency, thinking,
parsed_via=parsed.source))
log.append(f"{tool} -> {result.message}")
if result.done:
episode.done = True
episode.summary = result.summary
episode.stop_reason = "task_complete"
return episode
episode.stop_reason = f"step budget exhausted ({cfg.max_steps})"
return episode
benchmarks/edge_harness/android_agent/device.py
+354
View File
@@ -0,0 +1,354 @@
"""Where the agent's actions actually land.
Two backends behind one interface. `AdbDevice` talks to a phone over `adb`.
`SimulatedDevice` runs a small scripted phone in-process, and exists for a
reason beyond convenience: no Android device, emulator or `/dev/kvm` is
reachable from the host this was built on, so without it the loop, the tool
executor and the stopping rules would all be untested code shipped on a
promise.
The simulator emits **real uiautomator XML**, not pre-rendered text, so the
parser, the index-to-coordinate mapping and the screen-change detection are
exercised on the same path a phone would drive. What it cannot test is
whether a real Settings app is laid out the way the fixtures assume; that is
stated in the report rather than hidden.
"""
from __future__ import annotations
import shutil
import subprocess
from xml.sax.saxutils import escape as _xml_escape
import time
from dataclasses import dataclass, field
from typing import Callable, Protocol
KEYCODES = {
"back": "KEYCODE_BACK",
"home": "KEYCODE_HOME",
"enter": "KEYCODE_ENTER",
"delete": "KEYCODE_DEL",
"recent": "KEYCODE_APP_SWITCH",
}
class Device(Protocol):
def dump(self) -> str: ...
def tap(self, x: int, y: int) -> None: ...
def input_text(self, text: str) -> None: ...
def swipe(self, x1: int, y1: int, x2: int, y2: int, ms: int = 300) -> None: ...
def keyevent(self, key: str) -> None: ...
def launch(self, app: str) -> None: ...
def size(self) -> tuple[int, int]: ...
class AdbError(RuntimeError):
pass
class AdbDevice:
"""A real phone over `adb`."""
def __init__(self, serial: str | None = None, adb: str = "adb",
timeout: float = 30.0):
if shutil.which(adb) is None:
raise AdbError(
f"{adb!r} is not on PATH. Install platform-tools and connect a "
"device (`adb devices` should list it)."
)
self.adb, self.serial, self.timeout = adb, serial, timeout
self._size: tuple[int, int] | None = None
def _run(self, *args: str, binary: bool = False):
cmd = [self.adb] + (["-s", self.serial] if self.serial else []) + list(args)
proc = subprocess.run(cmd, capture_output=True, timeout=self.timeout)
if proc.returncode != 0:
raise AdbError(f"{' '.join(cmd)} failed: {proc.stderr.decode(errors='replace')[:400]}")
return proc.stdout if binary else proc.stdout.decode(errors="replace")
def dump(self) -> str:
# exec-out keeps the XML on stdout instead of round-tripping through
# /sdcard, which is both faster and avoids a writable-storage
# assumption that does not hold on every device.
out = self._run("exec-out", "uiautomator", "dump", "/dev/tty")
start = out.find("<?xml")
if start < 0:
start = out.find("<hierarchy")
if start < 0:
raise AdbError(f"no UI hierarchy in dump output: {out[:200]!r}")
end = out.rfind("</hierarchy>")
return out[start:end + len("</hierarchy>")] if end > 0 else out[start:]
def tap(self, x: int, y: int) -> None:
self._run("shell", "input", "tap", str(x), str(y))
def input_text(self, text: str) -> None:
# `input text` takes a single token: spaces become %s and the shell
# metacharacters have to go, or half the string is eaten.
escaped = text.replace(" ", "%s")
for ch in "()<>|;&*\\~\"'`$":
escaped = escaped.replace(ch, "\\" + ch)
self._run("shell", "input", "text", escaped)
def swipe(self, x1: int, y1: int, x2: int, y2: int, ms: int = 300) -> None:
self._run("shell", "input", "swipe", str(x1), str(y1), str(x2), str(y2), str(ms))
def keyevent(self, key: str) -> None:
code = KEYCODES.get(key, key)
self._run("shell", "input", "keyevent", code)
def launch(self, app: str) -> None:
"""Launch by package name, or by visible label via the launcher."""
if "." in app:
self._run("shell", "monkey", "-p", app, "-c",
"android.intent.category.LAUNCHER", "1")
return
raise AdbError(
f"open_app needs a package name, got {app!r}. Resolve the label to a "
"package first (`adb shell pm list packages`) or tap its icon."
)
def size(self) -> tuple[int, int]:
if self._size is None:
out = self._run("shell", "wm", "size")
part = out.strip().rsplit(":", 1)[-1].strip()
w, h = part.split("x")
self._size = (int(w), int(h))
return self._size
# --------------------------------------------------------------------------
# Simulated phone
# --------------------------------------------------------------------------
@dataclass
class SimElement:
label: str
cls: str = "android.widget.TextView"
clickable: bool = True
checkable: bool = False
checked: bool = False
selected: bool = False
editable: bool = False
goto: str | None = None
"""Screen this element navigates to when tapped."""
toggles: str | None = None
"""State key this element flips when tapped."""
saves: str | None = None
"""Record key. Tapping commits the current field values into
``records[saves]`` and clears them -- the difference between typing a time
and having an alarm, which the scored suite previously could not tell."""
label_of: str | None = None
"""State key whose value picks the label from ``labels``."""
labels: tuple[str, str] = ("", "")
"""(label when the state key is false, label when true)."""
def rendered_label(self, state: dict) -> str:
if self.label_of:
return self.labels[bool(state.get(self.label_of))]
return self.label
@dataclass
class SimScreen:
screen_id: str
elements: list[SimElement]
back: str | None = None
dynamic: "Callable[[SimulatedDevice], list[SimElement]] | None" = None
"""Rows derived from device state, appended after ``elements``. A saved
alarm has to appear *on the screen*, or an oracle that reads the screen --
which is the only kind a real phone can have -- cannot tell a committed
alarm from a typed one."""
@dataclass
class SimulatedDevice:
"""A scripted phone that speaks uiautomator XML.
Deliberately small. It is here to prove the loop terminates, recovers and
stops -- not to model Android.
"""
screens: dict[str, SimScreen]
current: str
state: dict[str, bool] = field(default_factory=dict)
width: int = 1080
height: int = 2400
typed: list[str] = field(default_factory=list)
actions: list[tuple[str, ...]] = field(default_factory=list)
home: str = "home"
focus: str | None = None
"""Label of the focused text field; typing lands here."""
fields: dict[str, str] = field(default_factory=dict)
records: dict[str, list[dict]] = field(default_factory=dict)
toggled: list[tuple[str, bool, bool]] = field(default_factory=list)
"""(key, before, after) per switch tap. A task whose starting state is
already correct is failed by a round trip, which final state cannot see."""
_ROW_H = 160
_TOP = 200
def _visible(self) -> list[SimElement]:
screen = self.screens[self.current]
rows = list(screen.elements)
if screen.dynamic:
rows = [r for r in rows if not getattr(r, "_placeholder", False)]
rows.extend(screen.dynamic(self))
return rows
def _bounds(self, row: int) -> tuple[int, int, int, int]:
y1 = self._TOP + row * self._ROW_H
return 0, y1, self.width, y1 + self._ROW_H
def dump(self) -> str:
rows = []
for i, el in enumerate(self._visible()):
x1, y1, x2, y2 = self._bounds(i)
checked = el.checked
if el.toggles:
checked = self.state.get(el.toggles, el.checked)
# uiautomator escapes attribute values, and labels really do
# contain them -- "Network & internet" is on every Android
# settings screen. Emitting raw breaks the parser on the second
# screen of the first task.
label = _xml_escape(el.rendered_label(self.state), {'"': "&quot;"})
rows.append(
f' <node index="{i}" text="{label}" resource-id="" '
f'class="{el.cls}" package="com.example.sim" content-desc="" '
f'checkable="{str(el.checkable).lower()}" '
f'checked="{str(checked).lower()}" '
f'clickable="{str(el.clickable).lower()}" enabled="true" '
f'focusable="{str(el.clickable).lower()}" focused="false" '
f'scrollable="false" long-clickable="false" password="false" '
f'selected="{str(el.selected).lower()}" '
f'bounds="[{x1},{y1}][{x2},{y2}]" />'
)
body = "\n".join(rows)
return (
'<?xml version="1.0" encoding="UTF-8"?>\n'
f'<hierarchy rotation="0">\n{body}\n</hierarchy>'
)
def _row_at(self, y: int) -> int | None:
row = (y - self._TOP) // self._ROW_H
return row if 0 <= row < len(self._visible()) else None
def tap(self, x: int, y: int) -> None:
self.actions.append(("tap", str(x), str(y)))
row = self._row_at(y)
if row is None:
return
el = self._visible()[row]
if el.editable:
self.focus = el.label
if el.toggles:
before = bool(self.state.get(el.toggles, el.checked))
self.state[el.toggles] = not before
self.toggled.append((el.toggles, before, not before))
if el.saves:
self.records.setdefault(el.saves, []).append(
{**self.fields, **{k: v for k, v in self.state.items()
if k.startswith(el.saves + "_")}})
self.fields = {}
self.focus = None
if el.goto and el.goto in self.screens:
self.current = el.goto
def input_text(self, text: str) -> None:
self.actions.append(("text", text))
self.typed.append(text)
if self.focus is not None:
self.fields[self.focus] = self.fields.get(self.focus, "") + text
def swipe(self, x1: int, y1: int, x2: int, y2: int, ms: int = 300) -> None:
self.actions.append(("swipe", str(x1), str(y1), str(x2), str(y2)))
def keyevent(self, key: str) -> None:
self.actions.append(("key", key))
if key == "back":
self.current = self.screens[self.current].back or self.home
elif key == "home":
self.current = self.home
def launch(self, app: str) -> None:
self.actions.append(("launch", app))
for screen in self.screens.values():
for el in screen.elements:
if el.label.lower() == app.lower() and el.goto:
self.current = el.goto
return
def size(self) -> tuple[int, int]:
return self.width, self.height
def _saved_alarm_rows(device: "SimulatedDevice") -> list[SimElement]:
"""Render committed alarms the way a clock app lists them."""
saved = device.records.get("alarm", [])
if not saved:
return [SimElement("No alarms", clickable=False)]
rows = []
for rec in saved:
hour = (rec.get("Hour") or "?").lstrip("0") or "0"
minute = (rec.get("Minute") or "?").rjust(2, "0")
meridiem = "PM" if rec.get("alarm_pm") else "AM"
rows.append(SimElement(f"{hour}:{minute} {meridiem}", clickable=False))
return rows
def demo_phone() -> SimulatedDevice:
"""A two-app phone: Settings (Wi-Fi) and Clock (alarms).
Enough to exercise a multi-screen task, a stateful toggle -- the case the
earlier single-step run got wrong, tapping a switch that was already on --
and a text field.
"""
screens = {
"home": SimScreen("home", [
SimElement("Phone"),
SimElement("Messages"),
SimElement("Clock", goto="clock"),
SimElement("Settings", goto="settings"),
SimElement("Camera"),
]),
"settings": SimScreen("settings", [
SimElement("Settings", clickable=False),
SimElement("Network & internet", goto="network"),
SimElement("Connected devices"),
SimElement("Apps"),
SimElement("Battery"),
], back="home"),
"network": SimScreen("network", [
SimElement("Network & internet", clickable=False),
SimElement("Wi-Fi", cls="android.widget.Switch", checkable=True,
toggles="wifi"),
SimElement("Mobile network"),
SimElement("Airplane mode", cls="android.widget.Switch",
checkable=True, toggles="airplane"),
], back="settings"),
"clock": SimScreen("clock", [
SimElement("Alarm", cls="android.widget.Button", goto="alarms"),
SimElement("Clock", cls="android.widget.Button", selected=True),
SimElement("Timer", cls="android.widget.Button"),
SimElement("Stopwatch", cls="android.widget.Button"),
], back="home"),
"alarms": SimScreen("alarms", [
SimElement("Alarm", cls="android.widget.Button", selected=True),
SimElement("Add alarm", cls="android.widget.ImageButton",
goto="new_alarm"),
], back="clock", dynamic=_saved_alarm_rows),
"new_alarm": SimScreen("new_alarm", [
SimElement("Set time", clickable=False),
SimElement("Hour", cls="android.widget.EditText", editable=True),
SimElement("Minute", cls="android.widget.EditText", editable=True),
# Android shows one AM/PM control whose label is its value, so
# "set an alarm for 7:30 AM" has a wrong answer as well as a right
# one. Without it the task could be passed by typing any digits.
SimElement("AM", cls="android.widget.Button",
label_of="alarm_pm", labels=("AM", "PM"),
toggles="alarm_pm"),
SimElement("Save", cls="android.widget.Button", goto="alarms",
saves="alarm"),
], back="alarms"),
}
return SimulatedDevice(screens=screens, current="home", state={"wifi": False})
benchmarks/edge_harness/android_agent/executor.py
+147
View File
@@ -0,0 +1,147 @@
"""Turn one tool call into one device action.
Every failure here is returned to the model as a sentence it can act on, not
raised. A 1.7B model asked for `index: 9` on a four-element screen has made a
recoverable mistake, and "index 9 is out of range, the screen has 4 elements
(0-3)" is the difference between the next step being right and the episode
ending. Exceptions are reserved for the device being gone.
"""
from __future__ import annotations
from dataclasses import dataclass
from device import KEYCODES, Device
from ui_tree import Element
@dataclass
class ActionResult:
ok: bool
message: str
done: bool = False
summary: str = ""
action: str = ""
def as_tool_reply(self) -> str:
return self.message
def _need(args: dict, key: str, kind: type, what: str) -> tuple[object, str | None]:
if key not in args:
return None, f"missing required argument {key!r} for {what}"
value = args[key]
if kind is int:
try:
value = int(value)
except (TypeError, ValueError):
return None, f"{key!r} must be a whole number, got {args[key]!r}"
elif not isinstance(value, kind):
return None, f"{key!r} must be {kind.__name__}, got {type(value).__name__}"
return value, None
def _resolve(index: object, elements: list[Element]) -> tuple[Element | None, str | None]:
i = int(index)
if not elements:
return None, "the screen has no addressable elements; try press_key(back)"
if not 0 <= i < len(elements):
return None, (
f"index {i} is out of range; the screen has {len(elements)} "
f"elements (0-{len(elements) - 1})"
)
return elements[i], None
def execute(name: str, args: dict, elements: list[Element],
device: Device) -> ActionResult:
"""Run one tool call. Returns what to tell the model."""
args = args or {}
if name == "task_complete":
summary = str(args.get("summary", "")).strip()
return ActionResult(True, "task marked complete", done=True,
summary=summary, action="task_complete")
if name == "tap":
index, err = _need(args, "index", int, "tap")
if err:
return ActionResult(False, err, action="tap")
el, err = _resolve(index, elements)
if err:
return ActionResult(False, err, action="tap")
x, y = el.center
device.tap(x, y)
state = ""
if "checked" in el.flags:
state = f" (was {'on' if el.flags['checked'] else 'off'})"
return ActionResult(True, f'tapped [{index}] "{el.text}"{state}',
action=f"tap:{index}")
if name == "type_text":
index, err = _need(args, "index", int, "type_text")
if err:
return ActionResult(False, err, action="type_text")
text, err = _need(args, "text", str, "type_text")
if err:
return ActionResult(False, err, action="type_text")
el, err = _resolve(index, elements)
if err:
return ActionResult(False, err, action="type_text")
# Focus first: `input text` goes to whatever holds focus, which on a
# freshly dumped screen is usually nothing.
x, y = el.center
device.tap(x, y)
device.input_text(str(text))
return ActionResult(True, f'typed "{text}" into [{index}] "{el.text}"',
action=f"type:{index}")
if name == "swipe":
direction = str(args.get("direction", "")).lower()
if direction not in ("up", "down", "left", "right"):
return ActionResult(
False, "direction must be one of up, down, left, right",
action="swipe")
w, h = device.size()
cx, cy = w // 2, h // 2
dx, dy = int(w * 0.35), int(h * 0.30)
# `direction` is where the *content* should move, i.e. what the user
# means by "scroll down to see more", so the finger travels the other
# way. Stated in the tool description too, because guessing this wrong
# is invisible in the transcript.
moves = {
"down": (cx, cy + dy, cx, cy - dy),
"up": (cx, cy - dy, cx, cy + dy),
"right": (cx + dx, cy, cx - dx, cy),
"left": (cx - dx, cy, cx + dx, cy),
}
device.swipe(*moves[direction])
return ActionResult(True, f"swiped to scroll {direction}",
action=f"swipe:{direction}")
if name == "press_key":
key = str(args.get("key", "")).lower()
if key not in KEYCODES:
return ActionResult(
False, f"unknown key {key!r}; use one of {', '.join(KEYCODES)}",
action="press_key")
device.keyevent(key)
return ActionResult(True, f"pressed {key}", action=f"key:{key}")
if name == "open_app":
app = str(args.get("name", "")).strip()
if not app:
return ActionResult(False, "open_app needs a name", action="open_app")
try:
device.launch(app)
except Exception as exc: # device-specific launch failures are recoverable
return ActionResult(False, f"could not launch {app!r}: {exc}",
action="open_app")
return ActionResult(True, f"launched {app}", action=f"open:{app}")
return ActionResult(
False,
f"unknown tool {name!r}; available: tap, type_text, swipe, press_key, "
"open_app, task_complete",
action="unknown",
)
benchmarks/edge_harness/android_agent/phone_state.py
+135
View File
@@ -0,0 +1,135 @@
"""What the phone is actually in, read the same way for both backends.
The scored suite used to take its predicates over ``SimulatedDevice``'s Python
attributes, so every scored run was a simulator run by construction -- the
real-device path existed for *driving* a phone and not for judging one. That
made "the agent suite" a claim about fixtures.
Success is therefore read two ways, and the order matters:
1. **Out of band, from the setting itself** -- ``device.state`` in the
simulator, ``settings get global wifi_on`` over adb. This is the ground
truth, and it does not depend on where the agent happened to leave the
phone. Scoring from the final screen alone fails an agent that turns Wi-Fi
on correctly and then presses Home, which is not a failure.
2. **From the UI dump**, for anything with no setting behind it -- a saved
alarm appears as a row in the clock app and nowhere else. The simulator
emits real uiautomator XML, so this goes down the same parse, the same
switch lookup and the same text matching a phone would.
Only the facts a phone genuinely cannot report after the fact -- which
switches were flipped on the way, and what was typed into which field -- come
from the simulator object, and any task needing them declares it, so a
real-device run refuses to score rather than scoring wrong.
"""
from __future__ import annotations
from dataclasses import dataclass, field
from ui_tree import parse_ui_dump
def _norm(s: str) -> str:
"""Match labels across the punctuation Android is inconsistent about."""
return "".join(c for c in s.lower() if c.isalnum())
@dataclass
class PhoneState:
source: str
"""``"sim"`` or ``"device"`` -- which facts below are trustworthy."""
settings: dict[str, bool] = field(default_factory=dict)
"""Authoritative setting values, read out of band. Preferred over
:attr:`switches`, which only shows what is currently on screen."""
screen: str | None = None
"""Simulator screen id. ``None`` on a real phone, where there is no such
thing; use :meth:`has_text` against something the screen actually shows."""
switches: dict[str, bool] = field(default_factory=dict)
texts: list[str] = field(default_factory=list)
typed: list[str] = field(default_factory=list)
fields: dict[str, str] = field(default_factory=dict)
records: dict[str, list[dict]] = field(default_factory=dict)
toggled: list[tuple[str, bool, bool]] = field(default_factory=list)
"""(key, before, after) for every switch the agent actually flipped."""
def switch(self, *names: str) -> bool | None:
"""State of the first control matching any of ``names``.
The setting wins over the screen: an agent that turned Wi-Fi on and
then navigated away has still turned Wi-Fi on. ``None`` means neither
source knows the control, which is not the same as off -- an agent that
never found the screen has not turned anything on, and predicates must
not read that as "off".
"""
wanted = {_norm(n) for n in names}
for key, value in self.settings.items():
if _norm(key) in wanted:
return bool(value)
for label, checked in self.switches.items():
if _norm(label) in wanted:
return checked
return None
def has_text(self, *needles: str) -> bool:
joined = " ".join(self.texts).lower()
return any(n.lower() in joined for n in needles)
def times_toggled(self, key: str) -> int:
return sum(1 for k, _, _ in self.toggled if k == key)
def read_state(device) -> PhoneState:
"""Read the phone through its screen, plus the simulator's own records."""
elements = parse_ui_dump(device.dump())
# `checkable` is deliberately not among the flags ui_tree parses, because
# that list is also what the model is shown; adding to it would change the
# agent's prompt. A switch is recognised by its class or by exposing a
# `checked` flag instead.
switches = {e.text: bool(e.flags.get("checked", False))
for e in elements
if "checked" in e.flags
or e.short_class in ("Switch", "CheckBox", "ToggleButton")}
texts = [e.text for e in elements if e.text]
if not hasattr(device, "screens"):
return PhoneState(source="device", switches=switches, texts=texts,
settings=adb_settings(device))
return PhoneState(
source="sim",
screen=device.current,
settings={k: bool(v) for k, v in device.state.items()},
switches=switches,
texts=texts,
typed=list(device.typed),
fields=dict(getattr(device, "fields", {})),
records={k: [dict(r) for r in v]
for k, v in getattr(device, "records", {}).items()},
toggled=list(getattr(device, "toggled", [])),
)
ADB_SETTINGS = {
"wifi": ("global", "wifi_on"),
"airplane": ("global", "airplane_mode_on"),
"bluetooth": ("global", "bluetooth_on"),
}
"""Settings readable out of band, so scoring does not depend on which screen
the agent left the phone on."""
def adb_settings(device) -> dict[str, bool]:
"""Read the authoritative values from a real phone. Missing keys are
omitted rather than defaulted -- an unreadable setting must not be scored
as "off"."""
out: dict[str, bool] = {}
for key, (namespace, name) in ADB_SETTINGS.items():
try:
raw = device._run("shell", "settings", "get", namespace, name)
except Exception:
continue
raw = raw.strip()
if raw in ("0", "1"):
out[key] = raw == "1"
return out
benchmarks/edge_harness/android_agent/run_agent.py
+98
View File
@@ -0,0 +1,98 @@
"""Drive Spark-X2.5 through Android control steps against an OpenAI endpoint.
Measures whether the model picks the right tool and the right element, and how
long a decision takes -- the two things that decide whether a 1.7B model can
actually run a phone agent locally.
"""
import argparse, json, time
from urllib import request as urlrequest
from action_space import SYSTEM_PROMPT, TOOLS
from scenarios import SCENARIOS
def chat(base_url, model, messages, tools, timeout=900, max_tokens=1024,
think=True):
payload = {
"model": model, "messages": messages, "tools": tools,
"tool_choice": "auto", "temperature": 0.0, "max_tokens": max_tokens,
}
if not think:
# Spark's template opens <think> by default; closing it up front makes
# the model answer directly, which is what a phone agent wants.
payload["chat_template_kwargs"] = {"enable_thinking": False}
body = json.dumps(payload).encode()
req = urlrequest.Request(f"{base_url}/chat/completions", data=body,
headers={"Content-Type": "application/json"})
t0 = time.perf_counter()
with urlrequest.urlopen(req, timeout=timeout) as r:
out = json.loads(r.read())
return out, time.perf_counter() - t0
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--base-url", default="http://127.0.0.1:8077/v1")
ap.add_argument("--model", default="XHToken/Spark-X2.5-1.7B")
ap.add_argument("--no-think", action="store_true")
ap.add_argument("--max-tokens", type=int, default=1024)
ap.add_argument("--out", default="android_agent_results.json")
a = ap.parse_args()
rows, correct_tool, correct_args = [], 0, 0
for sc in SCENARIOS:
msgs = [
{"role": "system", "content": SYSTEM_PROMPT},
{"role": "user", "content":
f"Task: {sc['task']}\n\nCurrent screen:\n{sc['screen']}\n\n"
"Choose the next action."},
]
try:
out, dt = chat(a.base_url, a.model, msgs, TOOLS,
max_tokens=a.max_tokens, think=not a.no_think)
except Exception as e:
rows.append({"id": sc["id"], "error": str(e)[:200]})
print(f"{sc['id']:<20} ERROR {str(e)[:90]}", flush=True)
continue
msg = out["choices"][0]["message"]
calls = msg.get("tool_calls") or []
got_name = calls[0]["function"]["name"] if calls else None
try:
got_args = json.loads(calls[0]["function"]["arguments"]) if calls else {}
except Exception:
got_args = {"_raw": calls[0]["function"]["arguments"]}
name_ok = got_name in sc["accept_names"]
args_ok = name_ok and all(
str(got_args.get(k)) == str(v) for k, v in sc["expect"]["args"].items()
) if got_name == sc["expect"]["name"] else name_ok and not sc["expect"]["args"]
correct_tool += int(name_ok)
correct_args += int(bool(args_ok))
rows.append({
"id": sc["id"], "latency_s": round(dt, 2),
"expected": sc["expect"], "got_name": got_name, "got_args": got_args,
"tool_ok": name_ok, "args_ok": bool(args_ok),
"reasoning_chars": len(msg.get("reasoning_content") or ""),
"completion_tokens": out.get("usage", {}).get("completion_tokens"),
})
flag = "OK " if args_ok else ("~ " if name_ok else "X ")
print(f"{flag}{sc['id']:<20} {got_name}({got_args}) {dt:.1f}s", flush=True)
n = len(SCENARIOS)
summary = {
"scenarios": n,
"thinking": not a.no_think,
"tool_accuracy": f"{correct_tool}/{n}",
"action_accuracy": f"{correct_args}/{n}",
"median_latency_s": round(
sorted(r["latency_s"] for r in rows if "latency_s" in r)[
max(0, len([r for r in rows if "latency_s" in r]) // 2)] , 2)
if any("latency_s" in r for r in rows) else None,
"rows": rows,
}
json.dump(summary, open(a.out, "w"), indent=2)
print(json.dumps({k: v for k, v in summary.items() if k != "rows"}, indent=2))
if __name__ == "__main__":
main()
benchmarks/edge_harness/android_agent/run_episode.py
+106
View File
@@ -0,0 +1,106 @@
"""Drive Spark-X2.5 through a whole phone task, not one decision.
# against the in-process simulated phone (no device needed)
python run_episode.py --task "Turn on Wi-Fi." --device sim
# against a real phone
adb devices # confirm it is listed
python run_episode.py --task "Turn on Wi-Fi." --device adb
The server side is vLLM's stock OpenAI endpoint; Spark emits GLM-4.5's tool
format, so no custom parser is needed:
vllm serve <model> --port 8077 --enable-auto-tool-choice \\
--tool-call-parser glm45 --reasoning-parser glm45
"""
from __future__ import annotations
import argparse
import json
import sys
import time
from agent import AgentConfig, run_episode
from device import AdbDevice, AdbError, demo_phone
TASKS = {
"wifi": "Turn on Wi-Fi.",
"alarm": "Set an alarm for 7:30 AM.",
}
def build_device(kind: str, serial: str | None):
if kind == "sim":
return demo_phone(), "simulated phone (settings + clock)"
try:
device = AdbDevice(serial=serial)
w, h = device.size()
return device, f"adb device {serial or '(default)'} {w}x{h}"
except AdbError as exc:
print(f"error: {exc}", file=sys.stderr)
raise SystemExit(2)
def main() -> None:
ap = argparse.ArgumentParser()
ap.add_argument("--task", default=TASKS["wifi"],
help=f"free text, or one of: {', '.join(TASKS)}")
ap.add_argument("--device", choices=("sim", "adb"), default="sim")
ap.add_argument("--serial", default=None)
ap.add_argument("--base-url", default="http://127.0.0.1:8077/v1")
ap.add_argument("--model", default="XHToken/Spark-X2.5-1.7B")
ap.add_argument("--max-steps", type=int, default=12)
ap.add_argument("--max-tokens", type=int, default=1024)
ap.add_argument("--no-think", action="store_true")
ap.add_argument("--no-think-fallback", action="store_true")
ap.add_argument("--out", default=None)
a = ap.parse_args()
task = TASKS.get(a.task, a.task)
device, described = build_device(a.device, a.serial)
cfg = AgentConfig(
base_url=a.base_url, model=a.model, max_steps=a.max_steps,
max_tokens=a.max_tokens, think=not a.no_think,
think_fallback=not a.no_think_fallback,
)
print(f"device: {described}")
print(f"task: {task}")
t0 = time.perf_counter()
episode = run_episode(device, task, cfg)
wall = time.perf_counter() - t0
print()
print(episode.transcript())
latencies = [s.latency_s for s in episode.steps]
print(f"\n{episode.n_steps} steps in {wall:.1f} s"
f" (median decision {sorted(latencies)[len(latencies) // 2]:.2f} s)"
if latencies else f"\n{episode.n_steps} steps in {wall:.1f} s")
if hasattr(device, "state"):
print(f"device state: {device.state}")
if a.out:
payload = {
"task": task, "device": described, "done": episode.done,
"summary": episode.summary, "stop_reason": episode.stop_reason,
"wall_s": round(wall, 2),
"think": cfg.think, "max_steps": cfg.max_steps,
"steps": [
{"n": s.n, "tool": s.tool, "args": s.args,
"ok": s.result.ok if s.result else False,
"message": s.result.message if s.result else s.note,
"latency_s": round(s.latency_s, 3), "thinking": s.thinking}
for s in episode.steps
],
"final_state": getattr(device, "state", None),
}
with open(a.out, "w") as f:
json.dump(payload, f, indent=2, ensure_ascii=False)
print(f"wrote {a.out}")
raise SystemExit(0 if episode.done else 1)
if __name__ == "__main__":
main()
benchmarks/edge_harness/android_agent/run_suite.py
+151
View File
@@ -0,0 +1,151 @@
"""Run the task suite and score it, so an agent change can be compared.
Success is decided by the phone's own settings and screen, not by the
transcript: an agent that says it turned Wi-Fi on and left it off fails here,
which is exactly the failure this model exhibits.
Two things changed after the plan review:
* ``--repeat N`` now walks the task's **variants** rather than re-running one
fixture. At temperature 0 against fixed fixtures, repeats measure
repeatability; they only become trials if the phone differs between them.
The summary reports a pass rate per task with the variants named, so a build
comparison rests on more than one cell.
* ``--device adb`` scores a real phone. Each run resets the device to the
task's declared starting state first, and a task that cannot be scored from
a phone (because its predicate needs flip history) is skipped and reported
as skipped, rather than silently scored from final state.
"""
from __future__ import annotations
import argparse
import json
import statistics
import time
from collections import defaultdict
from agent import AgentConfig, run_episode, wait_for_server
from phone_state import read_state
from tasks import SUITE, SUITE_BY_ID
def _reset_real_device(device, task) -> None:
if not task.reset_adb:
raise SystemExit(
f"task {task.task_id!r} declares no reset_adb, so a phone run would "
f"be scored against an unknown starting state. Add one or exclude "
f"the task with --tasks."
)
for cmd in task.reset_adb:
device._run(*cmd)
def main() -> None:
ap = argparse.ArgumentParser()
ap.add_argument("--base-url", default="http://127.0.0.1:8077/v1")
ap.add_argument("--model", default="spark")
ap.add_argument("--tasks", default="all",
help="comma-separated task ids, or 'all'")
ap.add_argument("--repeat", type=int, default=1,
help="runs per task; each walks the next variant")
ap.add_argument("--device", choices=("sim", "adb"), default="sim")
ap.add_argument("--serial", default=None, help="adb device serial")
ap.add_argument("--no-think", action="store_true")
ap.add_argument("--max-tokens", type=int, default=1024)
ap.add_argument("--label", default="")
ap.add_argument("--out", default=None)
a = ap.parse_args()
tasks = SUITE if a.tasks == "all" else [SUITE_BY_ID[t] for t in a.tasks.split(",")]
cfg = AgentConfig(base_url=a.base_url, model=a.model,
think=not a.no_think, max_tokens=a.max_tokens)
phone = None
if a.device == "adb":
from device import AdbDevice
phone = AdbDevice(serial=a.serial)
skipped = [t.task_id for t in tasks if t.needs]
if skipped:
print(f"skipping on a real phone (needs flip history): "
f"{', '.join(skipped)}")
tasks = [t for t in tasks if not t.needs]
if not wait_for_server(cfg, timeout_s=600):
raise SystemExit(f"no model server at {a.base_url}")
rows, t0 = [], time.perf_counter()
for task in tasks:
for run in range(a.repeat):
if phone is not None:
_reset_real_device(phone, task)
device = phone
else:
device = task.build(run)
cfg.max_steps = task.max_steps
episode = run_episode(device, task.prompt, cfg)
score = task.score(read_state(device), episode.done, run)
score.update(
run=run + 1, steps=episode.n_steps,
stop_reason=episode.stop_reason,
recovered_steps=sum(1 for s in episode.steps
if s.parsed_via not in ("tool_calls", "none")),
median_latency_s=round(
sorted(s.latency_s for s in episode.steps)[episode.n_steps // 2], 2)
if episode.n_steps else None,
)
rows.append(score)
mark = "PASS" if score["success"] else "fail"
detail = []
if not score.get("scorable", True):
mark, detail = "skip", [score.get("reason", "")]
else:
if score["reached_goal"] and not score["declared_done"]:
detail.append("goal reached but never stopped")
if score["declared_done"] and not score["reached_goal"]:
detail.append("claimed done, goal NOT reached")
if score["collateral_damage"]:
detail.append("collateral damage")
print(f" {mark} {task.task_id:18s} {score['variant']:10s} "
f"steps={score['steps']:2d} "
f"{score['stop_reason'][:34]:36s} {'; '.join(detail)}", flush=True)
wall = time.perf_counter() - t0
scored = [r for r in rows if r.get("scorable", True)]
n_pass = sum(1 for r in scored if r["success"])
reached = sum(1 for r in scored if r["reached_goal"])
claimed = sum(1 for r in scored if r["declared_done"] and not r["reached_goal"])
by_task: dict[str, list[dict]] = defaultdict(list)
for r in scored:
by_task[r["task_id"]].append(r)
print("\npass rate per task")
for task_id, rs in by_task.items():
passes = sum(1 for r in rs if r["success"])
failed = [r["variant"] for r in rs if not r["success"]]
print(f" {task_id:18s} {passes}/{len(rs)}"
+ (f" failed on: {', '.join(failed)}" if failed else ""))
rates = [sum(1 for r in rs if r["success"]) / len(rs) for rs in by_task.values()]
spread = (f"{min(rates):.0%}-{max(rates):.0%}" if rates else "-")
print(f"\n{n_pass}/{len(scored)} passed in {wall:.0f}s"
f" (goal reached {reached}/{len(scored)}; "
f"falsely claimed done {claimed}; per-task rate {spread}"
+ (f"; median {statistics.median(rates):.0%}" if rates else "") + ")")
if len(scored) and len(by_task) and max(len(rs) for rs in by_task.values()) < 3:
print(" NOTE: fewer than 3 runs per task -- enough to show a failure "
"reproduces, not to rank builds.")
if a.out:
with open(a.out, "w") as f:
json.dump({"label": a.label, "base_url": a.base_url,
"model": a.model, "think": cfg.think,
"device": a.device, "repeat": a.repeat,
"passed": n_pass, "total": len(scored),
"skipped": len(rows) - len(scored),
"wall_s": round(wall, 1), "rows": rows}, f, indent=2)
print(f"wrote {a.out}")
if __name__ == "__main__":
main()
benchmarks/edge_harness/android_agent/scenarios.py
+113
View File
@@ -0,0 +1,113 @@
"""Android screens as a text-only agent sees them, with the correct next action.
Each screen is a uiautomator-style accessibility dump flattened to one line per
element -- the same shape droidrun's Portal and AndroidWorld's a11y observation
hand to a text model. `expect` is the action a competent agent must take next;
`accept_names` allows a defensible alternative where one genuinely exists.
"""
SCENARIOS = [
{
"id": "launch_clock",
"task": "Set an alarm for 7:30 AM.",
"screen": """[0] TextView "Phone" clickable=true
[1] TextView "Messages" clickable=true
[2] TextView "Chrome" clickable=true
[3] TextView "Clock" clickable=true
[4] TextView "Settings" clickable=true
[5] TextView "Camera" clickable=true""",
"expect": {"name": "tap", "args": {"index": 3}},
"accept_names": ["tap", "open_app"],
},
{
"id": "clock_alarm_tab",
"task": "Set an alarm for 7:30 AM.",
"screen": """[0] Button "Alarm" clickable=true selected=false
[1] Button "Clock" clickable=true selected=true
[2] Button "Timer" clickable=true selected=false
[3] Button "Stopwatch" clickable=true selected=false
[4] TextView "14:32" clickable=false""",
"expect": {"name": "tap", "args": {"index": 0}},
"accept_names": ["tap"],
},
{
"id": "add_alarm",
"task": "Set an alarm for 7:30 AM.",
"screen": """[0] Button "Alarm" clickable=true selected=true
[1] TextView "No alarms" clickable=false
[2] ImageButton "Add alarm" clickable=true
[3] Button "Timer" clickable=true""",
"expect": {"name": "tap", "args": {"index": 2}},
"accept_names": ["tap"],
},
{
"id": "wifi_settings",
"task": "Turn on Wi-Fi.",
"screen": """[0] TextView "Settings" clickable=false
[1] TextView "Network & internet" clickable=true
[2] TextView "Connected devices" clickable=true
[3] TextView "Apps" clickable=true
[4] TextView "Battery" clickable=true
[5] TextView "Display" clickable=true""",
"expect": {"name": "tap", "args": {"index": 1}},
"accept_names": ["tap"],
},
{
"id": "wifi_toggle",
"task": "Turn on Wi-Fi.",
"screen": """[0] TextView "Network & internet" clickable=false
[1] Switch "Wi-Fi" checked=false clickable=true
[2] TextView "Mobile network" clickable=true
[3] Switch "Airplane mode" checked=false clickable=true""",
"expect": {"name": "tap", "args": {"index": 1}},
"accept_names": ["tap"],
},
{
"id": "search_contact",
"task": "Find the contact named Wei Zhang.",
"screen": """[0] EditText "Search contacts" clickable=true focused=true
[1] TextView "Alice Chen" clickable=true
[2] TextView "Bob Lee" clickable=true
[3] TextView "Carol Wu" clickable=true""",
"expect": {"name": "type_text", "args": {"index": 0, "text": "Wei Zhang"}},
"accept_names": ["type_text"],
},
{
"id": "scroll_for_item",
"task": "Open the Developer options screen in Settings.",
"screen": """[0] TextView "Settings" clickable=false
[1] TextView "Display" clickable=true
[2] TextView "Sound" clickable=true
[3] TextView "Storage" clickable=true
[4] TextView "Privacy" clickable=true
[5] TextView "Location" clickable=true""",
"expect": {"name": "swipe", "args": {"direction": "up"}},
"accept_names": ["swipe"],
},
{
"id": "go_back",
"task": "Return to the previous screen; you opened the wrong menu.",
"screen": """[0] TextView "About phone" clickable=false
[1] TextView "Model: Pixel 8" clickable=false
[2] TextView "Android version" clickable=true""",
"expect": {"name": "press_key", "args": {"key": "back"}},
"accept_names": ["press_key"],
},
{
"id": "open_app_by_name",
"task": "Open the Camera app.",
"screen": """[0] TextView "Swipe up to unlock" clickable=false
[1] TextView "10:04" clickable=false""",
"expect": {"name": "open_app", "args": {"name": "Camera"}},
"accept_names": ["open_app", "swipe"],
},
{
"id": "task_done",
"task": "Turn on Wi-Fi.",
"screen": """[0] TextView "Network & internet" clickable=false
[1] Switch "Wi-Fi" checked=true clickable=true
[2] TextView "Connected to HOME-5G" clickable=false""",
"expect": {"name": "task_complete", "args": {}},
"accept_names": ["task_complete"],
},
]
benchmarks/edge_harness/android_agent/tasks.py
+279
View File
@@ -0,0 +1,279 @@
"""A scored task suite, judged by what the phone shows rather than by what the
agent said -- or by what the simulator happens to remember.
Three things were wrong with the first version of this file, and all three
made it score more generously than it claimed to:
* **``add_alarm`` passed on any typed string containing "7".** Typing ``7``
and stopping scored a success on a task whose prompt is "Set an alarm for
7:30 AM": no minute, no meridiem, and no requirement that the alarm was ever
saved. It now requires a *committed* alarm reading 7:30 AM, which the clock
screen lists, so the check is one a real phone can also answer.
* **``wifi_on_already`` scored final state only.** Its note said "any tap on
the switch fails the task"; the code could not see a tap, so turning Wi-Fi
off and back on passed. Switch flips are now recorded and the task fails on
any of them -- which is the exact failure W4A8 exhibited (it turned Wi-Fi on
and then tapped the switch off again).
* **Predicates read ``SimulatedDevice`` attributes**, so every scored run was a
simulator run by construction. They read :class:`PhoneState` now, which comes
from the UI dump for both backends.
Repeats are only trials if something differs between them, so each task
carries **variants**: the same goal against a renamed control, a reordered
screen, or a different starting state. Five runs of one fixture at temperature
0 measure repeatability; five variants measure whether the agent can do the
task.
"""
from __future__ import annotations
import random
from dataclasses import dataclass, field
from typing import Callable
from device import SimElement, SimulatedDevice, demo_phone
from phone_state import PhoneState, read_state
Predicate = Callable[[PhoneState], bool]
CLOCK_PACKAGE = "com.google.android.deskclock"
"""Clock app to reset for the alarm tasks. AOSP builds and several OEMs ship a
different package (``com.android.deskclock``, ``com.sec.android.app.clockpackage``
on Samsung); override before a real-device run, because clearing the wrong
package leaves a stale alarm and the next run scores against it."""
# --------------------------------------------------------------- variations
def rename(screen: str, old: str, new: str) -> Callable[[SimulatedDevice], None]:
"""Same control, different word. Android OEMs really do ship "WLAN"."""
def apply(d: SimulatedDevice) -> None:
for el in d.screens[screen].elements:
if el.label == old:
el.label = new
return apply
def reorder(screen: str, seed: int) -> Callable[[SimulatedDevice], None]:
"""Shuffle the tappable rows. Index-based tapping must not be luck."""
def apply(d: SimulatedDevice) -> None:
els = d.screens[screen].elements
head = [e for e in els if not e.clickable]
body = [e for e in els if e.clickable]
random.Random(seed).shuffle(body)
d.screens[screen].elements = head + body
return apply
def decoy(screen: str, label: str) -> Callable[[SimulatedDevice], None]:
"""A plausible wrong row, e.g. "Wi-Fi calling" next to "Wi-Fi"."""
def apply(d: SimulatedDevice) -> None:
d.screens[screen].elements.append(SimElement(label))
return apply
def chain(*fns: Callable[[SimulatedDevice], None]) -> Callable[[SimulatedDevice], None]:
def apply(d: SimulatedDevice) -> None:
for fn in fns:
fn(d)
return apply
@dataclass
class Variant:
variant_id: str
mutate: Callable[[SimulatedDevice], None] = lambda d: None
# --------------------------------------------------------------------- task
@dataclass
class Task:
task_id: str
prompt: str
reaches: Predicate
must_not: Predicate = lambda s: False
setup: Callable[[SimulatedDevice], None] = lambda d: None
max_steps: int = 10
note: str = ""
variants: list[Variant] = field(default_factory=lambda: [Variant("base")])
needs: tuple[str, ...] = ()
"""PhoneState facts this task's predicates require that a screen cannot
show (``toggled``, ``fields``, ``records``). A task listing any of these
cannot be scored on a real phone, and says so instead of scoring wrong."""
reset_adb: tuple[tuple[str, ...], ...] = ()
"""Commands that put a real phone into this task's starting state. Without
them a real-device run is scored against an unknown start, so an empty
tuple means "refuse to run me on a phone"."""
def variant(self, n: int) -> Variant:
return self.variants[n % len(self.variants)]
def build(self, run: int = 0) -> SimulatedDevice:
device = demo_phone()
self.variant(run).mutate(device)
self.setup(device)
return device
def scorable_on(self, state: PhoneState) -> str | None:
"""Why this task cannot be scored from ``state``, or None if it can."""
if state.source == "device" and self.needs:
return (f"needs {', '.join(self.needs)}, which a screen cannot "
f"show; simulator only")
return None
def score(self, device_or_state, declared_done: bool, run: int = 0) -> dict:
state = (device_or_state if isinstance(device_or_state, PhoneState)
else read_state(device_or_state))
blocked = self.scorable_on(state)
if blocked:
return {"task_id": self.task_id, "variant": self.variant(run).variant_id,
"scorable": False, "reason": blocked, "success": False,
"reached_goal": None, "collateral_damage": None,
"declared_done": declared_done}
reached = bool(self.reaches(state))
damaged = bool(self.must_not(state))
return {
"task_id": self.task_id,
"variant": self.variant(run).variant_id,
"scorable": True,
"reached_goal": reached,
"collateral_damage": damaged,
"declared_done": declared_done,
# The only row that counts: the phone is in the requested state,
# nothing else was broken, and the agent knew to stop.
"success": reached and not damaged and declared_done,
}
# ----------------------------------------------------------------- oracles
# Both the on-screen labels and the setting key, because `switch()` matches
# either source: "Wi-Fi" normalises onto the `wifi` setting by luck, but
# "Airplane mode" does not normalise onto `airplane`, and a name the setting
# lookup misses silently falls through to "control not on screen".
WIFI = ("Wi-Fi", "WiFi", "WLAN", "wifi")
AIRPLANE = ("Airplane mode", "Flight mode", "airplane")
def _wifi_is(want: bool) -> Predicate:
# `switch` returns None when the control is not on screen, which is not
# the same as off: an agent that never navigated there has not turned
# anything on.
return lambda s: s.switch(*WIFI) is want
def _airplane_on(s: PhoneState) -> bool:
return s.switch(*AIRPLANE) is True
def _alarm_set(hour: int, minute: int, pm: bool) -> Predicate:
"""A *saved* alarm reading the requested time.
The clock screen lists committed alarms, so this reads the screen and works
on a real phone. Typing the digits without saving leaves the list empty.
"""
wanted = f"{hour}:{minute:02d} {'PM' if pm else 'AM'}"
return lambda s: s.has_text(wanted)
SUITE: list[Task] = [
Task(
task_id="wifi_on",
prompt="Turn on Wi-Fi.",
reaches=_wifi_is(True),
must_not=_airplane_on,
note="Two screens deep, then a switch. The switch's state is visible, "
"so stopping is decidable from the screen.",
variants=[
Variant("base"),
Variant("wlan", rename("network", "Wi-Fi", "WLAN")),
Variant("reordered", reorder("network", seed=7)),
Variant("decoy", decoy("network", "Wi-Fi calling")),
Variant("deep", chain(reorder("settings", seed=3),
decoy("network", "Wi-Fi preferences"))),
],
reset_adb=(("shell", "svc", "wifi", "disable"),),
),
Task(
task_id="wifi_on_already",
prompt="Turn on Wi-Fi.",
setup=lambda d: d.state.__setitem__("wifi", True),
reaches=lambda s: s.switch(*WIFI) is True and s.times_toggled("wifi") == 0,
must_not=_airplane_on,
note="Already satisfied. The only correct action is task_complete. "
"Any flip of the switch fails, including off-and-on again, which "
"final state cannot see and which W4A8 actually did.",
needs=("toggled",),
variants=[
Variant("base"),
Variant("wlan", rename("network", "Wi-Fi", "WLAN")),
Variant("reordered", reorder("network", seed=11)),
],
),
Task(
task_id="wifi_off",
prompt="Turn off Wi-Fi.",
setup=lambda d: d.state.__setitem__("wifi", True),
reaches=_wifi_is(False),
must_not=_airplane_on,
note="The mirror of wifi_on: guards against a model that has simply "
"learned to tap switches.",
variants=[
Variant("base"),
Variant("wlan", rename("network", "Wi-Fi", "WLAN")),
Variant("reordered", reorder("network", seed=5)),
],
reset_adb=(("shell", "svc", "wifi", "enable"),),
),
Task(
task_id="airplane_on",
prompt="Turn on airplane mode.",
reaches=_airplane_on,
note="Same screen as Wi-Fi, different row -- catches index confusion, "
"which is why one variant reorders that screen.",
variants=[
Variant("base"),
Variant("reordered", reorder("network", seed=2)),
Variant("decoy", decoy("network", "Airplane mode settings")),
],
# `settings put global` needs WRITE_SECURE_SETTINGS, which adb has on a
# debuggable build but not on every retail phone; the run fails loudly
# if it is refused rather than scoring against an unknown start.
reset_adb=(("shell", "settings", "put", "global", "airplane_mode_on", "0"),),
),
Task(
task_id="open_clock",
prompt="Open the Clock app and go to the Alarm tab.",
reaches=lambda s: s.has_text("Add alarm") or s.has_text("Set time"),
note="Navigation only, no state change. Isolates 'can it find an app' "
"from 'can it finish'.",
variants=[
Variant("base"),
Variant("reordered", reorder("home", seed=4)),
Variant("decoy", decoy("home", "Clock Widget")),
],
reset_adb=(("shell", "am", "force-stop", CLOCK_PACKAGE),
("shell", "input", "keyevent", "KEYCODE_HOME")),
),
Task(
task_id="add_alarm",
prompt="Set an alarm for 7:30 AM.",
reaches=_alarm_set(7, 30, pm=False),
must_not=lambda s: s.has_text("7:30 PM"),
max_steps=12,
note="The full flow: right app, right tab, add, type both fields, "
"leave the meridiem alone, save. Previously this passed on any "
"typed string containing '7'.",
variants=[
Variant("base"),
Variant("reordered", reorder("clock", seed=9)),
Variant("decoy", decoy("alarms", "Alarm settings")),
],
# `pm clear` removes saved alarms, so run N does not see run N-1's.
reset_adb=(("shell", "pm", "clear", CLOCK_PACKAGE),
("shell", "input", "keyevent", "KEYCODE_HOME")),
),
]
SUITE_BY_ID = {t.task_id: t for t in SUITE}
benchmarks/edge_harness/android_agent/test_android_agent.py
+525
View File
@@ -0,0 +1,525 @@
"""Tests for the Android agent toolkit.
No phone, emulator or model server is needed: the device is the in-process
simulator and the model is a scripted stand-in. That is deliberate -- this
host has no `adb`, no emulator and no `/dev/kvm`, so anything that could only
be checked against real hardware would otherwise ship unverified.
"""
import json
import pytest
import agent as agent_mod
from agent import AgentConfig, first_tool_call, run_episode
from device import demo_phone
from executor import execute
from ui_tree import parse_ui_dump, render_tree, screen_signature
# --------------------------------------------------------------------- ui_tree
REAL_DUMP = """<?xml version="1.0" encoding="UTF-8"?>
<hierarchy rotation="0">
<node index="0" text="" resource-id="" class="android.widget.FrameLayout" package="com.android.settings" content-desc="" checkable="false" checked="false" clickable="false" enabled="true" focusable="false" focused="false" scrollable="false" long-clickable="false" password="false" selected="false" bounds="[0,0][1080,2400]">
<node index="1" text="Network &amp; internet" resource-id="android:id/title" class="android.widget.TextView" package="com.android.settings" content-desc="" checkable="false" checked="false" clickable="true" enabled="true" focusable="true" focused="false" scrollable="false" long-clickable="false" password="false" selected="false" bounds="[0,300][1080,460]" />
<node index="2" text="" resource-id="" class="android.widget.LinearLayout" package="com.android.settings" content-desc="" checkable="false" checked="false" clickable="false" enabled="true" focusable="false" focused="false" scrollable="false" long-clickable="false" password="false" selected="false" bounds="[0,460][1080,620]" />
<node index="3" text="" resource-id="com.android.settings:id/switch_widget" class="android.widget.Switch" package="com.android.settings" content-desc="Wi-Fi" checkable="true" checked="true" clickable="true" enabled="true" focusable="true" focused="false" scrollable="false" long-clickable="false" password="false" selected="false" bounds="[800,460][1000,620]" />
</node>
</hierarchy>"""
def test_parses_a_real_dump_and_keeps_only_what_is_actionable():
els = parse_ui_dump(REAL_DUMP)
labels = [e.text for e in els]
assert "Network & internet" in labels # entity decoded
assert "Wi-Fi" in labels # fell back to content-desc
assert "" not in labels # bare layout node dropped
assert [e.index for e in els] == list(range(len(els)))
def test_switch_state_is_visible_to_the_model():
"""The one clearly wrong action measured earlier was tapping a switch that
was already on; that is only avoidable if the state is on the screen."""
rendered = render_tree(parse_ui_dump(REAL_DUMP))
assert "checked=true" in rendered
assert 'Switch "Wi-Fi"' in rendered
def test_index_maps_to_the_centre_of_the_element():
els = parse_ui_dump(REAL_DUMP)
wifi = next(e for e in els if e.text == "Wi-Fi")
assert wifi.center == (900, 540)
def test_signature_tracks_state_not_pixels():
a = parse_ui_dump(REAL_DUMP)
shifted = REAL_DUMP.replace("[0,300][1080,460]", "[0,310][1080,470]")
assert screen_signature(parse_ui_dump(shifted)) == screen_signature(a)
toggled = REAL_DUMP.replace('checkable="true" checked="true"',
'checkable="true" checked="false"')
assert screen_signature(parse_ui_dump(toggled)) != screen_signature(a)
def test_element_cap_bounds_the_context():
big = '<?xml version="1.0"?><hierarchy>' + "".join(
f'<node text="row{i}" class="android.widget.TextView" clickable="true" '
f'bounds="[0,{i},100,{i + 10}]" />' for i in range(200)
) + "</hierarchy>"
assert len(parse_ui_dump(big, max_elements=25)) == 25
# -------------------------------------------------------------------- executor
def test_tap_lands_on_the_right_row_and_reports_prior_state():
device = demo_phone()
device.current = "network"
els = parse_ui_dump(device.dump())
wifi = next(e for e in els if e.text == "Wi-Fi")
result = execute("tap", {"index": wifi.index}, els, device)
assert result.ok and device.state["wifi"] is True
assert "was off" in result.message
def test_out_of_range_index_is_a_message_not_an_exception():
device = demo_phone()
els = parse_ui_dump(device.dump())
result = execute("tap", {"index": 99}, els, device)
assert not result.ok
assert "out of range" in result.message and f"0-{len(els) - 1}" in result.message
def test_bad_arguments_explain_themselves():
device = demo_phone()
els = parse_ui_dump(device.dump())
assert "missing required argument" in execute("tap", {}, els, device).message
assert "whole number" in execute("tap", {"index": "banana"}, els, device).message
assert "unknown tool" in execute("frobnicate", {}, els, device).message
assert "direction must be" in execute("swipe", {"direction": "sideways"}, els, device).message
def test_type_text_focuses_the_field_first():
device = demo_phone()
device.current = "new_alarm"
els = parse_ui_dump(device.dump())
field = next(e for e in els if e.text == "Hour")
execute("type_text", {"index": field.index, "text": "07"}, els, device)
kinds = [a[0] for a in device.actions]
assert kinds == ["tap", "text"], kinds
assert device.typed == ["07"]
def test_swipe_direction_is_content_movement():
device = demo_phone()
els = parse_ui_dump(device.dump())
execute("swipe", {"direction": "down"}, els, device)
_, x1, y1, x2, y2 = device.actions[-1]
assert int(y1) > int(y2), "scrolling down must move the finger up"
def test_task_complete_carries_its_summary():
device = demo_phone()
result = execute("task_complete", {"summary": "Wi-Fi is on"}, [], device)
assert result.done and result.summary == "Wi-Fi is on"
# ----------------------------------------------------------------------- agent
def _scripted(calls):
"""A stand-in model that plays a fixed list of tool calls."""
seq = list(calls)
def fake_chat(cfg, messages, *, think):
call = seq.pop(0) if seq else None
if call is None:
return {"choices": [{"message": {"content": "", "tool_calls": []}}]}, 0.01
name, args = call
return {"choices": [{"message": {"tool_calls": [
{"function": {"name": name, "arguments": json.dumps(args)}}
]}}]}, 0.01
return fake_chat
def test_first_tool_call_handles_both_argument_encodings():
as_str = {"choices": [{"message": {"tool_calls": [
{"function": {"name": "tap", "arguments": '{"index": 2}'}}]}}]}
as_dict = {"choices": [{"message": {"tool_calls": [
{"function": {"name": "tap", "arguments": {"index": 2}}}]}}]}
assert first_tool_call(as_str) == ("tap", {"index": 2})
assert first_tool_call(as_dict) == ("tap", {"index": 2})
assert first_tool_call({"choices": [{"message": {"content": "hi"}}]}) == (None, {})
# Malformed JSON must not kill the episode.
broken = {"choices": [{"message": {"tool_calls": [
{"function": {"name": "tap", "arguments": "{not json"}}]}}]}
assert first_tool_call(broken) == ("tap", {})
def test_completes_a_multi_screen_task(monkeypatch):
device = demo_phone()
monkeypatch.setattr(agent_mod, "chat", _scripted([
("tap", {"index": 3}), # home -> Settings
("tap", {"index": 1}), # -> Network & internet
("tap", {"index": 1}), # Wi-Fi switch
("task_complete", {"summary": "Wi-Fi enabled"}),
]))
ep = run_episode(device, "Turn on Wi-Fi.", AgentConfig(max_steps=8))
assert ep.done and device.state["wifi"] is True
assert ep.summary == "Wi-Fi enabled"
assert ep.stop_reason == "task_complete"
assert ep.n_steps == 4
def test_says_so_when_an_action_changes_nothing(monkeypatch):
device = demo_phone()
monkeypatch.setattr(agent_mod, "chat", _scripted([
("tap", {"index": 0}), # "Phone" on home: goes nowhere
("task_complete", {"summary": "done"}),
]))
ep = run_episode(device, "Do something.", AgentConfig(max_steps=4))
assert "did not change" in ep.steps[0].result.message
def test_stops_when_it_repeats_a_no_op(monkeypatch):
device = demo_phone()
monkeypatch.setattr(agent_mod, "chat", _scripted([("tap", {"index": 0})] * 6))
ep = run_episode(device, "Do something.", AgentConfig(max_steps=6))
assert not ep.done
assert "stuck" in ep.stop_reason
assert ep.n_steps < 6, "should give up before the budget"
def test_empty_reply_retries_without_thinking(monkeypatch):
"""The measured failure: reasoning never closes `<think>`, so the parser
returns nothing. A non-thinking retry is worth ~0.6 s."""
seen = []
def fake_chat(cfg, messages, *, think):
seen.append(think)
if think:
return {"choices": [{"message": {"content": "", "tool_calls": []}}]}, 8.9
return {"choices": [{"message": {"tool_calls": [
{"function": {"name": "task_complete",
"arguments": '{"summary": "fallback"}'}}]}}]}, 0.61
monkeypatch.setattr(agent_mod, "chat", fake_chat)
ep = run_episode(device := demo_phone(), "Task.", AgentConfig(max_steps=3))
assert seen == [True, False], seen
assert ep.done and ep.summary == "fallback"
assert ep.steps[0].thinking is False
assert device is not None
def test_gives_up_rather_than_looping_when_no_tool_call_ever_comes(monkeypatch):
monkeypatch.setattr(agent_mod, "chat", _scripted([]))
ep = run_episode(demo_phone(), "Task.", AgentConfig(max_steps=5))
assert not ep.done
assert "no tool call" in ep.stop_reason
assert ep.n_steps == 1
def test_step_budget_is_enforced(monkeypatch):
device = demo_phone()
# Alternate two real navigations so nothing is a repeat or a no-op.
monkeypatch.setattr(agent_mod, "chat", _scripted(
[("tap", {"index": 3}), ("press_key", {"key": "back"})] * 5))
ep = run_episode(device, "Wander.", AgentConfig(max_steps=4))
assert not ep.done
assert ep.n_steps == 4
assert "budget exhausted" in ep.stop_reason
def test_transcript_is_readable(monkeypatch):
monkeypatch.setattr(agent_mod, "chat", _scripted([
("tap", {"index": 3}), ("task_complete", {"summary": "ok"})]))
ep = run_episode(demo_phone(), "Open settings.", AgentConfig(max_steps=4))
text = ep.transcript()
assert "task: Open settings." in text
assert "tap(" in text and "task_complete" in text
assert "stopped: task_complete" in text
def test_detects_undoing_its_own_action(monkeypatch):
"""The measured failure: tapping a switch that is already on, turning it
back off, forever. Each tap really does change the screen, so the
no-change check cannot see it."""
device = demo_phone()
device.current = "network"
monkeypatch.setattr(agent_mod, "chat", _scripted([("tap", {"index": 1})] * 6))
ep = run_episode(device, "Turn on Wi-Fi.", AgentConfig(max_steps=6))
assert not ep.done
assert "undoing its own action" in ep.stop_reason
assert any("undid your previous action" in s.result.message
for s in ep.steps if s.result)
def test_normal_navigation_is_not_mistaken_for_oscillation(monkeypatch):
"""Opening a submenu and pressing back revisits a screen; that is how
phones work and must not trip the guard."""
device = demo_phone()
monkeypatch.setattr(agent_mod, "chat", _scripted([
("tap", {"index": 3}), # home -> settings
("press_key", {"key": "back"}), # -> home
("tap", {"index": 2}), # home -> clock
("press_key", {"key": "back"}), # -> home
("task_complete", {"summary": "looked around"}),
]))
ep = run_episode(device, "Look around.", AgentConfig(max_steps=8))
assert ep.done, ep.stop_reason
assert ep.n_steps == 5
# ------------------------------------------------------------------ task suite
def test_suite_scores_on_device_state_not_on_what_the_agent_said():
"""The failure this model exhibits is claiming completion it did not
achieve, so the score must not read the transcript."""
from tasks import SUITE_BY_ID
task = SUITE_BY_ID["wifi_on"]
device = task.build()
assert task.score(device, declared_done=True)["success"] is False
device.state["wifi"] = True
assert task.score(device, declared_done=True)["success"] is True
# Reaching the goal without stopping is not success either.
assert task.score(device, declared_done=False)["success"] is False
def test_collateral_damage_fails_the_task():
from tasks import SUITE_BY_ID
task = SUITE_BY_ID["wifi_on"]
device = task.build()
device.state.update(wifi=True, airplane=True)
score = task.score(device, declared_done=True)
assert score["reached_goal"] and score["collateral_damage"]
assert score["success"] is False
def test_already_satisfied_task_starts_satisfied():
"""Its only correct action is task_complete; tapping the switch fails."""
from phone_state import read_state
from tasks import SUITE_BY_ID
task = SUITE_BY_ID["wifi_on_already"]
device = task.build()
assert task.reaches(read_state(device))
assert task.score(device, declared_done=True)["success"] is True
def test_already_satisfied_task_fails_on_a_switch_round_trip():
"""W4A8 turned Wi-Fi on and then tapped the switch off again. Final state
cannot see that, so this task scored it a pass until the flips were
recorded."""
from tasks import SUITE_BY_ID
task = SUITE_BY_ID["wifi_on_already"]
device = task.build()
device.launch("Settings")
device.tap(540, 360) # Network & internet
device.tap(540, 360) # Wi-Fi off
device.tap(540, 360) # ... and on again
assert device.state["wifi"] is True # final state looks correct
assert task.score(device, declared_done=True)["success"] is False
def test_goal_is_read_from_the_setting_not_from_the_final_screen():
"""An agent that turns Wi-Fi on and then goes Home has still done it."""
from tasks import SUITE_BY_ID
task = SUITE_BY_ID["wifi_on"]
device = task.build()
device.launch("Settings")
device.tap(540, 360)
device.tap(540, 360)
device.keyevent("home")
assert task.score(device, declared_done=True)["success"] is True
def test_never_finding_the_screen_is_not_scored_as_off():
"""`switch()` returns None when neither the setting nor the screen knows
the control; a predicate must not read that as 'off' and pass a task."""
from phone_state import PhoneState
from tasks import SUITE_BY_ID
empty = PhoneState(source="device")
assert empty.switch("Wi-Fi") is None
assert SUITE_BY_ID["wifi_off"].score(empty, declared_done=True)["success"] is False
def test_suite_includes_a_mirror_task_for_switch_tapping():
"""Guards against a model that has simply learned 'tap the switch'."""
from tasks import SUITE_BY_ID
from phone_state import read_state
off = SUITE_BY_ID["wifi_off"]
device = off.build()
assert device.state["wifi"] is True
assert not off.reaches(read_state(device))
device.state["wifi"] = False
assert off.reaches(read_state(device))
def test_chat_retries_transient_failures(monkeypatch):
"""A phone task is a dozen round trips; one dropped connection should cost
a retry, not the episode."""
import agent as mod
from urllib import error as urlerror
calls = {"n": 0}
def flaky(req, timeout=None):
calls["n"] += 1
if calls["n"] < 3:
raise urlerror.URLError("connection reset")
class _R:
def __enter__(self): return self
def __exit__(self, *a): return False
def read(self): return b'{"choices":[{"message":{"tool_calls":[]}}]}'
return _R()
monkeypatch.setattr(mod.urlrequest, "urlopen", flaky)
monkeypatch.setattr(mod.time, "sleep", lambda s: None)
cfg = AgentConfig(retries=3)
body, _ = mod.chat(cfg, [{"role": "user", "content": "x"}], think=False)
assert calls["n"] == 3 and "choices" in body
def test_chat_gives_up_with_a_clear_error(monkeypatch):
import agent as mod
from urllib import error as urlerror
monkeypatch.setattr(mod.urlrequest, "urlopen",
lambda *a, **k: (_ for _ in ()).throw(urlerror.URLError("down")))
monkeypatch.setattr(mod.time, "sleep", lambda s: None)
with pytest.raises(mod.ChatError, match="after 2 attempts"):
mod.chat(AgentConfig(retries=2), [{"role": "user", "content": "x"}], think=False)
# --------------------------------------------------- oracles that were wrong
def test_alarm_task_needs_a_saved_alarm_not_a_typed_digit():
"""This oracle used to pass on any typed string containing "7", so typing
`7` and stopping scored a success on "Set an alarm for 7:30 AM"."""
from tasks import SUITE_BY_ID
task = SUITE_BY_ID["add_alarm"]
typed_only = task.build()
typed_only.launch("Clock")
typed_only.tap(540, 200) # Alarm tab
typed_only.tap(540, 360) # Add alarm
typed_only.tap(540, 360)
typed_only.input_text("7")
assert task.score(typed_only, declared_done=True)["success"] is False
both_fields = task.build()
both_fields.launch("Clock")
both_fields.tap(540, 200)
both_fields.tap(540, 360)
both_fields.tap(540, 360)
both_fields.input_text("7")
both_fields.tap(540, 520)
both_fields.input_text("30")
assert task.score(both_fields, declared_done=True)["success"] is False, \
"typing both fields without saving is not a set alarm"
both_fields.tap(540, 840) # Save
assert task.score(both_fields, declared_done=True)["success"] is True
def test_alarm_task_rejects_the_wrong_meridiem():
from tasks import SUITE_BY_ID
task = SUITE_BY_ID["add_alarm"]
d = task.build()
d.launch("Clock"); d.tap(540, 200); d.tap(540, 360)
d.tap(540, 360); d.input_text("7")
d.tap(540, 520); d.input_text("30")
d.tap(540, 680) # AM -> PM
d.tap(540, 840) # Save
score = task.score(d, declared_done=True)
assert score["success"] is False and score["collateral_damage"] is True
def test_a_saved_alarm_is_visible_on_the_screen():
"""The oracle reads the clock app's list, so it works on a real phone."""
from phone_state import read_state
from tasks import SUITE_BY_ID
d = SUITE_BY_ID["add_alarm"].build()
d.launch("Clock"); d.tap(540, 200); d.tap(540, 360)
d.tap(540, 360); d.input_text("7")
d.tap(540, 520); d.input_text("30")
d.tap(540, 840)
assert read_state(d).has_text("7:30 AM")
# ------------------------------------------------------------------ variants
def test_every_task_has_variants_so_repeats_are_trials():
"""Five runs of one fixture at temperature 0 measure repeatability. The
suite has to vary something for a repeat to be a trial."""
from tasks import SUITE
for task in SUITE:
assert len(task.variants) >= 1
ids = [v.variant_id for v in task.variants]
assert len(ids) == len(set(ids)), task.task_id
assert any(len(t.variants) > 1 for t in SUITE)
def test_variants_change_the_phone_but_not_the_goal():
"""A renamed or reordered control is still the same task."""
from phone_state import read_state
from tasks import SUITE_BY_ID
task = SUITE_BY_ID["wifi_on"]
labels_seen = set()
for run in range(len(task.variants)):
d = task.build(run)
d.launch("Settings")
d.tap(540, 360)
labels_seen.add(read_state(d).source)
# The goal predicate is unchanged; only the route differs.
assert task.score(d, declared_done=True)["success"] is False
d.state["wifi"] = True
assert task.score(d, declared_done=True)["success"] is True
assert labels_seen == {"sim"}
def test_wlan_variant_renames_the_control_the_agent_must_find():
from phone_state import read_state
from tasks import SUITE_BY_ID
task = SUITE_BY_ID["wifi_on"]
variant = next(i for i, v in enumerate(task.variants) if v.variant_id == "wlan")
d = task.build(variant)
d.launch("Settings"); d.tap(540, 360)
assert read_state(d).has_text("WLAN")
assert not read_state(d).has_text("Wi-Fi")
# -------------------------------------------------- real device honesty
def test_history_dependent_task_refuses_to_score_on_a_real_phone():
"""A phone cannot report which switches were flipped, so the task says so
instead of scoring from final state and calling a round trip a pass."""
from phone_state import PhoneState
from tasks import SUITE_BY_ID
state = PhoneState(source="device", settings={"wifi": True})
score = SUITE_BY_ID["wifi_on_already"].score(state, declared_done=True)
assert score["scorable"] is False
assert "simulator only" in score["reason"]
assert score["success"] is False
def test_tasks_runnable_on_a_phone_declare_how_to_reset_it():
"""Without a reset the run is scored against an unknown starting state."""
from tasks import SUITE_BY_ID
assert SUITE_BY_ID["wifi_on"].reset_adb
assert SUITE_BY_ID["wifi_off"].reset_adb
# and the sim-only one is not pretending otherwise
assert SUITE_BY_ID["wifi_on_already"].needs == ("toggled",)
benchmarks/edge_harness/android_agent/tool_parse.py
+155
View File
@@ -0,0 +1,155 @@
"""Recover a tool call from what the model actually emitted.
vLLM's stock `glm45` parser handles Spark's tool output most of the time --
4 of 5 screens in a spot check -- and then does not. On the Wi-Fi toggle
screen, the one where a switch's state decides the right action, the model
dropped out of GLM-4.5 format entirely and emitted
[tap]
{"tool": "tap", "parameters": {"index": 1}
with the JSON unterminated. `tool_calls` comes back empty, and an agent that
only reads `tool_calls` sees a step with no action at all.
For a phone agent that is the difference between finishing a task and
stopping in the middle of it, so this module tries, in order:
1. the parsed `tool_calls` (the normal path -- nothing here second-guesses it),
2. a GLM-4.5 `<tool_call>` block left in the content,
3. the bracketed `[name]` + JSON form above, repairing unbalanced braces,
4. a bare JSON object naming a tool.
Every recovery is reported, because an agent that quietly repairs its model's
output will hide a regression in that model.
"""
from __future__ import annotations
import json
import re
from dataclasses import dataclass
_BRACKET = re.compile(r"\[([a-z_][a-z0-9_]*)\]", re.I)
_GLM_CALL = re.compile(
r"<tool_call>\s*([a-z_][a-z0-9_]*)(.*?)</tool_call>", re.I | re.S)
_GLM_ARG = re.compile(
r"<arg_key>(.*?)</arg_key>\s*<arg_value>(.*?)</arg_value>", re.I | re.S)
_ARG_CONTAINERS = ("parameters", "arguments", "args")
_NAME_KEYS = ("tool", "name", "tool_name", "function")
@dataclass
class ParsedCall:
name: str | None
args: dict
source: str
"""Which path produced it: tool_calls | glm_block | bracketed | bare_json."""
@property
def recovered(self) -> bool:
return self.name is not None and self.source != "tool_calls"
def _coerce(value: str):
"""Tool arguments arrive as strings; indices must become integers."""
text = value.strip()
if re.fullmatch(r"-?\d+", text):
return int(text)
if text.lower() in ("true", "false"):
return text.lower() == "true"
return text
def _balance(text: str) -> str:
"""Close a JSON object the model stopped short of finishing."""
depth_curly = text.count("{") - text.count("}")
depth_square = text.count("[") - text.count("]")
if depth_curly < 0 or depth_square < 0:
return text
return text + "]" * depth_square + "}" * depth_curly
def _load(text: str) -> dict | None:
for candidate in (text, _balance(text)):
try:
value = json.loads(candidate)
except json.JSONDecodeError:
continue
if isinstance(value, dict):
return value
return None
def _unwrap(payload: dict, fallback_name: str | None) -> tuple[str | None, dict]:
"""Pull a (name, args) pair out of assorted self-describing shapes."""
name = fallback_name
for key in _NAME_KEYS:
value = payload.get(key)
if isinstance(value, str) and value:
name = value
break
if isinstance(value, dict) and isinstance(value.get("name"), str):
name = value["name"]
payload = {**payload, **value}
break
args: dict = {}
for key in _ARG_CONTAINERS:
inner = payload.get(key)
if isinstance(inner, dict):
args = dict(inner)
break
if isinstance(inner, str):
parsed = _load(inner)
if parsed is not None:
args = parsed
break
if not args:
args = {k: v for k, v in payload.items()
if k not in _NAME_KEYS + _ARG_CONTAINERS}
return name, args
def _first_json_object(text: str) -> str | None:
start = text.find("{")
return text[start:] if start >= 0 else None
def extract_tool_call(message: dict) -> ParsedCall:
"""The one call to act on, from whichever shape the model used."""
calls = message.get("tool_calls") or []
if calls:
fn = calls[0].get("function", {}) or {}
raw = fn.get("arguments")
if isinstance(raw, dict):
args = raw
else:
args = _load(raw or "{}") or {}
if set(args) & set(_ARG_CONTAINERS):
_, args = _unwrap(args, fn.get("name"))
return ParsedCall(fn.get("name"), args, "tool_calls")
content = message.get("content") or ""
reasoning = message.get("reasoning_content") or ""
for text in (content, reasoning):
if not text.strip():
continue
if m := _GLM_CALL.search(text):
args = {k.strip(): _coerce(v) for k, v in _GLM_ARG.findall(m.group(2))}
return ParsedCall(m.group(1), args, "glm_block")
bracket = _BRACKET.search(text)
blob = _first_json_object(text)
if blob is not None:
payload = _load(blob)
if payload is not None:
name, args = _unwrap(payload, bracket.group(1) if bracket else None)
if name:
return ParsedCall(name, args, "bracketed" if bracket else "bare_json")
if bracket and not blob:
# A naked "[press_key]" with no arguments is still an instruction.
return ParsedCall(bracket.group(1), {}, "bracketed")
return ParsedCall(None, {}, "none")
benchmarks/edge_harness/android_agent/ui_tree.py
+146
View File
@@ -0,0 +1,146 @@
"""Turn a uiautomator dump into what a text-only agent can act on.
Spark-X2.5 never sees a screenshot, so the whole agent hinges on this file:
a screen has to arrive as a short, stable, indexed list of elements, and an
index the model returns has to map back to a point on the glass.
Two decisions shape the format, both forced by the model rather than by taste:
* **Filter hard.** A real settings screen dumps 150+ nodes, most of them
layout containers. At 56 KiB/token of KV cache and a 1.7B model that both
costs context and buries the answer, so only nodes that are interactive or
carry text survive.
* **Index, do not coordinate.** The model is asked for `index: 3`, never for
pixels -- a text model has no way to know where 3 is on the glass, and
coordinates in the transcript would be a source of silent nonsense. The
mapping back to a centre point happens here, where the bounds are known.
The rendered form is the one the scenario fixtures were written in, so the
decision benchmark and the live agent observe identically.
"""
from __future__ import annotations
import re
import xml.etree.ElementTree as ET
from dataclasses import dataclass, field
_BOUNDS = re.compile(r"\[(-?\d+),(-?\d+)\]\[(-?\d+),(-?\d+)\]")
# Attributes worth showing when they are true, in the order a reader wants
# them. `checked` earns its place: the one clearly wrong action in the earlier
# measurements was tapping a Wi-Fi switch that was already on, which is only
# avoidable if the state is on the screen the model sees.
_FLAGS = ("clickable", "checked", "selected", "scrollable", "focused")
@dataclass
class Element:
index: int
cls: str
text: str
resource_id: str = ""
package: str = ""
bounds: tuple[int, int, int, int] = (0, 0, 0, 0)
flags: dict[str, bool] = field(default_factory=dict)
editable: bool = False
@property
def center(self) -> tuple[int, int]:
x1, y1, x2, y2 = self.bounds
return (x1 + x2) // 2, (y1 + y2) // 2
@property
def short_class(self) -> str:
return self.cls.rsplit(".", 1)[-1] if self.cls else "View"
def render(self) -> str:
parts = [f"[{self.index}]", self.short_class, f'"{self.text}"']
for flag in _FLAGS:
if flag in self.flags:
parts.append(f"{flag}={str(self.flags[flag]).lower()}")
return " ".join(parts)
def _truthy(node: ET.Element, name: str) -> bool:
return node.get(name, "false") == "true"
def _label(node: ET.Element) -> str:
"""What to call this element, preferring what a person would read."""
for attr in ("text", "content-desc"):
value = (node.get(attr) or "").strip()
if value:
return value.replace("\n", " ")
# A bare resource id is a poor label but far better than an empty string
# when the element is the only way forward (unlabelled icon buttons).
rid = node.get("resource-id") or ""
return rid.rsplit("/", 1)[-1] if rid else ""
def _interesting(node: ET.Element) -> bool:
if any(_truthy(node, a) for a in ("clickable", "checkable", "scrollable",
"long-clickable", "focusable")):
return True
return bool(_label(node)) and not _truthy(node, "password")
def parse_bounds(raw: str | None) -> tuple[int, int, int, int]:
m = _BOUNDS.search(raw or "")
if not m:
return (0, 0, 0, 0)
return tuple(int(g) for g in m.groups()) # type: ignore[return-value]
def parse_ui_dump(xml_text: str, *, max_elements: int = 60) -> list[Element]:
"""uiautomator XML -> the elements worth showing, in reading order."""
root = ET.fromstring(xml_text)
elements: list[Element] = []
for node in root.iter("node"):
if not _interesting(node):
continue
label = _label(node)
clickable = _truthy(node, "clickable")
# An unlabelled, non-interactive node tells the model nothing it can
# use and costs it context.
if not label and not clickable and not _truthy(node, "checkable"):
continue
flags = {"clickable": clickable}
if _truthy(node, "checkable"):
flags["checked"] = _truthy(node, "checked")
if _truthy(node, "selected"):
flags["selected"] = True
if _truthy(node, "scrollable"):
flags["scrollable"] = True
if _truthy(node, "focused"):
flags["focused"] = True
elements.append(
Element(
index=len(elements),
cls=node.get("class", ""),
text=label,
resource_id=node.get("resource-id", ""),
package=node.get("package", ""),
bounds=parse_bounds(node.get("bounds")),
flags=flags,
editable="EditText" in (node.get("class") or ""),
)
)
if len(elements) >= max_elements:
break
return elements
def render_tree(elements: list[Element]) -> str:
"""The exact text the model is shown."""
return "\n".join(e.render() for e in elements)
def screen_signature(elements: list[Element]) -> str:
"""A cheap identity for a screen, to notice when an action changed nothing.
Labels and flags only: coordinates shift with animations and scroll offsets,
and treating those as a new screen would hide a stuck agent.
"""
return "|".join(f"{e.short_class}:{e.text}:{e.flags.get('checked', '')}"
for e in elements)
benchmarks/edge_harness/bench_cpu_vllm.py
+106
View File
@@ -0,0 +1,106 @@
"""pp512 / tg128 for vLLM on CPU, mirroring `llama-bench` semantics.
pp512: time to process a 512-token prompt (prefill throughput).
tg128: time to generate 128 tokens (decode throughput), prefill subtracted.
"""
import argparse, json, os, sys, time
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from bench_harness import provenance # noqa: E402
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--model", default="XHToken/Spark-X2.5-1.7B")
ap.add_argument("--dtype", default="bfloat16")
ap.add_argument("--pp", type=int, default=512)
ap.add_argument("--tg", type=int, default=128)
ap.add_argument("--reps", type=int, default=3)
ap.add_argument("--max-model-len", type=int, default=2048)
ap.add_argument("--enforce-eager", action="store_true")
ap.add_argument("--quantization", default=None)
ap.add_argument("--tag", default="vllm")
ap.add_argument("--no-detokenize", action="store_true",
help="skip incremental detokenization, as llama-bench does")
ap.add_argument("--compilation-config", default=None,
help="JSON passed straight to LLM(compilation_config=...)")
ap.add_argument("--out", default="vllm_cpu_bench.json")
a = ap.parse_args()
import vllm_omni # noqa: F401
from vllm import LLM, SamplingParams
t0 = time.perf_counter()
# Prefix caching must be off: the pp measurement replays one prompt, and a
# cache hit would report a prefill that never happened. llama-bench has no
# such cache, so leaving it on would not be a comparison.
extra = {}
if a.compilation_config:
extra["compilation_config"] = json.loads(a.compilation_config)
llm = LLM(model=a.model, dtype=a.dtype, max_model_len=a.max_model_len,
enforce_eager=a.enforce_eager, max_num_seqs=1,
quantization=a.quantization, enable_prefix_caching=False,
trust_remote_code=False, disable_log_stats=True, **extra)
init_s = time.perf_counter() - t0
def rss_mb(field):
with open("/proc/self/status") as f:
for line in f:
if line.startswith(field):
return round(int(line.split()[1]) / 1024, 1)
return None
prompt_pp = {"prompt_token_ids": list(range(10, 10 + a.pp))}
prompt_tg = {"prompt_token_ids": [10]}
det = not a.no_detokenize
sp_pre = SamplingParams(temperature=0.0, max_tokens=1, ignore_eos=True,
detokenize=det)
sp_gen = SamplingParams(temperature=0.0, max_tokens=a.tg, ignore_eos=True,
detokenize=det)
def timed(prompt, sp, reps):
best = []
for _ in range(reps):
s = time.perf_counter()
llm.generate([prompt], sp, use_tqdm=False)
best.append(time.perf_counter() - s)
return best
timed(prompt_pp, sp_pre, 1) # warmup
pp_t = timed(prompt_pp, sp_pre, a.reps)
tg_t = timed(prompt_tg, sp_gen, a.reps)
tg1_t = timed(prompt_tg, sp_pre, a.reps) # 1-token prefill overhead
pp_best, tg_best, tg1_best = min(pp_t), min(tg_t), min(tg1_t)
res = {
"tag": a.tag, "model": a.model, "dtype": a.dtype,
"compilation_config": a.compilation_config,
"detokenize": not a.no_detokenize,
"quantization": a.quantization,
"enforce_eager": a.enforce_eager, "init_s": round(init_s, 2),
# Which build, which threads, which commits -- a throughput number
# without these cannot be invalidated later, only doubted.
"provenance": provenance({
"model_path": os.path.realpath(a.model) if os.path.exists(a.model)
else a.model,
"pp": a.pp, "tg": a.tg, "reps": a.reps,
"max_model_len": a.max_model_len,
}),
"rss_mb": rss_mb("VmRSS:"), "rss_peak_mb": rss_mb("VmHWM:"),
"threads_bind": os.environ.get("VLLM_CPU_OMP_THREADS_BIND", ""),
"omp_num_threads": os.environ.get("OMP_NUM_THREADS", ""),
"pp_tokens": a.pp, "pp_s": round(pp_best, 4),
"pp_tps": round(a.pp / pp_best, 2),
"tg_tokens": a.tg,
"tg_s": round(tg_best - tg1_best, 4),
"tg_tps": round((a.tg - 1) / max(tg_best - tg1_best, 1e-9), 2),
"raw": {"pp": pp_t, "tg": tg_t, "tg1": tg1_t},
}
print(json.dumps({k: v for k, v in res.items() if k != "raw"}, indent=2))
with open(a.out, "a") as f:
f.write(json.dumps(res) + "\n")
if __name__ == "__main__":
main()
benchmarks/edge_harness/bench_harness.py
+634
View File
@@ -0,0 +1,634 @@
"""Run benchmark arms so the numbers survive a shared machine.
Every wrong conclusion this study produced came from measurement, not from
code: a sequential grid that swung the llama.cpp control 86.7 -> 60.4 tok/s
mid-run, a 3.1 GB "saving" that was really `VmRSS` counting the mmap'd
checkpoint's pages, and a W4A8/W4A16 ordering that flipped between grids. The
discipline that fixed each of them is the same, so it lives here instead of
being re-typed into the next shell script:
* **Rotate the arm order every pass.** Drift hits arm 1 and arm 5 differently,
and interleaving alone does not fix it: if every pass runs the arms in the
same order, arm 1 is always the one measured first and carries whatever the
machine was doing at the start of each pass. Pass ``p`` starts at arm
``(p - 1) mod n``.
* **Repeat, and report the spread.** A single number from a shared host is a
claim about the host, not the code. An arm whose passes disagree by more
than `--noise-threshold` is flagged, and the summary says so.
* **Record how busy the machine was, and refuse when it is.** A pass taken
under load is kept and labelled rather than silently averaged in -- but load
average is a whole-host number, and this host has 224 cores, so it says
nothing about the 24 this benchmark is pinned to. The pre-flight reads
per-core utilisation for the *target* cores, and a host lock keeps two
benchmarks off the same cores in the first place. Two sessions benchmarking
at once wrecked several measurements here, and one grid's control drifted
30% mid-run.
* **Prefer allocation over memory.** `RssAnon` repeats to within 0.5 MB where
`VmRSS` moved 1-2 GB between identical runs.
* **Know which direction is better.** ``best`` is the maximum for a rate and
the minimum for memory, init time or latency. Reporting the maximum for every
metric published the *worst* pass as the headline memory figure.
* **Stamp provenance.** A number whose model, build and thread set were not
recorded cannot be retracted later, only doubted -- and this study has
retracted four.
* **Give an arm its own compile cache when it changes the graph.** vLLM keys
its AOT compile cache on the *config*, so two arms that differ only by an
environment variable share a cached artifact -- and if that variable changes
the traced graph, later passes load the wrong one. Three passes of a prefill
grid died with ``KeyError: 'weight_halves'`` inside an AOT-loaded forward
after the first pass saved an artifact for the other path. Set
``VLLM_CACHE_ROOT`` per arm, and note that the first pass of each arm then
pays a cold compile, which is why ``init_s`` is not comparable across passes.
* **Gate on what the host can actually resolve.** ``--check`` compares against
a checked-in expectations file, with a band per metric taken from that
metric's *observed* spread: 0.2% for ``RssAnon``, ~5% for per-op self time,
and nothing for end-to-end throughput, which swings ~30% here and cannot
carry a gate at all. A band is an argument about what is measurable, not a
tolerance someone picked.
The arms are subprocess invocations of the existing single-run benchmarks, so
the measurement code stays the one that was verified; this only decides what
to run, when, and how to read it.
"""
from __future__ import annotations
import argparse
import contextlib
import json
import os
import platform
import shlex
import statistics
import subprocess
import time
from dataclasses import dataclass, field
def load_average() -> float:
return os.getloadavg()[0]
def parse_cores(spec: str) -> list[int]:
"""'0-23,48-71' -> [0..23, 48..71]"""
cores: list[int] = []
for part in spec.split(","):
part = part.strip()
if not part:
continue
if "-" in part:
lo, hi = part.split("-")
cores.extend(range(int(lo), int(hi) + 1))
else:
cores.append(int(part))
return sorted(set(cores))
@dataclass
class Arm:
tag: str
command: list[str]
env: dict[str, str] = field(default_factory=dict)
@dataclass
class Observation:
tag: str
pass_no: int
metrics: dict[str, float]
load_before: float
load_after: float
wall_s: float
ok: bool
error: str = ""
@property
def busy(self) -> bool:
return max(self.load_before, self.load_after) > BUSY_LOAD
BUSY_LOAD = 8.0
"""Above this 1-minute load average a pass is labelled contended.
A poor default on its own: load average is whole-host, and on a 224-core
machine a 24-thread benchmark contributes ~24 to it by itself, so a fixed 8
labels every pass contended and the label stops meaning anything. ``main()``
derives it from the target core count instead -- our own threads plus slack --
and the per-core pre-flight is the guard that actually decides whether to run.
"""
LOWER_IS_BETTER = (
"rss", "mem", "mb", "init", "load_s", "latency", "ms", "_s", "time", "peak",
)
"""Substrings that mark a metric where smaller wins. Rates (`tg_tps`, `pp_tps`)
take the maximum; memory, init time and latency take the minimum."""
HOST_LOCK_PATH = os.environ.get("SPARK_EDGE_LOCK", "/tmp/spark_edge_bench.lock")
"""Host-level lock. It has to live somewhere every session can see, which a
session-private scratchpad by definition is not -- that is the whole point of
a host lock. Override with ``SPARK_EDGE_LOCK``."""
CORE_BUSY_PCT = 50.0
"""A target core sustained above this is a straggler. OpenMP barriers mean the
slowest thread gates the whole region, so one heavily-loaded core costs far
more than its share -- which is why this is a separate test from the aggregate
below rather than a single averaged number."""
STOLEN_PCT = 5.0
"""Refuse when this much of the target cores' aggregate capacity is going to
someone else, which bounds the bias on the measurement.
A flat per-core threshold does not work on this host. With 224 cores and
several users there is always a transient blip somewhere: a 20% rule over a
single 0.5 s window refused three consecutive launches on cores that were
idle seconds earlier, which makes the guard useless rather than strict. So
utilisation is sampled twice and the per-core minimum is taken -- a spike in
one window does not count as load -- and the decision is made on aggregate
stolen capacity plus the straggler rule."""
def _cpu_times() -> dict[int, tuple[int, int]]:
"""(busy, total) jiffies per core, from /proc/stat."""
out: dict[int, tuple[int, int]] = {}
with open("/proc/stat") as f:
for line in f:
if not line.startswith("cpu") or line.startswith("cpu "):
continue
parts = line.split()
try:
cpu = int(parts[0][3:])
except ValueError:
continue
values = [int(v) for v in parts[1:11]]
idle = values[3] + values[4] # idle + iowait
total = sum(values)
out[cpu] = (total - idle, total)
return out
def core_utilization(cores: list[int], window_s: float = 0.5) -> dict[int, float]:
"""Percentage busy per core over ``window_s``."""
before = _cpu_times()
time.sleep(window_s)
after = _cpu_times()
usage: dict[int, float] = {}
for cpu in cores:
if cpu not in before or cpu not in after:
continue
busy = after[cpu][0] - before[cpu][0]
total = after[cpu][1] - before[cpu][1]
usage[cpu] = 0.0 if total <= 0 else busy / total * 100.0
return usage
def preflight(cores: list[int], *, max_busy_pct: float = CORE_BUSY_PCT,
max_stolen_pct: float = STOLEN_PCT,
window_s: float = 0.75, samples: int = 2) -> dict:
"""Refuse to start when contention would bias the measurement.
Two independent windows; a core's load is the *minimum* of them, so a
transient spike is not mistaken for a busy core. Then two tests: aggregate
stolen capacity (bounds the average bias) and a per-core straggler rule
(one slow core gates every OpenMP barrier).
"""
windows = [core_utilization(cores, window_s) for _ in range(max(samples, 1))]
usage = {c: min(w.get(c, 0.0) for w in windows) for c in cores
if any(c in w for w in windows)}
stolen = (sum(usage.values()) / len(usage)) if usage else 0.0
stragglers = {c: round(u, 1) for c, u in usage.items() if u > max_busy_pct}
return {"cores": cores,
"max_busy_pct": max_busy_pct, "max_stolen_pct": max_stolen_pct,
"utilization": {c: round(u, 1) for c, u in usage.items()},
"stolen_pct": round(stolen, 2),
"busy_cores": stragglers,
"clear": stolen <= max_stolen_pct and not stragglers}
class HostLock:
"""An advisory lock so two sessions do not benchmark the same cores.
Cooperative, and says so: it excludes other runs that take it, not the
unrelated job someone starts by hand. That is what the pre-flight is for.
"""
def __init__(self, path: str = HOST_LOCK_PATH, note: str = ""):
self.path, self.note, self._fd = path, note, None
def __enter__(self) -> "HostLock":
import fcntl
self._fd = open(self.path, "a+")
try:
fcntl.flock(self._fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
except BlockingIOError:
self._fd.seek(0)
holder = self._fd.read().strip()[-200:]
self._fd.close()
self._fd = None
raise SystemExit(
f"another benchmark holds {self.path}: {holder or 'unknown'}. "
f"Wait for it, or set SPARK_EDGE_LOCK to a different path if "
f"you really are on disjoint cores."
)
self._fd.seek(0)
self._fd.truncate()
self._fd.write(f"pid={os.getpid()} {time.strftime('%H:%M:%S')} {self.note}\n")
self._fd.flush()
return self
def __exit__(self, *exc) -> None:
if self._fd is not None:
import fcntl
fcntl.flock(self._fd, fcntl.LOCK_UN)
self._fd.close()
self._fd = None
DEFAULT_BANDS_PCT = {
"rss_anon_mb": 1.5,
"peak_rss_anon_mb": 2.0,
"rss_anon_at_context_mb": 2.0,
"init_s": 12.0,
"ready_s": 12.0,
"import_s": 12.0,
"step_ms": 6.0,
"gemm_ms": 6.0,
}
"""Regression bands, per metric, from the spread each one actually shows on
this host. ``RssAnon`` repeats to 0.2% and is given 1.5% for safety; init time
is dominated by page cache and repeats far less tightly; per-op self time
reproduces to ~5%.
Throughput is deliberately absent. ``tg_tps`` swings ~30% run to run on this
shared host, so any band wide enough not to fire constantly is too wide to
catch a real regression. Gate on memory and per-op time instead; that is the
whole reason the per-op profiler exists."""
UNGATEABLE = {"tg_tps", "pp_tps", "tps"}
"""Metrics this host cannot resolve well enough to gate on."""
def check_expectations(summaries: list["ArmSummary"], expected: dict) -> list[dict]:
"""Compare measured medians against a checked-in expectations file.
The median, not the best: `best` is one pass and a gate that fires on the
luckiest run of a noisy metric is a coin toss.
"""
bands = {**DEFAULT_BANDS_PCT, **(expected.get("bands_pct") or {})}
arms = expected.get("arms") or {}
results = []
for s in summaries:
want = (arms.get(s.tag) or {}).get(s.metric)
if want is None:
results.append({"arm": s.tag, "metric": s.metric, "status": "unlisted",
"measured": s.median, "expected": None})
continue
if s.metric in UNGATEABLE:
results.append({"arm": s.tag, "metric": s.metric, "status": "ungateable",
"measured": s.median, "expected": want,
"why": "this host cannot resolve it"})
continue
band = bands.get(s.metric, 5.0)
drift = abs(s.median - want) / want * 100.0 if want else 0.0
worse = (s.median > want) if s.lower_is_better else (s.median < want)
results.append({
"arm": s.tag, "metric": s.metric,
"measured": s.median, "expected": want,
"drift_pct": round(drift, 2), "band_pct": band,
"status": ("ok" if drift <= band else
"REGRESSION" if worse else "improved"),
})
return results
def format_check(results: list[dict]) -> str:
if not results:
return "nothing to check"
width = max(len(r["arm"]) for r in results)
lines = []
for r in results:
if r["status"] in ("unlisted", "ungateable"):
lines.append(f" {r['arm'].ljust(width)} {r['metric']:22s} "
f"{r['measured']:9.2f} {r['status']}"
+ (f" ({r['why']})" if r.get("why") else ""))
else:
lines.append(f" {r['arm'].ljust(width)} {r['metric']:22s} "
f"{r['measured']:9.2f} vs {r['expected']:9.2f} "
f"{r['drift_pct']:+5.2f}% (band {r['band_pct']:.1f}%) "
f"{r['status']}")
return "\n".join(lines)
def metric_is_lower_better(metric: str) -> bool:
name = metric.lower()
if name.endswith(("_tps", "_tok_s", "tps")):
return False
return any(tok in name for tok in LOWER_IS_BETTER)
def provenance(extra: dict | None = None) -> dict:
"""Everything needed to decide later whether a number is still valid.
Saved profiles in this tree carry no model path, context length or build,
so two of them are indistinguishable from each other and one set silently
describes a checkpoint that was later found wrong. This is the fix.
"""
root = os.environ.get("EMBEDDING_INFER_ROOT", "/data/zhoutaichang/embedding_infer")
def _sha(rel: str) -> str | None:
try:
out = subprocess.run(["git", "-C", os.path.join(root, rel),
"rev-parse", "HEAD"],
capture_output=True, text=True, timeout=10)
return out.stdout.strip() or None
except Exception:
return None
try:
affinity = sorted(os.sched_getaffinity(0))
except AttributeError:
affinity = []
return {
"recorded_at": time.strftime("%Y-%m-%dT%H:%M:%S%z"),
"host": platform.node(),
"python": platform.python_version(),
"cpu_count": os.cpu_count(),
"n_affinity": len(affinity),
"affinity": affinity,
"commits": {
"embedding_infer": _sha("."),
"vllm_omni": _sha("vllm-omni"),
"vllm": _sha("vllm"),
},
"env": {k: v for k, v in sorted(os.environ.items())
if k.startswith(("OMP_", "KMP_", "MKL_", "VLLM_", "LD_PRELOAD",
"CUDA_VISIBLE_DEVICES"))},
"loadavg": os.getloadavg(),
**(extra or {}),
}
def run_once(arm: Arm, pass_no: int, timeout: float, keys: list[str]) -> Observation:
env = {**os.environ, **arm.env}
load_before = load_average()
t0 = time.perf_counter()
proc = subprocess.run(arm.command, env=env, capture_output=True,
text=True, timeout=timeout)
wall = time.perf_counter() - t0
load_after = load_average()
if proc.returncode != 0:
return Observation(arm.tag, pass_no, {}, load_before, load_after, wall,
ok=False, error=proc.stderr.strip()[-400:])
metrics = _scrape(proc.stdout, keys)
return Observation(arm.tag, pass_no, metrics, load_before, load_after, wall,
ok=bool(metrics))
def _scrape(stdout: str, keys: list[str]) -> dict[str, float]:
"""Pull metrics out of the last JSON object the run printed."""
start = stdout.rfind("{")
while start >= 0:
try:
payload = json.loads(stdout[start:stdout.rfind("}") + 1])
except json.JSONDecodeError:
start = stdout.rfind("{", 0, start)
continue
return {k: float(payload[k]) for k in keys
if isinstance(payload.get(k), (int, float))}
return {}
@dataclass
class ArmSummary:
tag: str
metric: str
best: float
median: float
spread_pct: float
n: int
n_busy: int
lower_is_better: bool = False
noise_threshold: float = 8.0
@property
def noisy(self) -> bool:
return self.spread_pct > self.noise_threshold
NOISE_THRESHOLD = 8.0
def summarize(observations: list[Observation], metric: str,
noise_threshold: float = NOISE_THRESHOLD,
lower_is_better: bool | None = None) -> list[ArmSummary]:
"""Collapse passes into one row per arm.
``lower_is_better`` defaults to reading the metric name, so `rss_anon_mb`
and `init_s` take their minimum while `tg_tps` takes its maximum. Taking
the maximum for every metric is how the *worst* memory pass came to be
published as the best one.
"""
if lower_is_better is None:
lower_is_better = metric_is_lower_better(metric)
by_tag: dict[str, list[Observation]] = {}
for obs in observations:
if obs.ok and metric in obs.metrics:
by_tag.setdefault(obs.tag, []).append(obs)
summaries = []
for tag, rows in by_tag.items():
values = [r.metrics[metric] for r in rows]
lo, hi = min(values), max(values)
spread = 0.0 if hi == 0 else (hi - lo) / hi * 100.0
summaries.append(ArmSummary(
tag=tag, metric=metric, best=lo if lower_is_better else hi,
median=statistics.median(values), spread_pct=spread,
n=len(values), n_busy=sum(1 for r in rows if r.busy),
lower_is_better=lower_is_better, noise_threshold=noise_threshold,
))
return sorted(summaries, key=lambda s: s.best if lower_is_better else -s.best)
def indistinguishable(summaries: list[ArmSummary]) -> str | None:
"""Say so when the gap between the top two arms is inside the noise.
Per-op self time resolves ~5% on this host; end-to-end throughput resolves
~30%. An 8% difference in tg_tps is not a result, and publishing it as one
is how the W4A8/W4A16 ordering came to flip between grids.
"""
if len(summaries) < 2:
return None
first, second = summaries[0], summaries[1]
if not first.best:
return None
gap = abs(first.best - second.best) / abs(first.best) * 100.0
band = max(first.spread_pct, second.spread_pct,
DEFAULT_BANDS_PCT.get(first.metric, first.noise_threshold))
if gap > band:
return None
return (f"{first.tag} and {second.tag} differ by {gap:.1f}% on "
f"{first.metric}, inside this measurement's own spread "
f"({band:.1f}%). That is not a ranking. For a difference this "
f"size use per-op self time, which reproduces to ~5%: "
f"profile_cpu_step.py --model <build> --prompt-len <ctx>")
def format_summary(summaries: list[ArmSummary]) -> str:
if not summaries:
return "no successful observations"
width = max(len(s.tag) for s in summaries)
direction = "min" if summaries[0].lower_is_better else "max"
lines = [f"{'arm'.ljust(width)} {f'best({direction})':>9s} {'median':>9s} "
f"{'spread':>7s} n notes"]
for s in summaries:
notes = []
if s.noisy:
notes.append(f"NOISY (>{s.noise_threshold:.0f}%)")
if s.n_busy:
notes.append(f"{s.n_busy} pass(es) under load")
if s.n < 2:
notes.append("single pass — not a result")
lines.append(f"{s.tag.ljust(width)} {s.best:9.2f} {s.median:9.2f} "
f"{s.spread_pct:6.1f}% {s.n} {'; '.join(notes)}")
warning = indistinguishable(summaries)
if warning:
lines.append("")
lines.append(f" NOT A RANKING: {warning}")
return "\n".join(lines)
def run_grid(arms: list[Arm], *, passes: int, metric: str, timeout: float,
keys: list[str], settle_s: float = 0.0,
busy_load: float = BUSY_LOAD) -> list[Observation]:
"""Run every arm each pass, rotating the starting arm between passes.
Interleaving alone leaves arm 0 permanently first, so whatever the machine
does at the start of a pass always lands on the same arm. Pass ``p`` starts
at arm ``(p - 1) mod n``.
"""
globals()["BUSY_LOAD"] = busy_load
observations: list[Observation] = []
for p in range(1, passes + 1):
shift = (p - 1) % len(arms) if arms else 0
for arm in arms[shift:] + arms[:shift]:
if settle_s:
time.sleep(settle_s)
obs = run_once(arm, p, timeout, keys)
observations.append(obs)
status = "ok" if obs.ok else f"FAILED: {obs.error[:80]}"
value = obs.metrics.get(metric)
shown = f"{value:.2f}" if value is not None else "-"
flag = " [busy]" if obs.busy else ""
print(f" pass {p} {arm.tag:24s} {metric}={shown:>9s} "
f"load {obs.load_before:.1f}->{obs.load_after:.1f}{flag} {status}",
flush=True)
return observations
def main() -> None:
ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
ap.add_argument("--spec", required=True,
help="JSON file: {\"arms\": [{\"tag\":..., \"command\":[...], \"env\":{...}}]}")
ap.add_argument("--passes", type=int, default=3)
ap.add_argument("--metric", default="tg_tps")
ap.add_argument("--keys", default="tg_tps,pp_tps,init_s,rss_anon_mb,rss_mb")
ap.add_argument("--timeout", type=float, default=2400)
ap.add_argument("--settle-s", type=float, default=2.0)
ap.add_argument("--busy-load", type=float, default=None,
help="load average above which a pass is labelled "
"contended; defaults to (target cores + 8), since "
"the run itself contributes about one per core")
ap.add_argument("--noise-threshold", type=float, default=NOISE_THRESHOLD)
ap.add_argument("--out", default=None)
ap.add_argument("--cores", default="",
help="target cores, e.g. '0-23'. Checked for other "
"people's work before the run starts.")
ap.add_argument("--allow-busy", action="store_true",
help="run even if the target cores are already loaded "
"(the result is labelled, not trusted)")
ap.add_argument("--check", default=None,
help="expectations JSON; exit non-zero on a regression "
"outside the metric's measured band")
ap.add_argument("--no-lock", action="store_true",
help="skip the host lock (only if you are on cores no "
"other benchmark uses)")
a = ap.parse_args()
spec = json.loads(open(a.spec).read())
arms = [Arm(tag=x["tag"],
command=x["command"] if isinstance(x["command"], list)
else shlex.split(x["command"]),
env={k: str(v) for k, v in (x.get("env") or {}).items()})
for x in spec["arms"]]
keys = [k.strip() for k in a.keys.split(",") if k.strip()]
cores = parse_cores(a.cores) if a.cores else sorted(os.sched_getaffinity(0))
# Our own run puts ~len(cores) on the load average, so a threshold below
# that labels every pass contended and tells us nothing.
busy_load = a.busy_load if a.busy_load is not None else len(cores) + 8.0
check = preflight(cores)
if not check["clear"]:
detail = []
if check["stolen_pct"] > check["max_stolen_pct"]:
detail.append(f"{check['stolen_pct']}% of the cores' capacity is "
f"someone else's (limit {check['max_stolen_pct']}%)")
if check["busy_cores"]:
detail.append("sustained stragglers: "
+ ", ".join(f"{c}={p}%"
for c, p in check["busy_cores"].items()))
msg = "contended target cores: " + "; ".join(detail)
if not a.allow_busy:
raise SystemExit(
msg + "\n Another job is on these cores; this run would "
"measure the contention, not the code. Wait, pick other "
"cores with --cores, or pass --allow-busy to label the "
"result instead of trusting it.")
print(f"WARNING: {msg} (--allow-busy)")
print(f"{len(arms)} arms x {a.passes} passes, rotating; metric={a.metric}; "
f"{len(cores)} cores, pre-flight "
f"{'clear' if check['clear'] else 'BUSY'} "
f"({check['stolen_pct']}% stolen), "
f"busy-label above load {busy_load:.0f}")
lock = contextlib.nullcontext() if a.no_lock else HostLock(note=a.metric)
with lock:
observations = run_grid(arms, passes=a.passes, metric=a.metric,
timeout=a.timeout, keys=keys,
settle_s=a.settle_s, busy_load=busy_load)
summaries = summarize(observations, a.metric, a.noise_threshold)
print("\n" + format_summary(summaries))
checks = []
if a.check:
checks = check_expectations(summaries, json.loads(open(a.check).read()))
print("\nregression check\n" + format_check(checks))
if a.out:
with open(a.out, "w") as f:
json.dump({
"metric": a.metric, "passes": a.passes,
"lower_is_better": metric_is_lower_better(a.metric),
"provenance": provenance({
"arms": [{"tag": x.tag, "command": x.command, "env": x.env}
for x in arms],
"preflight": check,
"busy_load_threshold": busy_load,
}),
"observations": [vars(o) for o in observations],
"summary": [vars(s) for s in summaries],
"checks": checks,
}, f, indent=2)
print(f"\nwrote {a.out}")
regressions = [c for c in checks if c["status"] == "REGRESSION"]
if regressions:
raise SystemExit(
f"{len(regressions)} regression(s): "
+ ", ".join(f"{c['arm']}/{c['metric']} {c['drift_pct']:+.2f}%"
for c in regressions))
if __name__ == "__main__":
main()
benchmarks/edge_harness/bench_init_memory.py
+297
View File
@@ -0,0 +1,297 @@
"""Init time and resident memory, which decide whether a model can be served
from an edge box at all.
Three quantities, because the plan's items are defined in terms of them and
the earlier version of this file measured none of them:
* **ready_s -- process start to serving.** The old ``init_s`` started its timer
*after* ``import vllm_omni``, so the 14.4 s of import cost the plan wants to
cut was invisible to the benchmark that was supposed to prove the cut. The
clock now starts at process creation, read from ``/proc/self/stat``, so
interpreter start-up and imports are inside it. ``init_s`` is still reported,
unchanged in meaning, so older rows stay comparable.
* **peak_rss_anon_mb -- the high-water mark, sampled.** What an edge device has
to survive is the *peak*, not the steady state: this model loads at 7338 MB
and settles at 3510 MB only after the allocator is asked for its arenas back,
and a device that cannot hold the peak never reaches the release. There is no
kernel high-water mark for RssAnon (``VmHWM`` and ``ru_maxrss`` both track
VmRSS, which counts the mmap'd checkpoint and moved 1-2 GB between identical
runs), so a sampler thread takes it.
* **rss_anon_at_context_mb -- memory once the cache is actually populated.**
The old default generated 8 tokens, so the KV cache was allocated and never
touched; "KV sizing saves only 188 MB" was measured there. ``--context-len``
fills it.
And one correction that changes what the older rows mean. Everything here used
to read ``/proc/self/status``, i.e. **this process only**. That is the right
number only when the model is in-process, which happens when
``VLLM_ENABLE_V1_MULTIPROCESSING=0`` keeps the executor at ``uni``
(``vllm/platforms/cpu.py:273`` promotes it to ``mp`` otherwise). Every earlier
run set that variable, so the published figures are in-process figures. Run the
same benchmark at vLLM's defaults and the weights land in a worker process the
benchmark never looked at: 796 MB "resident" for a model whose weights alone
are 1.4 GB. Resident memory is now summed over the **process tree**, which is
what an edge device actually has to hold, and both numbers are reported so the
older rows remain comparable.
Every row carries provenance, so it can be invalidated later instead of merely
doubted.
"""
import time
_T_IMPORT = time.time()
import argparse, json, os, resource, sys, threading # noqa: E402
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from bench_harness import provenance # noqa: E402
def _rss(field: str, pid: str | int = "self") -> float | None:
try:
with open(f"/proc/{pid}/status") as f:
for line in f:
if line.startswith(field):
return round(int(line.split()[1]) / 1024, 1)
except OSError:
return None
return None
def _descendants(root: int) -> list[int]:
"""Every process under ``root``. vLLM's default CPU path runs the model in
a worker process, so a self-only reading misses the weights entirely."""
children: dict[int, list[int]] = {}
for entry in os.listdir("/proc"):
if not entry.isdigit():
continue
try:
with open(f"/proc/{entry}/stat") as f:
fields = f.read().rsplit(")", 1)[1].split()
ppid = int(fields[1])
except (OSError, IndexError, ValueError):
continue
children.setdefault(ppid, []).append(int(entry))
out, stack = [], [root]
while stack:
pid = stack.pop()
for child in children.get(pid, []):
out.append(child)
stack.append(child)
return out
def _tree_anon_mb(pids: list[int]) -> float:
return round(sum(_rss("RssAnon:", p) or 0.0 for p in pids), 1)
def _process_start_epoch() -> float | None:
"""When this process was created, so imports are inside the measurement."""
try:
with open("/proc/self/stat") as f:
fields = f.read().rsplit(")", 1)[1].split()
starttime_ticks = int(fields[19]) # /proc/self/stat field 22
hz = os.sysconf("SC_CLK_TCK")
with open("/proc/stat") as f:
for line in f:
if line.startswith("btime"):
return int(line.split()[1]) + starttime_ticks / hz
except Exception:
return None
return None
class AnonSampler(threading.Thread):
"""Peak RssAnon has no kernel counter, so poll for it -- across the whole
process tree, because the model may not be in this process."""
def __init__(self, interval: float = 0.02, rescan_every: float = 0.5):
super().__init__(daemon=True)
self.interval, self.rescan_every = interval, rescan_every
self.peak = self.peak_self = 0.0
self.marks: list[dict] = []
self.pids = [os.getpid()]
self.max_procs = 1
# not `_stop`: threading.Thread already defines a private _stop().
self._done = threading.Event()
def _refresh(self) -> None:
me = os.getpid()
self.pids = [me] + _descendants(me)
self.max_procs = max(self.max_procs, len(self.pids))
def _sample(self) -> tuple[float, float]:
mine = _rss("RssAnon:") or 0.0
tree = _tree_anon_mb(self.pids)
self.peak_self = max(self.peak_self, mine)
self.peak = max(self.peak, tree)
return mine, tree
def run(self) -> None:
last_scan = 0.0
while not self._done.is_set():
now = time.monotonic()
if now - last_scan >= self.rescan_every:
self._refresh()
last_scan = now
self._sample()
self._done.wait(self.interval)
def mark(self, name: str) -> float | None:
"""Record a phase boundary with the peak reached so far."""
self._refresh()
mine, tree = self._sample()
self.marks.append({"phase": name, "anon_mb": mine,
"tree_anon_mb": tree, "n_procs": len(self.pids),
"peak_tree_so_far_mb": round(self.peak, 1)})
return mine
def stop(self) -> None:
self._done.set()
self.join(timeout=2.0)
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--model", required=True)
ap.add_argument("--tag", required=True)
ap.add_argument("--max-model-len", type=int, default=2048)
ap.add_argument("--quantization", default=None)
ap.add_argument("--kv-cache-memory-bytes", type=int, default=None,
help="exact KV budget; VLLM_CPU_KVCACHE_SPACE is integer GiB "
"so it cannot express an edge-sized cache")
ap.add_argument("--max-num-batched-tokens", type=int, default=None,
help="vLLM sizes its warm-up forward from this; an edge "
"deployment serving one short stream does not need "
"the server default")
ap.add_argument("--trim", action="store_true",
help="return freed allocator arenas to the OS after init")
ap.add_argument("--generate", type=int, default=8,
help="tokens to generate, to prove the engine really works")
ap.add_argument("--context-len", type=int, default=0,
help="after init, run a request with this many prompt "
"tokens so the KV cache is populated, and report "
"resident memory there. 0 skips it -- and a memory "
"number taken without it says nothing about cache "
"sizing.")
ap.add_argument("--sample-interval", type=float, default=0.02)
ap.add_argument("--out", required=True)
a = ap.parse_args()
proc_start = _process_start_epoch()
sampler = AnonSampler(a.sample_interval)
sampler.start()
anon_start = sampler.mark("before_import")
import vllm_omni # noqa: F401
from vllm import LLM, SamplingParams
anon_after_import = sampler.mark("after_import")
import_s = time.time() - _T_IMPORT
t0 = time.perf_counter()
extra = {}
if a.kv_cache_memory_bytes:
extra["kv_cache_memory_bytes"] = a.kv_cache_memory_bytes
if a.max_num_batched_tokens:
extra["max_num_batched_tokens"] = a.max_num_batched_tokens
llm = LLM(model=a.model, dtype="bfloat16", max_model_len=a.max_model_len,
max_num_seqs=1, quantization=a.quantization,
enable_prefix_caching=False, trust_remote_code=False,
disable_log_stats=True, **extra)
init_s = time.perf_counter() - t0
ready_s = (time.time() - proc_start) if proc_start else None
anon_after_engine = sampler.mark("after_engine")
peak_during_load = round(sampler.peak, 1)
trimmed = None
if a.trim:
# The GPTQ repack unpacks 4-bit weights to int32 (4 bytes per weight)
# and throws the intermediates away, but a caching allocator keeps the
# arenas, so they stay resident. Ask for them back.
import ctypes
released = []
for lib, sym in ((None, "malloc_trim"),
("libtcmalloc_minimal.so.4",
"MallocExtension_ReleaseFreeMemory")):
try:
handle = ctypes.CDLL(lib) if lib else ctypes.CDLL("libc.so.6")
fn = getattr(handle, sym)
fn(0) if sym == "malloc_trim" else fn()
released.append(sym)
except Exception:
pass
trimmed = released
anon_after_trim = sampler.mark("after_trim")
out = llm.generate([{"prompt_token_ids": [10]}],
SamplingParams(temperature=0.0, max_tokens=a.generate,
ignore_eos=True), use_tqdm=False)
n_out = len(out[0].outputs[0].token_ids)
sampler.mark("after_short_generate")
anon_at_context = context_tokens = None
if a.context_len:
# Fill the cache. Without this the cache is allocated and never
# written, so resident memory says nothing about how it was sized.
context_tokens = min(a.context_len, a.max_model_len - 16)
llm.generate([{"prompt_token_ids": [10] * context_tokens}],
SamplingParams(temperature=0.0, max_tokens=16,
ignore_eos=True), use_tqdm=False)
anon_at_context = sampler.mark("after_context_generate")
sampler.stop()
rec = {
"tag": a.tag, "model": a.model,
"ready_s": round(ready_s, 2) if ready_s else None,
"import_s": round(import_s, 2),
"init_s": round(init_s, 2),
# VmRSS counts the mmap'd checkpoint's file pages, which come and go
# with system page-cache pressure: identical configurations measured
# 1.2-2.0 GB apart. RssAnon is the allocated memory -- weights copied
# into tensors, the KV cache, activations -- and is what a memory
# budget is actually about.
# self-only, for comparability with rows taken before the tree fix
"rss_anon_mb": _rss("RssAnon:"), "rss_file_mb": _rss("RssFile:"),
# what the device actually has to hold
"rss_anon_tree_mb": _tree_anon_mb([os.getpid()] + _descendants(os.getpid())),
"peak_rss_anon_mb": round(sampler.peak, 1),
"peak_rss_anon_self_mb": round(sampler.peak_self, 1),
"peak_rss_anon_during_load_mb": peak_during_load,
"n_procs_max": sampler.max_procs,
"anon_before_import_mb": anon_start,
"anon_after_import_mb": anon_after_import,
"anon_after_engine_mb": anon_after_engine,
"anon_after_trim_mb": anon_after_trim,
"rss_anon_at_context_mb": anon_at_context,
"context_tokens": context_tokens,
"phase_marks": sampler.marks,
"trim_calls": trimmed,
"rss_mb": _rss("VmRSS:"), "rss_peak_mb": _rss("VmHWM:"),
"maxrss_mb": round(resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024, 1),
"kvcache_space_gib": os.environ.get("VLLM_CPU_KVCACHE_SPACE", ""),
"kv_cache_memory_bytes": a.kv_cache_memory_bytes,
"max_num_batched_tokens": a.max_num_batched_tokens,
"cache_root": os.environ.get("VLLM_CACHE_ROOT", "(default)"),
"aot_shim": os.environ.get("VLLM_OMNI_AOT_CACHE_SHIM", "1"),
"max_model_len": a.max_model_len,
"tokens_generated": n_out,
"provenance": provenance({
"quantization": a.quantization,
"sample_interval_s": a.sample_interval,
# uni vs mp decides whether the model is in this process at all
"v1_multiprocessing": os.environ.get(
"VLLM_ENABLE_V1_MULTIPROCESSING", "1"),
"memory_scope": "process tree",
}),
}
print(json.dumps(rec, indent=2))
with open(a.out, "a") as f:
f.write(json.dumps(rec) + "\n")
if __name__ == "__main__":
main()
benchmarks/edge_harness/fidelity_kl.py
+179
View File
@@ -0,0 +1,179 @@
"""A fidelity metric that does not saturate.
Greedy divergence over 12 prompts -- "did the quantized build reproduce the
bf16 output exactly?" -- answers 2/12 for W4A16 and 1/12 for W4A8 and cannot
say whether the gap between them is small or enormous. Everything past the
first divergent token is thrown away, and one unlucky token near position 3
scores the same as a build that is wrong everywhere.
This measures the distribution instead, teacher-forced over a corpus, so every
position contributes:
* **top-1 agreement** -- how often the quantized model's argmax matches the
reference's. This is what actually decides whether greedy decoding diverges.
* **KL(reference || candidate)** -- the same quantity ``llama-perplexity
--kl-divergence`` reports, so llama.cpp's Q4_K_M stays comparable against
our builds rather than being compared on a different metric.
* **Delta log-prob of the reference token**, which is the cross-entropy
difference and is what perplexity is built from.
One honest limitation, stated because it changes how the number should be
read: vLLM returns the top-k log-probs per position, not the full
distribution, so KL is computed over the union of the two top-k sets with the
remaining mass as a single tail bucket. ``coverage`` reports the reference
probability mass that fell inside the top-k, and a run whose coverage is low
is reporting a lower bound, not a KL. llama.cpp stores full logits and has no
such limitation; raise ``--topk`` (and vLLM's ``--max-logprobs``) to close the
gap when the two are compared directly.
"""
from __future__ import annotations
import argparse, json, math, os, sys
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from bench_harness import provenance # noqa: E402
def chunk_tokens(token_ids: list[int], size: int, limit: int | None) -> list[list[int]]:
chunks = [token_ids[i:i + size] for i in range(0, len(token_ids), size)]
chunks = [c for c in chunks if len(c) >= 16]
return chunks[:limit] if limit else chunks
def collect(llm, chunks, topk: int):
"""Teacher-forced top-k log-probs at every position of every chunk."""
from vllm import SamplingParams
sp = SamplingParams(temperature=0.0, max_tokens=1, prompt_logprobs=topk)
outs = llm.generate([{"prompt_token_ids": c} for c in chunks], sp,
use_tqdm=False)
positions = []
for out, chunk in zip(outs, chunks):
for i, entry in enumerate(out.prompt_logprobs or []):
if entry is None: # first token has no prediction
continue
positions.append({
"target": chunk[i],
"logprobs": {int(tid): float(lp.logprob) for tid, lp in entry.items()},
})
return positions
def _top1(d: dict[int, float]) -> int:
return max(d.items(), key=lambda kv: kv[1])[0]
def compare(ref: list[dict], cand: list[dict]) -> dict:
if len(ref) != len(cand):
raise SystemExit(f"position count differs: {len(ref)} vs {len(cand)}; "
f"the two runs did not see the same corpus")
n = len(ref)
agree = 0
kls, dlps, coverages = [], [], []
for r, c in zip(ref, cand):
rl, cl = r["logprobs"], c["logprobs"]
if _top1(rl) == _top1(cl):
agree += 1
# Reference distribution over its own top-k, plus a tail bucket.
rp = {t: math.exp(lp) for t, lp in rl.items()}
covered = sum(rp.values())
coverages.append(min(covered, 1.0))
tail_r = max(0.0, 1.0 - covered)
# Candidate probability for the same tokens. A token outside the
# candidate's top-k is bounded by its smallest returned probability,
# which overestimates q and therefore underestimates KL -- the reason
# this is reported as a lower bound when coverage is poor.
floor = math.exp(min(cl.values())) if cl else 1e-12
cq = {t: math.exp(cl[t]) if t in cl else floor for t in rp}
covered_c = sum(cq.values())
tail_c = max(1e-12, 1.0 - min(covered_c, 1.0 - 1e-12))
kl = sum(p * math.log(p / max(cq[t], 1e-12)) for t, p in rp.items() if p > 0)
if tail_r > 0:
kl += tail_r * math.log(tail_r / tail_c)
kls.append(max(kl, 0.0))
target = r["target"]
if target in rl and target in cl:
dlps.append(cl[target] - rl[target])
kls.sort()
return {
"positions": n,
"top1_agreement": round(agree / n, 4) if n else None,
"kl_mean": round(sum(kls) / len(kls), 6) if kls else None,
"kl_median": round(kls[len(kls) // 2], 6) if kls else None,
"kl_p99": round(kls[int(len(kls) * 0.99)], 6) if kls else None,
"delta_logprob_mean": round(sum(dlps) / len(dlps), 6) if dlps else None,
"delta_logprob_n": len(dlps),
"coverage_mean": round(sum(coverages) / len(coverages), 4) if coverages else None,
}
def main() -> None:
ap = argparse.ArgumentParser()
ap.add_argument("--model", required=True)
ap.add_argument("--quantization", default=None)
ap.add_argument("--corpus", required=True, help="UTF-8 text file")
ap.add_argument("--tokenizer", default=None,
help="tokenizer to chunk with; defaults to --model")
ap.add_argument("--chunk", type=int, default=512)
ap.add_argument("--max-chunks", type=int, default=8)
ap.add_argument("--topk", type=int, default=64)
ap.add_argument("--max-model-len", type=int, default=2048)
ap.add_argument("--save-ref", default=None,
help="write this run's distributions as the reference")
ap.add_argument("--ref", default=None,
help="reference file to compare against")
ap.add_argument("--label", default="")
ap.add_argument("--out", required=True)
a = ap.parse_args()
if not a.save_ref and not a.ref:
ap.error("give --save-ref (to record a reference) or --ref (to compare)")
import vllm_omni # noqa: F401
from transformers import AutoTokenizer
from vllm import LLM
tok = AutoTokenizer.from_pretrained(a.tokenizer or a.model,
trust_remote_code=False)
text = open(a.corpus, encoding="utf-8").read()
chunks = chunk_tokens(tok(text).input_ids, a.chunk, a.max_chunks)
if not chunks:
raise SystemExit("corpus too short for one chunk")
llm = LLM(model=a.model, dtype="bfloat16", max_model_len=a.max_model_len,
max_num_seqs=1, quantization=a.quantization,
enable_prefix_caching=False, trust_remote_code=False,
disable_log_stats=True, max_logprobs=a.topk)
positions = collect(llm, chunks, a.topk)
meta = {"label": a.label, "model": a.model, "quantization": a.quantization,
"corpus": os.path.abspath(a.corpus), "chunk": a.chunk,
"n_chunks": len(chunks), "topk": a.topk,
"provenance": provenance({"metric": "fidelity_kl"})}
if a.save_ref:
with open(a.save_ref, "w") as f:
json.dump({**meta, "positions": positions}, f)
print(f"wrote reference: {len(positions)} positions -> {a.save_ref}")
result = {**meta, "role": "reference", "n_positions": len(positions)}
else:
ref = json.loads(open(a.ref).read())
if ref["chunk"] != a.chunk or ref["n_chunks"] != len(chunks):
raise SystemExit("reference was taken on a different corpus slicing")
result = {**meta, "role": "candidate",
"reference_model": ref["model"],
**compare(ref["positions"], positions)}
print(json.dumps({k: v for k, v in result.items()
if k != "provenance"}, indent=2))
with open(a.out, "w") as f:
json.dump(result, f, indent=2)
if __name__ == "__main__":
main()
benchmarks/edge_harness/memory_ledger.py
+224
View File
@@ -0,0 +1,224 @@
"""Where the resident memory actually is, by attribution rather than subtraction.
The report's breakdown was arithmetic: weights, imports, and "unexplained
~1.6 GB" left over after subtracting the first two from the total. Three
problems with that. It assigned a bf16 embedding to a checkpoint whose config
says ``quantize_embedding: true``; it left the prefill duplicate out of the
weight line entirely; and a remainder computed from two wrong terms is not a
measurement of anything.
So this counts instead:
* **Unique tensor storages.** Walking ``named_parameters`` double-counts tied
weights -- Spark ties its embedding to its head, and the quantized path ties
five tensors at once -- so storages are keyed by ``data_ptr`` and counted
once, and the tied ones are reported separately rather than silently
dropped.
* **The KV cache**, found by attribute rather than assumed from a formula --
and split into *allocated* and *resident*. A cache that has been allocated
but never written costs address space, not memory, so counting its size as
attributed memory would hide the remainder rather than explain it.
Residency comes from ``mincore(2)``, which reports per-page whether a
mapping is actually in core.
* **Every mapping in /proc/self/smaps**, bucketed into heap, anonymous, file
and the rest. ``smaps_rollup`` cannot do this: it is an aggregate, which is
the whole thing this is trying to break apart.
What it cannot do is attribute native allocations that no longer belong to a
live tensor -- freed allocator arenas, inductor's generated code, transient
buffers the caching allocator is holding. Those land in "unattributed anon",
and the point of the exercise is to see how big that really is once the
attributable part is counted properly.
"""
from __future__ import annotations
import argparse, collections, json, os, sys
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from bench_harness import provenance # noqa: E402
def resident_mb(ptr: int, nbytes: int) -> float:
"""How much of a mapping is actually in core, via mincore(2).
Allocated is not resident: an untouched KV cache is address space. Without
this the ledger would "explain" memory the process never used.
"""
import ctypes
if nbytes <= 0:
return 0.0
page = os.sysconf("SC_PAGE_SIZE")
start = ptr - (ptr % page) # mincore needs a page boundary
length = nbytes + (ptr - start)
pages = (length + page - 1) // page
vec = (ctypes.c_ubyte * pages)()
libc = ctypes.CDLL("libc.so.6", use_errno=True)
if libc.mincore(ctypes.c_void_p(start), ctypes.c_size_t(length), vec) != 0:
return -1.0 # unknown, rather than a guess
return round(sum(v & 1 for v in vec) * page / 2**20, 1)
def _rss(field: str) -> float:
with open("/proc/self/status") as f:
for line in f:
if line.startswith(field):
return round(int(line.split()[1]) / 1024, 1)
return 0.0
def smaps_breakdown() -> dict[str, float]:
"""Anonymous resident bytes per mapping kind. Needs smaps, not rollup."""
buckets: collections.Counter = collections.Counter()
name = "[anon]"
try:
with open("/proc/self/smaps") as f:
for line in f:
if not line[:1].isspace() and "-" in line.split()[0]:
parts = line.split()
path = parts[5] if len(parts) > 5 else "[anon]"
if path.startswith("["):
name = path
elif path.endswith(".so") or ".so." in path:
name = "shared libs"
elif path:
name = "file-backed"
else:
name = "[anon]"
elif line.startswith("Anonymous:"):
buckets[name] += int(line.split()[1])
except OSError:
return {}
return {k: round(v / 1024, 1) for k, v in buckets.most_common()}
def storage_ledger(model) -> dict:
"""Unique tensor storages, by role, counting each byte once."""
import torch
seen: dict[int, int] = {}
biggest: list = []
by_role: collections.Counter = collections.Counter()
counts: collections.Counter = collections.Counter()
tied_bytes = 0
def role_of(name: str) -> str:
low = name.lower()
if "embed" in low or "lm_head" in low:
return "embedding / head"
for key in ("int4pack", "halves", "scales_and_zeros", "group_scale",
"group_zero", "weight_scale", "weight_zero", "weight_packed"):
if key in low:
return f"quant:{key}"
if "weight" in low:
return "other weights"
return "other"
for name, tensor in list(model.named_parameters()) + list(model.named_buffers()):
if not isinstance(tensor, torch.Tensor) or tensor.device.type != "cpu":
continue
try:
ptr = tensor.untyped_storage().data_ptr()
nbytes = tensor.untyped_storage().nbytes()
except Exception:
ptr, nbytes = tensor.data_ptr(), tensor.numel() * tensor.element_size()
if ptr in seen:
tied_bytes += nbytes # tied / shared: counted once already
continue
seen[ptr] = nbytes
by_role[role_of(name)] += nbytes
counts[role_of(name)] += 1
if nbytes >= 16 * 2**20:
biggest.append((name, nbytes, resident_mb(ptr, nbytes),
tuple(tensor.shape), tensor.dtype))
biggest.sort(key=lambda t: -t[1])
return {
"unique_storage_mb": round(sum(seen.values()) / 2**20, 1),
"tied_duplicate_mb": round(tied_bytes / 2**20, 1),
"by_role_mb": {k: round(v / 2**20, 1) for k, v in by_role.most_common()},
"tensors_by_role": dict(counts),
# "other" is where an unclassified tensor hides, so name the largest
"largest_tensors": [
{"name": n, "mb": round(b / 2**20, 1), "resident_mb": r,
"shape": list(sh), "dtype": str(dt)}
for n, b, r, sh, dt in biggest[:15]
],
}
def main() -> None:
ap = argparse.ArgumentParser()
ap.add_argument("--model", required=True)
ap.add_argument("--quantization", default=None)
ap.add_argument("--max-model-len", type=int, default=2048)
ap.add_argument("--tag", default="ledger")
ap.add_argument("--out", required=True)
a = ap.parse_args()
anon_before = _rss("RssAnon:")
import vllm_omni # noqa: F401
import torch
from vllm import LLM
anon_after_import = _rss("RssAnon:")
llm = LLM(model=a.model, dtype="bfloat16", max_model_len=a.max_model_len,
max_num_seqs=1, quantization=a.quantization,
enable_prefix_caching=False, trust_remote_code=False,
disable_log_stats=True)
# The model only exists in this process when the executor stayed uniproc;
# say so rather than reporting a ledger of nothing.
try:
engine = llm.llm_engine
runner = engine.engine_core.engine_core.model_executor.driver_worker.model_runner
model = runner.model
except AttributeError as exc:
raise SystemExit(
f"cannot reach the model in this process ({exc}). Set "
f"VLLM_ENABLE_V1_MULTIPROCESSING=0 so the executor stays uniproc; "
f"otherwise the weights are in a worker and there is nothing here "
f"to attribute."
)
ledger = storage_ledger(model)
kv_mb = kv_resident = 0.0
try:
caches = runner.kv_caches
seen = set()
for c in caches:
if isinstance(c, torch.Tensor) and c.data_ptr() not in seen:
seen.add(c.data_ptr())
nbytes = c.numel() * c.element_size()
kv_mb += nbytes / 2**20
r = resident_mb(c.data_ptr(), nbytes)
kv_resident += max(r, 0.0)
except Exception:
kv_mb = kv_resident = -1.0
anon_total = _rss("RssAnon:")
# Only resident KV counts: allocated-but-untouched cache is address space.
attributed = (ledger["unique_storage_mb"] + max(kv_resident, 0.0)
+ anon_after_import)
rec = {
"tag": a.tag, "model": a.model, "quantization": a.quantization,
"anon_before_import_mb": anon_before,
"anon_after_import_mb": anon_after_import,
"anon_total_mb": anon_total,
"kv_cache_allocated_mb": round(kv_mb, 1),
"kv_cache_resident_mb": round(kv_resident, 1),
**ledger,
"attributed_mb": round(attributed, 1),
"unattributed_anon_mb": round(anon_total - attributed, 1),
"smaps_anon_mb": smaps_breakdown(),
"provenance": provenance({"scope": "in-process ledger"}),
}
print(json.dumps({k: v for k, v in rec.items() if k != "provenance"}, indent=2))
with open(a.out, "w") as f:
json.dump(rec, f, indent=2)
if __name__ == "__main__":
main()
benchmarks/edge_harness/profile_cpu_step.py
+158
View File
@@ -0,0 +1,158 @@
"""Where does a single-token CPU decode step go?
Runs the real engine and splits decode self-CPU time into linear GEMMs,
attention, elementwise and everything else.
Two things this deliberately does *not* do any more, because both produced
published numbers that meant something other than what they said:
* **It does not divide a whole ``generate`` by ``tg``.** That call contains the
prefill step and the sampling around it, so the quotient is not a decode
step. Decode is measured the way the wall clock measures it -- profile a
prefill-only run and a full run, subtract, divide by ``tg - 1`` -- so the
profile and the wall clock are finally the same quantity.
* **It does not profile a one-token prompt by default and call the result
representative.** Spark's sliding layers have a 512-token window; at a
1-token prompt nothing in the KV cache is being read, so attention cost here
says nothing about attention cost at context. ``--prompt-len`` sets it, and
the length is recorded in the artifact.
Every result carries provenance (model path, context, build SHAs, thread set),
because two saved profiles in this tree are indistinguishable from each other
and one set describes a checkpoint later found to be wrong.
"""
import argparse, json, os, sys, time
from collections import defaultdict
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from bench_harness import provenance # noqa: E402
GEMM = ("addmm", "mm", "linear", "matmul", "bmm", "packed_linear",
"int8_scaled_mm", "int4pack", "int8pack", "onednn", "dnnl")
ATTN = ("attention", "paged", "flash", "softmax", "reshape_and_cache")
def bucket_of(key: str) -> str:
n = key.lower()
if any(g in n for g in GEMM):
return "gemm"
if any(s in n for s in ATTN):
return "attn"
if n.startswith("aten::") or n.startswith("_c::"):
return "other_op"
return "framework"
def profile_once(llm, prompt, sp):
"""Self-CPU microseconds and call counts, per op key, for one generate."""
from torch.profiler import profile, ProfilerActivity
with profile(activities=[ProfilerActivity.CPU]) as prof:
llm.generate([prompt], sp, use_tqdm=False)
us, calls = defaultdict(float), defaultdict(int)
for e in prof.key_averages():
us[e.key] += e.self_cpu_time_total
calls[e.key] += e.count
return us, calls
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--model", default="XHToken/Spark-X2.5-1.7B")
ap.add_argument("--dtype", default="bfloat16")
ap.add_argument("--quantization", default=None)
ap.add_argument("--tg", type=int, default=64)
ap.add_argument("--prompt-len", type=int, default=1,
help="Prompt tokens. Use >=512 to fill Spark's sliding "
"window; the default of 1 measures an empty cache.")
ap.add_argument("--max-model-len", type=int, default=2048)
ap.add_argument("--enforce-eager", action="store_true")
ap.add_argument("--reps", type=int, default=3)
ap.add_argument("--out", default="cpu_step_profile.json")
a = ap.parse_args()
if a.tg < 2:
ap.error("--tg must be at least 2: decode is measured by difference")
import vllm_omni # noqa: F401
import torch
from vllm import LLM, SamplingParams
llm = LLM(model=a.model, dtype=a.dtype, max_model_len=a.max_model_len,
enforce_eager=a.enforce_eager, max_num_seqs=1,
quantization=a.quantization, enable_prefix_caching=False,
trust_remote_code=False, disable_log_stats=True)
prompt = {"prompt_token_ids": [10] * a.prompt_len}
sp1 = SamplingParams(temperature=0.0, max_tokens=1, ignore_eos=True)
spN = SamplingParams(temperature=0.0, max_tokens=a.tg, ignore_eos=True)
decode_steps = a.tg - 1
llm.generate([prompt], spN, use_tqdm=False) # warm
def wall(sp):
best = float("inf")
for _ in range(a.reps):
t = time.perf_counter()
llm.generate([prompt], sp, use_tqdm=False)
best = min(best, time.perf_counter() - t)
return best
t1, tN = wall(sp1), wall(spN)
step_ms = (tN - t1) / decode_steps * 1e3
# Same subtraction under the profiler, so the two agree by construction.
us1, calls1 = profile_once(llm, prompt, sp1)
usN, callsN = profile_once(llm, prompt, spN)
rows, buckets = [], defaultdict(float)
for key in set(usN) | set(us1):
delta_us = usN.get(key, 0.0) - us1.get(key, 0.0)
delta_calls = callsN.get(key, 0) - calls1.get(key, 0)
if delta_us <= 0 and delta_calls <= 0:
continue # prefill-only work; not part of a decode step
b = bucket_of(key)
buckets[b] += delta_us
rows.append((key, b, delta_us / 1e3, delta_calls))
rows.sort(key=lambda r: -r[2])
prefill_total_ms = sum(us1.values()) / 1e3
res = {
"phase": "decode",
"tag": f"{a.dtype}/{a.quantization}/eager={a.enforce_eager}",
"wall_step_ms": round(step_ms, 3),
"tps": round(1e3 / step_ms, 2),
"decode_steps": decode_steps,
"profiled_total_ms_per_step": round(
sum(buckets.values()) / 1e3 / decode_steps, 3),
"buckets_ms_per_step": {
k: round(v / 1e3 / decode_steps, 3)
for k, v in sorted(buckets.items(), key=lambda x: -x[1])},
"top": [{"op": r[0], "bucket": r[1],
"ms_per_step": round(r[2] / decode_steps, 4),
"calls_per_step": round(r[3] / decode_steps, 1)}
for r in rows[:35]],
"prefill": {
"wall_ms": round(t1 * 1e3, 3),
"profiled_self_ms": round(prefill_total_ms, 3),
"prompt_tokens": a.prompt_len,
},
"provenance": provenance({
"model": a.model,
"model_path": os.path.realpath(a.model) if os.path.exists(a.model)
else a.model,
"quantization": a.quantization,
"dtype": a.dtype,
"enforce_eager": a.enforce_eager,
"prompt_len": a.prompt_len,
"tg": a.tg,
"context_at_end": a.prompt_len + a.tg,
"max_model_len": a.max_model_len,
"torch_threads": torch.get_num_threads(),
}),
}
print(json.dumps(res, indent=2))
with open(a.out, "w") as f:
json.dump(res, f, indent=2)
if __name__ == "__main__":
main()
benchmarks/edge_harness/quantize_gptq_int4.py
+123
View File
@@ -0,0 +1,123 @@
"""Write a GPTQ-format W4A8 checkpoint for Spark-X2.5.
vLLM 0.28's CPU backend *does* have a 4-bit kernel -- it is just not reachable
under the names a quantization search turns up. `CPUWNA16LinearKernel`
(`vllm/model_executor/kernels/linear/mixed_precision/cpu.py`) dispatches GPTQ
and AWQ checkpoints to `int4_scaled_mm_cpu`, an AMX W4A8 GEMM, whenever the
activations are bf16, the ISA has AMX tiles and the checkpoint has no
activation reordering. Four-bit weights times eight-bit activations is
exactly the arithmetic llama.cpp's Q4_K_M does (Q4_K weights, Q8_K
activations), so this is the honest precision match.
The kernel pins the zero point at 8 regardless of what the checkpoint says
(`_process_gptq_weights_w4a8` builds its own), so the grid has to be
symmetric: w = (code - 8) * scale. `--search` sweeps the scale per group and
keeps the lowest squared error, which recovers most of what symmetry costs
against an asymmetric min/max grid.
Emitted per quantized linear, in the layout `AutoGPTQLinearMethod` creates:
<name>.qweight int32 [K // 8, N] 4-bit codes packed along the input dim
<name>.qzeros int32 [K // g, N // 8]
<name>.scales bf16 [K // g, N]
<name>.g_idx int32 [K]
Tied embeddings are untied here: the head is quantized (it is the single
largest read in a decode step) while the vocab table stays bf16, since an
embedding lookup reads one row and gains nothing from being 4-bit.
"""
import argparse, json, shutil
from pathlib import Path
from safetensors.torch import load_file, save_file
from vllm_omni.model_executor.layers.quantization.gptq_cpu_export import gptq_tensors
# Split the fused projection here rather than in the model: the GPTQ layout
# transposes N to dim 1, and the loader's stacked mapping already knows how to
# put q/k/v back together.
FUSED_QKV = "self_attn.q_k_v_proj.weight"
QUANT_SUFFIXES = (
"self_attn.out_proj.weight",
"mlp.gate_proj.weight",
"mlp.up_proj.weight",
"mlp.down_proj.weight",
)
EMBED_NAME = "model.embedding.weight"
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--src", required=True)
ap.add_argument("--dst", required=True)
ap.add_argument("--group", type=int, default=64)
ap.add_argument("--no-search", action="store_true")
ap.add_argument("--no-lm-head", action="store_true")
a = ap.parse_args()
src, dst = Path(a.src), Path(a.dst)
dst.mkdir(parents=True, exist_ok=True)
state = {}
for shard in sorted(src.glob("*.safetensors")):
state.update(load_file(str(shard)))
cfg = json.loads((src / "config.json").read_text())
q_dim = cfg["num_attention_heads"] * cfg["head_dim"]
kv_dim = cfg["num_key_value_heads"] * cfg["head_dim"]
def emit(out, base, w, report):
qw, qz, sc, gi, snr = gptq_tensors(w, a.group, not a.no_search)
out[f"{base}.qweight"] = qw
out[f"{base}.qzeros"] = qz
out[f"{base}.scales"] = sc
out[f"{base}.g_idx"] = gi
report.append({"name": base, "shape": list(w.shape), "snr_db": round(snr, 2)})
out, report = {}, []
for name, t in sorted(state.items()):
if name.endswith(FUSED_QKV):
base = name[: -len(".weight")].replace("q_k_v_proj", "")
for shard, part in zip(("q_proj", "k_proj", "v_proj"),
t.split([q_dim, kv_dim, kv_dim], dim=0)):
emit(out, base + shard, part, report)
elif name.endswith(QUANT_SUFFIXES):
emit(out, name[: -len(".weight")], t, report)
elif name == EMBED_NAME:
out[name] = t
if not a.no_lm_head:
emit(out, "lm_head", t, report)
else:
out[name] = t
save_file(out, str(dst / "model.safetensors"), metadata={"format": "pt"})
for f in ("tokenizer.json", "tokenizer_config.json", "special_tokens_map.json",
"vocab.json", "merges.txt", "generation_config.json",
"chat_template.jinja"):
if (src / f).exists():
shutil.copy(src / f, dst / f)
# The head is quantized separately from the table it used to be tied to,
# so the tie has to come off or vLLM would skip loading lm_head.*
cfg["tie_word_embeddings"] = a.no_lm_head
cfg["quantization_config"] = {
"quant_method": "gptq", "bits": 4, "group_size": a.group,
"desc_act": False, "sym": True, "lm_head": not a.no_lm_head,
"checkpoint_format": "gptq",
}
(dst / "config.json").write_text(json.dumps(cfg, indent=2))
snrs = [r["snr_db"] for r in report]
summary = {
"group": a.group, "search": not a.no_search,
"tensors_quantized": len(report),
"snr_db_min": round(min(snrs), 2), "snr_db_max": round(max(snrs), 2),
"snr_db_mean": round(sum(snrs) / len(snrs), 2),
"bits_per_weight": round(4 + 16 / a.group + 4 / a.group, 3),
"per_tensor": report,
}
(dst / "quantization_report.json").write_text(json.dumps(summary, indent=2))
print(json.dumps({k: v for k, v in summary.items() if k != "per_tensor"}, indent=2))
if __name__ == "__main__":
main()
benchmarks/edge_harness/quantize_int4.py
+138
View File
@@ -0,0 +1,138 @@
"""Write a `cpu_int4` W4A16 checkpoint for Spark-X2.5.
Group-wise 4-bit weights at the same bits-per-weight budget as llama.cpp's
Q4_K_M (4.5 at group 64), so the CPU decode comparison is precision-matched
rather than a quantization contest.
Round-to-nearest over the group's min/max is the obvious grid, but it is not
the best one: an outlier stretches the range and costs every other weight in
the group. `--search` sweeps a shrink factor per group and keeps the one with
the lowest squared error, which is the same idea as llama.cpp's
`make_qkx2_quants` and buys about a dB.
The head-wise gate (g_proj) stays in full precision: eight outputs that each
scale a whole attention head, 0.001% of the weights.
"""
import argparse, json, shutil
from pathlib import Path
import torch
from safetensors.torch import load_file, save_file
QUANT_SUFFIXES = (
"self_attn.q_k_v_proj.weight",
"self_attn.out_proj.weight",
"mlp.gate_proj.weight",
"mlp.up_proj.weight",
"mlp.down_proj.weight",
)
EMBED_NAME = "model.embedding.weight"
SHRINKS = (1.0, 0.98, 0.96, 0.94, 0.92, 0.90, 0.88, 0.85)
def _pack(q: torch.Tensor) -> torch.Tensor:
q = q.to(torch.uint8)
return (q[:, 0::2] | (q[:, 1::2] << 4)).contiguous()
def quantize(w: torch.Tensor, group: int, search: bool):
n, k = w.shape
wf = w.float().reshape(n, k // group, group)
wmin, wmax = wf.amin(dim=-1), wf.amax(dim=-1)
span = (wmax - wmin).clamp(min=1e-9)
best_err = None
best = None
for f in (SHRINKS if search else (1.0,)):
# Shrink the grid symmetrically about the group mid-point: the
# extremes lose a little, everything between them gains.
scale = span * f / 15.0
lo = wmin + span * (1.0 - f) * 0.5
q = ((wf - lo.unsqueeze(-1)) / scale.unsqueeze(-1)).round().clamp_(0, 15)
err = ((q * scale.unsqueeze(-1) + lo.unsqueeze(-1)) - wf).pow_(2).sum(-1)
if best_err is None:
best_err, best = err, (q, scale, lo)
else:
take = err < best_err
best_err = torch.where(take, err, best_err)
t = take.unsqueeze(-1)
best = (
torch.where(t, q, best[0]),
torch.where(take, scale, best[1]),
torch.where(take, lo, best[2]),
)
q, scale, lo = best
zero = lo + 8.0 * scale
deq = (q - 8.0) * scale.unsqueeze(-1) + zero.unsqueeze(-1)
ref = wf
snr = float(20 * torch.log10(ref.norm() / (ref - deq).norm().clamp(min=1e-12)))
return _pack(q.reshape(n, k)), scale.to(torch.bfloat16), zero.to(torch.bfloat16), snr
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--src", required=True)
ap.add_argument("--dst", required=True)
ap.add_argument("--group", type=int, default=64)
ap.add_argument("--no-search", action="store_true")
ap.add_argument("--keep-embedding", action="store_true",
help="leave the tied vocab table in bf16")
a = ap.parse_args()
src, dst = Path(a.src), Path(a.dst)
dst.mkdir(parents=True, exist_ok=True)
state = {}
for shard in sorted(src.glob("*.safetensors")):
state.update(load_file(str(shard)))
targets = list(QUANT_SUFFIXES)
out, report = {}, []
for name, t in sorted(state.items()):
is_embed = name == EMBED_NAME and not a.keep_embedding
if name.endswith(tuple(targets)) or is_embed:
packed, scale, zero, snr = quantize(t, a.group, not a.no_search)
base = name[: -len(".weight")]
out[f"{base}.weight_packed"] = packed
out[f"{base}.weight_scale"] = scale
out[f"{base}.weight_zero"] = zero
report.append({"name": name, "shape": list(t.shape), "snr_db": round(snr, 2)})
else:
out[name] = t
save_file(out, str(dst / "model.safetensors"), metadata={"format": "pt"})
for f in ("tokenizer.json", "tokenizer_config.json", "special_tokens_map.json",
"vocab.json", "merges.txt", "generation_config.json",
"chat_template.jinja"):
if (src / f).exists():
shutil.copy(src / f, dst / f)
cfg = json.loads((src / "config.json").read_text())
cfg["quantization_config"] = {
"quant_method": "cpu_int4",
"bits": 4,
"group_size": a.group,
"quantize_embedding": not a.keep_embedding,
"prefill_dequant": True,
# Only single-row steps may take the tinygemm kernel; see the AMX
# tile-state note in vllm_omni/.../cpu_int4.py.
"dequant_threshold": 2,
"ignore": [],
}
(dst / "config.json").write_text(json.dumps(cfg, indent=2))
snrs = [r["snr_db"] for r in report]
summary = {
"group": a.group, "search": not a.no_search,
"tensors_quantized": len(report),
"snr_db_min": round(min(snrs), 2), "snr_db_max": round(max(snrs), 2),
"snr_db_mean": round(sum(snrs) / len(snrs), 2),
"bits_per_weight": 4 + 2 * 16 / a.group,
"per_tensor": report,
}
(dst / "quantization_report.json").write_text(json.dumps(summary, indent=2))
print(json.dumps({k: v for k, v in summary.items() if k != "per_tensor"}, indent=2))
if __name__ == "__main__":
main()
benchmarks/edge_harness/test_bench_harness.py
+308
View File
@@ -0,0 +1,308 @@
"""Tests for the benchmark harness's judgement, not its plumbing.
The summariser is where a shared machine turns into a wrong claim, so that is
what is pinned here: spread reported, noisy arms flagged, single passes
refused as results, contended passes labelled.
"""
import json
import os
import pytest
from bench_harness import (
Arm,
HostLock,
check_expectations,
indistinguishable,
Observation,
core_utilization,
format_summary,
metric_is_lower_better,
parse_cores,
preflight,
provenance,
run_grid,
summarize,
)
def _obs(tag, p, value, load=1.0, ok=True):
return Observation(tag, p, {"tg_tps": value} if ok else {}, load, load, 1.0, ok)
def test_reports_best_median_and_spread():
s = summarize([_obs("a", 1, 80), _obs("a", 2, 100), _obs("a", 3, 90)], "tg_tps")[0]
assert s.best == 100 and s.median == 90
assert s.spread_pct == pytest.approx(20.0)
assert s.n == 3
def test_flags_an_arm_whose_passes_disagree():
"""A 30% swing between identical runs is the failure mode that produced
three wrong conclusions in this study."""
noisy = summarize([_obs("a", 1, 60), _obs("a", 2, 90)], "tg_tps")[0]
assert noisy.noisy
clean = summarize([_obs("b", 1, 88), _obs("b", 2, 90)], "tg_tps", noise_threshold=8.0)[0]
assert not clean.noisy
def test_single_pass_is_not_presented_as_a_result():
text = format_summary(summarize([_obs("a", 1, 80)], "tg_tps"))
assert "single pass" in text
def test_contended_passes_are_counted_not_hidden():
rows = [_obs("a", 1, 80, load=1.0), _obs("a", 2, 40, load=40.0)]
s = summarize(rows, "tg_tps")[0]
assert s.n_busy == 1
assert "under load" in format_summary([s])
def test_failed_runs_do_not_enter_the_summary():
rows = [_obs("a", 1, 80), _obs("a", 2, 0, ok=False)]
s = summarize(rows, "tg_tps")[0]
assert s.n == 1
def test_arm_order_rotates_between_passes(monkeypatch):
"""Interleaving alone leaves arm 0 permanently first, so whatever the host
is doing at the start of a pass always lands on the same arm. This test
previously asserted the un-rotated order, which is how that survived."""
order = []
def fake_run(arm, pass_no, timeout, keys):
order.append((pass_no, arm.tag))
return _obs(arm.tag, pass_no, 50.0)
monkeypatch.setattr("bench_harness.run_once", fake_run)
arms = [Arm("a", ["true"]), Arm("b", ["true"]), Arm("c", ["true"])]
run_grid(arms, passes=3, metric="tg_tps", timeout=1, keys=["tg_tps"])
assert order == [(1, "a"), (1, "b"), (1, "c"),
(2, "b"), (2, "c"), (2, "a"),
(3, "c"), (3, "a"), (3, "b")]
# every arm still runs once per pass
for p in (1, 2, 3):
assert sorted(t for q, t in order if q == p) == ["a", "b", "c"]
def test_best_is_the_minimum_for_a_memory_metric():
"""Taking the maximum for every metric published the worst memory pass as
the headline figure."""
rows = [Observation("a", 1, {"rss_anon_mb": 2802.0}, 1.0, 1.0, 1.0, True),
Observation("a", 2, {"rss_anon_mb": 2814.0}, 1.0, 1.0, 1.0, True)]
s = summarize(rows, "rss_anon_mb")[0]
assert s.lower_is_better is True
assert s.best == 2802.0
def test_best_arm_sorts_first_in_both_directions():
mem = [Observation("small", 1, {"rss_anon_mb": 2802.0}, 1.0, 1.0, 1.0, True),
Observation("big", 1, {"rss_anon_mb": 3542.0}, 1.0, 1.0, 1.0, True)]
assert [s.tag for s in summarize(mem, "rss_anon_mb")] == ["small", "big"]
assert [s.tag for s in summarize([_obs("slow", 1, 40.0), _obs("fast", 1, 86.6)],
"tg_tps")] == ["fast", "slow"]
@pytest.mark.parametrize("metric,lower", [
("tg_tps", False), ("pp_tps", False),
("rss_anon_mb", True), ("peak_rss_anon_mb", True),
("init_s", True), ("ready_s", True), ("step_ms", True),
])
def test_metric_direction_is_read_from_the_name(metric, lower):
assert metric_is_lower_better(metric) is lower
def test_a_custom_noise_threshold_does_not_leak_into_the_next_call():
"""The threshold used to be written into a module global after the rows
were built, so one call's setting silently changed the next one's."""
loose = summarize([_obs("a", 1, 60), _obs("a", 2, 90)], "tg_tps",
noise_threshold=50.0)[0]
assert not loose.noisy
strict = summarize([_obs("a", 1, 60), _obs("a", 2, 90)], "tg_tps")[0]
assert strict.noisy
def test_provenance_records_what_makes_a_number_retractable():
p = provenance({"model": "spark2_5-int4-g128"})
assert p["model"] == "spark2_5-int4-g128"
assert set(p) >= {"recorded_at", "host", "commits", "env", "n_affinity"}
assert set(p["commits"]) == {"embedding_infer", "vllm_omni", "vllm"}
def test_scrapes_the_last_json_object_a_run_printed():
from bench_harness import _scrape
stdout = 'INFO noise {"not": "it"}\nmore logs\n{"tg_tps": 86.6, "pp_tps": 3347}\n'
assert _scrape(stdout, ["tg_tps", "pp_tps"]) == {"tg_tps": 86.6, "pp_tps": 3347.0}
def test_end_to_end_against_a_real_subprocess(tmp_path):
"""The arm really is a subprocess, so a trivial one must round-trip."""
script = tmp_path / "emit.py"
script.write_text('print("noise"); print(\'{"tg_tps": 42.0}\')')
arm = Arm("trivial", ["python3", str(script)])
obs = run_grid([arm], passes=2, metric="tg_tps", timeout=60,
keys=["tg_tps"], settle_s=0)
assert all(o.ok for o in obs)
s = summarize(obs, "tg_tps")[0]
assert s.best == 42.0 and s.n == 2 and s.spread_pct == 0.0
def test_a_failing_arm_is_recorded_with_its_error(tmp_path):
arm = Arm("broken", ["python3", "-c", "import sys; sys.exit(3)"])
obs = run_grid([arm], passes=1, metric="tg_tps", timeout=60,
keys=["tg_tps"], settle_s=0)
assert obs[0].ok is False
assert summarize(obs, "tg_tps") == []
# ------------------------------------------------------- contention guard
def test_parse_cores_handles_ranges_and_singletons():
assert parse_cores("0-3,8") == [0, 1, 2, 3, 8]
assert parse_cores("5") == [5]
assert parse_cores("2-1") == [] # empty range, not an error
assert parse_cores("") == []
def test_core_utilization_reports_the_requested_cores():
"""Load average is a whole-host number and this host has 224 cores, so it
says nothing about the 24 a run is pinned to."""
usage = core_utilization([0, 1], window_s=0.05)
assert set(usage) <= {0, 1}
assert all(0.0 <= v <= 100.0 for v in usage.values())
def test_preflight_refuses_a_sustained_straggler(monkeypatch):
"""One heavily loaded core gates every OpenMP barrier, so it costs far more
than its 1/24 share of capacity."""
monkeypatch.setattr("bench_harness.core_utilization",
lambda cores, window_s=0.75: {0: 3.0, 1: 91.0})
check = preflight([0, 1])
assert check["clear"] is False
assert check["busy_cores"] == {1: 91.0}
def test_preflight_passes_on_idle_cores(monkeypatch):
monkeypatch.setattr("bench_harness.core_utilization",
lambda cores, window_s=0.75: {0: 1.0, 1: 2.0})
assert preflight([0, 1])["clear"] is True
def test_preflight_ignores_a_transient_spike(monkeypatch):
"""A flat per-core threshold over one window refused three consecutive
launches on cores that were idle seconds earlier. Load is the minimum of
two windows, so a blip in one of them is not load."""
windows = iter([{0: 1.0, 1: 30.0}, {0: 1.0, 1: 1.0}])
monkeypatch.setattr("bench_harness.core_utilization",
lambda cores, window_s=0.75: next(windows))
check = preflight([0, 1])
assert check["clear"] is True
assert check["utilization"][1] == 1.0
def test_preflight_refuses_when_aggregate_capacity_is_stolen(monkeypatch):
"""No single straggler, but a third of every core is someone else's."""
monkeypatch.setattr("bench_harness.core_utilization",
lambda cores, window_s=0.75: {c: 33.0 for c in cores})
check = preflight(list(range(24)))
assert check["clear"] is False
assert check["stolen_pct"] == 33.0
assert check["busy_cores"] == {} # none individually over 50%
def test_one_blipping_core_out_of_many_does_not_block_a_run(monkeypatch):
"""1 core of 24 at 30% is ~1.25% of capacity; refusing on that makes the
guard useless rather than strict."""
monkeypatch.setattr("bench_harness.core_utilization",
lambda cores, window_s=0.75: {c: (30.0 if c == 11 else 1.0)
for c in cores})
assert preflight(list(range(24)))["clear"] is True
def test_host_lock_excludes_a_second_benchmark(tmp_path):
"""Two sessions benchmarking the same cores wrecked several measurements."""
path = str(tmp_path / "bench.lock")
with HostLock(path, note="first"):
with pytest.raises(SystemExit) as excinfo:
with HostLock(path, note="second"):
pass
assert "another benchmark holds" in str(excinfo.value)
# released again afterwards
with HostLock(path, note="third"):
pass
def test_host_lock_records_who_holds_it(tmp_path):
path = str(tmp_path / "bench.lock")
with HostLock(path, note="peak_rss_anon_mb"):
held = open(path).read()
assert "peak_rss_anon_mb" in held and str(os.getpid()) in held
# ------------------------------------------------------- regression gates
def _summary(tag, metric, median, lower):
from bench_harness import ArmSummary
return ArmSummary(tag=tag, metric=metric, best=median, median=median,
spread_pct=0.0, n=3, n_busy=0, lower_is_better=lower)
EXPECT = {"arms": {"w4a16": {"rss_anon_mb": 3541.8, "tg_tps": 80.0}},
"bands_pct": {"rss_anon_mb": 1.5}}
def test_memory_regression_outside_the_band_fails():
rows = check_expectations([_summary("w4a16", "rss_anon_mb", 3900.0, True)], EXPECT)
assert rows[0]["status"] == "REGRESSION"
assert rows[0]["drift_pct"] > 1.5
def test_memory_inside_the_band_passes():
rows = check_expectations([_summary("w4a16", "rss_anon_mb", 3560.0, True)], EXPECT)
assert rows[0]["status"] == "ok"
def test_a_large_improvement_is_not_reported_as_a_regression():
rows = check_expectations([_summary("w4a16", "rss_anon_mb", 2800.0, True)], EXPECT)
assert rows[0]["status"] == "improved"
def test_throughput_is_not_gated_because_this_host_cannot_resolve_it():
"""tg_tps swings ~30% run to run here, so any band wide enough not to fire
constantly is too wide to catch a real regression."""
rows = check_expectations([_summary("w4a16", "tg_tps", 55.0, False)], EXPECT)
assert rows[0]["status"] == "ungateable"
def test_an_unlisted_arm_is_reported_not_silently_passed():
rows = check_expectations([_summary("brand-new", "rss_anon_mb", 100.0, True)], EXPECT)
assert rows[0]["status"] == "unlisted"
# ---------------------------------------- refusing to rank inside the noise
def test_a_difference_inside_the_noise_is_not_presented_as_a_ranking():
"""An 8% tg_tps gap on a host that swings 30% is not a result; publishing
one is how the W4A8/W4A16 ordering flipped between grids."""
rows = [_obs("a", 1, 86.0), _obs("a", 2, 80.0),
_obs("b", 1, 84.0), _obs("b", 2, 78.0)]
text = format_summary(summarize(rows, "tg_tps"))
assert "NOT A RANKING" in text
assert "profile_cpu_step.py" in text
def test_a_real_difference_is_still_reported_as_one():
rows = [_obs("fast", 1, 86.0), _obs("fast", 2, 86.2),
_obs("slow", 1, 41.0), _obs("slow", 2, 41.1)]
text = format_summary(summarize(rows, "tg_tps"))
assert "NOT A RANKING" not in text
def test_memory_arms_that_really_differ_are_ranked():
rows = [Observation("small", p, {"rss_anon_mb": v}, 1.0, 1.0, 1.0, True)
for p, v in ((1, 2802.0), (2, 2814.0))]
rows += [Observation("big", p, {"rss_anon_mb": v}, 1.0, 1.0, 1.0, True)
for p, v in ((1, 3541.0), (2, 3542.0))]
assert indistinguishable(summarize(rows, "rss_anon_mb")) is None
benchmarks/edge_harness/test_fidelity_kl.py
+70
View File
@@ -0,0 +1,70 @@
"""Tests for the fidelity metric's arithmetic, which is where it can lie.
Greedy divergence saturates; a KL that silently reports 0 for a broken build
would be worse than saturating, so identical / shifted / disagreeing cases are
pinned here.
"""
import math
import pytest
from fidelity_kl import chunk_tokens, compare
def _pos(target, dist):
"""dist: {token_id: probability} -> a position record in log space."""
return {"target": target,
"logprobs": {t: math.log(p) for t, p in dist.items()}}
def test_identical_distributions_score_zero_kl_and_full_agreement():
ref = [_pos(1, {1: 0.7, 2: 0.2, 3: 0.1})]
out = compare(ref, [dict(p) for p in ref])
assert out["top1_agreement"] == 1.0
assert out["kl_mean"] == pytest.approx(0.0, abs=1e-9)
assert out["delta_logprob_mean"] == pytest.approx(0.0, abs=1e-9)
def test_a_shifted_distribution_scores_positive_kl_without_changing_top1():
"""The case greedy divergence cannot see: same argmax, worse distribution."""
ref = [_pos(1, {1: 0.7, 2: 0.2, 3: 0.1})]
cand = [_pos(1, {1: 0.5, 2: 0.3, 3: 0.2})]
out = compare(ref, cand)
assert out["top1_agreement"] == 1.0
assert out["kl_mean"] > 0.01
assert out["delta_logprob_mean"] < 0 # reference token got less likely
def test_disagreement_is_counted():
ref = [_pos(1, {1: 0.7, 2: 0.3}), _pos(1, {1: 0.7, 2: 0.3})]
cand = [_pos(1, {1: 0.7, 2: 0.3}), _pos(1, {2: 0.8, 1: 0.2})]
assert compare(ref, cand)["top1_agreement"] == 0.5
def test_worse_builds_score_monotonically_worse():
"""The property greedy divergence lacks: it must rank, not just detect."""
ref = [_pos(1, {1: 0.8, 2: 0.15, 3: 0.05})]
near = compare(ref, [_pos(1, {1: 0.78, 2: 0.16, 3: 0.06})])["kl_mean"]
far = compare(ref, [_pos(1, {1: 0.40, 2: 0.35, 3: 0.25})])["kl_mean"]
assert 0 < near < far
def test_coverage_is_reported_so_a_lower_bound_is_not_read_as_a_kl():
"""Top-k truncation makes KL a lower bound; the run has to say so."""
ref = [_pos(1, {1: 0.3, 2: 0.2})] # only 0.5 of the mass covered
out = compare(ref, [_pos(1, {1: 0.3, 2: 0.2})])
assert out["coverage_mean"] == pytest.approx(0.5, abs=1e-6)
def test_mismatched_corpora_are_refused_not_averaged():
with pytest.raises(SystemExit):
compare([_pos(1, {1: 1.0})], [])
def test_chunking_drops_only_a_tail_too_short_to_score():
# 488 is a perfectly good chunk; both runs see the same slicing.
assert [len(c) for c in chunk_tokens(list(range(1000)), 512, None)] == [512, 488]
# 8 leftover tokens is not, so it is dropped rather than compared.
assert [len(c) for c in chunk_tokens(list(range(520)), 512, None)] == [512]
assert len(chunk_tokens(list(range(4096)), 512, 3)) == 3
docs/README.vllm-omni.md
+109
View File
@@ -0,0 +1,109 @@
<p align="center">
<picture>
<source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/logos/vllm-omni-logo.png">
<img alt="vllm-omni" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/logos/vllm-omni-logo.png" width=55%>
</picture>
</p>
<h3 align="center">
Easy, fast, and cheap omni-modality model serving for everyone
</h3>
<p align="center">
| <a href="https://vllm-omni.readthedocs.io/en/latest/"><b>Documentation</b></a> | <a href="https://deepwiki.com/vllm-project/vllm-omni"><b>DeepWiki</b></a> | <a href="https://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> | <a href="docs/assets/WeChat.jpg"><b>WeChat</b></a> | <a href="https://arxiv.org/abs/2602.02204"><b>Paper</b></a> | <a href="https://docs.google.com/presentation/d/1aPj0OGl_-ZVoib-Qne5dGDAlrRFB-PdHl6E-EE99g8E/edit?usp=sharing"><b>Slides</b></a> |
</p>
---
*Latest News* 🔥
- [2026/08] We released [0.28.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.28.0), featuring production-ready [MiniMax H3](recipes/MiniMaxAI/MiniMax-H3.md) serving on GPU and NPU, a unified AR/DiT paged KV cache runtime, and enhanced realtime full-duplex serving for the [MiniCPM-o series](recipes/OpenBMB/MiniCPM-o-4_5.md).
- [2026/08] [VeRL-Omni](https://github.com/verl-project/verl-omni) `v0.2.0` is released: faster diffusion RL powered by vLLM-Omni (request-level/step-wise batching with FA3), rebuilt Qwen3-Omni multimodal training (DPO & GSPO), plus LTX-2.3, Qwen-Image-Edit support and more. See the [release notes](https://github.com/verl-project/verl-omni/releases/tag/v0.2.0).
- [2026/08] We released [0.26.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.26.0) - aligned with the vLLM 0.26 release line, featuring [MiniMax H3](recipes/MiniMaxAI/MiniMax-H3.md) joint video/audio generation, an experimental full-duplex realtime runtime for [MiniCPM-o 4.5](recipes/OpenBMB/MiniCPM-o-4_5.md), distributed layerwise diffusion offload, and broader model, hardware, streaming, TTS, and quantization support.
- [2026/07] We released [0.24.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.24.0) - aligned with the vLLM 0.24 release line, expanding production-ready coverage across TTS, speech, diffusion, image/video generation, and robot-policy serving, with major Omni stage runtime refactoring, diffusion request-level batching, async output materialization, quantization/cache/memory improvements, and broad CUDA/ROCm/XPU/NPU support.
- [2026/06] Starting with [0.14.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.14.0), vLLM-Omni publishes a stable release aligned with every even-numbered upstream vLLM minor version. [0.16.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.16.0), [0.18.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.18.0), [0.20.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.20.0), and [0.22.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.22.0) continued this cadence, expanding omni and world-model support with [NVIDIA Cosmos3](recipes/cosmos3/Cosmos3-Nano.md) and DreamZero, adding models such as MiniCPM-o 4.5, MOSS-TTS, and Lance, and advancing TTS, diffusion, distributed execution, quantization, RL integration through [VeRL-Omni](https://github.com/verl-project/verl-omni), and CUDA/ROCm/MUSA/NPU/XPU coverage.
- [2026/03] Check out our first public [project deepdive](https://youtu.be/sgwNfsNnR9I) at the vLLM Hong Kong Meetup!
- [2025/11] vLLM community officially released [vllm-project/vllm-omni](https://github.com/vllm-project/vllm-omni) in order to support omni-modality models serving.
---
## About
[vLLM](https://github.com/vllm-project/vllm) was originally designed to support large language models for text-based autoregressive generation tasks. vLLM-Omni is a framework that extends its support for omni-modality model inference and serving:
- **Omni-modality**: Text, image, audio, video, and action data processing
- **Non-autoregressive Architectures**: extend the AR support of vLLM to Diffusion Transformers (DiT) and other parallel generation models
- **Heterogeneous outputs**: from traditional text generation to multimodal and action outputs
<p align="center">
<picture>
<img alt="vllm-omni" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/omni-modality-model-architecture.png" width=55%>
</picture>
</p>
vLLM-Omni is fast with:
- State-of-the-art AR support by leveraging efficient KV cache management from vLLM
- Pipelined stage execution overlapping for high throughput performance
- Fully disaggregation based on OmniConnector and dynamic resource allocation across stages
vLLM-Omni is flexible and easy to use with:
- Heterogeneous pipeline abstraction to manage complex model workflows
- Seamless integration with popular Hugging Face models
- Tensor, pipeline, data and expert parallelism support for distributed inference
- Streaming outputs
- OpenAI-compatible API server
- [Full-duplex realtime serving](docs/serving/realtime_duplex_api.md) with streaming audio input and output
vLLM-Omni seamlessly supports most popular open-source models on HuggingFace, including:
- **Omni-modality models** (e.g. Qwen3-Omni, MiniCPM-o 4.5, Cosmos3, HunyuanImage, BAGEL)
- **TTS models** (e.g. Qwen3-TTS, IndexTTS 2.5, dots.tts, CosyVoice3)
- **Diffusion models** — image, video, and audio generation (e.g. MiniMax H3, LTX-2.5, SANA-Video, Wan2.2)
- **Robot-policy and action models** (e.g. π0, GR00T-N1.7, DreamZero-DROID, InternVLA-A1)
## Getting Started
Visit our [documentation](https://vllm-omni.readthedocs.io/en/latest/) to learn more.
- [Installation](https://vllm-omni.readthedocs.io/en/latest/getting_started/installation/)
- [Quickstart](https://vllm-omni.readthedocs.io/en/latest/getting_started/quickstart/)
- [Nightly CUDA Docker images](https://hub.docker.com/r/vllm/vllm-omni/tags?name=nightly) built from the latest `main`
- [List of Supported Models](https://vllm-omni.readthedocs.io/en/latest/models/supported_models/)
- [Deployment Recipes](https://recipes.vllm.ai) for vLLM-Omni model serving
## Contributing
We welcome and value any contributions and collaborations.
Please check out [Contributing to vLLM-Omni](https://vllm-omni.readthedocs.io/en/latest/contributing/) for how to get involved.
## Citation
If you use vLLM-Omni for your research, please cite our [paper](https://arxiv.org/abs/2602.02204):
```bibtex
@article{yin2026vllmomni,
title={vLLM-Omni: Fully Disaggregated Serving for Any-to-Any Multimodal Models},
author={Peiqi Yin, Jiangyun Zhu, Han Gao, Chenguang Zheng, Yongxiang Huang, Taichang Zhou, Ruirui Yang, Weizhi Liu, Weiqing Chen, Canlin Guo, Didan Deng, Zifeng Mo, Cong Wang, James Cheng, Roger Wang, Hongsheng Liu},
journal={arXiv preprint arXiv:2602.02204},
year={2026}
}
```
## Join the Community
Feel free to ask questions, provide feedbacks and discuss with fellow users of vLLM-Omni in `#sig-omni` slack channel at [slack.vllm.ai](https://slack.vllm.ai) or vLLM user forum at [discuss.vllm.ai](https://discuss.vllm.ai).
## Star History
<a href="https://www.star-history.com/?repos=vllm-project%2Fvllm-omni&type=date&legend=top-left">
<picture>
<source media="(prefers-color-scheme: dark)" srcset="https://api.star-history.com/chart?repos=vllm-project/vllm-omni&type=date&theme=dark&legend=top-left&sealed_token=ExgLDZJoQEg27Zfhhut2LqN0GYO6Fw2PWLwPE6JYBUp2BgM3hmsYlwaIVopnUEfbRXidQ4nisumrTdKYydiKhy1SZXipw47qY2_tiUDhCpsPXeXtPuEVKVzBwKs3pw0tiHsJgtSfwXx5yjHXck0Y2SblzFWeJYCkTe1WLGTbUAOIETjXXQJjyCGZvKz5" />
<source media="(prefers-color-scheme: light)" srcset="https://api.star-history.com/chart?repos=vllm-project/vllm-omni&type=date&legend=top-left&sealed_token=ExgLDZJoQEg27Zfhhut2LqN0GYO6Fw2PWLwPE6JYBUp2BgM3hmsYlwaIVopnUEfbRXidQ4nisumrTdKYydiKhy1SZXipw47qY2_tiUDhCpsPXeXtPuEVKVzBwKs3pw0tiHsJgtSfwXx5yjHXck0Y2SblzFWeJYCkTe1WLGTbUAOIETjXXQJjyCGZvKz5" />
<img alt="Star History Chart" src="https://api.star-history.com/chart?repos=vllm-project/vllm-omni&type=date&legend=top-left&sealed_token=ExgLDZJoQEg27Zfhhut2LqN0GYO6Fw2PWLwPE6JYBUp2BgM3hmsYlwaIVopnUEfbRXidQ4nisumrTdKYydiKhy1SZXipw47qY2_tiUDhCpsPXeXtPuEVKVzBwKs3pw0tiHsJgtSfwXx5yjHXck0Y2SblzFWeJYCkTe1WLGTbUAOIETjXXQJjyCGZvKz5" />
</picture>
</a>
## License
Apache License 2.0, as found in the [LICENSE](./LICENSE) file.
docs/edge/cpu-4bit.md
+100
View File
@@ -0,0 +1,100 @@
# The 4-bit CPU path
Two 4-bit paths exist on an x86 CPU and they are not interchangeable. Which one
a device should use is a property of the device, and `vllm_omni/edge/weight_path.py`
picks it from a hardware probe.
## W4A8 — `_C::int4_scaled_mm_cpu`
Needs AMX tiles. Wants GPTQ-packed weights and quantizes activations to int8.
Fastest measured here (74.4 tok/s) and least faithful (1/12 greedy prompts
reproduced exactly).
One trap: a GPTQ checkpoint **must** quantize the output head. Spark ties its
embedding to its head, so the head has to be untied and emitted separately; if
it is not, that one projection runs in bf16 through `onednn_mm` and costs
2.97 ms/step by itself — more than the entire 4-bit saving. Every W4A8 number in
this project's first round was invalid for exactly this reason.
The symmetric grid is **our exporter's choice, not the kernel's limit**:
`_process_gptq_weights_w4a8` uses the checkpoint's zero points whenever
`config.zero_points` is set and synthesises the constant 8 only otherwise. An
asymmetric export is a config change and is the cheapest untried fidelity
experiment.
## W4A16 — PyTorch tinygemm
No AMX requirement, activations stay bf16, more faithful (2/12). 71.6 tok/s.
This is `quantization="cpu_int4"`.
The kernel dequantizes once per row of activations, so its cost is linear in
batch size: fine for decode, ruinous for prefill (453 tok/s if used for
everything). Large batches therefore take a second path — dequantize to bf16
once and hand it to the ordinary oneDNN/AMX GEMM.
### The weight layout, measured
That prefill path used to read a row-major *duplicate* of every weight, because
the tinygemm layout was treated as opaque. It is not. Packing an index pattern
and reading back where each nibble landed shows, on AVX-512:
```
[N/64][K][32 bytes]
```
K is the middle axis, so a K range is contiguous inside each n-block; and within
every 32-byte run, byte `d` carries output channel `d` in its low nibble and
`d + 32` in its high one — the same permutation for every `k` and every block.
Decoding it is a deinterleave, not a reverse-engineered ISA table.
So prefill can dequantize from the kernel's own layout and the duplicate is
never allocated: **615 MB**, bit-exact against the old path at 2048×2048,
3072×2048 and `gate_up`'s real 13312×2048.
It is not free. `to_split_halves` performed that deinterleave *once at load*, so
each prefill call read two contiguous planes; doing it per call costs **1.63×
prefill**. It is a dial, not a win:
- against `prefill_dequant: false` it is **strictly better** — same memory, 2.97×
the prefill, and multi-row batches never reach the tinygemm kernel, avoiding
the AMX tile-configuration SIGILL documented in `cpu_int4.py`. That setting is
dominated.
- against keeping both layouts it trades 615 MB for ~1.6× prefill. Default on,
because this path exists for devices chosen for their memory ceiling. Set
`direct_dequant: false` in the checkpoint's `quantization_config` to revert.
`supports_direct_dequant` gates the reshape on N tiling the 64-channel block,
which every Spark linear does. An N like 80 packs as 64 + a 16 tail the reshape
would mis-read, so those layers keep the duplicate rather than being quietly
wrong.
### Why the path choice is a config field
The two prefill paths build *different graphs* — one passes `weight_halves` into
the custom op, the other passes `None`. vLLM keys its AOT compile cache on the
config, so an env-var-only switch shares one cached artifact between them, and a
warm cache loads the wrong graph: three benchmark passes died with
`KeyError: 'weight_halves'` inside an AOT-loaded forward. `direct_dequant` is a
`quantization_config` field because `ModelConfig.compute_hash` hashes those.
## `cpu_gemm_wna16` — the fallback
vLLM selects this when `use_w4a8` is off. Measured at 32.3 tok/s, **2.2× slower
than tinygemm**, so the selector does not route to it. It is documented because
`weight_path.py` previously asserted it did not exist.
## Building a checkpoint
```bash
# W4A16 (cpu_int4)
python benchmarks/edge_harness/quantize_int4.py \
--model <bf16-model> --group-size 128 --out models/<name>
# W4A8 (GPTQ) — note the head must be untied and quantized
python benchmarks/edge_harness/quantize_gptq_int4.py \
--model <bf16-model> --group-size 128 --out models/<name>
```
`quantization_config` keys the `cpu_int4` path understands: `group_size`,
`quantize_embedding`, `prefill_dequant`, `direct_dequant`, `dequant_threshold`,
`use_custom_op`.
docs/edge/harness.md
+79
View File
@@ -0,0 +1,79 @@
# The measurement harness
`benchmarks/edge_harness/` — the tooling every number in this fork was produced
with. It exists because five earlier conclusions were wrong; see
[methodology.md](methodology.md).
## Running a comparison
Arms are subprocess invocations of a single-run benchmark, so the measurement
code stays the one that was verified; the harness decides what to run, when, and
how to read it.
```bash
python benchmarks/edge_harness/bench_harness.py \
--spec arms.json --passes 3 --metric rss_anon_tree_mb \
--keys ready_s,init_s,rss_anon_tree_mb,peak_rss_anon_mb \
--cores 56-79 --out summary.json
```
`arms.json` is `{"arms": [{"tag": ..., "command": [...], "env": {...}}]}`.
Give each arm its own `VLLM_CACHE_ROOT` if the arms differ in anything that
changes the traced graph.
It rotates arm order every pass, reports spread, labels contended passes, refuses to
start when the target cores are busy, prints `NOT A RANKING` when the top two
arms differ by less than their own spread, and stamps provenance into the
output. `--check expectations.json` turns it into a regression gate.
## The single-run benchmarks
| script | measures |
|---|---|
| `bench_init_memory.py` | `ready_s` (process start → serving), peak `RssAnon` sampled through load, resident memory over the **process tree**, memory at a filled context |
| `bench_cpu_vllm.py` | pp512 / tg128, mirroring `llama-bench` semantics |
| `profile_cpu_step.py` | per-op decode cost, measured by difference so it matches the wall clock |
| `memory_ledger.py` | attributes resident memory to tensor storages, KV cache and mappings, with `mincore(2)` residency |
| `fidelity_kl.py` | top-1 agreement and KL against a bf16 reference, comparable with `llama-perplexity --kl-divergence` |
Two arguments matter more than they look:
- `--context-len` on `bench_init_memory.py`. Without it the KV cache is
allocated and never touched, and a memory number taken that way says nothing
about cache sizing. That is how "KV sizing is not a lever" happened.
- `--prompt-len` on `profile_cpu_step.py`. The default of 1 measures an empty
cache; Spark's sliding layers have a 512-token window, so attention cost at a
1-token prompt is not attention cost at context.
## The agent harness
`benchmarks/edge_harness/android_agent/` drives a phone (real, over `adb`, or a
scripted simulator that emits real uiautomator XML) and scores tasks by the
phone's own settings and screen rather than by the transcript.
Scoring reads the *setting* where one exists and the UI dump otherwise, so the
same predicate works on both backends; an agent that turns Wi-Fi on and then
presses Home has still turned Wi-Fi on. Tasks carry variants — a renamed
control, reordered rows, a plausible decoy — because repeats at temperature 0
against one fixture measure repeatability, not whether the agent can do the task.
```bash
python benchmarks/edge_harness/android_agent/run_suite.py \
--base-url http://127.0.0.1:8077/v1 --repeat 5 # simulator
python benchmarks/edge_harness/android_agent/run_suite.py --device adb --repeat 5
```
Tasks needing flip history (a phone cannot report which switches were touched)
are skipped with a reason on a real device rather than scored from final state.
## Tests
```bash
pytest benchmarks/edge_harness/test_bench_harness.py \
benchmarks/edge_harness/test_fidelity_kl.py \
benchmarks/edge_harness/android_agent/test_android_agent.py
```
They pin judgement, not plumbing: that arm order rotates, that `best` is the
minimum for memory, that a difference inside the noise is not reported as a
ranking, that the alarm oracle needs a saved alarm rather than a typed digit.
docs/edge/methodology.md
+82
View File
@@ -0,0 +1,82 @@
# How to measure on a shared machine without fooling yourself
Five conclusions in this project were retracted. None was a coding error; every
one was a measurement error that looked like a result. This page is what was
left after each of them, and it is the part of this repository most likely to be
useful somewhere else.
## The five
**1. `VmRSS` counts pages you do not own.** A "3.1 GB saving" was the page cache
holding the mmap'd checkpoint. `VmRSS` moved 1–2 GB between *identical* runs;
`RssAnon` repeats to 0.2%. Measure allocation, not residency-of-everything.
**2. The build under test was not the build you think.** Every W4A8 timing was
taken on a checkpoint with `tie_word_embeddings: true` and no `lm_head` tensors,
so the output projection ran in bf16 and cost 2.97 ms/step by itself — more than
the entire 4-bit saving being measured. Artifacts now carry provenance: model
path, quantization, thread set, build SHAs. Two saved profiles in this tree were
indistinguishable from each other before that.
**3. You measured the wrong process.** The memory benchmark read
`/proc/self/status`, which is the whole story only when the model is in-process.
vLLM's CPU default promotes the executor to `mp`, putting the weights in a
worker: the benchmark reported 796 MB resident for a model whose weights alone
are 1.4 GB. Memory is summed over the process tree now.
**4. The knob was not connected.** "KV cache sizing saves only 188 MB, the pages
are never touched" sat on a *closed, do not retry* list. The comparison ran
through `CpuPlatform.check_and_update_config`, which overwrote
`kv_cache_memory_bytes` from an environment variable unconditionally — so both
arms ran the same cache. With the setting honoured the saving is **843 MB**, and
`mincore(2)` shows the cache fully in core rather than untouched. *A lever tested
through a path that silently ignores it measures nothing, and "no effect" is the
most dangerous result to accept without first checking the knob was connected.*
**5. One surviving pass is not a result.** A 1.9× prefill cost was quoted from
the single pass of a grid whose other passes had crashed — a number the harness
had already labelled `single pass — not a result`. With every pass completing it
is 1.63×, and only good to ±0.2× at that.
## What the harness enforces
- **Rotate arm order every pass.** Interleaving alone leaves arm 0 permanently
first. `run_grid` rotates; a test pins it, because the previous test asserted
the *un-rotated* order and that is how the bug survived.
- **Know which direction is better.** `best` is the maximum for a rate and the
minimum for memory or latency. Taking the maximum for everything published the
*worst* memory pass as the headline.
- **Refuse to rank inside the noise.** When the top two arms differ by less than
their own spread, the summary prints `NOT A RANKING` and points at the per-op
profiler. Per-op self time reproduces to ~5% here; end-to-end throughput swings
~30%, which is why throughput is explicitly ungateable in the regression check.
- **Guard the cores, not the host.** Load average is a whole-host number and this
machine has 224 cores, so it says nothing about the 24 a run is pinned to. The
pre-flight samples per-core utilisation over two windows, takes the minimum
(so a blip is not load), and refuses on either >5% aggregate stolen capacity or
any core sustained above 50% — one slow core gates every OpenMP barrier, so it
cannot be averaged away.
- **Record failures.** Three passes of one grid died with
`KeyError: 'weight_halves'`; the harness recorded the exit status and stderr
instead of dropping the runs, which is the only reason the bug was found.
- **Stamp provenance.** Build SHAs, host, thread affinity, environment. A number
whose origin was not recorded cannot be retracted later, only doubted.
## Two traps specific to this stack
**Compile-cache collisions.** vLLM keys its AOT compile cache on the *config*, so
two arms differing only by an environment variable share one cached artifact. If
that variable changes the traced graph, later passes load the wrong one. Give
each arm its own `VLLM_CACHE_ROOT`, and prefer expressing a graph-affecting
choice as a config field — `ModelConfig.compute_hash` hashes those.
**Allocated is not resident.** An untouched buffer costs address space. Use
`mincore(2)` before attributing it to memory: `memory_ledger.py` does, and that
is what proved the KV cache was real rather than reserved.
## The rule that generalises
Attribute, do not subtract. The "unexplained ~1.6 GB" in an earlier report was a
remainder computed from two uncertain terms, and it hid a 640 MB RoPE table in
plain sight. Counting unique tensor storages by `data_ptr`, naming the largest,
and checking residency per page turned "unexplained" into a one-line fix.
docs/edge/results.md
+116
View File
@@ -0,0 +1,116 @@
# Measured results
Spark-X2.5-1.7B unless stated. Xeon Platinum 8480C, 24 pinned physical cores,
memory bound to the same NUMA node, W4A16 g128, `max_model_len` 2048.
Each row is a median over at least three interleaved passes against a control
arm running in the same grid.
## Memory
Every lever below was measured as its own A/B, with the alternative reachable by
a flag so both arms run identical code.
| lever | before | after | saving | cost |
|---|---|---|---|---|
| KV budget 1 GiB → 192 MiB | 3264 MB | 2421 MB | **843 MB** | fewer cached tokens |
| RoPE table → `max_model_len` | 3891 MB | 3273 MB | **618 MB** | none |
| One weight layout, not two | 4591 MB | 3976 MB | **615 MB** | ~1.63× prefill |
| Blocked repack (peak) | 4502 MB | 4013 MB | **489 MB** peak | none |
**These do not add.** The single-layout change removed `to_split_halves`, which
was doing its own whole-matrix expansion, so part of the blocked-repack saving
was already gone by the time it was measured. The report each was taken from
states its baseline; quote them individually.
### Where the memory actually is
From `memory_ledger.py`, which counts unique tensor storages by `data_ptr` (so
tied weights are counted once) and checks residency with `mincore(2)`:
| | MB | note |
|---|---|---|
| KV cache | 1024 | fully resident — 1024.1 of 1023.8 allocated |
| int4 packed weights | 686 | |
| RoPE cos/sin tables | 640 | **before the cap** — sized for 1 048 576 positions |
| embedding / head | 408 | tied; includes a 128 MB row-major copy |
| imports (`import vllm_omni`) | 501 | |
The 640 MB RoPE entry is the clearest single finding in the project. Spark
advertises a 1 048 576-token context and the model passed
`max_position_embeddings` straight to `get_rope`, so vLLM built
`torch.arange(1048576)` cos/sin tables — two of them, because the sliding and
full layers use different theta and partial-rotary factors. At `max_model_len`
2048 they need 1.25 MB between them. It had been sitting inside an
"unexplained ~1.6 GB" remainder.
### Peak vs settled
`release_after_load()` runs only once loading and warm-up return, so a device
that cannot hold the *peak* never reaches the release. Sampling `RssAnon` on a
timer through the load, rather than reading the endpoint:
| | peak | settled | transient |
|---|---|---|---|
| whole-matrix repack | 4502 MB | 3905 MB | 550 MB |
| blocked repack | 4013 MB | 3892 MB | 105 MB |
`unpack_nibbles` expands 4-bit codes to int32 — eight times the packed form —
because the kernel's packer wants int32. Whole-matrix that is 104 MiB for
`gate_up` and **1024 MiB for the tied embedding**, against 24 MiB of packed
weights for a whole layer. Chunking is also *faster* (79 ms vs 208 ms for the
embedding) because the working set stays in cache.
## Speed
| path | decode | prefill (pp512) | spread |
|---|---|---|---|
| W4A8, AMX `int4_scaled_mm_cpu` | 74.4 tok/s | — | 1.1% |
| W4A16, tinygemm | 71.6 tok/s | 2190 tok/s | 9.7% |
| W4A16, one layout | 71.6 tok/s | 1344 tok/s | 18.3% |
| W4A16, `prefill_dequant: false` | 73.0 tok/s | 453 tok/s | 0.3% |
| `cpu_gemm_wna16` fallback | 32.3 tok/s | — | 1.5% |
| llama.cpp Q4_K_M | 90.0 tok/s | 1121 tok/s | — |
`cpu_gemm_wna16` is the kernel vLLM selects when `use_w4a8` is off. It exists —
`weight_path.py` used to call tinygemm "the only option" without AMX, which was
false — but at 2.2× slower than tinygemm the routing was right anyway. Caveat:
`_get_isa_hint` returns `"amx"` whenever AMX is present regardless of the W4A8
flag, so this measured its *best* path on this host, not the `"vec"` path a
genuinely AMX-less machine takes.
## Bandwidth headroom
A node-local streaming read on the same 24 cores sustains **222 GB/s**.
| | GB/s | of ceiling |
|---|---|---|
| W4A16 GEMM only | 132 | 60% |
| W4A8 GEMM only | 153 | 69% |
| W4A16 whole step | 73 | 33% |
| llama.cpp whole step | 86 | 39% |
Neither engine is bandwidth-bound. An earlier draft asserted the GEMM ran at
"85% of what this socket delivers" and concluded the remainder was framework
overhead; at 60% of the measured ceiling, kernel and scheduling efficiency are
still on the table. See [ROADMAP.md](../../ROADMAP.md) item I10.
## Fidelity
Against each build's own bf16 output over 12 greedy prompts:
| build | exact reproductions | bits/weight |
|---|---|---|
| llama.cpp Q4_K_M | 4/12 | 4.5 (groups of 32, mixed types) |
| W4A16 | 2/12 | 4.25 |
| W4A8 | 1/12 | 4.25 |
**Not a matched-budget comparison** — Q4_K_M stores more. And greedy divergence
saturates: it discards everything after the first divergent token, so one
unlucky token scores the same as a build that is wrong everywhere.
`fidelity_kl.py` replaces it with top-1 agreement and KL over a corpus.
On the Android agent harness, every build navigates a multi-screen task
correctly and only some notice they have arrived: W4A8 turned Wi-Fi on and then
tapped the switch off again. **Navigation survives quantization; recognising
completion does not.** That is one run per cell — enough to show a failure
reproduces, not to rank success rates.