From b36a3431bda13b53dfe74d9d2e5617ec9f86e5b7 Mon Sep 17 00:00:00 2001 From: tzhouam Date: Mon, 14 Sep 2026 13:25:32 +0800 Subject: [PATCH] Package this fork as vllm-embedding: README, docs, roadmap, harness 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 Claude-Session: https://claude.ai/code/session_01L1ax5wYLA6XbqEL1i8L9gZ --- NOTICE | 23 + README.md | 150 ++--- ROADMAP.md | 107 +++ .../edge_harness/android_agent/README.md | 92 +++ .../android_agent/action_space.py | 152 +++++ .../edge_harness/android_agent/agent.py | 271 ++++++++ .../edge_harness/android_agent/device.py | 354 ++++++++++ .../edge_harness/android_agent/executor.py | 147 ++++ .../edge_harness/android_agent/phone_state.py | 135 ++++ .../edge_harness/android_agent/run_agent.py | 98 +++ .../edge_harness/android_agent/run_episode.py | 106 +++ .../edge_harness/android_agent/run_suite.py | 151 +++++ .../edge_harness/android_agent/scenarios.py | 113 ++++ .../edge_harness/android_agent/tasks.py | 279 ++++++++ .../android_agent/test_android_agent.py | 525 +++++++++++++++ .../edge_harness/android_agent/tool_parse.py | 155 +++++ .../edge_harness/android_agent/ui_tree.py | 146 ++++ benchmarks/edge_harness/bench_cpu_vllm.py | 106 +++ benchmarks/edge_harness/bench_harness.py | 634 ++++++++++++++++++ benchmarks/edge_harness/bench_init_memory.py | 297 ++++++++ benchmarks/edge_harness/fidelity_kl.py | 179 +++++ benchmarks/edge_harness/memory_ledger.py | 224 +++++++ benchmarks/edge_harness/profile_cpu_step.py | 158 +++++ benchmarks/edge_harness/quantize_gptq_int4.py | 123 ++++ benchmarks/edge_harness/quantize_int4.py | 138 ++++ benchmarks/edge_harness/test_bench_harness.py | 308 +++++++++ benchmarks/edge_harness/test_fidelity_kl.py | 70 ++ docs/README.vllm-omni.md | 109 +++ docs/edge/cpu-4bit.md | 100 +++ docs/edge/harness.md | 79 +++ docs/edge/methodology.md | 82 +++ docs/edge/results.md | 116 ++++ 32 files changed, 5641 insertions(+), 86 deletions(-) create mode 100644 NOTICE create mode 100644 ROADMAP.md create mode 100644 benchmarks/edge_harness/android_agent/README.md create mode 100644 benchmarks/edge_harness/android_agent/action_space.py create mode 100644 benchmarks/edge_harness/android_agent/agent.py create mode 100644 benchmarks/edge_harness/android_agent/device.py create mode 100644 benchmarks/edge_harness/android_agent/executor.py create mode 100644 benchmarks/edge_harness/android_agent/phone_state.py create mode 100644 benchmarks/edge_harness/android_agent/run_agent.py create mode 100644 benchmarks/edge_harness/android_agent/run_episode.py create mode 100644 benchmarks/edge_harness/android_agent/run_suite.py create mode 100644 benchmarks/edge_harness/android_agent/scenarios.py create mode 100644 benchmarks/edge_harness/android_agent/tasks.py create mode 100644 benchmarks/edge_harness/android_agent/test_android_agent.py create mode 100644 benchmarks/edge_harness/android_agent/tool_parse.py create mode 100644 benchmarks/edge_harness/android_agent/ui_tree.py create mode 100644 benchmarks/edge_harness/bench_cpu_vllm.py create mode 100644 benchmarks/edge_harness/bench_harness.py create mode 100644 benchmarks/edge_harness/bench_init_memory.py create mode 100644 benchmarks/edge_harness/fidelity_kl.py create mode 100644 benchmarks/edge_harness/memory_ledger.py create mode 100644 benchmarks/edge_harness/profile_cpu_step.py create mode 100644 benchmarks/edge_harness/quantize_gptq_int4.py create mode 100644 benchmarks/edge_harness/quantize_int4.py create mode 100644 benchmarks/edge_harness/test_bench_harness.py create mode 100644 benchmarks/edge_harness/test_fidelity_kl.py create mode 100644 docs/README.vllm-omni.md create mode 100644 docs/edge/cpu-4bit.md create mode 100644 docs/edge/harness.md create mode 100644 docs/edge/methodology.md create mode 100644 docs/edge/results.md diff --git a/NOTICE b/NOTICE new file mode 100644 index 000000000..ae1967a26 --- /dev/null +++ b/NOTICE @@ -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. diff --git a/README.md b/README.md index 5063f9a3f..a52c985e8 100644 --- a/README.md +++ b/README.md @@ -1,109 +1,87 @@ -

- - - vllm-omni - -

-

-Easy, fast, and cheap omni-modality model serving for everyone -

+# vllm-embedding -

-| Documentation | DeepWiki | User Forum | Developer Slack | WeChat | Paper | Slides | -

+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: -

- - vllm-omni - -

+| 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='', 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 - - - - - - Star History Chart - - - -## 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). diff --git a/ROADMAP.md b/ROADMAP.md new file mode 100644 index 000000000..cf654b357 --- /dev/null +++ b/ROADMAP.md @@ -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. diff --git a/benchmarks/edge_harness/android_agent/README.md b/benchmarks/edge_harness/android_agent/README.md new file mode 100644 index 000000000..cfbf5a0f6 --- /dev/null +++ b/benchmarks/edge_harness/android_agent/README.md @@ -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 --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 ``; 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. diff --git a/benchmarks/edge_harness/android_agent/action_space.py b/benchmarks/edge_harness/android_agent/action_space.py new file mode 100644 index 000000000..191c7e25f --- /dev/null +++ b/benchmarks/edge_harness/android_agent/action_space.py @@ -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." +) diff --git a/benchmarks/edge_harness/android_agent/agent.py b/benchmarks/edge_harness/android_agent/agent.py new file mode 100644 index 000000000..145d8f83d --- /dev/null +++ b/benchmarks/edge_harness/android_agent/agent.py @@ -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 ``, 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 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 diff --git a/benchmarks/edge_harness/android_agent/device.py b/benchmarks/edge_harness/android_agent/device.py new file mode 100644 index 000000000..186829c46 --- /dev/null +++ b/benchmarks/edge_harness/android_agent/device.py @@ -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("") + return out[start:end + len("")] 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), {'"': """}) + rows.append( + f' ' + ) + body = "\n".join(rows) + return ( + '\n' + f'\n{body}\n' + ) + + 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}) diff --git a/benchmarks/edge_harness/android_agent/executor.py b/benchmarks/edge_harness/android_agent/executor.py new file mode 100644 index 000000000..c795201fe --- /dev/null +++ b/benchmarks/edge_harness/android_agent/executor.py @@ -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", + ) diff --git a/benchmarks/edge_harness/android_agent/phone_state.py b/benchmarks/edge_harness/android_agent/phone_state.py new file mode 100644 index 000000000..0f1ffbb00 --- /dev/null +++ b/benchmarks/edge_harness/android_agent/phone_state.py @@ -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 diff --git a/benchmarks/edge_harness/android_agent/run_agent.py b/benchmarks/edge_harness/android_agent/run_agent.py new file mode 100644 index 000000000..c848856d7 --- /dev/null +++ b/benchmarks/edge_harness/android_agent/run_agent.py @@ -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 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() diff --git a/benchmarks/edge_harness/android_agent/run_episode.py b/benchmarks/edge_harness/android_agent/run_episode.py new file mode 100644 index 000000000..f77291130 --- /dev/null +++ b/benchmarks/edge_harness/android_agent/run_episode.py @@ -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 --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() diff --git a/benchmarks/edge_harness/android_agent/run_suite.py b/benchmarks/edge_harness/android_agent/run_suite.py new file mode 100644 index 000000000..2c1196389 --- /dev/null +++ b/benchmarks/edge_harness/android_agent/run_suite.py @@ -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() diff --git a/benchmarks/edge_harness/android_agent/scenarios.py b/benchmarks/edge_harness/android_agent/scenarios.py new file mode 100644 index 000000000..3a7951248 --- /dev/null +++ b/benchmarks/edge_harness/android_agent/scenarios.py @@ -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"], + }, +] diff --git a/benchmarks/edge_harness/android_agent/tasks.py b/benchmarks/edge_harness/android_agent/tasks.py new file mode 100644 index 000000000..b0029c446 --- /dev/null +++ b/benchmarks/edge_harness/android_agent/tasks.py @@ -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} diff --git a/benchmarks/edge_harness/android_agent/test_android_agent.py b/benchmarks/edge_harness/android_agent/test_android_agent.py new file mode 100644 index 000000000..2ad919cec --- /dev/null +++ b/benchmarks/edge_harness/android_agent/test_android_agent.py @@ -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 = """ + + + + + + +""" + + +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 = '' + "".join( + f'' for i in range(200) + ) + "" + 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 ``, 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",) diff --git a/benchmarks/edge_harness/android_agent/tool_parse.py b/benchmarks/edge_harness/android_agent/tool_parse.py new file mode 100644 index 000000000..4fafaa56a --- /dev/null +++ b/benchmarks/edge_harness/android_agent/tool_parse.py @@ -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 `` 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"\s*([a-z_][a-z0-9_]*)(.*?)", re.I | re.S) +_GLM_ARG = re.compile( + r"(.*?)\s*(.*?)", 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") diff --git a/benchmarks/edge_harness/android_agent/ui_tree.py b/benchmarks/edge_harness/android_agent/ui_tree.py new file mode 100644 index 000000000..0862b96cc --- /dev/null +++ b/benchmarks/edge_harness/android_agent/ui_tree.py @@ -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) diff --git a/benchmarks/edge_harness/bench_cpu_vllm.py b/benchmarks/edge_harness/bench_cpu_vllm.py new file mode 100644 index 000000000..32b25eef2 --- /dev/null +++ b/benchmarks/edge_harness/bench_cpu_vllm.py @@ -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() diff --git a/benchmarks/edge_harness/bench_harness.py b/benchmarks/edge_harness/bench_harness.py new file mode 100644 index 000000000..7dd79ab67 --- /dev/null +++ b/benchmarks/edge_harness/bench_harness.py @@ -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 --prompt-len ") + + +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() diff --git a/benchmarks/edge_harness/bench_init_memory.py b/benchmarks/edge_harness/bench_init_memory.py new file mode 100644 index 000000000..3eb8fcf04 --- /dev/null +++ b/benchmarks/edge_harness/bench_init_memory.py @@ -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() diff --git a/benchmarks/edge_harness/fidelity_kl.py b/benchmarks/edge_harness/fidelity_kl.py new file mode 100644 index 000000000..1b151f0cb --- /dev/null +++ b/benchmarks/edge_harness/fidelity_kl.py @@ -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() diff --git a/benchmarks/edge_harness/memory_ledger.py b/benchmarks/edge_harness/memory_ledger.py new file mode 100644 index 000000000..e7fa12c5b --- /dev/null +++ b/benchmarks/edge_harness/memory_ledger.py @@ -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() diff --git a/benchmarks/edge_harness/profile_cpu_step.py b/benchmarks/edge_harness/profile_cpu_step.py new file mode 100644 index 000000000..74320c7f9 --- /dev/null +++ b/benchmarks/edge_harness/profile_cpu_step.py @@ -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() diff --git a/benchmarks/edge_harness/quantize_gptq_int4.py b/benchmarks/edge_harness/quantize_gptq_int4.py new file mode 100644 index 000000000..3c7cd237b --- /dev/null +++ b/benchmarks/edge_harness/quantize_gptq_int4.py @@ -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: + + .qweight int32 [K // 8, N] 4-bit codes packed along the input dim + .qzeros int32 [K // g, N // 8] + .scales bf16 [K // g, N] + .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() diff --git a/benchmarks/edge_harness/quantize_int4.py b/benchmarks/edge_harness/quantize_int4.py new file mode 100644 index 000000000..2dad6b249 --- /dev/null +++ b/benchmarks/edge_harness/quantize_int4.py @@ -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() diff --git a/benchmarks/edge_harness/test_bench_harness.py b/benchmarks/edge_harness/test_bench_harness.py new file mode 100644 index 000000000..c329ba1b8 --- /dev/null +++ b/benchmarks/edge_harness/test_bench_harness.py @@ -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 diff --git a/benchmarks/edge_harness/test_fidelity_kl.py b/benchmarks/edge_harness/test_fidelity_kl.py new file mode 100644 index 000000000..364dbf280 --- /dev/null +++ b/benchmarks/edge_harness/test_fidelity_kl.py @@ -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 diff --git a/docs/README.vllm-omni.md b/docs/README.vllm-omni.md new file mode 100644 index 000000000..5063f9a3f --- /dev/null +++ b/docs/README.vllm-omni.md @@ -0,0 +1,109 @@ +

+ + + vllm-omni + +

+

+Easy, fast, and cheap omni-modality model serving for everyone +

+ +

+| Documentation | DeepWiki | User Forum | Developer Slack | WeChat | Paper | Slides | +

+ +--- + +*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 + +

+ + vllm-omni + +

+ +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 + + + + + + Star History Chart + + + +## License + +Apache License 2.0, as found in the [LICENSE](./LICENSE) file. diff --git a/docs/edge/cpu-4bit.md b/docs/edge/cpu-4bit.md new file mode 100644 index 000000000..dd5f4f2d7 --- /dev/null +++ b/docs/edge/cpu-4bit.md @@ -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 --group-size 128 --out models/ + +# W4A8 (GPTQ) โ€” note the head must be untied and quantized +python benchmarks/edge_harness/quantize_gptq_int4.py \ + --model --group-size 128 --out models/ +``` + +`quantization_config` keys the `cpu_int4` path understands: `group_size`, +`quantize_embedding`, `prefill_dequant`, `direct_dequant`, `dequant_threshold`, +`use_custom_op`. diff --git a/docs/edge/harness.md b/docs/edge/harness.md new file mode 100644 index 000000000..82f22e66c --- /dev/null +++ b/docs/edge/harness.md @@ -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. diff --git a/docs/edge/methodology.md b/docs/edge/methodology.md new file mode 100644 index 000000000..bcec4d454 --- /dev/null +++ b/docs/edge/methodology.md @@ -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. diff --git a/docs/edge/results.md b/docs/edge/results.md new file mode 100644 index 000000000..df01f5d95 --- /dev/null +++ b/docs/edge/results.md @@ -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.