diff --git a/README.md b/README.md index a0115c3b0..ae6f0288d 100644 --- a/README.md +++ b/README.md @@ -27,6 +27,7 @@ Upstream's own README is preserved at [docs/README.vllm-omni.md](docs/README.vll | Memory | RoPE tables sized to the servable context; weights repacked in blocks; a KV budget that is actually honoured on CPU | | CPU platform + vLLM patches | `vllm_omni/platforms/cpu/`, `vllm_omni/patch.py` — vLLM is monkey-patched at import; its *source* is unmodified and no vLLM is redistributed. See [docs/edge/environment.md](docs/edge/environment.md) | | Device routing | `vllm_omni/edge/weight_path.py` picks the weight format per device class, from measurements rather than assumption | +| Device placement | `vllm_omni/edge/placement.py` decides which parameters leave the GPU when the checkpoint is bigger than the VRAM — cold ones first, because `embed_tokens` and `lm_head` are the same size and the opposite decision | | Measurement harness | `benchmarks/edge_harness/` — the grid runner, benchmarks, memory ledger, fidelity metric and agent suite every number here came from | | Raw results | `benchmarks/edge_harness/results/` — 100+ artifacts backing the published figures, each with provenance | | Baselines | `baselines/` — llama.cpp runs and the comparison scripts | @@ -104,7 +105,10 @@ print(llm.generate(['Hello'], SamplingParams(max_tokens=32))[0].outputs[0].text) | [docs/edge/cpu-4bit.md](docs/edge/cpu-4bit.md) | the two 4-bit paths, the packed layout as measured, building a checkpoint | | [docs/edge/comparisons.md](docs/edge/comparisons.md) | against llama.cpp and the other engines considered | | [docs/edge/environment.md](docs/edge/environment.md) | pins, machine, and why there is no vLLM patch series | +| [docs/edge/local-text-mode.md](docs/edge/local-text-mode.md) | the single-stage local text engine: device capabilities, memory admission, explicit refusals, cancellation | +| [docs/edge/external-stage-backends.md](docs/edge/external-stage-backends.md) | placing a stage on the Radeon 890M or the AMD NPU: out-of-process workers, proven placement, the shared-RAM ledger | | [docs/edge/harness.md](docs/edge/harness.md) | running a comparison | +| [docs/edge/laptop-heterogeneous.md](docs/edge/laptop-heterogeneous.md) | Qwen3.8-27B and Spark-X2.5-4B on a Ryzen AI laptop — what the CPU, NPU and GPU are each worth when decode is bandwidth-bound | | [ROADMAP.md](ROADMAP.md) | what is next, with expected gains and how each would be verified — and what was rejected | ## Licence diff --git a/docs/edge/external-stage-backends.md b/docs/edge/external-stage-backends.md new file mode 100644 index 000000000..ca9fb127c --- /dev/null +++ b/docs/edge/external-stage-backends.md @@ -0,0 +1,182 @@ +# External stage backends: the Radeon 890M and the AMD NPU + +`vllm_omni.edge.local.external` lets the local engine place one stage on a device +vLLM does not drive, prove it actually ran there, and charge its memory to the +same ledger as everything else. Status: the plumbing is done and measured +(M3.0); no model component runs on either device yet (M3.1–M3.4). + +## Why these devices need another process + +Neither can be opened from the engine's interpreter, and for different reasons. + +The **XDNA2 NPU** is an MCDM device. WSL2's GPU-PV forwards WDDM display +adapters only, so `/dev/accel` does not exist inside WSL and no amount of +installing there changes it; VitisAI is a Windows DLL besides. + +The **Radeon 890M** *is* reachable from WSL, over DirectML through `/dev/dxg` — +but `torch-directml` pins torch 2.4.1 and Omni runs 2.13. + +They cannot even share a Windows venv: `onnxruntime-directml` is pinned at +1.24.4, and the VitisAI EP needs onnxruntime ≥ 1.25. + +So: one protocol, four routes, each in its own interpreter. + +| route | interpreter | artifact | status on this laptop | +|---|---|---|---| +| `ort-vitisai` | native Windows, plain ORT ≥ 1.25 | `onnx:a16w8` | working | +| `torch-dml` | WSL `.venvs/dml` | `pt:fp16` / `pt:fp32` | working | +| `ort-dml` | native Windows, onnxruntime-directml | `onnx:fp16` / `onnx:fp32` | venv not installed | +| `ort-cpu` | anything with ORT | any ONNX | reference arm | + +`python -m vllm_omni.edge.local devices` prints which resolved; a route that did +not resolve says what is missing rather than disappearing. + +## Placement is verified, never assumed + +A VitisAI session lists the EP in `get_providers()` **whether or not it claimed +a single node**, and a graph it declined runs entirely on the CPU and returns +bit-identical numbers. Identical output cannot confirm NPU execution, and a +small difference cannot confirm it either. The only evidence is ORT's per-node +provider assignment, which exists only after something has run — so `open()` +always does one profiled warm-up run and the planner refuses on the node split. + +An *unmeasured* placement (`fraction_on_target is None`) is refused too. It is +not rounded to "fine". + +On the torch-directml route there is no partitioner, so the evidence is coarser +and equivalent: the device the output tensors came back from. A module that +fell back to the CPU returns CPU tensors. + +## Memory + +The 890M and the NPU report `memory_bytes == 0` on purpose. All three devices +allocate out of the CPU's RAM, so repeating the pool size on each row would let +a caller add three copies of the same memory together. Capacity comes from the +host row — the same number the CPU stage is being budgeted against. + +The worker is a different process, on Windows a different OS, so the engine's +psutil sampler cannot see it. It reports its own resident set, and that is what +closes the budget-versus-measurement loop. + +## Refusal codes + +| code | means | +|---|---| +| `no_exported_graph` | the device runs graphs and there is no export for this stage at this precision | +| `execution_provider_declined_graph` | the session ran and left the graph on the CPU | +| `numeric_gate_failed` | it ran on the device and disagreed with the reference | +| `device_unreachable_from_this_process` | the device works; no interpreter here drives it | +| `budget_exceeds_device` | the stage does not fit the shared host pool | + +All of them exit 2 from the CLI. A refusal is a result. + +## Using it + +```bash +python -m vllm_omni.edge.local external \ + --graph out/code2wav.onnx --format onnx:fp16 --component code2wav \ + --device gpu_integrated:amd --runs 20 --json record.json +``` + +`--device` requires that device, so "does the NPU run this?" is answered with a +refusal rather than quietly with the iGPU. `--prefer` falls through instead. +`--min-placement` sets how much of the graph the target EP must take. + +## First real model component: Qwen3.8-27B's vision tower + +`vllm_omni.edge.local.artifacts.vision_tower` exports it at a fixed image size. +It is the component that best fits the transport constraint — 27 blocks run once +per image, 4.8 MB in and 4.0 MB out — and it needs no quantization story for the +iGPU, because the FP8 checkpoint keeps all 333 visual tensors in bf16. + +The export wraps `transformers`' `Qwen3_5VisionModel` (whose parameters match the +checkpoint name for name, 333/333) and only fixes its shapes; the wrapper's +output is **bitwise identical** to the model's own `forward`. At a fixed size the +bilinear indices, position ids and `cu_seqlens` are constants, so folding them +into buffers removes every integer `Gather` from the graph. + +| | 890M (DirectML, fp32) | AMD NPU (A16W8) | +|---|---|---| +| placement | **85 of 113 nodes** | **0 of 3246** — declined | +| latency | **227.0 ms** device, 237.2 round trip, 11.4 transport | — | +| vs torch CPU (1363 ms) | **6.0×**, or 5.7× including transport | — | +| vs RTX 5090 (95.9 ms) | 2.4× slower | — | +| parity vs torch CPU fp32 | 9.10e-06 normalized, cosine 0.99999994 | — | + +Moving the tower off the discrete GPU frees **1862 MiB** of its VRAM, which +matters for a 21–29 GB model on a 23.89 GiB card. + +The NPU refusal is about **this artifact**, not the model or the precision: the +same A16W8 scheme applied to the merger subgraph alone *was* accepted (1 of 8 +nodes), and cutting requantization from 3058 to 984 Q/DQ nodes did not change the +whole-graph answer. The next step is a bounded bisect of what the partitioner +will take, not a broader claim. + +One number worth keeping straight when comparing devices: the 890M's fp32 output +is **closer** to the CPU reference (9.10e-06) than the RTX 5090's is (5.65e-04), +because CUDA runs matmul in TF32 by default. That is a statement about what each +device actually computed, not about which is more accurate. + +## Measured, 2026-09-15, this laptop + +Transport, WSL ↔ native Windows: **11 µs** half round trip, **330 MB/s**. + +| | NPU (A16W8 probe graph) | 890M (100.7 M-param torch.export graph) | +|---|---|---| +| device compute | 0.855 ms | 9.749 ms | +| same graph on CPU | 6.270 ms | 49.661 ms | +| transport | **2.797 ms** (cross-OS) | 0.523 ms (same OS) | +| placement | vitisai 1 / CPU 2 nodes | outputs on `privateuseone` | +| numerics vs CPU | 3.03e-05 normalized | — | + +**The transport column is the design constraint.** On the NPU it consumes more +than half of a 5.4 ms compute saving, because that graph moves 1 MiB in and +1 MiB out. An external stage has to be compute-heavy and I/O-light: a vision +tower that runs once per image qualifies, a decode step that runs once per token +does not. This is the measured basis for the proposal's "coarse stages only". + +The 890M's 4.8× includes transport, and deserves a caveat: the iGPU hit +1.32 TFLOP/s, matching its recorded 1.23 fp32 rate, while the CPU path reached +only 0.26 TFLOP/s against a 1.54 peak. It is as much a slow CPU path at that +shape as a fast iGPU. Separately measured this session: the 890M does +**2.85 TFLOP/s in fp16**, 2.3× the fp32-only figure the earlier notes carried. + +## Exporting a graph these backends will accept + +Run `sanitize_onnx` on every dynamo export. It is not a nicety. + +torch's dynamo exporter emits `Reshape` with `allowzero=1`. onnxruntime's +DirectML provider rejects that at **session initialisation** with a bare +`E_INVALIDARG` out of `MLOperatorAuthorImpl.cpp` — and onnxruntime then does not +raise. It prints `Falling back to ['CPUExecutionProvider'] and retrying` and +hands back a working CPU session. Measured here: the 27B vision tower placed +**0 of 1171 nodes** on the iGPU and computed correct features entirely on the +CPU. With 58 `allowzero` attributes cleared, the same graph placed **85 of 113**. + +`vllm_omni.edge.qnn_export.sanitize_onnx` was written for Qualcomm's converter +and fixes exactly this, along with `IsNaN`/`Where` guard pairs and rank-0 `axes` +initializers — all generic dynamo warts rather than vendor-specific ones. It +preserves external data, so a multi-gigabyte export stays a pair of files. + +`vllm_omni.edge.local.artifacts.vision_tower.export_onnx` always calls it. + +Two independent guards catch the fallback if a graph slips through: the worker +compares `session.get_providers()` against what was requested and reports the +swap by name, and the node-assignment gate refuses on the placement itself. + +## Known limits + +- **A16W8 is the only door on XDNA2.** A8W8 — what `quantize_static` produces by + default — and fp32 are declined with no error and every node left on the CPU. +- **Dequantize nodes stay on the CPU** by the partitioner's own account, so a + graph that is mostly requantization will not benefit. +- **TorchScript cannot be placed on DirectML** (torch 2.4.1 / torch-directml + 0.2.5): `map_location` raises `Cannot access storage of OpaqueTensorImpl`, and + load-then-`.to(device)` trips `assert isinstance(param, Parameter)` inside + `Module._apply`. Use `torch.export` `.pt2`, which is what + `vllm_omni.edge.decoder_export` already emits — and do not call `.eval()` on + the loaded module, which 2.4.1 does not support. +- Small graphs cannot amortise dispatch; the NPU's advantage grows with GEMM + size (2.6× at 256×1024×1024, 5.3× at 1024×2048×2048). + +Raw records: [`analysis/experiments/m3_amd_devices_20260915/`](../../../analysis/experiments/m3_amd_devices_20260915/README.md) diff --git a/docs/edge/laptop-heterogeneous.md b/docs/edge/laptop-heterogeneous.md new file mode 100644 index 000000000..8239ae3ae --- /dev/null +++ b/docs/edge/laptop-heterogeneous.md @@ -0,0 +1,466 @@ +# Qwen3.8-27B and Spark-X2.5-4B on a Ryzen AI laptop + +Target machine: **AMD Ryzen AI 9 HX 370** (Zen 5, 12 cores / 24 threads, +AVX-512 with VNNI and BF16, **no AMX**), 30.9 GiB LPDDR5X, **RTX 5090 Laptop** +(24463 MiB, sm_120, PCIe 4.0 x8), **Radeon 890M** iGPU and an **XDNA2 NPU**, +running under **WSL2** (kernel 6.18.33.2). + +Two models were added and both run: + +| model | status | evidence | +|---|---|---| +| `XHToken/Spark-X2.5-4B` | runs; **24/24 greedy tokens identical to the HF reference** | `analysis/experiments/spark_edge/parity_hf_vs_vllm.py` | +| `Qwen/Qwen3.8-27B-FP8` | runs; correct output; 1.05 tok/s at the tightest fit | offload sweep below | + +Neither needed a new model class. Spark-4B is the same `Spark2_5ForCausalLM` +the 1.7B uses and the implementation was already config-driven; Qwen3.8 is +`Qwen3_5ForConditionalGeneration`, which vLLM 0.28 already carries. What both +needed was the *sizing* and *placement* code around them, which is what this +document is about, and four environment faults, which are in +[§7](#7-what-was-actually-broken). + +--- + +## 1. Three accelerators, and what each one costs to reach + +The request was to use the CPU, the NPU and the GPU together. An earlier +version of this document said the iGPU and the NPU were simply unreachable +under WSL2. That was true of the default setup and false as a statement about +the machine: both have a route, the routes are not equivalent, and finding them +changed one of the two answers. + +| device | reachable from WSL? | route | state | +|---|---|---|---| +| RTX 5090 Laptop | yes | `/dev/dxg` → `torch.cuda` | working | +| Radeon 890M iGPU | **yes** | DirectML over `/dev/dxg` (`torch-directml`) | **working, and slower than the CPU** | +| XDNA2 NPU | no | Windows-side ONNX Runtime + **VitisAI** EP | **working** — 2.6–5.3× the CPU on GEMMs | + +### The iGPU: reachable, and measured not worth it + +There is no `/dev/dri` under WSL2, and there does not need to be. Mesa ships a +D3D12 Gallium driver and Microsoft ships `libd3d12.so` / `libdxcore.so` in +`/usr/lib/wsl/lib`, so any WDDM adapter is addressable through `/dev/dxg`. +`glxinfo` already reports `D3D12 (AMD Radeon(TM) 890M Graphics)`, accelerated, +OpenGL 4.6. For compute, `torch-directml` offers it as a torch device: + +``` +DirectML device_count: 2 + [0] NVIDIA GeForce RTX 5090 Laptop GPU + [1] AMD Radeon(TM) 890M Graphics +``` + +It computes correctly — a 512×512 matmul agrees with the CPU to **3.6e-07**. +Then it was benchmarked, which is why this section is short: + +| | iGPU via DirectML | CPU | ratio | +|---|---:|---:|---| +| fp32 GEMM 2048³ | 1.12 TFLOP/s | 1.54 TFLOP/s | **0.73×** | +| memory bandwidth | 5.0–7.4 GB/s | 20.7 GB/s | **0.30×** | +| host↔device, 64 MB | 39.2 ms (3.3 GB/s) | — | — | + +The iGPU is **slower than the CPU it shares memory with, on both axes**: the +D3D12 translation layer costs more than the hardware wins. It is genuinely +usable now, and the honest advice is to route work to it only when the point is +to free the CPU rather than to go faster. A native path — ROCm on bare-metal +Linux, or DirectML inside a Windows process — would not pay the translation and +is the version worth measuring. + +Note the pin conflict: `torch-directml` requires `torch==2.4.1`, so it can +never share a virtualenv with the engine (torch 2.13). It gets its own. + +### The NPU: working, at 2.6-5.3x the CPU + +Not from WSL, and no driver install changes that: it is an **MCDM** device +(`IpuMcdmDriver`, AMD 32.0.203.329), and WSL2's GPU-PV forwards WDDM *display* +adapters only, so it has no device node in Linux. PCI passthrough is not a WSL2 +feature. It has to be driven from a Windows-side process. + +From Windows it works. Three things have to be right **at once**, and each one +failing on its own produces exactly the same symptom — a session that builds, +runs, returns correct numbers, and executes every node on the CPU: + +1. **The right EP.** The appx package ships two. + `onnxruntime_providers_ryzenai.dll` registers cleanly, enumerates the NPU as + `OrtHardwareDeviceType.NPU`, opens a session — and claims **zero nodes** from + everything: fp32, INT8 QDQ symmetric and asymmetric, per-channel, + power-of-two scales, MatMul and Conv, and ahead-of-time `ModelCompiler` + (which emitted zero `EPContext` nodes). It is the *Light* EP. The general one + is `onnxruntime_vitisai_ep.dll`. +2. **The DLL search path.** That one fails to load from a bare + `register_execution_provider_library` with `Error loading … But no + dependency`. It needs its own directory added via `os.add_dll_directory` + first, because it pulls `dyn_dispatch_core.dll` and `vaiml.dll` from beside + it. +3. **A16W8, not A8W8.** The overlay names are embedded in the DLL and say so: + + ``` + 4x2_psf_model_a8w8_qdq.xclbin + 8x4_psu_model_a16w8_qdq.xclbin <- most of them + 2x4x4_bmm_model_a16w16.xclbin + ``` + + XDNA2 wants **16-bit activations with 8-bit weights**. A8W8 — the default + shape of `quantize_static` — is rejected as completely as fp32, silently. + +Also needed: **onnxruntime ≥ 1.25**, because the EP requests ORT API 25. +`onnxruntime-directml` stops at 1.24.4 and is refused, so the NPU and DirectML +cannot share a process. + +With all of it right the partitioner starts talking and names its target, +`AMD_AIE2P_4x8_CMC_Overlay`. Measured against the same INT8 graph on the CPU: + +| GEMM | CPU | NPU | speedup | NPU | +|---|---:|---:|---:|---:| +| 256 × 1024 × 1024 | 1.47 ms | 0.56 ms | 2.64× | 0.97 TOPS | +| 512 × 2048 × 2048 | 9.64 ms | 2.14 ms | 4.51× | 2.01 TOPS | +| 1024 × 2048 × 2048 | 19.89 ms | 3.76 ms | **5.29×** | **2.29 TOPS** | + +Max relative error against the CPU: **3.2e-05**. It scales the right way — a +small GEMM cannot amortise the dispatch. Dequantize nodes stay on the CPU (the +partitioner says so explicitly), so a graph that is mostly requantization gains +nothing. + +The recipe is code in `vllm_omni/edge/npu_ryzenai.py`, including the one check +worth trusting: **count nodes in an ORT profile.** A session reporting the EP in +`get_providers()` may still be running everything on the CPU — and when it does, +its outputs are bit-identical to the CPU, which is precisely what a *correct* +NPU run does not look like. A 0.000e+00 error is the bug signature, not the +success signature. + +**What it is worth here.** Not decode: 2.29 TOPS against the dGPU's +48.2 TFLOP/s, over the same 20.7 GB/s system memory the CPU uses, cannot help a +bandwidth-bound step (§3). It earns its place on the compute-bound work off the +per-token path — a vision tower, a draft model for speculative decoding — where +5× the CPU on large GEMMs is real. + +### Reaching a phone + +The same question for mobile has a cleaner answer. The harness's `AdbDevice` +backend needs only an `adb` binary and a reachable device, and WSL2 is in NAT +mode, so **outbound** connections to the LAN work even though inbound do not. +That makes wireless debugging the path of least resistance: + +```bash +sudo apt-get install -y adb # done; adb 1.0.41 (34.0.5) +adb connect :5555 # Android 11+ wireless debugging +``` + +No USB passthrough is involved, which matters because WSL2 has none by default +— the USB route needs `usbipd-win` on the Windows side, and is worth the extra +moving part only when the phone cannot be on the same network. `AdbDevice` also +takes an explicit binary path, so `adb="/mnt/c/.../adb.exe"` drives a +USB-attached phone through the Windows adb server with no passthrough at all, +at the cost of translating paths on every `push`/`pull`. + +The mobile **export** path is now working as well: the QNN export tests needed +`onnxscript` and `tabulate`, and the ExecuTorch lowering needed `flatc`. All +three are installed, and the one remaining ExecuTorch failure turned out not to +be a failure — pytest run from the repository root lets the bare `executorch/` +source checkout shadow the installed package. Run from `vllm-omni/` the suite +is **124 passed, 1 skipped**. + +## 2. The measurement that decides the design + +Taken on an **idle** machine (see §7 for why that word is load-bearing): + +| path | measured | +|---|---| +| GPU VRAM bandwidth | **472 GB/s** | +| GPU bf16 GEMM 8192³ | **48.2 TFLOP/s** (91 W, 1447 MHz) | +| PCIe H2D, pinned, bulk | **13.9 GB/s** | +| PCIe, as seen by offloaded weight reads | **10.6 GB/s** (§5) | +| CPU RAM bandwidth | **20.7 GB/s** | +| CPU bf16 GEMM 4096³ | **1.54 TFLOP/s** | + +Two ratios follow, and everything else in this document is a consequence: + +``` +VRAM : system RAM = 22.8 x +VRAM : PCIe link = 34 x (44x on the effective figure) +``` + +Autoregressive decode reads **every weight once per token** and does two FLOPs +per weight. At batch 1 it is bandwidth-bound with no room for argument, so +decode rate is set by *where the weights are*, and almost nothing else. + +--- + +## 3. Why "use all three devices for decode" is the wrong goal + +The intuition -- three processors, so split the model three ways -- inverts on +a bandwidth-bound workload. Splitting weights across pools does not add +bandwidth; the step ends when the *slowest* pool finishes, and the slow pool is +22.8x slower. Putting 10% of a model on the CPU makes the step +`0.9/472 + 0.1/20.7`, which is **3.2x slower** than putting all of it on the +GPU, not 10% slower. + +The iGPU and the NPU do not escape this: both read the same LPDDR5X the CPU +does, so all three share one 20.7 GB/s pool — and the iGPU, measured, reaches +only 5-7.4 GB/s of it (§1). Adding them to a decode split adds arithmetic to a +problem that is not short of arithmetic, over a pool that is already the +bottleneck. + +So the rule is: **for decode, one device, and it is the one with the memory +bandwidth.** The other devices earn their place on work that is *not* on the +per-token critical path: + +| device | job | why it fits there | +|---|---|---| +| dGPU | every decode step, all resident layers | the only 472 GB/s in the machine | +| CPU | `embed_tokens` gather, detokenisation, scheduling, prompt templating | a gather reads one 10 KiB row per token, not the 2.54 GB table | +| NPU | the **vision tower** on image prompts; a **draft model** for speculative decoding | both are compute-bound and off the decode path — *pending the compiler in §1* | +| iGPU | the same class of work, when the goal is to free the CPU | reachable today, but measured at 0.73x the CPU's compute and 0.30x its bandwidth (§1) | + +The vision tower is the clean case: 0.92 GB, run **once per image**, never per +token. Moving it off the GPU frees 0.92 GB of VRAM -- which by §5 is worth +86 ms/token -- and costs one prefill-time hop. That is a real heterogeneous win +and it does not touch the decode loop. + +Speculative decoding is the other one, and it is the only way a second device +raises *decode* throughput: a draft model on the NPU proposes k tokens, the GPU +verifies them in a single batched pass, and the GPU's per-token weight read is +amortised over k tokens. Qwen3.8-27B **ships an MTP head for exactly this** +(`mtp_num_hidden_layers: 1`, 0.48 GB, in `mtp.safetensors`). + +--- + +## 4. What the two models actually cost + +`vllm_omni/edge/kv_budget.py` now prices both. It could price neither before: +it read `num_hidden_layers` off the top-level config, which `Qwen3_5Config` +does not have, and it charged every layer for a KV cache. + +**Qwen3.8-27B-FP8** is a hybrid of a kind the old code had no concept of: of its +64 layers, **48 are gated DeltaNet**, which keep a fixed-size recurrent state +and *no KV cache at any context length*. Only the 16 full-attention layers +cache. + +``` +flat over 64 layers 256 KiB/token <- what a layer-blind budget charges +16 full layers only 64 KiB/token <- the truth, 4x smaller +GDN recurrent state 147 MiB/sequence <- constant, and previously invisible +``` + +The state is constant and the cache is linear, so they cross at ~2350 tokens: +below that, the thing the old budget did not count at all is the larger half. +At a 4096-token context the whole cache plus state is **467 MiB** (147 MiB of +state, 320 MiB of cache), and **787 MiB** at 8192 -- against 30.87 GB of +weights. For this model, **memory is a weights problem and not a +cache problem**, and long context is nearly free. That is the architectural +gift of the 3:1 linear hybrid and it is the strongest reason to want this model +on an edge box. + +**Spark-X2.5-4B** is the other hybrid: 27 sliding-window(512) layers to 9 full +ones, 144 KiB/token flat, and a **3.37x** discount if the backend accounts per +layer type. Both sizes now have their published geometry pinned in +`tests/model_executor/models/spark2_5/test_spark2_5.py`, so a change that +assumes the 1.7B's 8 query heads or 28 layers fails in CI rather than on a +checkpoint. + +--- + +## 5. The cost of not fitting, measured + +Qwen3.8-27B-FP8 is **30.87 GB** of weights against a 24 GiB card. It only runs +with part of it on the host. Sweeping how much (4096-token context, batch 1, +greedy, `enforce_eager`, idle machine): + +| offloaded | tok/s | s/token | marginal | +|---|---:|---:|---| +| 10 GB (the tightest fit; 9 GB fails) | **1.05** | 0.952 | | +| 12 GB | 0.89 | 1.124 | +0.086 s/GB | +| 14 GB | 0.76 | 1.316 | +0.096 s/GB | +| 16 GB | 0.66 | 1.515 | +0.100 s/GB | + +Straight line, slope **+94 ms per offloaded GB**, i.e. an effective +**10.6 GB/s** -- the 13.9 GB/s bulk figure less ~20% of overhead. The model + +``` +s/token = resident_GB / 472 + offloaded_GB / 10.6 +``` + +predicts the 10 GB point at 987 ms against 952 ms measured, and the whole curve +within 5%. At the tightest fit, **the 10 GB on the host is 94% of the step and +the 20.9 GB in VRAM is 4%.** + +This is the number to design against: **a gigabyte that does not fit costs +44x what a gigabyte that fits costs** (94 ms against 2.1 ms). Fitting is not +one optimisation among several; it is the only one that matters until it is +solved. + +### Which gigabytes to spend + +vLLM's `cpu_offload_gb` offloads "non-selectively until the limit is reached", +which is the expensive way. Not every parameter is read once per token, and +`vllm_omni/edge/placement.py` sorts them: + +``` +lm_head 2.543 GB HOT one 248320 x 5120 GEMM per token +embed_tokens 2.543 GB COLD a gather: one row, ~10 KiB per token +visual.* 0.922 GB COLD unless the prompt carries an image +mtp.* 0.478 GB COLD unless speculative decoding is on + --------- +cold set 3.94 GB free to evict +``` + +`embed_tokens` and `lm_head` are the same shape and the same dtype and the +opposite decision. A size-ordered, name-blind policy cannot tell them apart; +`cpu_offload_params`, which matches name segments, can. + +For FP8 the cold set is 3.94 GB against a ~10.6 GB deficit, so it helps but +cannot close it -- **selective placement does not rescue this checkpoint**. + +And a larger caveat, found after this section was first written: *the cold set +cannot currently be offloaded at all*. `cpu_offload_params` is real and matches +name segments as documented, but the offloader is only ever offered decoder +layers (`make_layers`, `utils.py:812`), so `embed_tokens`, `lm_head`, the vision +tower and the MTP head are out of its reach. The classification below is right +about which bytes are cheap to move; moving them needs an upstream change. See +§6. + +--- + +## 6. The design: make it fit — and the retraction that follows + +The previous version of this section said: use a 4-bit checkpoint, it is worth +~20×, everything else is worth less than 2×. That was arithmetic — +`resident_GB / 472 GB/s` — and the arithmetic was **wrong by roughly 30×**. The +4-bit build was then measured, and it is *slower* than the 8-bit one. + +| build | weights | offloaded | decode | +|---|---:|---:|---:| +| `Qwen/Qwen3.8-27B-FP8` | 30.87 GB | 10 GB | **1.05 tok/s** | +| `nvidia/Qwen3.8-27B-NVFP4` | 21.95 GB | 3 GB | **0.67 tok/s** | +| the same, torch.compile + CUDA graphs | 21.95 GB | 3 GB | 0.72 tok/s | +| the same, 4 GB offloaded | 21.95 GB | 4 GB | 0.56 tok/s | + +Three things went wrong at once, and each is worth stating separately. + +**It does not fit, even at 21.95 GB.** The weights load to 19.92 GiB and the +card has 22.58 GiB free (the Windows desktop holds the rest). vLLM's runtime +overhead — non-torch allocations plus the activation peak — is about **2.8 GiB** +on this model, which leaves the KV cache at **−0.31 GiB**. Shrinking +`max_model_len` from 4096 to 1024 moved that to −0.30: the overhead is not the +cache, so context length cannot buy the fit. + +**The lever that would have made it fit does not exist.** §5 identified 4.31 GB +of cold parameters and prescribed `cpu_offload_params={embed_tokens, visual, +mtp}`. The flag is real and matches name segments as documented. But the +offloader is only ever *offered decoder layers*: `wrap_modules` is called from +inside `make_layers` (`vllm/model_executor/models/utils.py:812`) and nowhere +else. `embed_tokens`, `lm_head`, the vision tower and the MTP head are built +outside it. Asking for the cold set moved **~0.5 GiB** and the run failed +identically to one with no offload at all. The classification was right about +which bytes are cheap to move; the mechanism to move them is missing, and +closing it is an upstream change. + +**And fewer bytes turned out not to mean faster.** This is the part that +invalidates the projection rather than merely complicating it. The traffic model +`resident/472 + offloaded/10.6` fits the FP8 offload curve to within 4.3% across +four points. Against NVFP4 it is **4.7× optimistic**: predicted 3.1 tok/s at +3 GB offloaded, measured 0.67. Its own offload slope is 0.293 s/GB — an +effective **3.4 GB/s**, against FP8's 10.6 — and extrapolating to zero offload +gives ~**1.6 tok/s** even if it fit. Enabling torch.compile and CUDA graphs +changed 0.67 to 0.72, so this is not an eager-mode artifact. + +The conclusion is that **NVFP4 decode on this stack is compute-bound, not +bandwidth-bound**, and a model that counts bytes cannot see that. Why it is +compute-bound is not established here: the checkpoint is a `modelopt` mix of +FP8 and NVFP4 with group-16 scales, and this fork already carries a correctness +patch for the W4A4 NVFP4 linear path, which is a hint that the path is not a +well-trodden one on sm_120. Diagnosing it is the obvious next experiment, and +until someone does, the honest statement is the measurement. + +### What the design actually is, then + +1. **Run the FP8 checkpoint with the smallest offload that fits** — 10 GB, 1.05 + tok/s. This is the fastest configuration measured on this machine, and it is + the 8-bit one. +2. **Do not change quantization format on an estimate.** The one estimate this + document made without measuring is the one it got wrong, by 30×. +3. **Size the cache from the model** — 467 MiB at 4096 tokens (§4). Still right, + and still worth little here, because weights dominate. +4. **Speculative decoding via the shipped MTP head** is now the largest + *unexplored* lever, and the only one that attacks the per-token weight read + rather than working within it. It is untested here. +5. **Getting the cold set offloadable upstream** would free 3.94 GB on FP8, worth + ~370 ms/token by the traffic model — which, on the FP8 path where that model + demonstrably holds, is the difference between 1.05 and roughly 1.6 tok/s. + +The blunt summary: a 27B model on a 24 GB laptop GPU runs at **about one token +per second**, and the things that would change that by an order of magnitude — +a checkpoint that genuinely fits, or speculative decoding — are not yet in hand. + +## 7. Where this beats llama.cpp, and where it does not + +llama.cpp at the pinned revision **does** support this architecture -- `QWEN35`, +`QWEN35MOE` and `QWEN3NEXT` are in its arch registry and it has gated-DeltaNet +kernels. It is a real competitor here, and on a fitted 4-bit model on one +stream it is a strong one: this repository's own CPU measurements have +llama.cpp's Q4_K_M decode at 90 tok/s against our 74, and `docs/edge/` explains +that the gap is the execution model rather than the kernels. + +So the case for vLLM-Omni on this machine is not raw single-stream decode. It is: + +| axis | why | +|---|---| +| **Prefill** | measured ~2x llama.cpp on the CPU path (2190 vs 1121 tok/s, pp512). Prompt processing is compute-bound, which is the one regime where the hardware has headroom. | +| **Concurrency** | continuous batching and paged KV. At batch 1 they cost overhead; past ~4 streams they are the whole game, and llama.cpp has no equivalent. | +| **Speculative decoding** | the shipped MTP head is wired in vLLM and is not a GGUF concept. | +| **Multimodal pipelines** | the reason this fork exists: staged vision/audio/codec execution with an OpenAI-compatible server. | +| **Hybrid cache accounting** | 64 KiB/token instead of 256, once §4 is honoured. | + +A design that claims to beat llama.cpp at batch-1 decode of a fitted model on +one GPU is claiming something this project has already measured to be false. +The defensible claim is the four rows above, and the way to establish it is +`benchmarks/edge_harness/` against `baselines/`, not assertion. + +--- + +## 8. What was actually broken + +Four faults stood between "the models are supported" and "the models run", none +of them in the model code: + +1. **`RuntimeError: UVA is not available`.** vLLM disables pinned memory on WSL + by default; its own gate is kernel ≥ 4.19.121 and this kernel is 6.18. Fixed + by `VLLM_WSL2_ENABLE_PIN_MEMORY=1`. Without it the CUDA runner does not + start **at all**, for any model. +2. **No C compiler**, so Triton could not JIT and the GDN kernels could not + build. `gcc` plus `libc6-dev` (gcc alone did not pull the headers). +3. **FlashInfer's sampler wants `nvcc`**, which a runtime-only CUDA install does + not have. `VLLM_USE_FLASHINFER_SAMPLER=0`. Note this surfaced as a failure in + *warmup*, long after the model had loaded and run a forward pass. +4. **transformers 5 renamed `input_embeds` → `inputs_embeds`**, which broke the + vendored Spark modeling file used as the parity *reference*. Shimmed in + `parity_hf_vs_vllm.py` rather than by editing the cached module, so the + reference stays the code the checkpoint ships. + +The required environment, in full: + +```bash +export VLLM_WSL2_ENABLE_PIN_MEMORY=1 +export VLLM_USE_FLASHINFER_SAMPLER=0 +sudo apt-get install -y gcc libc6-dev +``` + +### And one measurement error, which is the real lesson + +The first roofline taken on this machine read **3.67 GB/s** H2D, **38 GB/s** +VRAM and **9 TFLOP/s**, with `nvidia-smi` reporting a 33 W power cap and a +220 MHz SM clock. The conclusion drawn from it -- that the dGPU had no +bandwidth advantage over the CPU and the design should lean on the CPU -- was +wrong in every part. + +Those numbers were taken **while two model downloads were saturating the disk, +the network and the CPU**. Re-measured on an idle machine the same calls give +13.9 GB/s, 472 GB/s and 48.2 TFLOP/s, at 91 W and 1447 MHz, with the power +limit reading 100.98 W. The GPU was never capped; it was idling because the +benchmark could not feed it. + +This repository already has [a methodology +document](methodology.md) that exists because of five retractions of exactly +this shape. This is the sixth, and the rule it re-teaches is the one already +written down: **a number taken on a contended machine is not a number.** The +harness has a contention guard for this; it was not in the path of an +interactive measurement, which is precisely where the mistake got made. diff --git a/docs/edge/local-text-mode.md b/docs/edge/local-text-mode.md new file mode 100644 index 000000000..7ab00781d --- /dev/null +++ b/docs/edge/local-text-mode.md @@ -0,0 +1,286 @@ +# The local text mode: one stage, one device, and a plan you can argue with + +`vllm_omni.edge.local` runs a single autoregressive text stage on whichever +device *this process* can actually drive, and writes down why. It is the M0 +milestone of the [single-device engine +proposal](../../../analysis/architecture_local_engine_20260915.md): the smallest +thing that turns "a configuration that happens to run" into "a configuration +that can be explained, reproduced and refused". + +It adds no kernels, no scheduler and no graph compiler. vLLM keeps doing all of +that inside the Omni stage. What this adds is the layer above: which device, +which artifact, how many bytes, and what happens when the answer is no. + +```bash +python -m vllm_omni.edge.local devices # what this process can drive +python -m vllm_omni.edge.local plan --model # exit 2 == explicit refusal +python -m vllm_omni.edge.local run --model --prompt "..." +python -m vllm_omni.edge.local accept --model --json out.json +``` + +--- + +## 1. "It fits" is not "it runs" + +`models/Spark-X2.5-1.7B-int8` is 1.84 GiB of compressed-tensors int8 W8A8. The +RTX 5090 Laptop in the reference machine has 23.89 GiB. Every capacity check +passes, the engine starts, the weights load — and the first forward dies: + +``` +dispatch_scaled_mm, csrc/.../c3x/scaled_mm_helper.hpp:34, +Int8 not supported on SM120. Use FP8 quantization instead, or run on older arch (SM < 100). +``` + +cutlass' `c3x` has no int8 `scaled_mm` for SM >= 100, and sm_120 is inside that +range. Roughly forty seconds and a stack trace from a subprocess to learn +something that is written in the checkpoint's `quantization_config` and the +device's compute capability. + +The planner reads both and refuses in about a second: + +``` +REFUSED: no device in this process can run this checkpoint as asked. + [weight_format_unsupported] cuda:0: the checkpoint is compressed-tensors:int8 and + cuda:0 (NVIDIA GeForce RTX 5090 Laptop GPU, sm_120) executes only + ['awq', 'compressed-tensors:fp8', 'dense', 'gptq', 'modelopt_fp4'] + remedy: this device has fp8 tensor cores but no int8 scaled_mm path + (cutlass c3x has none for SM >= 100). Use an fp8 or a dense build here, + or run the int8 build on the CPU backend (.venvs/omni-cpu). +``` + +The remedy is not decoration. The same checkpoint completes the full acceptance +run on the CPU backend, at 16.3 tok/s median — so this is a property of the +device, not a defect in the artifact, and the refusal says which. + +## 2. What a device *is*, here + +`hardware_probe` answers "what is in the box". That is the wrong question at +load time, because a vLLM wheel is built for one platform: the CUDA wheel +cannot run a CPU stage however many cores the probe finds, and the CPU wheel +cannot reach the card beside it. So a device is runnable only when the probe +finds it **and** `vllm.platforms.current_platform` is the platform that drives +it. Everything else is enumerated with the reason it is not available: + +``` + [runnable] cuda:0 NVIDIA GeForce RTX 5090 Laptop GPU (23.9GiB vram) backend=vllm:cuda + [not here] npu:amd NPU Compute Accelerator Device (host_ram) backend=None + An MCDM device; WSL2's GPU-PV forwards WDDM display adapters only... + [not here] gpu_integrated:amd AMD Radeon(TM) 890M Graphics (host_ram) backend=None + No /dev/dri under WSL2, but reachable through DirectML over /dev/dxg... + [not here] cpu AMD Ryzen AI 9 HX 370 (30.9GiB host_ram) backend=None + vllm.platforms.current_platform is 'cuda', not 'cpu'... +``` + +The iGPU and the NPU are reported rather than dropped, because "the machine has +an NPU" and "a stage can run on it" are different claims and a plan has to be +able to show which one it relied on. They share `host_ram` with the CPU and +carry zero bytes, so nothing can accidentally add three copies of the same RAM +together. + +`--mask gpu_discrete` hides the discrete GPU to exercise the no-NVIDIA +deployment class. It is a policy, not amnesia: the device stays in the record +with `masked: true` and its real memory. + +## 3. The budget is a sum, and every term is visible + +``` +ADMITTED on cuda:0 (NVIDIA GeForce RTX 5090 Laptop GPU) via vllm:cuda + artifact 43e867b180f1a3d2 dense 7.66GiB + budget 12.51GiB of 23.89GiB vram (total) + weights 7843 MiB vram session_weights D 5 safetensors shard(s), on-disk size + kv_cache 720 MiB vram session E flat budget at 1x context; 248 MiB if + the backend accounts 27 sliding layers separately + activations 120 MiB vram request E 2048-token prefill chunk x 12 live tensors + backend_workspace 512 MiB vram session E attention scratch, driver/allocator context + load_transient 1569 MiB vram load E freed once the weights are resident + external_reserve 1024 MiB vram session E memory this engine does not own + safety_margin 1024 MiB vram session E fragmentation and allocator slack +``` + +`D` is read off the artifact; `E` is estimated. The peak is a **max**, not a +sum: the loader's transient is gone before vLLM sizes the KV pool, so adding +both would refuse plans that fit. + +The KV line comes from `vllm_omni.edge.kv_budget`, which reads the real +`layer_types`. Spark-4B is 9 full-attention and 27 sliding@512 layers, so a +backend that accounts them separately needs 248 MiB where a flat charge needs +720. The budget takes the flat number (it is the safe one) and reports the +other. + +**The estimates are checked, not believed.** After every run the engine reports +the budgeted peak against the sampled one: + +| | CUDA / Spark-4B bf16 | CPU / Spark-1.7B int8 | +|---|---:|---:| +| counter | `nvml_device_used` (whole card) | `tree_rss` (process tree) | +| budgeted peak | 12.51 GiB | 6.55 GiB | +| measured attributable peak | 8.76 GiB | 5.48 GiB | +| budget / measured | 1.43x | 1.20x | + +A budget above the measurement is the intended direction — admission needs a +provable upper bound, not a prediction — but 1.40x is loose enough to be worth +narrowing, and the number is in the record so somebody can. + +### The flag that is not GPU-only + +The plan sets `gpu_memory_utilization` on **both** pools, derived from its own +budget. On the CPU backend vLLM reads that same field as the fraction of *host +RAM* to reserve, so leaving it at its default there is not "no opinion" — it is +asking for 92% of the machine. A 6.55 GiB CPU plan failed to start with +`desired CPU memory utilization (0.92, 28.43 GiB)` until this was fixed. + +The KV pool is set as an absolute `kv_cache_memory_bytes` rather than left as +whatever remains inside that fraction, so the cache is the size the request +limits imply. Both backends honour it exactly: CUDA logs 754974720 bytes, CPU +logs `Explicitly set (0.14/30.91) GiB for KV cache`. + +## 4. Streaming, state and cancellation + +Three contracts, which is what the proposal says the first version needs: + +* **`ChunkEvent`** carries the *delta*, never the output so far, plus a `seq` + that detects loss and an `epoch` that retires it. +* **`StateHandle`** is an opaque reference to state the backend owns. It is + tied to the artifact that produced it (the same shapes from a different + quantization are not the same state) and is never `migratable` — moving a KV + cache between backends needs a measured layout conversion for the specific + pair. +* **`BoundedEventStream`** bounds chunks *and* bytes and blocks the producer, + which is the backpressure. A single event larger than the whole byte bound + still passes on an empty queue, or it would deadlock. + +Cancelling is not "stop calling `next()`". The backend has work in flight and +its tokens arrive after the cancel. `cancel()` retires the epoch first, then +aborts in the backend, then closes the stream — *before* waiting on the +producer, so a backend that hangs while unwinding delays cleanup rather than +the caller. In the recorded runs the fence discarded 3 in-flight events on CUDA +and 2 on CPU, and zero events from the retired epoch reached the consumer. + +The device counter does **not** fall on cancel, because vLLM's KV pool is +pre-allocated. What falls is the in-flight count. The report says so, so that +"memory did not drop" is not read as "nothing was released". + +## 5. Measuring memory without fooling yourself + +Three counters, three different accountings, never summed: + +* `nvml_device_used` is the whole card, every process. Under WSL2 + `nvmlDeviceGetComputeRunningProcesses` returns an empty list, so GPU bytes + **cannot** be attributed to this engine's PIDs on this host. "Attributable + peak" means peak minus a pre-load baseline (2.79 GiB here), not a private + total. +* `tree_pss` includes a proportional share of file-backed pages, which on an + mmap'd safetensors load is partly page cache. +* `host_available` moves for reasons that are not us. + +Sampling is every 0.25 s, so every peak is a lower bound, and the record says +that rather than assuming it away. + +## 6. What Spark needed before any of this worked + +Two gaps that made "the model is in the registry" and "the model runs under +Omni" different statements: + +1. **No `PipelineConfig`.** Spark's model definition had been registered for a + while, but with no pipeline `Omni(model=)` had nothing to + resolve, so the model was reachable only through plain `vllm.LLM`. Added as + a single `LLM_AR` text stage in + `vllm_omni/model_executor/models/spark2_5/pipeline.py`. +2. **A stage-call mismatch.** Omni's AR model runner hands every stage model + `sampling_metadata`, `logits_index` and `sampler`. Spark's + `ForCausalLM.forward` caught them in `**kwargs` and forwarded them into + `Spark2_5Model.forward`, which takes four arguments — `TypeError` on the + first forward. They now stop at the `ForCausalLM` boundary, where the + decoder stack has no use for them anyway. + +One stage is the design, not a placeholder. An autoregressive session stays on +one backend by default: prefill and decode share the KV cache, and a per-token +device boundary buys nothing and costs a synchronization. Multi-stage is M2, +and it needs the tensor/buffer contract that a single text stage — which hands +a buffer to nobody — cannot be used to design. + +## 7. Does the Omni path change the answer? + +This is the one claim the rest of M0 cannot make for itself. Everything above +is about the engine's own behaviour; this asks whether routing a decode through +the Omni stage changed the model's output. The validation plan makes it a +*hard* initial acceptance condition. + +Same checkpoint, same device, same request limits, same environment, greedy +with `ignore_eos` -- the only difference is the path: plain `vllm.LLM` straight +into vLLM against `AsyncOmni` -> StageRuntime -> the stage's EngineCore. + +**12 of 12 prompts, 128 greedy tokens each, token-for-token identical -- on +both paths**: CUDA with the bf16 build and CPU with the int8 one. The quantized +build gets its own run rather than inheriting the result, because quantization +is exactly what makes a near-tie likelier to flip. + +The comparison is on **token ids, not text**: two different id sequences can +detokenize to the same string, and a text comparison would hide exactly the +drift worth catching. The record keeps the first divergent position rather than +a yes/no, because greedy decoding is not bitwise stable across batch shapes -- +a different prefill chunking reorders reductions and can flip an argmax at a +near-tie, so "split at token 120" and "split at token 3" are different +findings. Neither occurred. + +The two engines load sequentially in one process, reference first and released +before the second starts; two copies of a 7.66 GB checkpoint would not fit the +budget this milestone is about respecting. + +```bash +python -m vllm_omni.edge.local parity --model --json out.json # exit 3 on divergence +``` + +### The two paths do not need the same environment + +The reference cannot start on this host with a default environment, and both +reasons are worth knowing: + +* under WSL2, vLLM's own GPU model runner allocates a `UvaBuffer` for its input + batch and raises `RuntimeError: UVA is not available` without pinned memory. + Omni's AR model runner does not take that path and starts either way. +* FlashInfer's sampler JIT-compiles on first use and needs `nvcc`, which is not + installed -- `nvidia-smi` reporting a CUDA version is not a toolkit. This one + kills the engine at the *first token*, not at load. + +Both now live in one `RUNTIME_ENV` applied by every entry point that starts an +engine, because a comparison whose reference is configured differently from its +subject is measuring the configuration. + +## 8. Limits + +* One model, one stage; no multi-stage, no tensor/buffer contract. +* `--mask` is a process-level policy, not a physical absence. A full no-card + verification needs another machine. +* The iGPU and NPU are enumerated only; neither has a stage backend (M3). +* Reference consistency covers **one decode mode**: single sequence, greedy, + 128 tokens, on both paths. It does not cover concurrency above 1 (which + changes batch shapes, and greedy is not bitwise stable across those), longer + outputs, or non-greedy sampling -- and it is not a quality evaluation. It + says the Omni path did not change the output, not that the output is good. +* Performance figures are one sample per prompt, not p50/p95, with no fixed + power profile and no cache-cold guarantee. + +## 9. Tests + +```bash +# hermetic; no GPU, no weights, ~3 s +pytest tests/edge/local --ignore=tests/edge/local/test_e2e_local_text.py + +# end to end against a real checkpoint; the venv decides which artifact +pytest tests/edge/local/test_e2e_local_text.py +``` + +The e2e tests run the acceptance criteria themselves: both startups reach a +decision, twelve fixed prompts each reach 128 tokens, the budget bounds the +measured peak, cancellation leaks nothing, and the Omni path decodes what plain +vLLM decodes. They skip rather than fail when no checkpoint is present. +**78 pass in each venv**; about 4 min on this machine either way. + +The twelve prompts are the same twelve +`analysis/experiments/spark_edge/quality_int4.py` measures the quantized builds +against, so these numbers stay comparable with the existing W4A16 / W4A8 / +Q4_K_M fidelity work. + +Full run records: +[`analysis/experiments/m0_local_text_20260915/`](../../../analysis/experiments/m0_local_text_20260915/README.md). diff --git a/tests/config/test_vllm_native_pipelines.py b/tests/config/test_vllm_native_pipelines.py new file mode 100644 index 000000000..9ef13a66e --- /dev/null +++ b/tests/config/test_vllm_native_pipelines.py @@ -0,0 +1,50 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Pipelines for text models vLLM implements and Omni does not own. + +A model being in vLLM's registry does not make it reachable through +``Omni(model=...)``: the orchestrator resolves a pipeline by ``model_type`` +before any weight is read, so a ``model_type`` with no entry fails early and +opaquely. Spark-X2.5 hit exactly that; these tests pin the general case. +""" + +from __future__ import annotations + +import pytest + +from vllm_omni.config.pipeline_registry import OMNI_PIPELINES +from vllm_omni.config.stage_config import StageExecutionType +from vllm_omni.config.vllm_native_pipelines import LLAMA_PIPELINE + + +def test_llama_resolves_from_the_registry(): + """MiniCPM5-2B declares model_type "llama"; without this entry the + orchestrator cannot resolve it at all.""" + assert OMNI_PIPELINES["llama"] is LLAMA_PIPELINE + + +def test_it_is_one_autoregressive_text_stage(): + assert LLAMA_PIPELINE.get_validation_errors() == [] + assert len(LLAMA_PIPELINE.stages) == 1 + stage = LLAMA_PIPELINE.stages[0] + assert stage.execution_type is StageExecutionType.LLM_AR + assert stage.input_sources == () + assert stage.final_output and stage.final_output_type == "text" + assert stage.owns_tokenizer + assert stage.model_arch == "LlamaForCausalLM" + + +def test_the_stage_declares_no_cross_chunk_state(): + """A causal LM's state is its KV cache, which the backend owns. Claiming + otherwise would make the scheduler park requests to hold capacity.""" + assert not LLAMA_PIPELINE.stages[0].retains_state_across_chunks + + +def test_registering_llama_is_keyed_on_architecture_not_checkpoint(): + """Worth pinning because it is the surprising part: this entry covers every + checkpoint whose config says "llama", not only the one it was added for. + That is the right granularity — pipeline shape is a property of the + architecture — but it is a claim about shape, not a verification of any + particular weights.""" + assert LLAMA_PIPELINE.model_type == "llama" + assert LLAMA_PIPELINE.hf_architectures == ("LlamaForCausalLM",) diff --git a/tests/edge/local/__init__.py b/tests/edge/local/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/edge/local/conftest.py b/tests/edge/local/conftest.py new file mode 100644 index 000000000..2035c12c4 --- /dev/null +++ b/tests/edge/local/conftest.py @@ -0,0 +1,154 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Fixtures for the M0 local text mode. + +Everything here is synthetic on purpose: a hardware profile built from +literals, and a checkpoint directory with a real ``config.json`` and +zero-filled weight files. The planner reads sizes and geometry, never tensors, +so this exercises the real code path without a GPU or 7.7 GB of weights. The +one test module that does need real weights says so and skips. +""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + +from vllm_omni.edge.hardware_probe import ( + ACCEL_GPU_DISCRETE, + ACCEL_GPU_INTEGRATED, + ACCEL_NPU, + HardwareProfile, +) + +# A 28-layer Spark-shaped decoder: 7 full-attention and 21 sliding@512, which +# is the 1.7B layout. Small hidden/vocab so nothing here is slow. +SPARK_SHAPED_CONFIG = { + "architectures": ["Spark2_5ForCausalLM"], + "model_type": "spark2_5", + "dtype": "bfloat16", + "hidden_size": 2048, + "intermediate_size": 8192, + "num_hidden_layers": 28, + "num_attention_heads": 8, + "num_key_value_heads": 2, + "head_dim": 256, + "hidden_act": "gelu", + "vocab_size": 131072, + "sliding_window": 512, + "max_position_embeddings": 1048576, + "tie_word_embeddings": True, + "layer_types": [ + "full_attention" if (i + 1) % 4 == 0 else "sliding_attention" + for i in range(28) + ], + "rope_parameters": { + "full_attention": {"rope_theta": 5000000, "partial_rotary_factor": 0.25}, + "sliding_attention": {"rope_theta": 10000, "partial_rotary_factor": 1.0}, + }, +} + +INT8_QUANT_CONFIG = { + "quant_method": "compressed-tensors", + "format": "int-quantized", + "quantization_status": "compressed", + "ignore": ["lm_head"], + "config_groups": { + "group_0": { + "targets": ["Linear"], + "weights": {"num_bits": 8, "type": "int", "symmetric": True, "strategy": "channel"}, + } + }, +} + +FP8_QUANT_CONFIG = { + "quant_method": "compressed-tensors", + "format": "float-quantized", + "config_groups": {}, +} + + +def write_checkpoint( + root: Path, + *, + quantization: dict | None = None, + weight_bytes: int = 4 * 1024 * 1024, + shards: int = 1, +) -> Path: + """A checkpoint directory the manifest and planner can read end to end.""" + root.mkdir(parents=True, exist_ok=True) + config = dict(SPARK_SHAPED_CONFIG) + if quantization is not None: + config["quantization_config"] = quantization + (root / "config.json").write_text(json.dumps(config), encoding="utf-8") + (root / "tokenizer_config.json").write_text('{"model_max_length": 4096}', encoding="utf-8") + (root / "chat_template.jinja").write_text("{{ messages }}", encoding="utf-8") + per_shard = max(weight_bytes // shards, 1) + for i in range(shards): + name = "model.safetensors" if shards == 1 else f"model-{i + 1:05d}-of-{shards:05d}.safetensors" + with (root / name).open("wb") as fh: + fh.truncate(per_shard) + return root + + +@pytest.fixture +def dense_checkpoint(tmp_path: Path) -> Path: + return write_checkpoint(tmp_path / "dense", weight_bytes=8 * 1024 * 1024, shards=2) + + +@pytest.fixture +def int8_checkpoint(tmp_path: Path) -> Path: + return write_checkpoint(tmp_path / "int8", quantization=INT8_QUANT_CONFIG) + + +def blackwell_laptop(ram_available_bytes: int = 28 * 2**30) -> HardwareProfile: + """The reference machine: sm_120 discrete GPU, an iGPU and an NPU beside it. + + Modelled on the probe output recorded in the architecture proposal, so the + tests exercise the same shape the real host produces. + """ + return HardwareProfile( + arch="x86_64", + cpu_model="AMD Ryzen AI 9 HX 370 w/ Radeon 890M", + cpu_flags=["avx2", "avx512", "avx512_bf16", "avx512_vnni", "fma"], + n_logical=24, + n_physical=12, + clusters=[{"max_khz": 5000000, "cpus": list(range(24))}], + ram_total_bytes=int(30.91 * 2**30), + ram_available_bytes=ram_available_bytes, + accelerator="cuda_discrete", + gpu_name="NVIDIA GeForce RTX 5090 Laptop GPU", + gpu_mem_bytes=24463 * 2**20, + gpu_sm_count=110, + gpu_capability=[12, 0], + supports_bf16=True, + host_os="wsl2", + accelerators=[ + {"kind": ACCEL_GPU_DISCRETE, "vendor": "nvidia", + "name": "NVIDIA GeForce RTX 5090 Laptop GPU", "usable": True, + "source": "torch.cuda", "mem_bytes": 24463 * 2**20}, + {"kind": ACCEL_NPU, "vendor": "amd", "name": "NPU Compute Accelerator Device", + "usable": False, "source": "wsl-host", "reason": "no device node under WSL2"}, + {"kind": ACCEL_GPU_INTEGRATED, "vendor": "amd", "name": "AMD Radeon(TM) 890M Graphics", + "usable": False, "source": "wsl-host", "reason": "no /dev/dri under WSL2"}, + ], + source="test-fixture", + ) + + +def cpu_only_box() -> HardwareProfile: + """The no-NVIDIA deployment class, as a machine rather than as a mask.""" + profile = blackwell_laptop() + profile.accelerator = "none" + profile.gpu_name = "" + profile.gpu_mem_bytes = 0 + profile.gpu_capability = [] + profile.accelerators = [a for a in profile.accelerators if a["vendor"] != "nvidia"] + return profile + + +@pytest.fixture +def laptop_profile() -> HardwareProfile: + return blackwell_laptop() diff --git a/tests/edge/local/external/__init__.py b/tests/edge/local/external/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/edge/local/external/conftest.py b/tests/edge/local/external/conftest.py new file mode 100644 index 000000000..f78da4755 --- /dev/null +++ b/tests/edge/local/external/conftest.py @@ -0,0 +1,96 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Fixtures that stand in for the two AMD devices. + +Nothing here needs a Radeon, an NPU, onnxruntime or Windows: the worker is +replaced by ``fake_worker.py`` and the devices by hand-built capabilities. The +paths that matter -- framing, the placement gate, the refusal codes -- are +logic, and logic that only runs on one laptop is logic nobody can change +safely. +""" + +from __future__ import annotations + +import json +import sys +from pathlib import Path + +import pytest + +from vllm_omni.edge.local.capabilities import ( + FORMAT_ONNX_A16W8, + FORMAT_ONNX_FP16, + FORMAT_ONNX_FP32, + DeviceCapability, +) +from vllm_omni.edge.local.external import launch + +FAKE_WORKER = str(Path(__file__).resolve().parent / "fake_worker.py") + + +@pytest.fixture +def fake_route(monkeypatch, tmp_path): + """Point the ``ort-cpu`` route at the fake worker and this interpreter. + + Returns a callable that writes the worker's reply script and hands back the + resolved route. + """ + + def configure(**config) -> launch.Route: + config_path = tmp_path / "fake_worker.json" + config_path.write_text(json.dumps(config)) + monkeypatch.setenv("VLLM_OMNI_FAKE_WORKER_JSON", str(config_path)) + monkeypatch.setenv("VLLM_OMNI_EXTERNAL_WORKER_ORT_CPU", FAKE_WORKER) + monkeypatch.setenv("VLLM_OMNI_EXTERNAL_PYTHON_ORT_CPU", sys.executable) + return launch.resolve(launch.ROUTE_CPU) + + return configure + + +def _device(device_id: str, kind: str, formats, route: str, *, runnable: bool = True, + reason: str = "") -> DeviceCapability: + return DeviceCapability( + device_id=device_id, + kind=kind, + vendor="amd", + name=f"test {kind}", + backend="ort:cpu" if runnable else None, + runnable=runnable, + memory_pool="host_ram", + memory_bytes=32 * 2**30, + compute_capability=None, + weight_formats=frozenset(formats), + evidence="B" if runnable else "D", + reason=reason, + extra={"worker_route": route, "routes": [{"name": route, "available": runnable, + "reason": reason}]}, + ) + + +@pytest.fixture +def npu_device() -> DeviceCapability: + """A16W8 only -- the real constraint, not a simplification.""" + return _device("npu:amd", "npu", {FORMAT_ONNX_A16W8}, launch.ROUTE_CPU) + + +@pytest.fixture +def igpu_device() -> DeviceCapability: + return _device("igpu:amd", "gpu_integrated", {FORMAT_ONNX_FP16, FORMAT_ONNX_FP32}, + launch.ROUTE_CPU) + + +@pytest.fixture +def unreachable_npu() -> DeviceCapability: + return _device( + "npu:amd", "npu", {FORMAT_ONNX_A16W8}, launch.ROUTE_VITISAI, + runnable=False, + reason="an MCDM device; WSL2 forwards display adapters only, so there is no device node here", + ) + + +@pytest.fixture +def a16w8_graph(tmp_path) -> Path: + """A file that stands in for an exported graph. Never loaded by ORT here.""" + path = tmp_path / "stage.onnx" + path.write_bytes(b"not really onnx, and nothing in these tests parses it" * 64) + return path diff --git a/tests/edge/local/external/fake_worker.py b/tests/edge/local/external/fake_worker.py new file mode 100644 index 000000000..7d422ac1f --- /dev/null +++ b/tests/edge/local/external/fake_worker.py @@ -0,0 +1,110 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""A stage worker that speaks the protocol and never touches a device. + +Selected with ``VLLM_OMNI_EXTERNAL_WORKER_ORT_CPU``, so the client, the +framing, the placement gate and every refusal path are exercised on a machine +with no AMD hardware, no onnxruntime and no Windows. What it replies is read +from the JSON file named by ``VLLM_OMNI_FAKE_WORKER_JSON``: + +``load`` body merged into the load reply -- set ``node_counts`` and + ``fraction_on_target`` to stage a declined graph. +``run`` ``{"scale": float}``; outputs are the inputs times it. +``die_on`` op name after which the worker exits without replying, for the + "worker vanished mid-request" path. +""" + +from __future__ import annotations + +import importlib.util +import json +import os +import sys +from pathlib import Path + +import numpy as np + +_HERE = Path(__file__).resolve() +_PROTOCOL = ( + _HERE.parents[4] / "vllm_omni" / "edge" / "local" / "external" / "protocol.py" +) + +spec = importlib.util.spec_from_file_location("_fake_protocol", _PROTOCOL) +proto = importlib.util.module_from_spec(spec) +sys.modules["_fake_protocol"] = proto +spec.loader.exec_module(proto) + + +def main() -> int: + host, port, token = sys.argv[2], int(sys.argv[4]), sys.argv[6] + config = {} + config_path = os.environ.get("VLLM_OMNI_FAKE_WORKER_JSON") + if config_path and Path(config_path).is_file(): + config = json.loads(Path(config_path).read_text()) + + import socket + + sock = socket.create_connection((host, port), timeout=30) + proto.send_message( + sock, + proto.OP_HELLO, + { + "token": config.get("token", token), + "executable": sys.executable, + "python": sys.version, + "platform": sys.platform, + "pid": os.getpid(), + "onnxruntime": "fake", + "numpy": np.__version__, + "available_providers": ["CPUExecutionProvider"], + "rss_bytes": 1024, + }, + ) + + while True: + op, body, tensors = proto.recv_message(sock) + if config.get("die_on") == op: + sock.close() + return 0 + if op == proto.OP_LOAD: + reply = { + "ep": body.get("ep", "cpu"), + "session_providers": ["CPUExecutionProvider"], + "available_providers": ["CPUExecutionProvider"], + "node_counts": {"CPUExecutionProvider": 4}, + "total_nodes": 4, + "target_nodes": 4, + "fraction_on_target": 1.0, + "session_create_s": 0.01, + "warmup_s": 0.001, + "rss_bytes": 111 * 2**20, + "inputs": [{"name": n, "type": "tensor(float)", "shape": list(t.shape)} + for n, t in tensors.items()], + "outputs": [{"name": "y", "type": "tensor(float)", "shape": []}], + } + reply.update(config.get("load", {})) + proto.send_message(sock, proto.OP_OK, reply) + elif op == proto.OP_RUN: + scale = float(config.get("run", {}).get("scale", 2.0)) + proto.send_message( + sock, + proto.OP_OK, + {"run_s": 0.001}, + {f"{k}_out": (v * scale) for k, v in tensors.items()}, + ) + elif op == proto.OP_STATS: + proto.send_message( + sock, proto.OP_OK, + {"runs": 1, "total_run_s": 0.001, "rss_bytes": 111 * 2**20, + "peak_rss_bytes": 122 * 2**20, "report": {}}, + ) + elif op == proto.OP_CLOSE: + proto.send_message(sock, proto.OP_OK, {}) + return 0 + else: + proto.send_message(sock, proto.OP_ERR, + {"message": f"unknown op {op}", "code": "ValueError"}) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/edge/local/external/test_client.py b/tests/edge/local/external/test_client.py new file mode 100644 index 000000000..1441f9f77 --- /dev/null +++ b/tests/edge/local/external/test_client.py @@ -0,0 +1,103 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Starting, driving and losing a worker process.""" + +from __future__ import annotations + +import numpy as np +import pytest + +from vllm_omni.edge.local.external import launch +from vllm_omni.edge.local.external.client import ExternalWorker, WorkerStartError + + +def test_handshake_and_round_trip(fake_route): + with ExternalWorker(fake_route()) as worker: + assert worker.hello["onnxruntime"] == "fake" + report = worker.load("/nonexistent/stage.onnx", example_inputs={"x": np.ones(4, np.float32)}) + assert report.fraction_on_target == 1.0 + assert report.rss_bytes == 111 * 2**20 + + outputs, timing = worker.run({"x": np.ones(4, np.float32)}) + assert np.array_equal(outputs["x_out"], np.full(4, 2.0, np.float32)) + assert timing.round_trip_s >= timing.worker_s + assert timing.transport_s == pytest.approx(timing.round_trip_s - timing.worker_s) + + assert worker.stats()["peak_rss_bytes"] == 122 * 2**20 + + +def test_wrong_token_is_refused(fake_route): + """The listener is on a host-visible address; a stray peer must not be served.""" + worker = ExternalWorker(fake_route(token="not-the-token")) + with pytest.raises(WorkerStartError, match="wrong token"): + worker.start() + worker.close() + + +def test_worker_that_vanishes_mid_request_reports_rather_than_hangs(fake_route): + worker = ExternalWorker(fake_route(die_on="run")) + worker.start() + worker.load("/nonexistent/stage.onnx", example_inputs={"x": np.ones(2, np.float32)}) + with pytest.raises(WorkerStartError, match="worker connection failed"): + worker.run({"x": np.ones(2, np.float32)}) + worker.close() + + +def test_unavailable_route_refuses_at_construction(): + route = launch.Route( + name="ort-vitisai", interpreter=None, worker="w.py", ep="vitisai", + is_windows=True, available=False, reason="no interpreter installed", + ) + with pytest.raises(WorkerStartError, match="no interpreter installed"): + ExternalWorker(route) + + +def test_close_is_idempotent(fake_route): + worker = ExternalWorker(fake_route()) + worker.start() + worker.close() + worker.close() + + +def test_bind_address_is_loopback_for_a_same_os_worker(): + from vllm_omni.edge.local.external import client + + assert client._bind_address(is_windows=False) == "127.0.0.1" + + +def test_bind_address_for_a_windows_worker_is_not_loopback_under_wsl(): + """The regression this exists for: probing the WSL nameserver instead of the + gateway yields 10.255.255.254, an address only this distro can reach, and + the Windows worker then dies with ECONNREFUSED after a silent start.""" + from vllm_omni.edge.local.external import client + + if not launch.is_wsl(): + pytest.skip("not running under WSL interop") + assert not client._bind_address(is_windows=True).startswith(("127.", "10.255.")) + + +def test_external_data_counts_toward_the_artifact_size(tmp_path): + """An ONNX graph with external weights is a few MB of topology pointing at + gigabytes beside it. Charging only the .onnx told the admission ledger that + a 1.8 GB vision tower costs 2 MiB.""" + import onnx + from onnx import helper, numpy_helper + + from vllm_omni.edge.local.manifest import build_graph_artifact + + weight = numpy_helper.from_array(np.zeros((256, 256), np.float32), name="w") + node = helper.make_node("MatMul", ["x", "w"], ["y"]) + graph = helper.make_graph( + [node], "g", + [helper.make_tensor_value_info("x", onnx.TensorProto.FLOAT, [1, 256])], + [helper.make_tensor_value_info("y", onnx.TensorProto.FLOAT, [1, 256])], + [weight], + ) + path = tmp_path / "m.onnx" + onnx.save(helper.make_model(graph), str(path), save_as_external_data=True, + all_tensors_to_one_file=True, location="m.onnx.data") + + artifact = build_graph_artifact(path, fmt="onnx:fp32", opset=21, source_model="t", + component="c", exporter="e") + assert artifact.bytes > path.stat().st_size + assert artifact.bytes >= 256 * 256 * 4 diff --git a/tests/edge/local/external/test_protocol.py b/tests/edge/local/external/test_protocol.py new file mode 100644 index 000000000..6dbe37445 --- /dev/null +++ b/tests/edge/local/external/test_protocol.py @@ -0,0 +1,94 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""The wire format, which has to survive a peer that is not us.""" + +from __future__ import annotations + +import socket +import threading + +import numpy as np +import pytest + +from vllm_omni.edge.local.external import protocol + + +def _round_trip(op, body=None, tensors=None): + near, far = socket.socketpair() + try: + threading.Thread( + target=lambda: protocol.send_message(near, op, body, tensors), daemon=True + ).start() + return protocol.recv_message(far) + finally: + near.close() + far.close() + + +def test_tensors_survive_dtype_and_shape(): + sent = { + "f32": np.arange(12, dtype=np.float32).reshape(3, 4), + "i64": np.array([[1, 2], [3, 4]], dtype=np.int64), + "u16": np.array([7, 8, 9], dtype=np.uint16), + "scalar": np.array(3.5, dtype=np.float64), + } + op, body, got = _round_trip(protocol.OP_RUN, {"n": 1}, sent) + assert op == protocol.OP_RUN and body == {"n": 1} + for name, array in sent.items(): + assert got[name].dtype == array.dtype + assert got[name].shape == array.shape + assert np.array_equal(got[name], array) + + +def test_received_tensors_are_writable(): + """``frombuffer`` hands back a read-only view aliasing the socket buffer. + + ORT's IOBinding writes into the arrays it is given, so a read-only result + would fail at the far end of a working connection rather than here. + """ + _, _, got = _round_trip(protocol.OP_RUN, {}, {"x": np.ones(4, dtype=np.float32)}) + got["x"][0] = 5.0 # must not raise + assert got["x"][0] == 5.0 + + +def test_payload_larger_than_one_recv(): + """Stream sockets short-read; a naive recv would misparse the next tensor.""" + big = np.arange(4 * 1024 * 1024, dtype=np.float32) + _, _, got = _round_trip(protocol.OP_RUN, {}, {"big": big, "after": np.array([1.0])}) + assert np.array_equal(got["big"], big) + assert np.array_equal(got["after"], np.array([1.0])) + + +def test_empty_message_and_no_tensors(): + op, body, got = _round_trip(protocol.OP_STATS) + assert (op, body, got) == (protocol.OP_STATS, {}, {}) + + +def test_error_reply_raises_with_remote_detail(): + with pytest.raises(protocol.WorkerError) as excinfo: + _round_trip(protocol.OP_ERR, {"message": "boom", "code": "ValueError", + "traceback": "far side traceback"}) + assert excinfo.value.code == "ValueError" + assert "far side traceback" in excinfo.value.remote_traceback + + +def test_closed_peer_is_a_protocol_error_not_a_hang(): + near, far = socket.socketpair() + near.close() + with pytest.raises(protocol.ProtocolError): + protocol.recv_message(far) + far.close() + + +def test_oversized_header_is_refused_before_allocation(): + """A corrupt length must not become a multi-gigabyte allocation.""" + near, far = socket.socketpair() + try: + import struct + + near.sendall(struct.pack(">I", protocol.MAX_HEADER_BYTES + 1)) + with pytest.raises(protocol.ProtocolError, match="over the limit"): + protocol.recv_message(far) + finally: + near.close() + far.close() diff --git a/tests/edge/local/external/test_stage.py b/tests/edge/local/external/test_stage.py new file mode 100644 index 000000000..c08425648 --- /dev/null +++ b/tests/edge/local/external/test_stage.py @@ -0,0 +1,209 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Placing a stage on an external device, and refusing to. + +The gate tested here is the one that keeps a CPU fallback from being reported +as NPU execution. It cannot be replaced by a cheaper check: a VitisAI session +lists the EP in ``get_providers()`` whether or not it claimed a node, and a +graph it declined runs on the CPU and returns bit-identical numbers. So these +tests assert on node assignments, never on output values. +""" + +from __future__ import annotations + +import numpy as np +import pytest + +from vllm_omni.edge.local.capabilities import FORMAT_ONNX_A16W8, FORMAT_ONNX_FP16 +from vllm_omni.edge.local.external.stage import ( + ExternalStage, + PlacementRefused, + plan_external_stage, + verify_placement, +) +from vllm_omni.edge.local.manifest import build_graph_artifact +from vllm_omni.edge.local.plan import ( + REFUSE_EP_PLACEMENT, + REFUSE_NO_ARTIFACT, + REFUSE_NO_DEVICE, + REFUSE_ROUTE, +) + + +def _artifact(path, fmt=FORMAT_ONNX_A16W8, component="stage"): + return build_graph_artifact( + path, fmt=fmt, opset=21, source_model="test", component=component, + exporter="tests.edge.local.external", + ) + + +def test_a16w8_graph_is_placed_on_the_npu(a16w8_graph, npu_device, igpu_device): + plan = plan_external_stage(_artifact(a16w8_graph), [igpu_device, npu_device]) + assert plan.admitted + assert plan.device is not None and plan.device.device_id == "npu:amd" + # The iGPU was considered and rejected on format; the record says so rather + # than only naming the winner. + assert [r.code for r in plan.refusals] == [REFUSE_NO_ARTIFACT] + + +def test_fp16_graph_is_refused_by_the_npu_not_quietly_downgraded(a16w8_graph, npu_device): + """A16W8 is the only door on XDNA2, and the failure mode for anything else + is silent: no error, every node on the CPU. So it is refused at plan time.""" + plan = plan_external_stage(_artifact(a16w8_graph, fmt=FORMAT_ONNX_FP16), [npu_device]) + assert not plan.admitted + assert plan.refusals[0].code == REFUSE_NO_ARTIFACT + assert "a16w8" in plan.refusals[0].remedy.lower() + + +def test_require_refuses_rather_than_substituting(a16w8_graph, npu_device, igpu_device): + """"Does the NPU run this" must not be answered with the iGPU.""" + plan = plan_external_stage( + _artifact(a16w8_graph, fmt=FORMAT_ONNX_FP16), [npu_device, igpu_device], + require="npu:amd", + ) + assert not plan.admitted + assert {r.device_id for r in plan.refusals} == {"npu:amd"} + + +def test_prefer_falls_through(a16w8_graph, npu_device, igpu_device): + plan = plan_external_stage( + _artifact(a16w8_graph, fmt=FORMAT_ONNX_FP16), [npu_device, igpu_device], + prefer="npu:amd", + ) + assert plan.admitted and plan.device.device_id == "igpu:amd" + + +def test_unreachable_device_refuses_with_the_route_reason(a16w8_graph, unreachable_npu): + plan = plan_external_stage(_artifact(a16w8_graph), [unreachable_npu]) + assert not plan.admitted + assert plan.refusals[0].code == REFUSE_ROUTE + assert "MCDM" in plan.refusals[0].message + assert "VLLM_OMNI_EXTERNAL_PYTHON_ORT_VITISAI" in plan.refusals[0].remedy + + +def test_no_external_device_at_all(a16w8_graph): + plan = plan_external_stage(_artifact(a16w8_graph), []) + assert not plan.admitted and plan.refusals[0].code == REFUSE_NO_DEVICE + + +def test_required_device_absent_from_enumeration(a16w8_graph, igpu_device): + plan = plan_external_stage(_artifact(a16w8_graph), [igpu_device], require="npu:amd") + assert not plan.admitted and plan.refusals[0].code == REFUSE_NO_DEVICE + + +def test_budget_is_on_the_shared_pool(a16w8_graph, npu_device): + plan = plan_external_stage(_artifact(a16w8_graph), [npu_device]) + assert {r.pool for r in plan.reservations} == {"host_ram"} + assert plan.budget_bytes > 0 + + +def test_open_verifies_placement_and_closes_the_ledger(a16w8_graph, npu_device, fake_route, tmp_path): + fake_route() + plan = plan_external_stage(_artifact(a16w8_graph), [npu_device]) + with ExternalStage(plan) as stage: + report = stage.open({"x": np.ones(4, np.float32)}, profile_dir=tmp_path) + assert report.fraction_on_target == 1.0 + runtime = next(r for r in plan.reservations if r.purpose == "worker runtime") + assert runtime.actual_upper_bound_bytes == 111 * 2**20 + assert runtime.bytes >= runtime.actual_upper_bound_bytes, "budget must bound the measurement" + + +def test_declined_graph_is_refused_even_though_it_ran(a16w8_graph, npu_device, fake_route, tmp_path): + """The whole point: it produced correct output, on the CPU.""" + fake_route(load={"node_counts": {"CPUExecutionProvider": 9, "vitisai": 1}, + "total_nodes": 10, "target_nodes": 1, "fraction_on_target": 0.1}) + plan = plan_external_stage(_artifact(a16w8_graph), [npu_device]) + stage = ExternalStage(plan) + with pytest.raises(PlacementRefused) as excinfo: + stage.open({"x": np.ones(4, np.float32)}, profile_dir=tmp_path) + assert excinfo.value.refusal.code == REFUSE_EP_PLACEMENT + assert "1 of 10" in excinfo.value.refusal.message + assert "Light" in excinfo.value.refusal.remedy + + +def test_unverified_placement_is_refused_not_assumed(a16w8_graph, npu_device, fake_route, tmp_path): + """``None`` means nobody looked. It must not be rounded to "fine".""" + fake_route(load={"fraction_on_target": None, "node_counts": {}, + "placement_unknown_reason": "no profiled run"}) + plan = plan_external_stage(_artifact(a16w8_graph), [npu_device]) + stage = ExternalStage(plan) + with pytest.raises(PlacementRefused) as excinfo: + stage.open({"x": np.ones(4, np.float32)}, profile_dir=tmp_path) + assert excinfo.value.refusal.code == REFUSE_EP_PLACEMENT + assert "not measured" in excinfo.value.refusal.message + + +def test_open_without_example_inputs_is_a_usage_error(a16w8_graph, npu_device, fake_route): + fake_route() + plan = plan_external_stage(_artifact(a16w8_graph), [npu_device]) + stage = ExternalStage(plan) + with pytest.raises(ValueError, match="example_inputs are required"): + stage.open({}) + stage.close() + + +def test_verify_placement_accepts_a_real_split(a16w8_graph, npu_device): + from vllm_omni.edge.local.external.client import LoadReport + + plan = plan_external_stage(_artifact(a16w8_graph), [npu_device]) + plan.min_fraction_on_target = 0.3 + report = LoadReport(ep="vitisai", graph_path="g.onnx", total_nodes=3, target_nodes=1, + fraction_on_target=1 / 3, node_counts={"vitisai": 1, "CPUExecutionProvider": 2}) + assert verify_placement(plan, report) is None + + +def test_refusing_a_plan_cannot_be_opened(a16w8_graph, npu_device): + plan = plan_external_stage(_artifact(a16w8_graph, fmt=FORMAT_ONNX_FP16), [npu_device]) + with pytest.raises(PlacementRefused): + ExternalStage(plan) + + +def test_a_masked_device_is_told_to_drop_the_mask(a16w8_graph, npu_device): + """The remedy has to name the cause the caller actually has. + + A masked device and a device WSL cannot forward are both "unreachable", and + explaining MCDM forwarding to somebody who simply passed ``--mask npu`` + sends them after a problem they do not have. + """ + import dataclasses + + masked = dataclasses.replace(npu_device, runnable=False, masked=True, + backend=None, reason="masked out by --mask (npu)") + plan = plan_external_stage(_artifact(a16w8_graph), [masked]) + assert not plan.admitted + assert plan.refusals[0].code == REFUSE_ROUTE + assert "--mask" in plan.refusals[0].remedy + assert "MCDM" not in plan.refusals[0].remedy + + +def test_a_stage_larger_than_the_shared_pool_is_refused(a16w8_graph, npu_device): + plan = plan_external_stage(_artifact(a16w8_graph), [npu_device], pool_capacity_bytes=1) + assert not plan.admitted + from vllm_omni.edge.local.plan import REFUSE_CAPACITY + + assert plan.refusals[0].code == REFUSE_CAPACITY + assert "same pool the CPU stage is budgeted against" in plan.refusals[0].message + + +def test_a_silently_swapped_provider_is_refused_by_name(a16w8_graph, npu_device, fake_route, tmp_path): + """onnxruntime does not raise when a provider fails to initialise. + + It prints "Falling back to ['CPUExecutionProvider'] and retrying" and hands + back a working CPU session — which is how a 27B vision tower ran end to end + on the CPU here while the caller believed it was on the iGPU. The report + carries the swap, and it outranks the node count because "the provider + never initialised" says what to fix and "0 of 1171 nodes" does not. + """ + fake_route(load={ + "session_providers": ["CPUExecutionProvider"], + "requested_provider_missing": "DmlExecutionProvider", + "fallback_note": "DmlExecutionProvider failed to initialise and onnxruntime silently fell back", + "node_counts": {"CPUExecutionProvider": 113}, + "total_nodes": 113, "target_nodes": 0, "fraction_on_target": 0.0, + }) + plan = plan_external_stage(_artifact(a16w8_graph), [npu_device]) + stage = ExternalStage(plan) + with pytest.raises(PlacementRefused) as excinfo: + stage.open({"x": np.ones(4, np.float32)}, profile_dir=tmp_path) + assert excinfo.value.refusal.code == REFUSE_EP_PLACEMENT + assert "silently fell back" in excinfo.value.refusal.message diff --git a/tests/edge/local/external/test_vision_tower.py b/tests/edge/local/external/test_vision_tower.py new file mode 100644 index 000000000..a6e94bff7 --- /dev/null +++ b/tests/edge/local/external/test_vision_tower.py @@ -0,0 +1,64 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Shape arithmetic for the exported vision tower. + +The tower itself needs a 29 GB checkpoint, so what is covered here is the part +that decides whether an export is even well-formed — and that is where a silent +mistake would land: a grid that does not divide by the merge size produces a +graph that exports, runs, and returns the wrong number of tokens to the LLM. +""" + +from __future__ import annotations + +import json + +import pytest + +from vllm_omni.edge.local.artifacts.vision_tower import tower_shape + +VISION_CONFIG = { + "depth": 27, "hidden_size": 1152, "in_channels": 3, "intermediate_size": 4304, + "num_heads": 16, "num_position_embeddings": 2304, "out_hidden_size": 5120, + "patch_size": 16, "spatial_merge_size": 2, "temporal_patch_size": 2, +} + + +@pytest.fixture +def checkpoint(tmp_path): + (tmp_path / "config.json").write_text(json.dumps({"vision_config": VISION_CONFIG})) + return tmp_path + + +def test_448_gives_784_patches_and_196_tokens(checkpoint): + shape = tower_shape(checkpoint, 448) + assert (shape.grid_h, shape.grid_w) == (28, 28) + assert shape.patches == 784 + # 3 channels x 2 temporal x 16 x 16 + assert shape.patch_dim == 1536 + # 784 patches merged 2x2 + assert shape.merged_tokens == 196 + assert shape.out_hidden_size == 5120 + + +def test_image_size_must_divide_by_the_patch_size(checkpoint): + with pytest.raises(ValueError, match="multiple of patch_size"): + tower_shape(checkpoint, 450) + + +def test_grid_must_divide_by_the_spatial_merge(checkpoint): + """A 3x3 grid of patches cannot be merged 2x2. Exporting it anyway would + hand the language model a silently wrong number of image tokens.""" + with pytest.raises(ValueError, match="spatial_merge_size"): + tower_shape(checkpoint, 48) + + +def test_a_checkpoint_without_a_tower_is_refused(tmp_path): + (tmp_path / "config.json").write_text(json.dumps({"model_type": "spark2_5"})) + with pytest.raises(ValueError, match="no vision_config"): + tower_shape(tmp_path, 448) + + +@pytest.mark.parametrize("size,patches,tokens", [(224, 196, 49), (448, 784, 196), (896, 3136, 784)]) +def test_token_count_scales_with_area(checkpoint, size, patches, tokens): + shape = tower_shape(checkpoint, size) + assert (shape.patches, shape.merged_tokens) == (patches, tokens) diff --git a/tests/edge/local/external/test_win_cuda_route.py b/tests/edge/local/external/test_win_cuda_route.py new file mode 100644 index 000000000..0c34f2046 --- /dev/null +++ b/tests/edge/local/external/test_win_cuda_route.py @@ -0,0 +1,106 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""[W4] The native-Windows vLLM route, and why it is refused. + +``win-cuda`` is the first route whose far side is torch rather than +onnxruntime, and the first that exists to leave WSL rather than to reach a +device WSL cannot see. Declaring it makes a Windows engine a row in the route +table instead of a parallel universe. + +It is also the first route that **cannot pass the placement gate**, and these +tests pin that down so nobody later mistakes the refusal for a bug. M3 made a +measured device attribution a hard admission condition after a VitisAI session +listed its EP, claimed zero nodes, and returned output bit-identical to the +CPU's. vLLM emits no per-operation device map, so a ``win-cuda`` worker has +nothing of that kind to report -- and the rule is about evidence, not about +whose code it is, so our own engine is refused on it exactly as the NPU is. +""" + +from __future__ import annotations + +import numpy as np +import pytest + +from vllm_omni.edge.local.capabilities import FORMAT_ONNX_A16W8 +from vllm_omni.edge.local.external import launch +from vllm_omni.edge.local.external.stage import ( + ExternalStage, + PlacementRefused, + plan_external_stage, +) +from vllm_omni.edge.local.manifest import build_graph_artifact +from vllm_omni.edge.local.plan import REFUSE_EP_PLACEMENT, REFUSE_NO_DEVICE + +from .conftest import _device + + +def _artifact(path): + return build_graph_artifact( + path, fmt=FORMAT_ONNX_A16W8, opset=21, source_model="test", + component="stage", exporter="tests.edge.local.external", + ) + + +def test_win_cuda_is_a_known_route(): + """It resolves through the same machinery as every other route.""" + route = launch.resolve(launch.ROUTE_WIN_CUDA) + assert route.name == "win-cuda" + assert route.ep == "cuda" + # It needs a native-Windows interpreter, like the two AMD Windows routes. + assert route.is_windows is True + # Available or not is environment-dependent; a reason is not. + assert route.available or route.reason + + +def test_win_cuda_reports_no_placement_evidence(): + """The constant is the contract: ``None`` means nobody can look yet. + + Asserted rather than assumed, because the tempting edit -- setting it to + 1.0 because the worker "obviously" runs on the GPU -- is exactly the + reasoning M3 disproved. + """ + assert launch.WIN_CUDA_PLACEMENT_EVIDENCE is None + + +def test_unverified_placement_is_refused_for_our_own_engine(a16w8_graph, fake_route, tmp_path): + """A worker that cannot say where it ran is refused, not trusted. + + The fake worker stands in for a ``win-cuda`` one: it runs, it returns + correct-looking numbers, and it reports ``fraction_on_target=None`` because + vLLM has no per-op device map to derive one from. That is precisely the + shape of a silent CPU fallback, so it must not be admitted. + """ + fake_route(load={ + "fraction_on_target": None, + "node_counts": {}, + "placement_unknown_reason": "vllm emits no per-operation device map", + }) + device = _device("npu:amd", "npu", [FORMAT_ONNX_A16W8], launch.ROUTE_CPU) + plan = plan_external_stage(_artifact(a16w8_graph), [device]) + stage = ExternalStage(plan) + with pytest.raises(PlacementRefused) as excinfo: + stage.open({"x": np.ones(4, np.float32)}, profile_dir=tmp_path) + assert excinfo.value.refusal.code == REFUSE_EP_PLACEMENT + assert "not measured" in excinfo.value.refusal.message + + +def test_external_planner_does_not_yet_admit_a_discrete_gpu(a16w8_graph, fake_route): + """A documented gap, pinned so the next person does not rediscover it. + + ``plan_external_stage`` filters candidates to ``("gpu_integrated", "npu")`` + (``stage.py``). Every external route so far reaches a device WSL cannot see + *and* that is integrated or an NPU, so the filter was exactly right. A + ``win-cuda`` route breaks that coupling: its device is a **discrete** GPU + reached out-of-process for an OS reason, not a visibility one. + + This test asserts the current refusal rather than a desired behaviour. When + the planner learns about externally-routed discrete GPUs, this test should + fail and be rewritten -- that failure is the signal, not a regression. + """ + fake_route() + device = _device("cuda:0", "gpu_discrete", [FORMAT_ONNX_A16W8], launch.ROUTE_CPU) + plan = plan_external_stage(_artifact(a16w8_graph), [device]) + assert plan.device is None + assert plan.refusals + assert plan.refusals[0].code == REFUSE_NO_DEVICE + assert "integrated GPU or NPU" in plan.refusals[0].message diff --git a/tests/edge/local/test_capabilities.py b/tests/edge/local/test_capabilities.py new file mode 100644 index 000000000..2e5fe7d9b --- /dev/null +++ b/tests/edge/local/test_capabilities.py @@ -0,0 +1,174 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""What this process can run on, as opposed to what is plugged in.""" + +import pytest + +from vllm_omni.edge.hardware_probe import ACCEL_GPU_DISCRETE +from vllm_omni.edge.local.capabilities import ( + FORMAT_DENSE, + FORMAT_FP8, + FORMAT_GPTQ, + FORMAT_INT8, + FORMAT_NVFP4, + MASK_ENV, + cuda_weight_formats, + enumerate_devices, + parse_mask, + runs_exported_graphs, +) + +from .conftest import blackwell_laptop, cpu_only_box + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _by_id(devices): + return {d.device_id: d for d in devices} + + +# ---------------------------------------------------------- format gating +def test_blackwell_has_no_int8_path(): + """sm_120 is inside the SM >= 100 range where cutlass c3x has no int8 + scaled_mm. This is the single fact that turns a 40-second crash into a + 30-millisecond refusal, so it gets its own test.""" + formats = cuda_weight_formats((12, 0)) + assert FORMAT_INT8 not in formats + assert FORMAT_FP8 in formats + assert FORMAT_NVFP4 in formats + + +def test_ampere_keeps_int8_and_has_no_nvfp4(): + formats = cuda_weight_formats((8, 6)) + assert FORMAT_INT8 in formats + assert FORMAT_NVFP4 not in formats + assert FORMAT_FP8 not in formats + + +def test_dense_runs_on_every_cuda_generation(): + for capability in ((7, 5), (8, 0), (9, 0), (12, 0)): + assert FORMAT_DENSE in cuda_weight_formats(capability) + + +# ------------------------------------------------------------ runnability +def test_a_cuda_wheel_cannot_drive_the_cpu_and_says_so(): + devices = _by_id(enumerate_devices(blackwell_laptop(), platform_device_type="cuda")) + assert devices["cuda:0"].runnable + assert not devices["cpu"].runnable + assert "omni-cpu" in devices["cpu"].reason + + +def test_a_cpu_wheel_cannot_drive_the_card_and_says_so(): + devices = _by_id(enumerate_devices(blackwell_laptop(), platform_device_type="cpu")) + assert devices["cpu"].runnable + assert not devices["cuda:0"].runnable + assert "omni-cuda" in devices["cuda:0"].reason + + +def test_the_igpu_and_npu_are_reported_not_dropped(): + """'The machine has an NPU' and 'a stage can run on it' are different + claims; the plan has to be able to show which one it relied on. + + Both are now reachable (M3), through a worker in another interpreter, so + the invariant is no longer "never runnable" -- it is that the record always + says which of the two claims it is making, and never leaves a caller to + guess. + """ + devices = _by_id(enumerate_devices(blackwell_laptop(), platform_device_type="cuda")) + assert "npu:amd" in devices + assert "gpu_integrated:amd" in devices + for key in ("npu:amd", "gpu_integrated:amd"): + device = devices[key] + assert device.evidence in ("B", "D") + if device.runnable: + assert device.backend, "a runnable device must name the backend that drives it" + assert device.extra["worker_route"], "and the route it was resolved through" + else: + assert device.reason, "and an unrunnable one must say why" + assert device.backend is None + # Either way these execute exported graphs, never a checkpoint. + assert runs_exported_graphs(device) + + +def test_the_amd_devices_refuse_with_the_missing_interpreter_named(monkeypatch): + """A device that works and a venv that is not installed are different + problems; the reason has to distinguish them well enough to act on.""" + for route in ("ORT_VITISAI", "ORT_DML", "TORCH_DML"): + monkeypatch.setenv(f"VLLM_OMNI_EXTERNAL_PYTHON_{route}", "/nonexistent/python") + devices = _by_id(enumerate_devices(blackwell_laptop(), platform_device_type="cuda")) + for key in ("npu:amd", "gpu_integrated:amd"): + device = devices[key] + assert not device.runnable and device.backend is None + assert "/nonexistent/python" in device.reason or "native-Windows" in device.reason + + +def test_the_amd_devices_get_no_memory_budget_of_their_own(): + """Section 6: the iGPU and the NPU allocate out of the CPU's RAM. + + Carrying the pool size on their rows would let a caller add three copies of + the same memory together, so the number lives on the host row only and the + stage planner reads it from there. + """ + devices = _by_id(enumerate_devices(blackwell_laptop(), platform_device_type="cpu")) + for key in ("npu:amd", "gpu_integrated:amd"): + assert devices[key].memory_bytes == 0 + assert devices[key].memory_pool == "host_ram" + assert devices[key].extra["host_pool_bytes"] > 0 + + +def test_a_cpu_only_box_reports_no_cuda_row_at_all(): + """A masked GPU and an absent GPU are different situations and must not + produce the same record.""" + devices = _by_id(enumerate_devices(cpu_only_box(), platform_device_type="cpu")) + assert "cuda:0" not in devices + assert devices["cpu"].runnable + + +# ------------------------------------------------------------------ masking +def test_masking_disables_without_hiding(): + devices = _by_id(enumerate_devices( + blackwell_laptop(), mask=frozenset({ACCEL_GPU_DISCRETE}), platform_device_type="cuda" + )) + gpu = devices["cuda:0"] + assert gpu.masked and not gpu.runnable + # Still enumerated, with its real memory: a mask is a policy, not an + # amnesia, and the record has to show the card was there. + assert gpu.memory_bytes > 0 + assert "mask" in gpu.reason + + +def test_mask_is_read_from_the_environment(monkeypatch): + monkeypatch.setenv(MASK_ENV, "gpu_discrete, cpu") + assert parse_mask() == frozenset({"gpu_discrete", "cpu"}) + monkeypatch.delenv(MASK_ENV) + assert parse_mask() == frozenset() + + +def test_an_explicit_mask_beats_the_environment(monkeypatch): + monkeypatch.setenv(MASK_ENV, "gpu_discrete") + assert parse_mask("") == frozenset() + + +# ---------------------------------------------------------------- pooling +def test_cpu_and_accelerators_share_one_host_pool(): + """Section 6 forbids giving the iGPU and the NPU each their own RAM.""" + devices = enumerate_devices(blackwell_laptop(), platform_device_type="cpu") + host_pool = [d for d in devices if d.memory_pool == "host_ram"] + assert {d.device_id for d in host_pool} >= {"cpu", "npu:amd", "gpu_integrated:amd"} + # Only the CPU row carries the bytes; the others are 0 so nothing can add + # three copies of the same RAM together. + assert sum(d.memory_bytes for d in host_pool) == next( + d.memory_bytes for d in host_pool if d.device_id == "cpu" + ) + + +def test_the_discrete_gpu_has_its_own_pool(): + devices = _by_id(enumerate_devices(blackwell_laptop(), platform_device_type="cuda")) + assert devices["cuda:0"].memory_pool == "vram" + assert devices["cpu"].memory_pool == "host_ram" + + +def test_cpu_weight_formats_include_int8_on_x86(): + devices = _by_id(enumerate_devices(blackwell_laptop(), platform_device_type="cpu")) + assert FORMAT_INT8 in devices["cpu"].weight_formats + assert FORMAT_GPTQ in devices["cpu"].weight_formats diff --git a/tests/edge/local/test_e2e_local_text.py b/tests/edge/local/test_e2e_local_text.py new file mode 100644 index 000000000..12bece39d --- /dev/null +++ b/tests/edge/local/test_e2e_local_text.py @@ -0,0 +1,279 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""End to end: the M0 acceptance criteria, executed against a real checkpoint. + +This is the test that says the local text mode works, so it tests the thing the +roadmap actually asks for rather than a proxy for it: + +* both startups -- device-visible and capability-masked -- reach a decision, + and a refusal is a pass, because "选择兼容方案或明确拒绝" is an either/or; +* the same fixed set of twelve prompts, each generating at least 128 tokens; +* TTFT and decode rate recorded per request; +* the budgeted peak is an upper bound on the measured one, and nothing OOMs; +* cancellation releases, and no event from the retired epoch is delivered. + +It needs real weights and a real backend, so it skips rather than fails when +there is no checkpoint. Which one it runs is whatever this interpreter can +drive: the CUDA venv gets the bf16 build, the CPU venv gets the int8 one -- +which is itself the point, since the int8 build is the one the CUDA plan +refuses. Override with ``VLLM_OMNI_LOCAL_TEST_MODEL``. + +Roughly a minute on the reference laptop's GPU and three on its CPU. Marked +``local_model`` because it is not a unit test and should not run in a PR gate. +""" + +from __future__ import annotations + +import os +from pathlib import Path + +import pytest +import pytest_asyncio + +from vllm_omni.edge.local import ( + LocalTextEngine, + PlanNotAdmitted, + acceptance_prompts, + plan_text_session, +) +from vllm_omni.edge.local.capabilities import FORMAT_INT8 + +# One event loop for the whole module. The engine fixture is module-scoped -- +# loading 7.7 GB per test would be absurd -- and the asyncio primitives inside +# AsyncOmni belong to whichever loop created them. Leaving the tests on +# pytest-asyncio's default per-function loop means a cancel awaits a task whose +# backend cleanup lives on a different, no-longer-running loop, and it hangs +# there rather than failing. +pytestmark = [pytest.mark.local_model, pytest.mark.asyncio(loop_scope="module")] + +MIN_TOKENS = 128 +MAX_MODEL_LEN = 4096 + +# Checkpoints under the repo's models/ directory, by the platform that can run +# them. The int8 build is deliberately the CPU entry: Blackwell refuses it. +_CANDIDATES = { + "cuda": ("Spark-X2.5-4B", "Spark-X2.5-1.7B-int8"), + "cpu": ("Spark-X2.5-1.7B-int8", "Spark-X2.5-4B"), +} + + +def _models_root() -> Path: + # tests/edge/local/ -> vllm-omni/ -> the workspace that holds models/ + return Path(__file__).resolve().parents[3].parent / "models" + + +def _platform() -> str: + from vllm.platforms import current_platform + + return str(getattr(current_platform, "device_type", "") or "") + + +def _find_model() -> str: + override = os.environ.get("VLLM_OMNI_LOCAL_TEST_MODEL") + if override: + return override + root = _models_root() + for name in _CANDIDATES.get(_platform(), ()): + if (root / name / "config.json").is_file(): + return str(root / name) + pytest.skip( + f"no Spark checkpoint for platform {_platform()!r} under {root}; " + "set VLLM_OMNI_LOCAL_TEST_MODEL to run this" + ) + + +@pytest.fixture(scope="module") +def model_dir() -> str: + return _find_model() + + +@pytest.fixture(scope="module") +def plan(model_dir: str): + p = plan_text_session(model_dir, max_model_len=MAX_MODEL_LEN, max_num_seqs=1) + if not p.admitted: + pytest.skip(f"this host cannot run {model_dir}:\n{p.summary()}") + return p + + +@pytest_asyncio.fixture(scope="module", loop_scope="module") +async def engine(plan): + eng = LocalTextEngine(plan) + await eng.start() + try: + yield eng + finally: + await eng.close() + + +# ------------------------------------------------------------------ planning +def test_the_masked_startup_reaches_a_decision(model_dir): + """The no-NVIDIA class must produce a plan or an explicit refusal -- never + a crash and never a silent substitution.""" + masked = plan_text_session( + model_dir, max_model_len=MAX_MODEL_LEN, mask=frozenset({"gpu_discrete"}) + ) + if masked.admitted: + assert masked.selected.memory_pool == "host_ram" + assert masked.selected.kind != "gpu_discrete" + else: + assert masked.refusals + assert all(r.message and r.remedy for r in masked.refusals) + assert any("mask" in r.message for r in masked.refusals) or _platform() != "cuda" + + +def test_a_refused_plan_cannot_be_started(model_dir): + refused = plan_text_session( + model_dir, max_model_len=MAX_MODEL_LEN, mask=frozenset({"gpu_discrete", "cpu"}) + ) + assert not refused.admitted + with pytest.raises(PlanNotAdmitted): + LocalTextEngine(refused) + + +@pytest.mark.skipif( + not (_models_root() / "Spark-X2.5-1.7B-int8" / "config.json").is_file(), + reason="the int8 checkpoint is not present", +) +def test_int8_is_refused_on_cuda_and_admitted_on_cpu(): + """The refusal and its remedy, checked against the real checkpoint on + whichever backend this interpreter has.""" + int8 = str(_models_root() / "Spark-X2.5-1.7B-int8") + p = plan_text_session(int8, max_model_len=MAX_MODEL_LEN) + assert p.manifest.weight_format == FORMAT_INT8 + if _platform() == "cuda": + assert not p.admitted + formats = [r for r in p.refusals if r.code == "weight_format_unsupported"] + assert formats and "omni-cpu" in formats[0].remedy + else: + assert p.admitted + assert p.selected.device_id == "cpu" + + +# ------------------------------------------------------------------- loading +async def test_the_engine_loads_and_reports_where_it_landed(engine): + assert engine.load_seconds and engine.load_seconds > 0 + placement = engine.report_placement() + configured = placement["configured"] + assert "error" not in configured + # Read back from the stage's own config, not echoed from the plan. + assert configured["device_type"] == _platform() + assert configured["max_model_len"] == MAX_MODEL_LEN + assert configured["enforce_eager"] is True + + +async def test_the_kv_pool_is_the_size_the_plan_asked_for(engine, plan): + """An absolute budget, honoured -- not whatever was left inside a + utilisation fraction.""" + configured = engine.report_placement()["configured"] + assert configured["kv_cache_memory_bytes"] == plan.kv_budget.bytes_total + + +# ---------------------------------------------------------- the twelve prompts +async def test_twelve_prompts_each_reach_128_tokens(engine): + prompts = acceptance_prompts() + assert len(prompts) == 12 + + records = [] + for name, text in prompts: + session = engine.open_session() + request_id, stream = await engine.submit( + session, text, max_tokens=MIN_TOKENS, temperature=0.0, ignore_eos=True + ) + seqs = [] + async for event in stream: + if event.kind == "token": + seqs.append(event.seq) + assert event.epoch == session.epoch + assert event.error is None, f"{name}: {event.error}" + record = engine.records[request_id] + engine.close_session(session.session_id) + + assert record.finished, f"{name} did not finish" + assert record.output_tokens >= MIN_TOKENS, f"{name} produced {record.output_tokens}" + assert record.ttft_s is not None and record.ttft_s > 0 + assert record.decode_tok_per_s is not None and record.decode_tok_per_s > 0 + # Incremental, in order, no repeats: seq is what detects a lost chunk. + assert seqs == sorted(seqs) + assert len(seqs) == len(set(seqs)) + records.append(record) + + # The long prompt crosses the 512 sliding window many times; if window + # rollover were broken it would be broken here and not on the short ones. + long_record = next(r for r in records if r.prompt_tokens and r.prompt_tokens > 2000) + assert long_record.output_tokens >= MIN_TOKENS + + +async def test_the_budget_bounded_the_measured_peak(engine, plan): + """Admission's whole claim: the plan's number is an upper bound, proven + before anything allocated.""" + measured = engine.measured_peak() + assert measured["budgeted_peak_bytes"] == plan.peak_bytes + attributable = measured["attributable_peak_bytes"] + assert attributable is not None and attributable > 0 + assert attributable <= measured["budgeted_peak_bytes"], ( + f"measured {attributable / 2**30:.2f} GiB exceeded the budgeted " + f"{measured['budgeted_peak_bytes'] / 2**30:.2f} GiB" + ) + + +async def test_usage_separates_what_was_measured_from_what_was_estimated(engine): + usage = engine.report_usage() + by_purpose = {r["purpose"]: r for r in usage["reservations"]} + assert by_purpose["weights"]["evidence"] == "D" + assert by_purpose["kv_cache"]["evidence"] == "E" + # Nothing gets an actual it cannot support: only the weights line is + # bounded by a counter on this host. + bounded = [p for p, r in by_purpose.items() if r["actual_upper_bound_bytes"]] + assert bounded == ["weights"] + + +# ------------------------------------------------------------- cancellation +async def test_cancel_releases_and_delivers_nothing_from_the_retired_epoch(engine): + session = engine.open_session() + request_id, stream = await engine.submit( + session, acceptance_prompts()[0][1], max_tokens=4096, ignore_eos=True + ) + seen = 0 + async for event in stream: + if event.kind == "token": + seen += 1 + if seen >= 8: + break + assert seen >= 8 + + report = await engine.cancel(request_id) + assert report["aborted_in_backend"] + assert report["producer_stopped"], "the producer task did not unwind" + assert report["inflight_after"] == 0 + # The session's state is retired, not reused. + assert report["session_epoch_now"] == session.epoch + 1 + + leaked = [] + while True: + event = await stream.get() + if event is None: + break + leaked.append(event) + assert leaked == [], f"{len(leaked)} events escaped the retired epoch" + + +async def test_a_retired_handle_cannot_be_resubmitted(engine): + session = engine.open_session() + request_id, stream = await engine.submit(session, "Hello", max_tokens=32) + async for _ in stream: + break + await engine.cancel(request_id) + with pytest.raises(RuntimeError, match="stale"): + await engine.submit(session, "Hello again", max_tokens=8) + + +async def test_the_engine_still_generates_after_a_cancel(engine): + """A cancelled request must not poison the engine for the next one.""" + session = engine.open_session() + request_id, stream = await engine.submit( + session, "The capital of Australia is", max_tokens=16, ignore_eos=True + ) + async for _ in stream: + pass + record = engine.records[request_id] + assert record.finished and record.output_tokens >= 16 diff --git a/tests/edge/local/test_gpu_attribution.py b/tests/edge/local/test_gpu_attribution.py new file mode 100644 index 000000000..d1c90d975 --- /dev/null +++ b/tests/edge/local/test_gpu_attribution.py @@ -0,0 +1,71 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""The placement report must measure per-process GPU attribution, not assume it. + +[edge-infer W4] The W1 acceptance on native Windows carried a hardcoded "under +WSL2" explanation for an empty NVML process list. These tests pin the repaired +behaviour: the reason names the host it actually ran on, and a successful +attribution (NVML lists one of our PIDs) is reported as data, not a reason. +""" + +from __future__ import annotations + +import builtins +import sys +import types + +import pytest + +from vllm_omni.edge.local import engine + + +def test_host_kind_is_one_of_the_three_hosts(): + assert engine._host_kind() in {"windows", "wsl2", "linux"} + + +def test_reason_names_the_host_when_nvml_is_unavailable(monkeypatch): + real_import = builtins.__import__ + + def no_pynvml(name, *args, **kwargs): + if name == "pynvml": + raise ImportError("pynvml is not installed in this venv") + return real_import(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", no_pynvml) + out = engine.gpu_attribution_for_tree() + assert out["per_process_gpu_attribution"] is None + assert engine._host_kind() in out["per_process_gpu_attribution_reason"] + assert "WSL2" not in out["per_process_gpu_attribution_reason"] or engine._host_kind() == "wsl2" + + +def _fake_pynvml(pids): + m = types.ModuleType("pynvml") + m.nvmlInit = lambda: None + m.nvmlShutdown = lambda: None + m.nvmlDeviceGetHandleByIndex = lambda i: object() + m.nvmlDeviceGetComputeRunningProcesses = lambda h: [ + types.SimpleNamespace(pid=p, usedGpuMemory=1234) for p in pids + ] + return m + + +def test_attribution_is_data_when_nvml_lists_our_pid(monkeypatch): + import os + + monkeypatch.setitem(sys.modules, "pynvml", _fake_pynvml([os.getpid(), 1])) + out = engine.gpu_attribution_for_tree() + assert out["per_process_gpu_attribution"] == [ + {"pid": os.getpid(), "used_gpu_memory": 1234} + ] + assert out["per_process_gpu_attribution_reason"] is None + + +def test_reason_states_the_measured_count_when_none_are_ours(monkeypatch): + monkeypatch.setitem(sys.modules, "pynvml", _fake_pynvml([1, 2])) + out = engine.gpu_attribution_for_tree() + assert out["per_process_gpu_attribution"] is None + reason = out["per_process_gpu_attribution_reason"] + assert reason.startswith("nvmlDeviceGetComputeRunningProcesses returned 2 process(es)") + host = engine._host_kind() + expected = {"wsl2": "WSL2", "windows": "WDDM", "linux": "no compute process"}[host] + assert expected in reason diff --git a/tests/edge/local/test_manifest.py b/tests/edge/local/test_manifest.py new file mode 100644 index 000000000..3065d980b --- /dev/null +++ b/tests/edge/local/test_manifest.py @@ -0,0 +1,99 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Identifying a checkpoint well enough to tell it from its siblings.""" + +import json + +import pytest + +from vllm_omni.edge.local.capabilities import FORMAT_DENSE, FORMAT_FP8, FORMAT_INT8 +from vllm_omni.edge.local.manifest import build_manifest, runtime_versions + +from .conftest import FP8_QUANT_CONFIG, INT8_QUANT_CONFIG, write_checkpoint + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_quantization_format_is_read_from_the_config(tmp_path): + assert build_manifest(write_checkpoint(tmp_path / "d")).weight_format == FORMAT_DENSE + assert build_manifest( + write_checkpoint(tmp_path / "i", quantization=INT8_QUANT_CONFIG) + ).weight_format == FORMAT_INT8 + assert build_manifest( + write_checkpoint(tmp_path / "f", quantization=FP8_QUANT_CONFIG) + ).weight_format == FORMAT_FP8 + + +def test_compressed_tensors_int8_and_fp8_are_distinguished_by_format_not_method(tmp_path): + """Both say ``compressed-tensors``; only ``format`` separates the one + Blackwell can execute from the one it cannot.""" + int8 = build_manifest(write_checkpoint(tmp_path / "i", quantization=INT8_QUANT_CONFIG)) + fp8 = build_manifest(write_checkpoint(tmp_path / "f", quantization=FP8_QUANT_CONFIG)) + assert int8.quantization["quant_method"] == fp8.quantization["quant_method"] + assert int8.weight_format != fp8.weight_format + + +def test_two_quantizations_of_one_model_get_different_ids(tmp_path): + dense = build_manifest(write_checkpoint(tmp_path / "d")) + int8 = build_manifest(write_checkpoint(tmp_path / "i", quantization=INT8_QUANT_CONFIG)) + assert dense.artifact_id != int8.artifact_id + + +def test_the_id_is_stable_across_reads(tmp_path): + path = write_checkpoint(tmp_path / "d") + assert build_manifest(path).artifact_id == build_manifest(path).artifact_id + + +def test_a_changed_tokenizer_changes_the_id(tmp_path): + """The tokenizer decides what the model is asked, and moves independently + of the weights.""" + path = write_checkpoint(tmp_path / "d") + before = build_manifest(path).artifact_id + (path / "chat_template.jinja").write_text("{{ different }}", encoding="utf-8") + assert build_manifest(path).artifact_id != before + + +def test_weights_are_not_hashed_unless_asked(tmp_path): + path = write_checkpoint(tmp_path / "d") + assert build_manifest(path).weight_sha256 is None + assert any("digest_weights" in n for n in build_manifest(path).notes) + assert build_manifest(path, digest_weights=True).weight_sha256 is not None + + +def test_weight_bytes_and_shards_are_recorded(tmp_path): + path = write_checkpoint(tmp_path / "d", weight_bytes=8192, shards=4) + manifest = build_manifest(path) + assert len(manifest.weight_files) == 4 + assert manifest.weight_bytes == 8192 + + +def test_a_directory_without_a_config_is_an_error(tmp_path): + (tmp_path / "empty").mkdir() + with pytest.raises(FileNotFoundError): + build_manifest(tmp_path / "empty") + + +def test_an_hf_snapshot_path_yields_its_revision(tmp_path): + root = tmp_path / "models--org--name" / "snapshots" / "deadbeef1234" + write_checkpoint(root) + assert build_manifest(root).revision == "deadbeef1234" + + +def test_a_local_checkpoint_reports_no_revision_rather_than_guessing(tmp_path): + assert build_manifest(write_checkpoint(tmp_path / "d")).revision is None + + +def test_the_runtime_records_both_the_wheel_and_the_checkout(): + versions = runtime_versions() + assert versions.python + # The import path, not just the version string: an editable install and a + # site-packages copy report the same version and are different code. + assert versions.vllm_omni_file.endswith("__init__.py") + assert isinstance(versions.checkouts, dict) + + +def test_the_manifest_is_json_serialisable(tmp_path): + manifest = build_manifest(write_checkpoint(tmp_path / "d")) + back = json.loads(json.dumps(manifest.to_dict(), default=str)) + assert back["artifact_id"] == manifest.artifact_id + assert back["architectures"] == ["Spark2_5ForCausalLM"] diff --git a/tests/edge/local/test_parity.py b/tests/edge/local/test_parity.py new file mode 100644 index 000000000..a7e22b83f --- /dev/null +++ b/tests/edge/local/test_parity.py @@ -0,0 +1,129 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Reference consistency: does the Omni stage path change the model's answer? + +The comparison logic is unit-tested here; the run that actually loads two +engines is at the bottom and skips without a checkpoint. It is separated from +``test_e2e_local_text`` because it cannot share that module's engine fixture -- +it needs the reference engine loaded and released *before* the Omni one, and +two copies of a 7.7 GB checkpoint do not fit the budget this milestone is about. +""" + +from __future__ import annotations + +import os +from pathlib import Path + +import pytest + +from vllm_omni.edge.local.parity import compare_ids, summarize +from vllm_omni.edge.local.plan import plan_text_session + +MAX_MODEL_LEN = 4096 + + +# ------------------------------------------------------------------- unit +@pytest.mark.core_model +@pytest.mark.cpu +class TestCompare: + def test_identical_sequences_are_exact(self): + assert compare_ids([1, 2, 3], [1, 2, 3]) == (True, None, 3) + + def test_a_divergence_reports_where(self): + assert compare_ids([1, 2, 3, 4], [1, 2, 9, 4]) == (False, 2, 2) + + def test_a_short_engine_sequence_is_not_exact(self): + """A truncation is a divergence at the end, not agreement.""" + assert compare_ids([1, 2, 3], [1, 2]) == (False, 2, 2) + + def test_a_long_engine_sequence_is_not_exact(self): + assert compare_ids([1, 2], [1, 2, 3]) == (False, 2, 2) + + def test_divergence_at_the_first_token(self): + assert compare_ids([5, 6], [7, 6]) == (False, 0, 0) + + def test_empty_sequences_agree_vacuously(self): + assert compare_ids([], []) == (True, None, 0) + + def test_summary_counts_and_reports_the_worst_prefix(self): + from vllm_omni.edge.local.parity import PromptParity + + rows = [ + PromptParity("a", 4, 8, 8, True, None, 8), + PromptParity("b", 4, 8, 8, False, 5, 5), + PromptParity("c", 4, 8, 8, False, 1, 1), + ] + stats = summarize(rows) + assert stats["exact"] == 1 + assert stats["prompts"] == 3 + assert stats["min_agreement_prefix"] == 1 + assert stats["mean_agreement_prefix"] == 3.0 + + def test_an_all_exact_summary_has_no_divergence_stats(self): + from vllm_omni.edge.local.parity import PromptParity + + stats = summarize([PromptParity("a", 4, 8, 8, True, None, 8)]) + assert stats["exact_fraction"] == 1.0 + assert stats["min_agreement_prefix"] is None + + +# -------------------------------------------------------------------- e2e +def _models_root() -> Path: + return Path(__file__).resolve().parents[3].parent / "models" + + +def _platform() -> str: + from vllm.platforms import current_platform + + return str(getattr(current_platform, "device_type", "") or "") + + +def _find_model() -> str: + override = os.environ.get("VLLM_OMNI_LOCAL_TEST_MODEL") + if override: + return override + candidates = ("Spark-X2.5-4B",) if _platform() == "cuda" else ("Spark-X2.5-1.7B-int8",) + for name in candidates: + if (_models_root() / name / "config.json").is_file(): + return str(_models_root() / name) + pytest.skip(f"no Spark checkpoint for platform {_platform()!r}") + + +@pytest.mark.local_model +@pytest.mark.asyncio(loop_scope="module") +async def test_the_omni_path_decodes_what_plain_vllm_decodes(): + """The hard acceptance condition: same checkpoint, same limits, same + greedy decode -- the only difference is the path through the Omni stage. + + Four prompts rather than the full twelve: this loads two engines and the + point is to detect a systematic difference in the stage path, which four + prompts of different shapes (short chat, tool-formatted, long, plain) show + as well as twelve. The full set runs from the CLI's ``parity`` command. + """ + from vllm_omni.edge.local.cli import _engine_token_ids + from vllm_omni.edge.local.parity import reference_token_ids + from vllm_omni.edge.local.prompts import acceptance_prompts + + model_dir = _find_model() + plan = plan_text_session(model_dir, max_model_len=MAX_MODEL_LEN, max_num_seqs=1) + if not plan.admitted: + pytest.skip(f"this host cannot run {model_dir}") + + all_prompts = dict(acceptance_prompts()) + prompts = [(n, all_prompts[n]) for n in ("chat_short", "chat_tools", "chat_long", "plain_0")] + max_tokens = 64 + + reference = reference_token_ids( + model_dir, prompts, engine_kwargs=plan.engine_kwargs, max_tokens=max_tokens + ) + engine = await _engine_token_ids(plan, prompts, max_tokens) + + mismatches = [] + for name, _ in prompts: + ref = reference[name]["token_ids"] + eng = engine[name]["token_ids"] + assert len(ref) == max_tokens, f"{name}: reference produced {len(ref)}" + exact, first, prefix = compare_ids(ref, eng) + if not exact: + mismatches.append(f"{name}: diverged at {first}/{len(ref)} (prefix {prefix})") + assert not mismatches, "the Omni stage path changed the decode:\n" + "\n".join(mismatches) diff --git a/tests/edge/local/test_plan.py b/tests/edge/local/test_plan.py new file mode 100644 index 000000000..fe2fa4170 --- /dev/null +++ b/tests/edge/local/test_plan.py @@ -0,0 +1,200 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Admission: say yes with a budget, or no with a reason and a remedy.""" + +import pytest + +from vllm_omni.edge.local.plan import ( + GiB, + REFUSE_CAPACITY, + REFUSE_FORMAT, + REFUSE_NO_DEVICE, + plan_text_session, +) + +from .conftest import FP8_QUANT_CONFIG, blackwell_laptop, cpu_only_box, write_checkpoint + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _plan(model_dir, profile=None, platform="cuda", **kw): + return plan_text_session( + str(model_dir), + profile=profile if profile is not None else blackwell_laptop(), + platform_device_type=platform, + **kw, + ) + + +def _codes(plan): + return {r.code for r in plan.refusals} + + +# ------------------------------------------------------------------ the yes +def test_a_dense_checkpoint_is_admitted_to_the_gpu(dense_checkpoint): + plan = _plan(dense_checkpoint) + assert plan.admitted + assert plan.selected.device_id == "cuda:0" + assert plan.backend == "vllm:cuda" + + +def test_an_admitted_plan_prices_every_line_of_section_6_2(dense_checkpoint): + """The budget is not 'the weights'. It is the sum the design specifies, + and each term has to be visible so a reviewer can argue with one.""" + plan = _plan(dense_checkpoint) + purposes = {r.purpose for r in plan.reservations} + assert purposes == { + "weights", "kv_cache", "activations", "backend_workspace", + "load_transient", "external_reserve", "safety_margin", + } + assert all(r.pool == "vram" for r in plan.reservations) + assert all(r.evidence in {"D", "E"} for r in plan.reservations) + # The weights line is read off the files, not estimated. + weights = next(r for r in plan.reservations if r.purpose == "weights") + assert weights.evidence == "D" + assert weights.bytes == plan.manifest.weight_bytes + + +def test_graph_capture_is_charged_only_when_graphs_are_on(dense_checkpoint): + eager = _plan(dense_checkpoint, enforce_eager=True) + graphs = _plan(dense_checkpoint, enforce_eager=False) + assert "graph_capture" not in {r.purpose for r in eager.reservations} + assert "graph_capture" in {r.purpose for r in graphs.reservations} + assert graphs.peak_bytes > eager.peak_bytes + + +def test_the_peak_is_a_max_not_a_sum(dense_checkpoint): + """The loader's transient is gone before the KV pool is sized, so adding + both would refuse plans that fit.""" + plan = _plan(dense_checkpoint) + everything = sum(r.bytes for r in plan.reservations) + assert plan.peak_bytes <= everything + + +def test_kv_reservation_tracks_the_context(dense_checkpoint): + short = _plan(dense_checkpoint, max_model_len=1024) + long = _plan(dense_checkpoint, max_model_len=8192) + kv_short = next(r for r in short.reservations if r.purpose == "kv_cache").bytes + kv_long = next(r for r in long.reservations if r.purpose == "kv_cache").bytes + assert kv_long > kv_short + + +def test_sliding_layers_are_not_charged_full_length(dense_checkpoint): + """21 of 28 layers hold a 512 window, not the whole context. Charging them + all at full length is the arithmetic the kv_budget module exists to fix.""" + plan = _plan(dense_checkpoint, max_model_len=8192) + assert plan.kv_budget.sliding_layers == 21 + assert plan.kv_budget.full_layers == 7 + assert plan.kv_budget.bytes_total_hybrid < plan.kv_budget.bytes_total + + +# ------------------------------------------------------------------- the no +def test_int8_is_refused_on_blackwell_although_it_fits(int8_checkpoint): + """4 MiB of weights against 24 GiB of VRAM: every capacity check passes and + the answer is still no.""" + plan = _plan(int8_checkpoint) + assert not plan.admitted + assert REFUSE_FORMAT in _codes(plan) + refusal = next(r for r in plan.refusals if r.code == REFUSE_FORMAT) + assert refusal.device_id == "cuda:0" + assert "int8" in refusal.message + assert "omni-cpu" in refusal.remedy + + +def test_the_same_int8_checkpoint_is_admitted_on_the_cpu(int8_checkpoint): + """The remedy the refusal names has to actually be a plan.""" + plan = _plan(int8_checkpoint, profile=cpu_only_box(), platform="cpu") + assert plan.admitted + assert plan.selected.device_id == "cpu" + assert plan.backend == "vllm:cpu" + + +def test_fp8_is_admitted_where_int8_is_not(tmp_path): + fp8 = write_checkpoint(tmp_path / "fp8", quantization=FP8_QUANT_CONFIG) + assert _plan(fp8).admitted + + +def test_a_checkpoint_too_large_for_the_pool_is_refused_with_a_number(tmp_path): + huge = write_checkpoint(tmp_path / "huge", weight_bytes=40 * GiB, shards=8) + plan = _plan(huge) + assert not plan.admitted + assert REFUSE_CAPACITY in _codes(plan) + refusal = next(r for r in plan.refusals if r.code == REFUSE_CAPACITY) + assert "GiB" in refusal.message + + +def test_a_capacity_refusal_never_silently_shortens_the_context(tmp_path): + """Trimming max_model_len is offered as a remedy the operator applies, not + something the planner does on its behalf.""" + # Sized so the overrun is smaller than the KV cache it asked for: that is + # the case where a shorter context would in fact fit, which is exactly when + # an engine that "helpfully" truncates would do damage. + big = write_checkpoint(tmp_path / "big", weight_bytes=16 * GiB, shards=8) + plan = _plan(big, max_model_len=131072) + assert not plan.admitted + refusal = next(r for r in plan.refusals if r.code == REFUSE_CAPACITY) + assert "max-model-len" in refusal.remedy + assert "forbids" in refusal.remedy + # And the plan still reports the context that was asked for, not a + # quietly reduced one. + assert plan.request_limits["max_model_len"] == 131072 + + +def test_masking_the_gpu_leaves_nothing_runnable_in_a_cuda_process(dense_checkpoint): + """The capability-masked startup of the acceptance run: a refusal here is + the correct outcome, not a failure.""" + plan = _plan(dense_checkpoint, mask=frozenset({"gpu_discrete"})) + assert not plan.admitted + assert REFUSE_NO_DEVICE in _codes(plan) + assert any("mask" in r.message for r in plan.refusals) + + +def test_every_unusable_device_is_explained(dense_checkpoint): + plan = _plan(dense_checkpoint) + explained = {r.device_id for r in plan.refusals} + assert {"cpu", "npu:amd", "gpu_integrated:amd"} <= explained + assert all(r.message and r.remedy for r in plan.refusals) + + +# ------------------------------------------------------------ engine kwargs +def test_the_memory_fraction_is_derived_not_defaulted(dense_checkpoint): + """vLLM reads this field as host RAM on the CPU backend, so leaving it at + its default there asks for 92% of the machine. It is always derived.""" + for platform, profile in (("cuda", blackwell_laptop()), ("cpu", cpu_only_box())): + plan = _plan(dense_checkpoint, profile=profile, platform=platform) + assert plan.admitted + fraction = plan.engine_kwargs["gpu_memory_utilization"] + assert 0 < fraction <= 0.95 + assert fraction >= plan.peak_bytes / plan.selected.memory_bytes * 0.5 + + +def test_the_kv_pool_is_absolute_not_leftover(dense_checkpoint): + plan = _plan(dense_checkpoint) + assert plan.engine_kwargs["kv_cache_memory_bytes"] == plan.kv_budget.bytes_total + + +def test_request_limits_reach_the_backend(dense_checkpoint): + plan = _plan(dense_checkpoint, max_model_len=2048, max_num_seqs=2) + assert plan.engine_kwargs["max_model_len"] == 2048 + assert plan.engine_kwargs["max_num_seqs"] == 2 + assert plan.engine_kwargs["enforce_eager"] is True + + +# ---------------------------------------------------------------- the stage +def test_the_plan_declares_exactly_one_text_stage(dense_checkpoint): + plan = _plan(dense_checkpoint) + assert plan.stage["stage_id"] == 0 + assert plan.stage["execution_type"] == "LLM_AR" + assert plan.stage["final_output_type"] == "text" + assert plan.stage["input_sources"] == [] + + +def test_the_record_round_trips_to_json(dense_checkpoint): + import json + + plan = _plan(dense_checkpoint) + text = json.dumps(plan.to_dict(), default=str) + back = json.loads(text) + assert back["admitted"] + assert back["selected_device"]["device_id"] == "cuda:0" + assert back["manifest"]["artifact_id"] == plan.manifest.artifact_id diff --git a/tests/edge/local/test_session.py b/tests/edge/local/test_session.py new file mode 100644 index 000000000..4b6e6182c --- /dev/null +++ b/tests/edge/local/test_session.py @@ -0,0 +1,156 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""The event stream: backpressure that blocks, and an epoch that fences.""" + +import asyncio + +import pytest + +from vllm_omni.edge.local.session import ( + STATE_LAYOUT_VERSION, + BoundedEventStream, + ChunkEvent, + StateHandle, + StreamClosed, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _event(seq: int, *, epoch: int = 0, text: str = "x", kind: str = "token") -> ChunkEvent: + return ChunkEvent( + request_id="r", stage_id=0, seq=seq, epoch=epoch, kind=kind, + payload={"text": text}, started_unix=0.0, emitted_unix=0.0, + ) + + +# -------------------------------------------------------------- state handle +def test_a_handle_carries_the_artifact_it_was_made_for(): + handle = StateHandle(session_id="s", backend="vllm:cuda", artifact_id="abc") + assert handle.layout_version == STATE_LAYOUT_VERSION + assert handle.accepts(handle) + + +def test_state_from_another_checkpoint_is_not_accepted(): + """Same shapes from a different quantization are not the same state.""" + a = StateHandle(session_id="s", backend="vllm:cuda", artifact_id="abc") + b = StateHandle(session_id="s", backend="vllm:cuda", artifact_id="def") + assert not a.accepts(b) + + +def test_state_from_a_retired_epoch_is_not_accepted(): + handle = StateHandle(session_id="s", backend="vllm:cuda", artifact_id="abc") + assert not handle.next_epoch().accepts(handle) + + +def test_state_is_never_migratable_in_m0(): + handle = StateHandle(session_id="s", backend="vllm:cuda", artifact_id="abc") + assert handle.migratable is False + assert handle.replayable is True + + +# ---------------------------------------------------------------- ordering +async def test_events_arrive_in_order(): + stream = BoundedEventStream() + for i in range(5): + await stream.put(_event(i)) + await stream.close() + assert [e.seq async for e in stream] == [0, 1, 2, 3, 4] + + +async def test_closing_ends_the_iteration(): + stream = BoundedEventStream() + await stream.put(_event(0)) + await stream.close() + seen = [e async for e in stream] + assert len(seen) == 1 + + +async def test_producing_after_close_raises(): + stream = BoundedEventStream() + await stream.close() + with pytest.raises(StreamClosed): + await stream.put(_event(0)) + + +# ------------------------------------------------------------ backpressure +async def test_a_full_queue_blocks_the_producer(): + """The bound has to actually stop the producer. A queue that grows instead + is how a slow consumer turns into an OOM.""" + stream = BoundedEventStream(max_chunks=2, max_bytes=1 << 20) + await stream.put(_event(0)) + await stream.put(_event(1)) + blocked = asyncio.create_task(stream.put(_event(2))) + await asyncio.sleep(0.05) + assert not blocked.done() + await stream.get() # consumer drains one; credit returns + await asyncio.wait_for(blocked, timeout=1.0) + assert stream.high_water_chunks == 2 + + +async def test_the_byte_bound_also_blocks(): + stream = BoundedEventStream(max_chunks=100, max_bytes=64) + await stream.put(_event(0, text="a" * 40)) + blocked = asyncio.create_task(stream.put(_event(1, text="b" * 40))) + await asyncio.sleep(0.05) + assert not blocked.done() + await stream.get() + await asyncio.wait_for(blocked, timeout=1.0) + + +async def test_an_oversized_event_still_gets_through_on_an_empty_queue(): + """Otherwise one chunk larger than the whole bound deadlocks forever.""" + stream = BoundedEventStream(max_chunks=4, max_bytes=8) + await asyncio.wait_for(stream.put(_event(0, text="x" * 400)), timeout=1.0) + assert (await stream.get()).seq == 0 + + +async def test_closing_wakes_a_blocked_producer(): + stream = BoundedEventStream(max_chunks=1, max_bytes=1 << 20) + await stream.put(_event(0)) + blocked = asyncio.create_task(stream.put(_event(1))) + await asyncio.sleep(0.05) + await stream.close() + with pytest.raises(StreamClosed): + await asyncio.wait_for(blocked, timeout=1.0) + + +# ------------------------------------------------------------------- epochs +async def test_a_retired_epoch_cannot_be_produced_into(): + stream = BoundedEventStream() + stream.retire_epoch(0) + await stream.put(_event(0, epoch=0)) + await stream.close() + assert [e async for e in stream] == [] + assert stream.dropped_stale == 1 + + +async def test_events_already_queued_from_a_retired_epoch_are_not_delivered(): + """The real cancellation race: the backend had tokens in flight when the + cancel landed, and they must not reach the consumer.""" + stream = BoundedEventStream() + for i in range(3): + await stream.put(_event(i, epoch=0)) + stream.retire_epoch(0) + await stream.put(_event(0, epoch=1, text="new")) + await stream.close() + delivered = [e async for e in stream] + assert [e.epoch for e in delivered] == [1] + assert stream.dropped_stale == 3 + + +async def test_retiring_is_monotonic(): + stream = BoundedEventStream() + stream.retire_epoch(3) + stream.retire_epoch(0) + assert stream.epoch == 4 + + +async def test_stats_report_the_bounds_and_what_was_dropped(): + stream = BoundedEventStream(max_chunks=8, max_bytes=256) + await stream.put(_event(0)) + stats = stream.stats() + assert stats["max_chunks"] == 8 + assert stats["max_bytes"] == 256 + assert stats["high_water_chunks"] == 1 + assert stats["dropped_stale"] == 0 diff --git a/tests/edge/test_hardware_probe.py b/tests/edge/test_hardware_probe.py index 114eec882..79b13a853 100644 --- a/tests/edge/test_hardware_probe.py +++ b/tests/edge/test_hardware_probe.py @@ -115,3 +115,141 @@ def test_live_probe_runs_on_this_host(): p = hp.probe(use_torch=False) assert p.n_logical >= 1 and p.ram_total_bytes > 0 assert hp.hardware_class(p) in hp.HW_CLASSES + + +# --------------------------------------------------------- accelerators +def _ryzen_ai_root(tmp_path: Path, *, wsl: bool, driver: bool = True) -> Path: + """A Ryzen AI laptop: XDNA2 NPU, Radeon iGPU, and one discrete card. + + Under WSL none of it has a device node; on native Linux the NPU shows up + at ``/dev/accel/accel0`` and the GPUs at ``/dev/dri/renderD*``. + """ + root = _snapshot(tmp_path, X86_CPUINFO, cpus={0: (5100, 0), 1: (5100, 1)}) + (root / "proc" / "version").write_text( + "Linux version 6.18.33.2-microsoft-standard-WSL2" if wsl + else "Linux version 6.18.0-generic (gcc 14)" + ) + npu = root / "sys" / "bus" / "pci" / "devices" / "0000:c5:00.1" + npu.mkdir(parents=True) + (npu / "vendor").write_text("0x1022\n") + (npu / "device").write_text("0x17f0\n") + if wsl: + return root + if driver: + (npu / "driver").mkdir() + node = root / "sys" / "class" / "accel" / "accel0" / "device" + node.mkdir(parents=True) + (node / "vendor").write_text("0x1022\n") + (node / "device").write_text("0x17f0\n") + (node / "driver").mkdir() + (root / "dev" / "accel").mkdir(parents=True) + (root / "dev" / "accel" / "accel0").write_text("") + (root / "dev" / "dri").mkdir(parents=True) + for name, vendor in (("renderD128", "0x1002"), ("renderD129", "0x10de")): + (root / "dev" / "dri" / name).write_text("") + d = root / "sys" / "class" / "drm" / name / "device" + d.mkdir(parents=True) + (d / "vendor").write_text(vendor + "\n") + (d / "device").write_text("0x150e\n") + return root + + +def test_native_linux_finds_the_npu_and_both_gpus(tmp_path): + p = hp.probe(_ryzen_ai_root(tmp_path, wsl=False), use_torch=False) + assert p.host_os == hp.HOST_OS_LINUX + (npu,) = p.npus + assert npu["usable"] and npu["name"].startswith("AMD XDNA2") + assert [a["kind"] for a in p.accelerators].count(hp.ACCEL_GPU_INTEGRATED) == 1 + assert p.accelerators_of(hp.ACCEL_GPU_DISCRETE, usable=True) + assert p.unusable_accelerators == [] + + +def test_an_npu_with_no_driver_is_reported_present_and_unusable(tmp_path): + """"No /dev/accel node" and "no such hardware" need different answers.""" + p = hp.probe(_ryzen_ai_root(tmp_path, wsl=False, driver=False), use_torch=False) + (npu,) = p.npus + assert npu["usable"] is False + assert "amdxdna" in npu["reason"] + + +def test_wsl_is_detected_and_sees_none_of_the_host_devices(tmp_path): + """The iGPU and the NPU are in the machine and unreachable from Linux; + the probe must not report either as absent or as available.""" + p = hp.probe(_ryzen_ai_root(tmp_path, wsl=True), use_torch=False, probe_host=False) + assert p.host_os == hp.HOST_OS_WSL2 + # The PCI sweep still finds the NPU; there is simply no way to open it. + (npu,) = p.npus + assert npu["usable"] is False + assert p.accelerators_of(hp.ACCEL_GPU_INTEGRATED, usable=True) == [] + + +def test_a_snapshot_root_never_shells_out_to_the_host(tmp_path): + """probe_host defaults on, but a captured snapshot must stay hermetic.""" + p = hp.probe(_ryzen_ai_root(tmp_path, wsl=True), use_torch=False, probe_host=True) + assert all(a["source"] != "wsl-interop" for a in p.accelerators) + + +def test_accelerators_survive_the_json_roundtrip(tmp_path): + p = hp.probe(_ryzen_ai_root(tmp_path, wsl=False), use_torch=False) + back = hp.HardwareProfile.from_json(p.to_json()) + assert back.accelerators == p.accelerators + assert back.host_os == p.host_os + + +def test_describe_accelerators_says_why_each_one_is_out_of_reach(tmp_path): + p = hp.probe(_ryzen_ai_root(tmp_path, wsl=True), use_torch=False, probe_host=False) + text = hp.describe_accelerators(p) + assert "not here" in text and "amdxdna" in text + + +# ------------------------------------------------------- reachability routes +def test_the_npu_reason_names_the_recipe_not_just_the_symptom(): + """The iGPU and the NPU are both off-limits to a default WSL process, but + for different reasons with different fixes. A blanket "unreachable" loses + the only part a reader can act on -- and for the NPU the actionable part is + a three-part recipe, each part of which fails silently on its own.""" + assert "MCDM" in hp._WSL_NPU_ROUTE # why WSL cannot see it + assert "onnxruntime_vitisai_ep" in hp._WSL_NPU_ROUTE # which of the two EPs + assert "DLL search path" in hp._WSL_NPU_ROUTE # how to make it load + assert "A16W8" in hp._WSL_NPU_ROUTE # what it will accept + assert "torch-directml" in hp._WSL_IGPU_ROUTE # the iGPU's other fix + + +def test_the_npu_recipe_module_states_the_three_preconditions(): + """Each of the three can fail on its own and the symptom is identical: + a session that runs correctly with every node on the CPU.""" + from vllm_omni.edge import npu_ryzenai as npu + + assert npu.EP_GENERAL == "onnxruntime_vitisai_ep.dll" + assert npu.EP_LIGHT != npu.EP_GENERAL + assert npu.PARTITION_TARGET == "AMD_AIE2P_4x8_CMC_Overlay" + kw = npu.quantization_kwargs() + # A16W8: 16-bit activations, 8-bit weights. A8W8 is silently rejected. + assert "16" in str(kw["activation_type"]) + assert "8" in str(kw["weight_type"]) and "16" not in str(kw["weight_type"]) + + +def test_the_igpu_is_usable_exactly_when_directml_is_importable(monkeypatch): + """torch-directml pins torch==2.4.1, so it never shares the engine's venv. + ``usable`` must answer for *this* process, not for the machine.""" + calls = {"n": 0} + + def fake_run(*a, **k): + calls["n"] += 1 + class R: stdout = "AMD Radeon(TM) 890M Graphics\n" + return R() + + import subprocess + monkeypatch.setattr(subprocess, "run", fake_run) + monkeypatch.setattr(hp, "directml_available", lambda: False) + (igpu,) = [a for a in hp.wsl_host_accelerators() if a["kind"] == hp.ACCEL_GPU_INTEGRATED] + assert igpu["usable"] is False and igpu["route"] == "directml" + + monkeypatch.setattr(hp, "directml_available", lambda: True) + (igpu,) = [a for a in hp.wsl_host_accelerators() if a["kind"] == hp.ACCEL_GPU_INTEGRATED] + assert igpu["usable"] is True + + +def test_directml_available_is_false_without_the_package(monkeypatch): + monkeypatch.setitem(__import__("sys").modules, "torch_directml", None) + assert hp.directml_available() is False diff --git a/tests/edge/test_kv_budget.py b/tests/edge/test_kv_budget.py index 85c1af22c..82f124e6a 100644 --- a/tests/edge/test_kv_budget.py +++ b/tests/edge/test_kv_budget.py @@ -105,3 +105,100 @@ def test_dtype_bytes_covers_the_kv_dtypes_vllm_accepts(): def test_summary_names_the_shape_it_priced(): s = kv_budget_for(_spark(), max_model_len=32768).summary() assert "flat" in s and "21 sliding@512" in s and "MiB" in s + + +def _qwen3_8(**kw) -> _Cfg: + """Qwen3.8-27B: 64 layers, 3 linear (gated DeltaNet) to 1 full, 4 KV + heads x 256. The linear layers hold a recurrent state and no KV.""" + layer_types = ["linear_attention"] * 3 + ["full_attention"] + base = dict( + num_hidden_layers=64, num_attention_heads=24, num_key_value_heads=4, + head_dim=256, hidden_size=5120, layer_types=layer_types * 16, + linear_num_key_heads=16, linear_num_value_heads=48, + linear_key_head_dim=128, linear_value_head_dim=128, + linear_conv_kernel_dim=4, mamba_ssm_dtype="float32", + ) + base.update(kw) + return _Cfg(**base) + + +class _Outer: + """A multimodal config: the decoder geometry lives one level down.""" + + def __init__(self, text): + self.text_config = text + self.model_type = "qwen3_5" + + def get_text_config(self): + return self.text_config + + +def test_linear_attention_layers_are_charged_no_kv_at_all(): + """Not "a shorter cache" -- none. A flat 64-layer price is 4x the truth.""" + b = kv_budget_for(_qwen3_8(), max_model_len=8192) + assert (b.linear_layers, b.full_layers, b.sliding_layers) == (48, 16, 0) + # 16 full layers x 4 KV heads x 256 x 2 (K and V) x 2 bytes = 64 KiB. + assert b.bytes_per_token_flat == 16 * 4 * 256 * 2 * 2 == 64 * 1024 + assert b.token_layers_flat == 16 * 8192 + + +def test_the_recurrent_state_is_sized_from_vllms_own_shapes(): + """conv (kernel-1, 2*k_heads*k_dim + v_heads*v_dim); temporal + (v_heads, v_dim, k_dim) in mamba_ssm_dtype.""" + b = kv_budget_for(_qwen3_8(), max_model_len=8192) + conv_dim = 128 * 16 * 2 + 128 * 48 + per_layer = (4 - 1) * conv_dim * 2 + 48 * 128 * 128 * 4 + assert b.state_bytes_per_seq == 48 * per_layer + assert b.state_bytes_per_seq / 2**20 == pytest.approx(147.0, abs=0.5) + assert b.state_bytes_total == b.state_bytes_per_seq # max_num_seqs=1 + + +def test_the_state_is_inside_the_budget_and_scales_with_sequences(): + """It comes out of the same pool as the KV cache, so a budget that + reports only the cache is short before the first token.""" + one = kv_budget_for(_qwen3_8(), max_model_len=8192, max_num_seqs=1) + two = kv_budget_for(_qwen3_8(), max_model_len=8192, max_num_seqs=2) + assert two.state_bytes_total == 2 * one.state_bytes_total + assert one.bytes_total > one.state_bytes_total + kv_only = int(one.token_layers_flat * one.bytes_per_token_flat / 16 * 1.25) + assert one.bytes_total == kv_only + one.state_bytes_total + + +def test_the_state_dominates_a_short_context(): + """147 MiB of state against 64 KiB/token: the cache only outgrows it past + ~2350 tokens, so over an edge deployment's range the constant is the + bigger half. The crossover is on raw bytes -- ``headroom`` is a safety + factor on the cache, not a cost the state has to beat.""" + assert kv_budget_for( + _qwen3_8(), max_model_len=2048 + ).state_crossover_tokens == pytest.approx(2352, abs=8) + short = kv_budget_for(_qwen3_8(), max_model_len=1024) + assert short.state_bytes_total > short.bytes_total - short.state_bytes_total + long = kv_budget_for(_qwen3_8(), max_model_len=32768) + assert long.state_bytes_total < long.bytes_total - long.state_bytes_total + + +def test_a_model_without_linear_layers_reports_no_state(): + b = kv_budget_for(_spark(), max_model_len=2048) + assert b.linear_layers == 0 + assert b.state_bytes_per_seq == 0 + assert b.state_crossover_tokens == float("inf") + + +def test_a_multimodal_config_is_resolved_to_its_text_config(): + """Qwen3_5Config has no num_hidden_layers; this used to raise.""" + nested = kv_budget_for(_Outer(_qwen3_8()), max_model_len=8192) + direct = kv_budget_for(_qwen3_8(), max_model_len=8192) + assert nested == direct + + +def test_linear_layers_without_their_geometry_raise_rather_than_guess(): + cfg = _qwen3_8() + del cfg.linear_num_value_heads + with pytest.raises(ValueError, match="linear_num_value_heads"): + kv_budget_for(cfg, max_model_len=8192) + + +def test_summary_names_the_state_it_priced(): + s = kv_budget_for(_qwen3_8(), max_model_len=8192).summary() + assert "recurrent state for 48 linear layers" in s diff --git a/tests/edge/test_placement.py b/tests/edge/test_placement.py new file mode 100644 index 000000000..e82b6d2d7 --- /dev/null +++ b/tests/edge/test_placement.py @@ -0,0 +1,214 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Placing a checkpoint that does not fit in the VRAM.""" + +import json +import struct + +import pytest + +from vllm_omni.edge.placement import ( + GiB, + ParamGroup, + checkpoint_param_groups, + group_key, + is_cold, + plan_placement, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +LINK = 13.9e9 # measured pinned H2D on the reference laptop, PCIe 4.0 x8 +VRAM = 472e9 # measured VRAM bandwidth on the same machine + + +def _plan(groups, avail_gib, **kw): + return plan_placement( + groups, device_bytes_available=int(avail_gib * GiB), + link_bytes_per_s=LINK, vram_bytes_per_s=VRAM, **kw) + + +# ------------------------------------------------------------ classification +def test_embed_tokens_is_cold_and_lm_head_is_not(): + """Same shape, same dtype, opposite decision: the embedding is a gather + and the head is a GEMM over the whole table.""" + assert is_cold("model.language_model.embed_tokens.weight") + assert not is_cold("lm_head.weight") + + +def test_cold_matching_is_by_segment_not_substring(): + assert is_cold("model.visual.blocks.0.attn.qkv.weight") + assert not is_cold("model.visualiser.weight") + + +def test_a_decoder_layer_is_never_folded_into_one_group(): + """A 3-segment prefix makes all 64 layers a single 24 GB group, and a + planner that moves groups whole then has no choice between none and all.""" + assert group_key("model.language_model.layers.17.mlp.down_proj.weight") == ( + "model.language_model.layers.17") + assert group_key("model.visual.blocks.3.attn.qkv.weight").startswith("model.visual") + + +# ------------------------------------------------------------------ planning +def test_a_checkpoint_that_fits_is_left_alone(): + groups = [ParamGroup("model.layers.0", 8 * GiB, False)] + p = _plan(groups, 20) + assert p.strategy == "resident" + assert p.host_bytes == 0 and p.est_link_s_per_token == 0.0 + assert p.engine_kwargs == {} + + +def test_the_deficit_comes_out_of_the_cold_set_first(): + """4 GiB over, and there is 6 GiB of cold weight to spend: nothing that a + decode step reads should end up on the host.""" + groups = [ + ParamGroup("model.layers.0", 18 * GiB, False), + ParamGroup("model.embed_tokens", 6 * GiB, True), + ] + p = _plan(groups, 20) + assert p.strategy == "cold_offload" + assert p.host_bytes == 6 * GiB + assert p.host_read_per_token == 0 + assert p.est_link_s_per_token == 0.0 + + +def test_an_offloaded_embedding_still_pays_one_row_per_token(): + groups = [ + ParamGroup("model.layers.0", 18 * GiB, False), + ParamGroup("model.embed_tokens", 6 * GiB, True), + ] + p = _plan(groups, 20, embedding_row_bytes=5120 * 2) + assert p.host_read_per_token == 5120 * 2 + # 10 KiB over a 13.9 GB/s link is well under a microsecond. + assert p.est_link_s_per_token < 1e-6 + + +def test_hot_weight_is_only_spent_once_the_cold_set_runs_out(): + groups = [ + ParamGroup("model.embed_tokens", 2 * GiB, True), + *[ParamGroup(f"model.layers.{i}", 1 * GiB, False) for i in range(20)], + ] + p = _plan(groups, 16) + assert p.strategy == "hot_offload" + # 6 GiB over; 2 GiB is free, so exactly 4 layers should cross. + assert p.host_read_per_token == 4 * GiB + assert p.host_bytes == 6 * GiB + + +def test_the_estimate_prices_host_bytes_at_the_link_and_the_rest_at_vram(): + groups = [ + ParamGroup("model.embed_tokens", 2 * GiB, True), + *[ParamGroup(f"model.layers.{i}", 1 * GiB, False) for i in range(20)], + ] + p = _plan(groups, 16) + assert p.est_link_s_per_token == pytest.approx(4 * GiB / LINK) + assert p.est_vram_s_per_token == pytest.approx(16 * GiB / VRAM) + # The link is 34x slower, so 4 GiB on the host outweighs 16 GiB resident. + assert p.est_link_s_per_token > 8 * p.est_vram_s_per_token + + +def test_a_checkpoint_too_big_for_host_and_device_says_so(): + groups = [ParamGroup("model.layers.0", 40 * GiB, False)] + p = _plan(groups, 8) + # Everything hot can still move to the host, so this one is feasible but + # miserable; infeasible needs the groups themselves to run out. + assert p.strategy == "hot_offload" + assert p.deficit_bytes == 32 * GiB + + +def test_engine_kwargs_name_only_segments_vllm_can_match(): + groups = [ + ParamGroup("model.embed_tokens", 6 * GiB, True), + ParamGroup("model.layers.0", 18 * GiB, False), + ] + kw = _plan(groups, 20).engine_kwargs + assert kw["cpu_offload_params"] == ["embed_tokens"] + assert kw["cpu_offload_gb"] >= 6.0 + assert "offload_group_size" not in kw # nothing hot was offloaded + + +# --------------------------------------------------------------- checkpoints +def _write_shard(path, tensors): + """A minimal safetensors file: 8-byte header length, JSON header, data.""" + header, off = {}, 0 + for name, nbytes in tensors: + header[name] = {"dtype": "BF16", "shape": [nbytes // 2], + "data_offsets": [off, off + nbytes]} + off += nbytes + blob = json.dumps(header).encode() + path.write_bytes(struct.pack("= {"model.language_model.layers.0", + "model.language_model.layers.1"} + assert groups["model.language_model.layers.0"].bytes == 2048 + + +def test_a_group_is_cold_only_if_every_tensor_in_it_is(tmp_path): + _write_shard(tmp_path / "a.safetensors", [ + ("model.visual.blocks.0.attn.qkv.weight", 1024), + ("model.visual.blocks.0.lm_head.weight", 1024), # contrived, but hot + ]) + (group,) = [g for g in checkpoint_param_groups(tmp_path) + if g.name.startswith("model.visual")] + assert not group.cold + + +def test_a_directory_with_no_shards_raises(tmp_path): + with pytest.raises(FileNotFoundError, match="safetensors"): + checkpoint_param_groups(tmp_path) + + +# ------------------------------------------- the cost model, against hardware +# Qwen3.8-27B-FP8, 4096-token context, batch 1, greedy, eager, idle machine. +# Four points from benchmarks/edge_harness; see docs/edge/laptop-heterogeneous.md. +MEASURED_S_PER_TOKEN = {10: 0.952, 12: 1.124, 14: 1.316, 16: 1.515} +QWEN38_FP8_GB = 30.87 + + +def test_the_estimate_reproduces_the_measured_offload_curve(): + """The whole design rests on `resident/472 + offloaded/10.6`. It is worth + one test that the arithmetic matches the hardware it claims to describe. + + It runs slightly *pessimistic* at every point, which is the right + direction for a number used to decide whether a build will fit. + """ + # 0.25 GB chunks, so the planner offloads close to exactly the target + # rather than having to take one oversized group whole -- which is also + # what the real checkpoint looks like at 0.38 GB per decoder layer. + CHUNK = 0.25 + n = round(QWEN38_FP8_GB / CHUNK) + for offloaded_gb, measured in MEASURED_S_PER_TOKEN.items(): + groups = [ParamGroup(f"layers.{i}", int(CHUNK * 1e9), False) for i in range(n)] + p = plan_placement( + groups, + device_bytes_available=int((QWEN38_FP8_GB - offloaded_gb) * 1e9), + link_bytes_per_s=10.6e9, vram_bytes_per_s=472e9, + ) + modelled = p.est_link_s_per_token + p.est_vram_s_per_token + assert modelled == pytest.approx(measured, rel=0.05), ( + f"{offloaded_gb} GB: modelled {modelled:.3f}s vs measured {measured:.3f}s") + assert modelled >= measured, "the estimate must not flatter the build" + + +def test_the_link_term_dominates_at_every_measured_point(): + """At the tightest fit the 10 GB on the host is ~94% of the step. If this + ever stops being true the advice in the design document changes.""" + groups = [ + ParamGroup("resident", int(20.87e9), False), + ParamGroup("onhost", int(10e9), False), + ] + p = plan_placement(groups, device_bytes_available=int(20.87e9), + link_bytes_per_s=10.6e9, vram_bytes_per_s=472e9) + share = p.est_link_s_per_token / (p.est_link_s_per_token + p.est_vram_s_per_token) + assert share > 0.9 diff --git a/tests/edge/test_qnn_export.py b/tests/edge/test_qnn_export.py index 2f7d641b2..add71ce58 100644 --- a/tests/edge/test_qnn_export.py +++ b/tests/edge/test_qnn_export.py @@ -105,6 +105,26 @@ def test_masked_cache_slots_do_not_change_the_result(): assert torch.allclose(first, module(*args)[0], atol=1e-5) +def test_llama_style_config_skips_qk_norm(): + """MiniCPM-o's TTS head is a Llama decoder: no per-head q/k RMSNorm, so those params must not be applied.""" + cfg = qx.DecodeStepConfig(**{**{f.name: getattr(CFG, f.name) for f in CFG.__dataclass_fields__.values()}, "qk_norm": False}) + module = qx.DecodeStep(cfg).eval() + for p in module.parameters(): + torch.nn.init.normal_(p, std=0.15) + args = qx.example_inputs(cfg, L) + before = module(*args)[0].clone() + with torch.no_grad(): # q_norm/k_norm must be dead weights when qk_norm is off + module.layers[0].q_norm.mul_(3.0) + module.layers[0].k_norm.mul_(0.1) + assert torch.equal(before, module(*args)[0]) + + +def test_minicpmo_config_maps_the_tts_head(): + cfg = qx.config_from_minicpmo({"tts_config": {"num_hidden_layers": 20, "hidden_size": 768, "num_attention_heads": 12, + "num_key_value_heads": 12, "intermediate_size": 3072, "num_audio_tokens": 6562}}) + assert cfg.qk_norm is False and cfg.head_dim == 64 and cfg.vocab_size == 6562 and cfg.group_size == 1 + + def test_placement_table_prefers_npu_for_the_step_and_gpu_for_the_vocoder(): p = qx.PLACEMENT_MS_PER_FRAME assert p["talker_step"]["npu"] < p["talker_step"]["gpu"] diff --git a/tests/model_executor/models/spark2_5/test_spark2_5.py b/tests/model_executor/models/spark2_5/test_spark2_5.py index 468f53a0c..eb938e062 100644 --- a/tests/model_executor/models/spark2_5/test_spark2_5.py +++ b/tests/model_executor/models/spark2_5/test_spark2_5.py @@ -243,3 +243,84 @@ def test_rope_table_is_capped_to_the_servable_context(monkeypatch): if os.environ.get("VLLM_OMNI_SPARK_ROPE_FULL", "0") == "0": max_position = min(max_position, 2048) assert max_position == 1048576, "the escape hatch must restore the full table" + + +# ------------------------------------------------------- published sizes +# Spark-X2.5 ships at two sizes and the vLLM implementation is meant to be +# size-generic. It is, but only one of them had ever been run: the constants +# below are the published configs, so a change that quietly assumes 1.7B's +# 8 query heads or 28 layers fails here instead of on a checkpoint. +PUBLISHED = { + "XHToken/Spark-X2.5-1.7B": dict( + hidden_size=2048, intermediate_size=6656, num_hidden_layers=28, + num_attention_heads=8, num_key_value_heads=2, + ), + "XHToken/Spark-X2.5-4B": dict( + hidden_size=2560, intermediate_size=10240, num_hidden_layers=36, + num_attention_heads=16, num_key_value_heads=4, + ), +} +_COMMON = dict( + head_dim=256, vocab_size=131072, sliding_window=512, hidden_act="gelu", + headwise_attn_output_gate=True, tie_word_embeddings=True, + max_position_embeddings=1048576, + rope_parameters={ + "full_attention": {"rope_theta": 5000000, "partial_rotary_factor": 0.25}, + "sliding_attention": {"rope_theta": 10000, "partial_rotary_factor": 1.0}, + }, +) + + +def _published(name: str) -> Spark2_5Config: + kw = dict(_COMMON, **PUBLISHED[name]) + n = kw["num_hidden_layers"] + kw["layer_types"] = (["sliding_attention"] * 3 + ["full_attention"]) * (n // 4) + return Spark2_5Config(**kw) + + +@pytest.mark.core_model +@pytest.mark.cpu +@pytest.mark.parametrize("name", sorted(PUBLISHED)) +def test_published_configs_validate_and_keep_the_3_to_1_hybrid(name): + config = _published(name) + assert len(config.layer_types) == config.num_hidden_layers + assert config.layer_types.count("full_attention") == config.num_hidden_layers // 4 + # Both sizes keep head_dim 256 and widen by head count, not by head width. + assert config.head_dim == 256 + assert config.num_attention_heads % config.num_key_value_heads == 0 + + +@pytest.mark.core_model +@pytest.mark.cpu +@pytest.mark.parametrize("name", sorted(PUBLISHED)) +def test_split_fused_qkv_holds_at_both_published_head_counts(name): + """4B is 16 q heads to 4 KV heads where 1.7B is 8 to 2. The split is by + dimension, not by a ratio, so getting one right does not imply the other.""" + config = _published(name) + q_dim = config.num_attention_heads * config.head_dim + kv_dim = config.num_key_value_heads * config.head_dim + fused = torch.cat([ + torch.full((q_dim, 8), 1.0), + torch.full((kv_dim, 8), 2.0), + torch.full((kv_dim, 8), 3.0), + ], dim=0) + out = dict(_split_fused_qkv( + [("model.layers.0.self_attn.q_k_v_proj.weight", fused)], config)) + assert out["model.layers.0.self_attn.q_proj.weight"].shape[0] == q_dim + assert out["model.layers.0.self_attn.k_proj.weight"].shape[0] == kv_dim + assert torch.all(out["model.layers.0.self_attn.q_proj.weight"] == 1.0) + assert torch.all(out["model.layers.0.self_attn.v_proj.weight"] == 3.0) + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_4b_is_priced_as_a_hybrid_not_as_36_full_layers(): + """The KV budget reads the same config the model does; 4B is 27 sliding + layers to 9 full ones, which is a 3.4x discount a flat price would miss.""" + from vllm_omni.edge.kv_budget import kv_budget_for + + b = kv_budget_for(_published("XHToken/Spark-X2.5-4B"), max_model_len=32768) + assert (b.full_layers, b.sliding_layers, b.linear_layers) == (9, 27, 0) + # 36 layers x 4 KV heads x 256 x 2 (K and V) x 2 bytes = 144 KiB/token. + assert b.bytes_per_token_flat == 36 * 4 * 256 * 2 * 2 == 144 * 1024 + assert b.saving_vs_flat > 3.3 diff --git a/tests/worker/test_forward_kwargs_filtering.py b/tests/worker/test_forward_kwargs_filtering.py new file mode 100644 index 000000000..e4db8483f --- /dev/null +++ b/tests/worker/test_forward_kwargs_filtering.py @@ -0,0 +1,69 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Omni-only forward kwargs must not be forced on a model that cannot take them. + +The AR runner hands every stage ``sampling_metadata``, ``logits_index`` and +``sampler``, which omni-native models absorb through ``**kwargs``. A model vLLM +owns does not: ``LlamaForCausalLM.forward`` is +``(input_ids, positions, intermediate_tensors, inputs_embeds)`` and nothing +else, so passing the extras is a ``TypeError`` on the very first forward. +""" + +from __future__ import annotations + +import torch + +from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner + + +class _OmniNativeModel(torch.nn.Module): + def forward(self, input_ids, positions, intermediate_tensors=None, inputs_embeds=None, **kwargs): + return input_ids + + +class _VLLMNativeModel(torch.nn.Module): + """The shape of vLLM's own causal LMs: no ``**kwargs``.""" + + def forward(self, input_ids, positions, intermediate_tensors=None, inputs_embeds=None): + return input_ids + + +class _Probe: + """Just enough of a runner to exercise the helper.""" + + def __init__(self, model): + self.model = model + + _accepted_forward_kwargs = OmniGPUModelRunner._accepted_forward_kwargs + + +def test_a_model_taking_kwargs_is_left_alone(): + assert _Probe(_OmniNativeModel())._accepted_forward_kwargs() is None + + +def test_a_vllm_native_model_reports_exactly_its_parameters(): + accepted = _Probe(_VLLMNativeModel())._accepted_forward_kwargs() + assert accepted == frozenset( + {"input_ids", "positions", "intermediate_tensors", "inputs_embeds"} + ) + for omni_only in ("sampling_metadata", "logits_index", "sampler"): + assert omni_only not in accepted + + +def test_the_answer_is_cached(): + """``inspect.signature`` on every decode step would be absurd.""" + probe = _Probe(_VLLMNativeModel()) + first = probe._accepted_forward_kwargs() + probe.model = _OmniNativeModel() # would change the answer if recomputed + assert probe._accepted_forward_kwargs() is first + + +def test_an_uninspectable_forward_falls_back_to_passing_everything(): + """Degrade to the previous behaviour rather than silently dropping + arguments a model may actually need.""" + + class Weird(torch.nn.Module): + forward = staticmethod(print) # no usable signature for our purposes + + accepted = _Probe(Weird())._accepted_forward_kwargs() + assert accepted is None or "sampling_metadata" not in accepted diff --git a/vllm_omni/_filelock_compat.py b/vllm_omni/_filelock_compat.py new file mode 100755 index 000000000..a94012d5c --- /dev/null +++ b/vllm_omni/_filelock_compat.py @@ -0,0 +1,69 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Non-blocking exclusive file locks, on POSIX and on Windows. + +[edge-infer W1] ``vllm_omni.engine.stage_init_utils`` imports ``fcntl`` at module +scope for the device-initialization lock. ``fcntl`` does not exist on Windows, and +because that module is in the import closure of ``AsyncOmni`` the import fails +before anything can be configured. The locking *logic* is sound and already +torch-free; only the primitive has to change. + +Two behavioural differences are deliberate and worth knowing: + +* ``flock`` is advisory and whole-file; ``msvcrt.locking`` is mandatory and + byte-range. We therefore lock **one byte at a fixed high offset**, far past the + PID text the caller writes, so the lock never overlaps the data. Locking byte 0 + instead would make the caller's own ``ftruncate``/``write`` fail on Windows. +* ``msvcrt.locking`` raises ``OSError`` when the region is already held, where + ``flock`` raises ``BlockingIOError``. This module normalizes to + ``BlockingIOError`` so the caller's existing ``except BlockingIOError`` branch + -- which is what triggers stale-lock cleanup -- keeps working untouched. + +Both flavours release on process death, which is what the stale-lock path relies +on: POSIX ``flock`` by the kernel dropping the fd, Windows by the file handle +closing when the process exits. +""" + +from __future__ import annotations + +import os + +_WINDOWS = os.name == "nt" + +# Far beyond any PID line the caller writes, so the locked byte and the file +# contents never overlap. +_LOCK_OFFSET = 1 << 30 + +if _WINDOWS: + import msvcrt + + def flock_exclusive_nb(fd: int) -> None: + """Take an exclusive, non-blocking lock. Raises BlockingIOError if held.""" + here = os.lseek(fd, 0, os.SEEK_CUR) + os.lseek(fd, _LOCK_OFFSET, os.SEEK_SET) + try: + msvcrt.locking(fd, msvcrt.LK_NBLCK, 1) + except OSError as exc: + raise BlockingIOError(str(exc)) from exc + finally: + os.lseek(fd, here, os.SEEK_SET) + + def funlock(fd: int) -> None: + here = os.lseek(fd, 0, os.SEEK_CUR) + os.lseek(fd, _LOCK_OFFSET, os.SEEK_SET) + try: + msvcrt.locking(fd, msvcrt.LK_UNLCK, 1) + finally: + os.lseek(fd, here, os.SEEK_SET) + +else: + import fcntl + + def flock_exclusive_nb(fd: int) -> None: + fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB) + + def funlock(fd: int) -> None: + fcntl.flock(fd, fcntl.LOCK_UN) + + +__all__ = ["flock_exclusive_nb", "funlock"] diff --git a/vllm_omni/config/pipeline_registry.py b/vllm_omni/config/pipeline_registry.py index d9751ae5f..2b25f582e 100644 --- a/vllm_omni/config/pipeline_registry.py +++ b/vllm_omni/config/pipeline_registry.py @@ -108,6 +108,8 @@ from vllm_omni.model_executor.models.qwen3_omni.pipeline import ( resolve_qwen3_omni_pipeline, ) from vllm_omni.model_executor.models.qwen3_tts.pipeline import QWEN3_TTS_PIPELINE, QWEN3_TTS_TALKER_ONLY_PIPELINE +from vllm_omni.config.vllm_native_pipelines import LLAMA_PIPELINE +from vllm_omni.model_executor.models.spark2_5.pipeline import SPARK2_5_PIPELINE from vllm_omni.model_executor.models.step_audio2.pipeline import ( STEP_AUDIO2_ASR_PIPELINE, STEP_AUDIO2_PIPELINE, @@ -135,6 +137,11 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = { "qwen3_omni_moe_thinker_only": QWEN3_OMNI_THINKER_ONLY_PIPELINE, "qwen3_tts": QWEN3_TTS_PIPELINE, "qwen3_tts_talker_only": QWEN3_TTS_TALKER_ONLY_PIPELINE, + # Single-stage text: an AR decoder with no companion stage. It lives in + # the same registry as the multi-stage pipelines so that one code path + # resolves every model (see the M0 local text mode). + "llama": LLAMA_PIPELINE, + "spark2_5": SPARK2_5_PIPELINE, "step_audio_2": STEP_AUDIO2_PIPELINE, "step_audio_2_asr": STEP_AUDIO2_ASR_PIPELINE, "covo_audio": COVO_AUDIO_PIPELINE, diff --git a/vllm_omni/config/vllm_native_pipelines.py b/vllm_omni/config/vllm_native_pipelines.py new file mode 100644 index 000000000..201585e78 --- /dev/null +++ b/vllm_omni/config/vllm_native_pipelines.py @@ -0,0 +1,73 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Pipelines for text models vLLM already implements and Omni does not own. + +Every other pipeline in the registry sits next to an Omni model definition, in +``vllm_omni/model_executor/models//pipeline.py``. These have no such +directory, because there is no Omni model to put there: the model executor is +vLLM's, and all that is missing is the one thing vLLM has no concept of -- how +many stages the pipeline has and what they produce. + +That gap is worth naming, because it is not obvious from either side. A model +being in vLLM's registry does not make it reachable through ``Omni(model=...)``: +the orchestrator resolves a pipeline by ``model_type`` first, and a +``model_type`` with no entry here fails before any weight is read. Spark-X2.5 +hit exactly this in M0 -- its model definition had been registered for a while +and ``Omni(model=)`` still could not resolve it. + +**One stage, and the model_type is the key.** These are autoregressive causal +language models with a single stage that owns the tokenizer and emits text. +Because the registry is keyed on ``model_type`` rather than on a checkpoint, +registering ``llama`` covers every checkpoint whose config says ``llama`` -- not +only the one this was added for. That is the correct granularity (the pipeline +shape really is a property of the architecture, not of the weights), but it +should be read as what it is: a claim about the *shape* of the pipeline, not a +statement that any particular checkpoint has been run and verified. What has +actually been run is recorded in ``analysis/experiments/``. +""" + +from __future__ import annotations + +from vllm_omni.config.stage_config import ( + PipelineConfig, + StageExecutionType, + StagePipelineConfig, +) + + +def _single_stage_text(model_type: str, model_arch: str, hf_architectures: tuple[str, ...]) -> PipelineConfig: + """One autoregressive stage that owns the tokenizer and emits text.""" + return PipelineConfig( + model_type=model_type, + model_arch=model_arch, + hf_architectures=hf_architectures, + stages=( + StagePipelineConfig( + stage_id=0, + model_stage=model_type, + execution_type=StageExecutionType.LLM_AR, + input_sources=(), + final_output=True, + final_output_type="text", + owns_tokenizer=True, + engine_output_type="text", + model_arch=model_arch, + ), + ), + ) + + +LLAMA_PIPELINE = _single_stage_text( + model_type="llama", + model_arch="LlamaForCausalLM", + hf_architectures=("LlamaForCausalLM",), +) +"""The Llama architecture as a single-stage Omni text pipeline. + +Added for MiniCPM5-2B, whose ``config.json`` declares ``model_type: "llama"`` +and ``architectures: ["LlamaForCausalLM"]`` -- it is a Llama-shaped 2B model +(42 layers, 16 heads over 2 KV heads, head_dim 128, 130560 vocab, rope_theta +5e6), not a new architecture, so the support it needed was this entry plus the +runner no longer forcing omni-only kwargs on a forward that cannot take them +(see ``OmniGPUModelRunner._accepted_forward_kwargs``). +""" diff --git a/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py b/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py index 3a51795c4..10999f2b0 100644 --- a/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py +++ b/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py @@ -1,7 +1,15 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import fcntl +# [edge-infer W1] This module is in the import closure of AsyncOmni, so a +# module-scope `import fcntl` makes the whole engine unimportable on Windows +# -- before anyone has chosen a connector. The transport itself is genuinely +# POSIX (it hardcodes /dev/shm lock paths), so the honest Windows behaviour +# is to stay importable and refuse at construction, not to fake a port. +try: + import fcntl +except ImportError: # Windows + fcntl = None import os from multiprocessing import shared_memory as shm_pkg from typing import Any @@ -25,6 +33,12 @@ class SharedMemoryConnector(OmniConnectorBase): """ def __init__(self, config: dict[str, Any]): + if fcntl is None: + raise RuntimeError( + "SharedMemoryConnector requires POSIX shared memory and fcntl " + "locking (it addresses lock files under /dev/shm); neither exists " + "on Windows. Select a different connector for this stage." + ) self.config = config self.stage_id = config.get("stage_id", -1) self._pending_keys: set[str] = set() diff --git a/vllm_omni/edge/hardware_probe.py b/vllm_omni/edge/hardware_probe.py index de1e2a4d5..539222bbe 100644 --- a/vllm_omni/edge/hardware_probe.py +++ b/vllm_omni/edge/hardware_probe.py @@ -8,6 +8,47 @@ overrides (see :mod:`vllm_omni.edge.adapt`). The probe reads from an injectable filesystem root so tests can feed captured snapshots of other devices, and ``VLLM_OMNI_HW_PROFILE=`` replaces the probe entirely (used for emulating a device class on a development host). + +**Reachability, not presence (2026-09-15).** The first version of this +inventory reported the iGPU and the NPU as simply unreachable under WSL2, which +was true of the *default* setup and false as a statement about the machine. +Both turned out to have a route, and the routes are not equivalent: + +* the **Radeon 890M** is reachable from inside WSL today, through DirectML over + ``/dev/dxg`` (``torch-directml``). It computes correctly -- a 512x512 matmul + agrees with the CPU to 3.6e-07. It is also, measured, **slower than the CPU it + shares memory with**: 1.12 TFLOP/s fp32 against the CPU's 1.54, and 5-7.4 GB/s + against the CPU's 20.7, because the D3D12 translation layer costs more than + the iGPU wins. Usable, and rarely worth using through this route. +* the **XDNA2 NPU** is *not* reachable from inside WSL at all: it is an MCDM + device (``IpuMcdmDriver``) and WSL's GPU-PV forwards WDDM display adapters + only, so no ``/dev/accel`` node exists and DirectML enumerates 2 adapters, not + 3. From a Windows-side process it **works**, at 2.6-5.3x the CPU on GEMMs, + once three things are right at once: the general VitisAI EP rather than the + *Light* one beside it, its directory on the DLL search path, and A16W8 + quantization. Each of those failing looks identical -- a session that runs + correctly with every node on the CPU. The recipe and the measurements are in + :mod:`vllm_omni.edge.npu_ryzenai`. + +So ``usable`` answers "can this process compute on it", and ``route`` says what +would have to be true. Reporting a device as absent because the default path +does not reach it is the mistake this field exists to stop. + +**Accelerator inventory (2026-09-14).** ``accelerator`` is one string, chosen +for the one device vLLM will run on. That was enough while the edge target was +a phone or a single-GPU box, and it is not enough on the machines this is now +aimed at: a Ryzen AI laptop has a discrete GPU, an integrated GPU and an XDNA2 +NPU at once, and which of the three a process can actually open depends on the +kernel it is running under, not on the silicon. So ``accelerators`` lists every +one that was found, each with ``usable`` and, when it is not, the ``reason``. + +The reason field earns its place immediately. Under WSL2 -- which is where this +project's own laptop measurements run -- the host passes through ``/dev/dxg`` +and nothing else: CUDA works, while ``/dev/dri`` and ``/dev/accel`` do not +exist, so the iGPU and the NPU are present in the machine and unreachable from +the process. Reporting them as absent would be false, and reporting them as +available would send a deploy planner after a device it cannot open. They are +reported as present, unusable, and told why. """ from __future__ import annotations @@ -29,6 +70,39 @@ HW_CLASS_X86_CPU = "x86_cpu" HW_CLASS_NPU_PHONE = "npu_phone" HW_CLASSES = (HW_CLASS_JETSON, HW_CLASS_CUDA_DISCRETE, HW_CLASS_ARM64_CPU, HW_CLASS_X86_CPU, HW_CLASS_NPU_PHONE) +HOST_OS_LINUX = "linux" +HOST_OS_WSL2 = "wsl2" + +ACCEL_GPU_DISCRETE = "gpu_discrete" +ACCEL_GPU_INTEGRATED = "gpu_integrated" +ACCEL_NPU = "npu" + +# PCI vendor ids, for naming a device that has no driver bound to it. +_PCI_VENDORS = {"0x10de": "nvidia", "0x1002": "amd", "0x8086": "intel", "0x1022": "amd"} +# PCI device ids of neural accelerators that present as a plain PCI function. +_NPU_PCI_IDS = { + ("0x1022", "0x17f0"): "AMD XDNA2 NPU (Strix Point)", + ("0x1022", "0x1502"): "AMD XDNA NPU (Phoenix)", + ("0x8086", "0x643e"): "Intel NPU (Meteor Lake)", + ("0x8086", "0x7d1d"): "Intel NPU (Lunar Lake)", +} +_WSL_IGPU_ROUTE = ( + "No /dev/dri under WSL2, but reachable through DirectML over /dev/dxg: " + "pip install torch-directml in its own venv (it pins torch==2.4.1). " + "Measured 1.12 TFLOP/s fp32 and 5-7.4 GB/s -- below the CPU on both, so " + "route work here only when the CPU is the thing you are trying to free." +) +_WSL_NPU_ROUTE = ( + "An MCDM device; WSL2's GPU-PV forwards WDDM display adapters only, so it " + "has no device node here and nothing installed inside WSL reaches it. It " + "does work from a Windows-side process -- onnxruntime>=1.25 plus the " + "WindowsWorkload.EP.AMD.VitisAI.Framework appx, loading " + "onnxruntime_vitisai_ep.dll (not the RyzenAI *Light* EP beside it) with " + "its own directory on the DLL search path, against an A16W8-quantized " + "opset-21 graph. Measured 5.3x the CPU on a 1024x2048x2048 GEMM. See " + "vllm_omni/edge/npu_ryzenai.py." +) + # Flags we care about, normalized across x86 (``flags``) and ARM (``Features``). _X86_FLAGS = ("avx2", "avx512f", "avx512_bf16", "avx512_vnni", "amx_bf16", "amx_tile", "amx_int8", "fma", "f16c") _ARM_FLAGS = ("asimd", "asimdhp", "asimddp", "i8mm", "bf16", "sve", "sve2", "sme") @@ -54,6 +128,9 @@ class HardwareProfile: supports_bf16: bool = False supports_fp16: bool = True thermal_zones: int = 0 + host_os: str = HOST_OS_LINUX + accelerators: list[dict[str, Any]] = field(default_factory=list) + """Every accelerator found, usable or not. See the module docstring.""" source: str = "probe" # ---------------------------------------------------------------- helpers @@ -64,6 +141,28 @@ class HardwareProfile: def big_cpus(self) -> list[int]: return list(self.clusters[0]["cpus"]) if self.clusters else list(range(self.n_logical)) + def accelerators_of(self, kind: str, *, usable: bool | None = None) -> list[dict[str, Any]]: + return [ + a + for a in self.accelerators + if a.get("kind") == kind and (usable is None or bool(a.get("usable")) is usable) + ] + + @property + def npus(self) -> list[dict[str, Any]]: + return self.accelerators_of(ACCEL_NPU) + + @property + def integrated_gpus(self) -> list[dict[str, Any]]: + return self.accelerators_of(ACCEL_GPU_INTEGRATED) + + @property + def unusable_accelerators(self) -> list[dict[str, Any]]: + """Present in the machine, not openable from this process. The reason + each one carries is the actionable part: usually a different kernel or + a helper process on the host OS.""" + return [a for a in self.accelerators if not a.get("usable")] + def to_dict(self) -> dict[str, Any]: return asdict(self) @@ -174,6 +273,154 @@ def _is_tegra(root: Path) -> bool: ) +def detect_host_os(root: Path) -> str: + """``wsl2`` when this kernel is Microsoft's, else ``linux``. + + It is read from ``/proc/version`` rather than from an env var so a captured + snapshot of another machine classifies the same way the machine would. + """ + return HOST_OS_WSL2 if "microsoft" in (_read(root / "proc" / "version") or "").lower() else HOST_OS_LINUX + + +def _accel(kind: str, vendor: str, name: str, usable: bool, source: str, reason: str = "", **extra: Any) -> dict[str, Any]: + entry = {"kind": kind, "vendor": vendor, "name": name, "usable": usable, "source": source} + if reason: + entry["reason"] = reason + entry.update(extra) + return entry + + +def _pci_ids(root: Path, device_dir: Path) -> tuple[str, str]: + vendor = (_read(device_dir / "vendor") or "").strip().lower() + device = (_read(device_dir / "device") or "").strip().lower() + return vendor, device + + +def linux_accelerators(root: Path) -> list[dict[str, Any]]: + """Accelerators this kernel actually exposes a device node for. + + Three sources, in the order of how much they prove: + + * ``/dev/dri/renderD*`` -- a GPU a process can submit to. Discrete vs + integrated is **inferred from the PCI vendor** (NVIDIA discrete, AMD and + Intel integrated), which is right on the laptop and edge-box parts this + targets and wrong for a discrete Radeon or Arc. Entries say so. + * ``/dev/accel/accel*`` -- the Linux accel subsystem, which is how both + ``amdxdna`` (XDNA/XDNA2) and ``intel_vpu`` present an NPU. + * a PCI sweep, which finds an NPU whose driver is not loaded. Those are + reported unusable with the driver named, because "no /dev/accel node" and + "no such hardware" are very different problems to be handed. + """ + found: list[dict[str, Any]] = [] + seen_npu_slots: set[str] = set() + + for node in sorted((root / "dev" / "dri").glob("renderD*")) if (root / "dev" / "dri").exists() else []: + dev_link = root / "sys" / "class" / "drm" / node.name / "device" + vendor_id, _ = _pci_ids(root, dev_link) + vendor = _PCI_VENDORS.get(vendor_id, "unknown") + kind = ACCEL_GPU_DISCRETE if vendor == "nvidia" else ACCEL_GPU_INTEGRATED + found.append( + _accel(kind, vendor, f"{vendor} GPU ({node.name})", True, f"/dev/dri/{node.name}", inferred_kind=vendor != "nvidia") + ) + + for node in sorted((root / "dev" / "accel").glob("accel*")) if (root / "dev" / "accel").exists() else []: + driver = "" + link = root / "sys" / "class" / "accel" / node.name / "device" / "driver" + try: + driver = link.resolve().name + except OSError: + pass + slot_dir = root / "sys" / "class" / "accel" / node.name / "device" + vendor_id, device_id = _pci_ids(root, slot_dir) + name = _NPU_PCI_IDS.get((vendor_id, device_id), f"NPU ({driver or node.name})") + seen_npu_slots.add(device_id) + found.append( + _accel(ACCEL_NPU, _PCI_VENDORS.get(vendor_id, "unknown"), name, True, f"/dev/accel/{node.name}", driver=driver) + ) + + pci_root = root / "sys" / "bus" / "pci" / "devices" + for slot in sorted(pci_root.glob("*")) if pci_root.exists() else []: + vendor_id, device_id = _pci_ids(root, slot) + name = _NPU_PCI_IDS.get((vendor_id, device_id)) + if not name or device_id in seen_npu_slots: + continue + bound = (slot / "driver").exists() + found.append( + _accel( + ACCEL_NPU, + _PCI_VENDORS.get(vendor_id, "unknown"), + name, + False, + f"pci:{slot.name}", + reason=( + "driver bound but no /dev/accel node" if bound else + "no driver bound: load amdxdna (Linux 6.14+) or intel_vpu" + ), + ) + ) + return found + + +def wsl_host_accelerators(timeout_s: float = 5.0) -> list[dict[str, Any]]: + """What the *Windows* host has, asked over WSL interop. + + Costs one ``powershell.exe`` round trip (~0.7 s) and is only worth making + on WSL, where the Linux-side probe is guaranteed to find nothing but the + paravirtualised CUDA device. Everything returned is marked unusable: the + point is to record that the silicon exists and name what would have to + change to reach it, not to suggest it can be opened from here. + """ + import subprocess + + try: + out = subprocess.run( + [ + "powershell.exe", "-NoProfile", "-Command", + "Get-PnpDevice -Class ComputeAccelerator,Display -Status OK | " + "Select-Object -ExpandProperty FriendlyName", + ], + capture_output=True, text=True, timeout=timeout_s, check=False, + ).stdout + except (OSError, subprocess.SubprocessError): + return [] + + found: list[dict[str, Any]] = [] + for line in (l.strip() for l in out.splitlines()): + if not line: + continue + low = line.lower() + if "nvidia" in low: + continue # reachable through /dev/dxg; the torch probe owns it + if "npu" in low or "neural" in low or "ai boost" in low: + vendor = "amd" if ("amd" in low or "npu compute" in low) else "intel" + found.append(_accel( + ACCEL_NPU, vendor, line, False, "wsl-interop", + reason=_WSL_NPU_ROUTE, route="ryzenai-ep-windows", + )) + elif "radeon" in low or "graphics" in low or "arc" in low: + vendor = "amd" if "radeon" in low else "intel" + found.append(_accel( + ACCEL_GPU_INTEGRATED, vendor, line, directml_available(), "wsl-interop", + reason=_WSL_IGPU_ROUTE, route="directml", + )) + return found + + +def directml_available() -> bool: + """Whether a DirectML torch device can actually be opened from here. + + ``torch-directml`` pins ``torch==2.4.1``, so it never shares a virtualenv + with the engine; this only reports True in a venv built for it, which is + the honest answer to "can *this* process compute on the iGPU". + """ + try: + import torch_directml + + return torch_directml.device_count() > 0 + except Exception: + return False + + def _torch_gpu_info() -> dict[str, Any]: info: dict[str, Any] = {"available": False} try: @@ -196,8 +443,13 @@ def _torch_gpu_info() -> dict[str, Any]: # -------------------------------------------------------------------- probe -def probe(root: str | Path = "/", use_torch: bool = True) -> HardwareProfile: - """Probe the current machine (or a snapshot rooted at ``root``).""" +def probe(root: str | Path = "/", use_torch: bool = True, probe_host: bool = True) -> HardwareProfile: + """Probe the current machine (or a snapshot rooted at ``root``). + + ``probe_host`` allows the one WSL-interop round trip that finds the host's + NPU and iGPU. It only fires on a live WSL probe; a snapshot root never + shells out, so tests stay hermetic. + """ root = Path(root) model, flags, n_logical = parse_cpuinfo(_read(root / "proc" / "cpuinfo") or "") ram_total, ram_avail = parse_meminfo(_read(root / "proc" / "meminfo") or "") @@ -213,6 +465,7 @@ def probe(root: str | Path = "/", use_torch: bool = True) -> HardwareProfile: ram_total_bytes=ram_total, ram_available_bytes=ram_avail, thermal_zones=_thermal_zones(root), + host_os=detect_host_os(root), ) gpu = _torch_gpu_info() if (use_torch and root == Path("/")) else {"available": False} tegra = _is_tegra(root) @@ -232,6 +485,21 @@ def probe(root: str | Path = "/", use_torch: bool = True) -> HardwareProfile: profile.supports_fp16 = ( True if profile.accelerator.startswith("cuda") else profile.has("fp16_arith") or arch == "x86_64" ) + + inventory = linux_accelerators(root) + if gpu.get("available"): + # torch owns the CUDA device; drop the DRM node that is the same card. + inventory = [a for a in inventory if a["vendor"] != "nvidia"] + inventory.insert( + 0, + _accel( + ACCEL_GPU_DISCRETE, "nvidia", gpu["name"], True, "torch.cuda", + mem_bytes=gpu["mem_bytes"], + ), + ) + if probe_host and root == Path("/") and profile.host_os == HOST_OS_WSL2: + inventory += wsl_host_accelerators() + profile.accelerators = inventory return profile @@ -264,14 +532,14 @@ def hardware_class(profile: HardwareProfile) -> str: return HW_CLASS_X86_CPU -def load_profile(path: str | Path | None = None, use_torch: bool = True) -> HardwareProfile: +def load_profile(path: str | Path | None = None, use_torch: bool = True, probe_host: bool = True) -> HardwareProfile: """Env override (``VLLM_OMNI_HW_PROFILE``) or explicit path wins over probing.""" override = path or os.environ.get(ENV_OVERRIDE) if override: p = HardwareProfile.from_json(Path(override).read_text(encoding="utf-8")) p.source = f"override:{override}" return p - return probe(use_torch=use_torch) + return probe(use_torch=use_torch, probe_host=probe_host) def describe(profile: HardwareProfile) -> str: @@ -282,11 +550,28 @@ def describe(profile: HardwareProfile) -> str: f"logical={profile.n_logical} physical={profile.n_physical} " f"big_cpus={len(big)} flags={','.join(profile.cpu_flags)} ram={profile.ram_total_bytes / 2**30:.1f}GiB " f"accel={profile.accelerator} gpu='{profile.gpu_name}' gpu_mem={profile.gpu_mem_bytes / 2**30:.1f}GiB " - f"bf16={profile.supports_bf16} source={profile.source}" + f"bf16={profile.supports_bf16} host_os={profile.host_os} source={profile.source}" ) +def describe_accelerators(profile: HardwareProfile) -> str: + """One line per accelerator, saying which can be opened from here.""" + if not profile.accelerators: + return " (none found)" + lines = [] + for a in profile.accelerators: + mark = "usable" if a.get("usable") else "not here" + route = f" route={a['route']}" if a.get("route") else "" + tail = f"\n {a['reason']}" if a.get("reason") else "" + lines.append( + f" [{mark:^8}] {a['kind']:<16} {a['vendor']:<8} {a['name']} " + f"(via {a['source']}){route}{tail}" + ) + return "\n".join(lines) + + if __name__ == "__main__": # pragma: no cover - manual use prof = load_profile() print(describe(prof)) + print(describe_accelerators(prof)) print(prof.to_json()) diff --git a/vllm_omni/edge/kv_budget.py b/vllm_omni/edge/kv_budget.py index 89d29b440..163b799d7 100644 --- a/vllm_omni/edge/kv_budget.py +++ b/vllm_omni/edge/kv_budget.py @@ -46,6 +46,38 @@ and the hybrid number is reported alongside it as ``bytes_total_hybrid``, for a backend that accounts per layer type. A budget that is too small is not a graceful degradation: vLLM preempts by recompute, which on a single-stream edge deployment is a stall. + +**Linear-attention layers (2026-09-14).** Spark's hybrid is sliding-vs-full, +which is still a KV cache in every layer, only a shorter one. Qwen3.5/3.6/3.8 +are a different hybrid: 3 of every 4 layers are *gated DeltaNet*, which keeps a +fixed-size recurrent state and **no KV cache at all**, at any context length. +Charging those layers per token is not conservative, it is wrong -- the cache +they need is zero. Qwen3.8-27B is 48 linear layers to 16 full ones, so a flat +64-layer price over-counts its KV by 4x: + + flat over 64 layers 256 KiB/token + 16 full layers only 64 KiB/token + +What the linear layers cost instead is a constant per *sequence*, and it is not +small -- vLLM allocates it from the same pool as the KV cache, so a budget that +omits it is short by that amount before the first token: + + conv state (kernel-1) x (2 x k_heads x k_dim + v_heads x v_dim) + temporal state v_heads x v_dim x k_dim, in ``mamba_ssm_dtype`` + +At Qwen3.8-27B's geometry that is 60 KiB + 3 MiB per layer per sequence, i.e. +**147 MiB per sequence** across the 48 linear layers, against 64 KiB/token of +KV. The two cross at about 2350 tokens: below that the state costs more than +the cache does. The shapes follow vLLM's own +``MambaStateShapeCalculator.gated_delta_net_state_shape`` +(``vllm/model_executor/layers/mamba/mamba_utils.py:262``) rather than a +re-derivation, so they stay right if that changes. + +The nesting matters too. A multimodal checkpoint puts all of this under +``text_config`` -- ``Qwen3_5Config`` has no ``num_hidden_layers`` of its own -- +so this module resolves the text config first. Passed the outer config it used +to raise "cannot read num_hidden_layers", which reads as "this model is +unsupported" when the real answer was one attribute deeper. """ from dataclasses import dataclass @@ -57,16 +89,64 @@ _DTYPE_BYTES = { "float8_e4m3fn": 1, "fp8": 1, "int8": 1, } +# Layer-type spellings that mean "recurrent state, no KV cache". vLLM's own +# configs use "linear_attention" (Qwen3.5/3.6/3.8, Qwen3-Next) and "mamba" +# (Bamba, Zamba); the substring test catches "mamba2" and friends. +_LINEAR_LAYER_MARKERS = ("linear_attention", "mamba", "recurrent", "gated_delta") + def dtype_bytes(dtype: Any) -> int: name = getattr(dtype, "name", None) or str(dtype).replace("torch.", "") return _DTYPE_BYTES.get(name, 2) +def text_config(hf_config: Any) -> Any: + """The sub-config that carries the decoder geometry. + + Transformers puts a multimodal model's decoder under ``text_config`` and + exposes ``get_text_config()`` to reach it. Both are tried, and a config + that is already the text config is returned unchanged, so callers can pass + whichever one they have. + """ + getter = getattr(hf_config, "get_text_config", None) + if callable(getter): + try: + resolved = getter() + except Exception: + resolved = None + if resolved is not None and getattr(resolved, "num_hidden_layers", None): + return resolved + nested = getattr(hf_config, "text_config", None) + if nested is not None and getattr(nested, "num_hidden_layers", None): + return nested + return hf_config + + +def _is_linear(layer_type: Any) -> bool: + name = str(layer_type).lower() + return any(marker in name for marker in _LINEAR_LAYER_MARKERS) + + +@dataclass(frozen=True) +class LinearState: + """Per-sequence recurrent state of the gated-DeltaNet layers. + + Zero on a model that has none, so callers never have to branch. + """ + + layers: int + conv_bytes_per_layer: int + temporal_bytes_per_layer: int + + @property + def bytes_per_seq(self) -> int: + return self.layers * (self.conv_bytes_per_layer + self.temporal_bytes_per_layer) + + @dataclass(frozen=True) class KVBudget: bytes_total: int - """Flat budget: safe to pass to vLLM today.""" + """Flat budget: safe to pass to vLLM today. Includes the recurrent state.""" bytes_total_hybrid: int """What the model's own cache needs, if the backend accounts per layer type.""" bytes_per_token_flat: int @@ -75,6 +155,12 @@ class KVBudget: sliding_layers: int full_layers: int window: int | None + linear_layers: int = 0 + """Layers with a recurrent state and no KV cache. They are excluded from + both token-layer counts, because their cache is zero at every context.""" + state_bytes_per_seq: int = 0 + state_bytes_total: int = 0 + """``state_bytes_per_seq * max_num_seqs``, already inside both totals.""" @property def saving_vs_flat(self) -> float: @@ -85,6 +171,18 @@ class KVBudget: else 1.0 ) + @property + def state_crossover_tokens(self) -> float: + """Context at which the KV cache first costs more than the fixed state. + + ``inf`` when there is no state, 0 when there are no attention layers. + """ + if not self.state_bytes_per_seq: + return float("inf") + if not self.bytes_per_token_flat: + return 0.0 + return self.state_bytes_per_seq / self.bytes_per_token_flat + def summary(self) -> str: extra = "" if self.sliding_layers: @@ -93,7 +191,54 @@ class KVBudget: f"accounts {self.sliding_layers} sliding@{self.window} layers " f"separately ({self.saving_vs_flat:.2f}x)" ) - return f"{self.bytes_total / 2**20:.0f} MiB flat{extra}" + state = "" + if self.linear_layers: + state = ( + f"; includes {self.state_bytes_total / 2**20:.0f} MiB of " + f"recurrent state for {self.linear_layers} linear layers, " + f"which hold no KV at any context" + ) + return f"{self.bytes_total / 2**20:.0f} MiB flat{extra}{state}" + + +def linear_state_for( + hf_config: Any, + *, + model_dtype: Any = "bfloat16", + n_linear_layers: int, +) -> LinearState: + """Per-sequence gated-DeltaNet state, in vLLM's own shapes. + + Mirrors ``MambaStateShapeCalculator.gated_delta_net_state_shape`` at + ``tp_world_size=1``: a conv state of ``(kernel - 1, conv_dim)`` and a + temporal state of ``(v_heads, v_dim, k_dim)``. The temporal state is the + one that matters -- it is quadratic in head dim and, by default, fp32. + """ + cfg = text_config(hf_config) + if n_linear_layers <= 0: + return LinearState(0, 0, 0) + k_heads = int(getattr(cfg, "linear_num_key_heads", 0) or 0) + v_heads = int(getattr(cfg, "linear_num_value_heads", 0) or 0) + k_dim = int(getattr(cfg, "linear_key_head_dim", 0) or 0) + v_dim = int(getattr(cfg, "linear_value_head_dim", 0) or 0) + kernel = int(getattr(cfg, "linear_conv_kernel_dim", 0) or 0) + if not all((k_heads, v_heads, k_dim, v_dim, kernel)): + raise ValueError( + f"{n_linear_layers} linear-attention layers, but the config does " + "not carry the linear_* geometry needed to size their state " + "(linear_num_key_heads, linear_num_value_heads, " + "linear_key_head_dim, linear_value_head_dim, " + "linear_conv_kernel_dim)" + ) + conv_dim = k_dim * k_heads * 2 + v_dim * v_heads + conv_elems = (kernel - 1) * conv_dim + temporal_elems = v_heads * v_dim * k_dim + ssm_dtype = getattr(cfg, "mamba_ssm_dtype", None) or model_dtype + return LinearState( + layers=n_linear_layers, + conv_bytes_per_layer=conv_elems * dtype_bytes(model_dtype), + temporal_bytes_per_layer=temporal_elems * dtype_bytes(ssm_dtype), + ) def kv_budget_for( @@ -108,27 +253,36 @@ def kv_budget_for( """Bytes of KV cache a deployment of this model actually needs. Reads ``layer_types``/``sliding_window`` when present, so hybrid-attention - models are charged for their real cache and not for the longest layer. + models are charged for their real cache and not for the longest layer, and + linear-attention layers are charged a per-sequence state instead of a + per-token cache. Accepts a multimodal config and resolves its text config. """ + cfg = text_config(hf_config) n_layers = int( - getattr(hf_config, "num_hidden_layers", None) - or getattr(hf_config, "n_layer", 0) + getattr(cfg, "num_hidden_layers", None) + or getattr(cfg, "n_layer", 0) ) if n_layers <= 0: raise ValueError("cannot read num_hidden_layers from the config") - n_heads = int(getattr(hf_config, "num_attention_heads", 0) or 0) - kv_heads = int(getattr(hf_config, "num_key_value_heads", 0) or n_heads or 1) + n_heads = int(getattr(cfg, "num_attention_heads", 0) or 0) + kv_heads = int(getattr(cfg, "num_key_value_heads", 0) or n_heads or 1) head_dim = int( - getattr(hf_config, "head_dim", 0) - or (getattr(hf_config, "hidden_size", 0) // max(n_heads, 1)) + getattr(cfg, "head_dim", 0) + or (getattr(cfg, "hidden_size", 0) // max(n_heads, 1)) ) if head_dim <= 0: raise ValueError("cannot read head_dim from the config") - layer_types = list(getattr(hf_config, "layer_types", None) or []) - window = getattr(hf_config, "sliding_window", None) + layer_types = list(getattr(cfg, "layer_types", None) or []) + window = getattr(cfg, "sliding_window", None) + linear = sum(1 for t in layer_types if _is_linear(t)) sliding = sum(1 for t in layer_types if "sliding" in str(t)) if window else 0 - full = n_layers - sliding + full = n_layers - sliding - linear + + state = linear_state_for( + hf_config, model_dtype=kv_dtype, n_linear_layers=linear + ) + state_total = state.bytes_per_seq * max_num_seqs def rounded(tokens: int) -> int: return ((tokens + block_size - 1) // block_size) * block_size @@ -141,16 +295,24 @@ def kv_budget_for( token_layers_hybrid = max_num_seqs * ( full * per_seq_full + sliding * per_seq_sliding ) - token_layers_flat = max_num_seqs * n_layers * per_seq_full + # "Flat" means every *attention* layer at full length. A linear layer is + # not a shorter cache, it is no cache, so it is out of both counts rather + # than conservatively included. + token_layers_flat = max_num_seqs * (n_layers - linear) * per_seq_full bytes_per_token_layer = kv_heads * head_dim * 2 * dtype_bytes(kv_dtype) return KVBudget( - bytes_total=int(token_layers_flat * bytes_per_token_layer * headroom), - bytes_total_hybrid=int(token_layers_hybrid * bytes_per_token_layer * headroom), - bytes_per_token_flat=n_layers * bytes_per_token_layer, + bytes_total=int(token_layers_flat * bytes_per_token_layer * headroom) + + state_total, + bytes_total_hybrid=int(token_layers_hybrid * bytes_per_token_layer * headroom) + + state_total, + bytes_per_token_flat=(n_layers - linear) * bytes_per_token_layer, token_layers_hybrid=token_layers_hybrid, token_layers_flat=token_layers_flat, sliding_layers=sliding, full_layers=full, window=int(window) if sliding else None, + linear_layers=linear, + state_bytes_per_seq=state.bytes_per_seq, + state_bytes_total=state_total, ) diff --git a/vllm_omni/edge/local/__init__.py b/vllm_omni/edge/local/__init__.py new file mode 100644 index 000000000..a13a48192 --- /dev/null +++ b/vllm_omni/edge/local/__init__.py @@ -0,0 +1,52 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""M0: an auditable local text mode on the Omni stage runtime. + +One model, one autoregressive stage, on the device this process can actually +drive -- with the execution plan, the memory budget, the state handles and the +cancellation semantics made explicit instead of implicit. + + from vllm_omni.edge.local import plan_text_session, LocalTextEngine + + plan = plan_text_session("models/Spark-X2.5-4B", max_model_len=4096) + print(plan.summary()) # says yes with a budget, or no with a reason + async with LocalTextEngine(plan) as engine: + session = engine.open_session() + rid, stream = await engine.submit(session, "Hello", max_tokens=128) + async for event in stream: + ... + +The command line is ``python -m vllm_omni.edge.local`` (``devices``, ``plan``, +``run``, ``accept``). +""" + +from vllm_omni.edge.local.capabilities import ( + DeviceCapability, + describe as describe_devices, + enumerate_devices, +) +from vllm_omni.edge.local.engine import LocalTextEngine, PlanNotAdmitted, RequestRecord +from vllm_omni.edge.local.manifest import ArtifactManifest, build_manifest, runtime_versions +from vllm_omni.edge.local.plan import ExecutionPlan, MemoryReservation, Refusal, plan_text_session +from vllm_omni.edge.local.prompts import acceptance_prompts +from vllm_omni.edge.local.session import BoundedEventStream, ChunkEvent, StateHandle + +__all__ = [ + "ArtifactManifest", + "BoundedEventStream", + "ChunkEvent", + "DeviceCapability", + "ExecutionPlan", + "LocalTextEngine", + "MemoryReservation", + "PlanNotAdmitted", + "Refusal", + "RequestRecord", + "StateHandle", + "acceptance_prompts", + "build_manifest", + "describe_devices", + "enumerate_devices", + "plan_text_session", + "runtime_versions", +] diff --git a/vllm_omni/edge/local/__main__.py b/vllm_omni/edge/local/__main__.py new file mode 100644 index 000000000..5a34e2788 --- /dev/null +++ b/vllm_omni/edge/local/__main__.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +from vllm_omni.edge.local.cli import main + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/vllm_omni/edge/local/artifacts/__init__.py b/vllm_omni/edge/local/artifacts/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/vllm_omni/edge/local/artifacts/quantize.py b/vllm_omni/edge/local/artifacts/quantize.py new file mode 100644 index 000000000..61b4bee8f --- /dev/null +++ b/vllm_omni/edge/local/artifacts/quantize.py @@ -0,0 +1,180 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""A16W8 quantization, which is the only door into the XDNA2 NPU. + +The overlays shipped in the VitisAI EP are named for what they take -- +``8x4_psu_model_a16w8_qdq.xclbin`` and friends -- and they mean it. A8W8, which +is what :func:`onnxruntime.quantization.quantize_static` produces by default, +and plain fp32 are both declined **with no error at all**: the session builds, +runs, returns correct numbers, and every node executed on the CPU. So the +precision is not a tuning knob here, it is the entry condition, and +:func:`vllm_omni.edge.npu_ryzenai.quantization_kwargs` is where it is written +down. + +**Calibration is part of the artifact.** AGENTS.md requires a quantized export +to carry the data it was calibrated on, and this project has the receipts for +why: the Qwen3-TTS vocoder calibrated to int8 measured -4.3 dB with correlation +0.00 against its fp16 export, and uncalibrated w4a16 Spark builds scored +-1.6 dB -- uncorrelated with the reference, while looking like a normal model. +So :class:`ImageCalibrationReader` records what it fed the quantizer, and +:func:`quantize_a16w8` returns that record for the manifest. +""" + +from __future__ import annotations + +import hashlib +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Iterator + +import numpy as np + +CALIBRATION_IMAGES = 16 +"""Small on purpose: the quantizer walks the whole graph per sample, and this +is a ranging pass, not training. Recorded either way -- the number matters less +than being able to say which images they were.""" + + +@dataclass +class CalibrationRecord: + """What a quantized artifact was calibrated on.""" + + kind: str + count: int + sources: list[str] = field(default_factory=list) + digest: str = "" + shape: list[int] = field(default_factory=list) + note: str = "" + + def to_dict(self) -> dict[str, Any]: + return { + "kind": self.kind, + "count": self.count, + "sources": self.sources, + "digest": self.digest, + "shape": self.shape, + "note": self.note, + } + + +class ArrayCalibrationReader: + """Feeds pre-built tensors to ``quantize_static`` and remembers them.""" + + def __init__(self, input_name: str, samples: list[np.ndarray], record: CalibrationRecord) -> None: + self.input_name = input_name + self.samples = samples + self.record = record + self._iter: Iterator[np.ndarray] | None = None + + def get_next(self) -> dict[str, np.ndarray] | None: + if self._iter is None: + self._iter = iter(self.samples) + sample = next(self._iter, None) + return None if sample is None else {self.input_name: sample} + + def rewind(self) -> None: + self._iter = None + + +def image_calibration( + model_dir: str | Path, + image_paths: list[str | Path], + *, + image_size: int, + input_name: str = "patches", + dtype: Any = np.float32, +) -> ArrayCalibrationReader: + """Real images, put through the checkpoint's own image processor. + + Using the shipped ``Qwen2VLImageProcessor`` rather than a hand-rolled + patchify is the point: the tower's input layout (temporal pairs, 16x16 + patches, channel-major within a patch) is exactly the kind of thing that is + easy to get subtly wrong, and a calibration set in the wrong layout ranges + the activations of a model nobody is going to run. + """ + from PIL import Image + from transformers import AutoImageProcessor + + processor = AutoImageProcessor.from_pretrained(str(model_dir)) + pixels = image_size * image_size + samples: list[np.ndarray] = [] + used: list[str] = [] + hasher = hashlib.sha256() + + for path in image_paths: + try: + image = Image.open(path).convert("RGB").resize((image_size, image_size)) + except Exception: + continue # an unreadable file is skipped, and absent from `sources` + encoded = processor(images=image, return_tensors="np", + min_pixels=pixels, max_pixels=pixels) + sample = np.asarray(encoded["pixel_values"], dtype=dtype) + samples.append(sample) + used.append(str(path)) + hasher.update(sample.tobytes()) + + if not samples: + raise ValueError(f"no readable images among {len(image_paths)} paths") + + record = CalibrationRecord( + kind="real images, resized to a square and run through the checkpoint's " + "Qwen2VLImageProcessor", + count=len(samples), + sources=used, + digest=hasher.hexdigest()[:32], + shape=list(samples[0].shape), + note=( + "a ranging set, not a quality claim: these are whatever images this " + "repository happens to contain, not a task-representative sample" + ), + ) + return ArrayCalibrationReader(input_name, samples, record) + + +def quantize_a16w8( + source: str | Path, + target: str | Path, + reader: ArrayCalibrationReader, + *, + op_types_to_quantize: tuple[str, ...] | None = None, +) -> dict[str, Any]: + """Quantize an ONNX graph to A16W8 QDQ, the shape XDNA2 partitions. + + The quantization settings come from :mod:`vllm_omni.edge.npu_ryzenai` so + that there is one statement of them in the tree; this function only adds + the calibration and the record. + """ + from onnxruntime.quantization import QuantFormat, quantize_static + + from vllm_omni.edge.npu_ryzenai import quantization_kwargs + + source, target = Path(source), Path(target) + target.parent.mkdir(parents=True, exist_ok=True) + kwargs = quantization_kwargs() + + quantize_static( + str(source), + str(target), + reader, + quant_format=QuantFormat.QDQ, + op_types_to_quantize=list(op_types_to_quantize) if op_types_to_quantize else None, + **kwargs, + ) + reader.rewind() + + return { + "activation_type": str(kwargs["activation_type"]), + "weight_type": str(kwargs["weight_type"]), + "per_channel": kwargs["per_channel"], + "quant_format": "QDQ", + "calibration": reader.record.to_dict(), + "source_sha256_prefix": _digest(source)[:16], + } + + +def _digest(path: Path) -> str: + hasher = hashlib.sha256() + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1 << 20), b""): + hasher.update(chunk) + return hasher.hexdigest() diff --git a/vllm_omni/edge/local/artifacts/vision_tower.py b/vllm_omni/edge/local/artifacts/vision_tower.py new file mode 100644 index 000000000..322c961fd --- /dev/null +++ b/vllm_omni/edge/local/artifacts/vision_tower.py @@ -0,0 +1,304 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Exporting Qwen3.8-27B's vision tower as a stage for the iGPU or the NPU. + +Of everything in the three models this project targets, this tower is the +component that best fits what an external device can actually do. The +constraint measured in M3.0 is that leaving the process costs real time -- +2.8 ms across the WSL/Windows boundary for a megabyte each way -- so a stage +only pays for itself when it is compute-heavy and I/O-light. The tower is 27 +transformer blocks run **once per image**, against 4.8 MB in and 4.0 MB out; +a decode step, run once per token, is the opposite shape and does not qualify. + +It is also free of the usual obstacles. The 461 M parameters are already bf16 +in the FP8 checkpoint -- ``quantization_config.modules_to_not_convert`` excludes +every visual module -- so the DirectML path needs no quantization story at all. +And the tower carries no KV and no cross-request state: pixels in, one +``[num_merged_tokens, out_hidden_size]`` tensor out. + +**Which definition this exports.** ``transformers``' +:class:`Qwen3_5VisionModel`, whose 333 parameters match the checkpoint's +``model.visual.*`` tensors exactly, name for name. No third definition of the +model is written here -- open question 8 in the proposal is precisely about +export code drifting from the implementation. :class:`FixedImageVisionTower` +holds the real module and only fixes its shapes, the same way +:class:`~vllm_omni.edge.decoder_export.WindowedCode2Wav` does for the vocoder. + +**What fixing the shape buys.** ``Qwen3_5VisionModel.forward`` derives bilinear +interpolation indices, position ids and ``cu_seqlens`` from ``grid_thw`` on +every call. At a fixed image size all three are constants, so they are computed +once here and kept as buffers -- and with them go every integer ``Gather`` in +the graph, which is the op most likely to have no kernel on a DirectML or NPU +backend. +""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import torch +from torch import nn + +DEFAULT_IMAGE_SIZE = 448 +"""One image, 448x448: a 28x28 patch grid, 784 patches in, 196 merged tokens +out. Large enough that the 27 blocks dominate, small enough to stay a single +attention segment.""" + + +@dataclass(frozen=True) +class TowerShape: + """The geometry an exported tower is fixed at.""" + + image_size: int + grid_t: int + grid_h: int + grid_w: int + patches: int + patch_dim: int + merged_tokens: int + out_hidden_size: int + + def to_dict(self) -> dict[str, Any]: + return { + "image_size": self.image_size, + "grid_thw": [self.grid_t, self.grid_h, self.grid_w], + "patches": self.patches, + "patch_dim": self.patch_dim, + "merged_tokens": self.merged_tokens, + "out_hidden_size": self.out_hidden_size, + } + + +def _vision_config(model_dir: str | Path) -> Any: + from transformers.models.qwen3_5 import Qwen3_5VisionConfig + + config = json.loads((Path(model_dir) / "config.json").read_text()) + if "vision_config" not in config: + raise ValueError( + f"{model_dir} has no vision_config; this checkpoint carries no tower " + "and there is nothing here to place on an external device" + ) + return Qwen3_5VisionConfig(**config["vision_config"]) + + +def tower_shape(model_dir: str | Path, image_size: int = DEFAULT_IMAGE_SIZE) -> TowerShape: + config = _vision_config(model_dir) + if image_size % config.patch_size: + raise ValueError(f"image_size {image_size} is not a multiple of patch_size {config.patch_size}") + grid = image_size // config.patch_size + merge = config.spatial_merge_size + if (grid * grid) % (merge * merge): + raise ValueError(f"a {grid}x{grid} grid does not divide by spatial_merge_size {merge}") + return TowerShape( + image_size=image_size, + grid_t=1, + grid_h=grid, + grid_w=grid, + patches=grid * grid, + patch_dim=config.in_channels * config.temporal_patch_size * config.patch_size**2, + merged_tokens=(grid * grid) // (merge * merge), + out_hidden_size=config.out_hidden_size, + ) + + +def load_tower( + model_dir: str | Path, + *, + dtype: torch.dtype = torch.float32, + attn_implementation: str = "sdpa", +) -> nn.Module: + """The real tower, with the checkpoint's weights, on the CPU. + + Weights are read from whichever shard holds ``model.visual.*`` -- in + Qwen3.8-27B-FP8 that is ``outside.safetensors``, the same file the hot + ``lm_head`` and the cold ``embed_tokens`` live in. ``strict=True``: a + silently half-loaded tower would export cleanly and produce plausible + garbage, which is the failure this project has the most scars from. + """ + from safetensors.torch import load_file + from transformers.models.qwen3_5 import Qwen3_5VisionModel + + config = _vision_config(model_dir) + config._attn_implementation = attn_implementation + + weights: dict[str, torch.Tensor] = {} + for shard in sorted(Path(model_dir).glob("*.safetensors")): + for name, tensor in load_file(str(shard)).items(): + if ".visual." in name: + weights[name.split("model.visual.", 1)[1]] = tensor.to(dtype) + if not weights: + raise ValueError(f"no model.visual.* tensors in {model_dir}") + + tower = Qwen3_5VisionModel(config).to(dtype) + tower.load_state_dict(weights, strict=True) + return tower.eval() + + +class FixedImageVisionTower(nn.Module): + """The tower at one image size, with its index arithmetic folded away. + + Everything ``forward`` would recompute from ``grid_thw`` -- interpolated + position embeddings, rotary cos/sin, the attention segment boundaries -- is + a function of the grid alone, so at a fixed size it is a constant. Folding + them into buffers leaves a graph of convolution, GEMM, softmax and + normalisation, with no integer gather anywhere in it. + """ + + def __init__(self, tower: nn.Module, shape: TowerShape) -> None: + super().__init__() + self.tower = tower + self.shape = shape + from transformers.models.qwen3_5.modeling_qwen3_5 import ( + get_vision_bilinear_indices_and_weights, + get_vision_cu_seqlens, + get_vision_position_ids, + ) + + grid = torch.tensor([[shape.grid_t, shape.grid_h, shape.grid_w]], dtype=torch.long) + with torch.no_grad(): + indices, bilinear_weights = get_vision_bilinear_indices_and_weights( + grid, + num_grid_per_side=tower.num_grid_per_side, + spatial_merge_size=tower.config.spatial_merge_size, + kwargs={}, + ) + position_ids = get_vision_position_ids(grid, tower.spatial_merge_size, kwargs={}) + cu_seqlens = get_vision_cu_seqlens(grid, kwargs={}) + + # The interpolated position embedding is a weighted gather over a + # constant table at constant indices: constant. + pos_embeds = (tower.pos_embed(indices) * bilinear_weights[:, :, None]).sum(0) + rotary = tower.rotary_pos_emb(position_ids).reshape(shape.patches, -1) + emb = torch.cat((rotary, rotary), dim=-1) + + dtype = next(tower.parameters()).dtype + self.register_buffer("pos_embeds", pos_embeds.to(dtype), persistent=False) + self.register_buffer("rope_cos", emb.cos().to(dtype), persistent=False) + self.register_buffer("rope_sin", emb.sin().to(dtype), persistent=False) + # Kept as a tensor because the block signature takes one; for a single + # image it is [0, patches], so the attention runs as one segment. + self.register_buffer("cu_seqlens", cu_seqlens, persistent=False) + + def forward(self, patches: torch.Tensor) -> torch.Tensor: + """``[patches, patch_dim]`` -> ``[merged_tokens, out_hidden_size]``.""" + hidden = self.tower.patch_embed(patches) + hidden = hidden + self.pos_embeds.to(hidden.dtype) + hidden = hidden.reshape(self.shape.patches, -1) + position_embeddings = (self.rope_cos, self.rope_sin) + for block in self.tower.blocks: + hidden = block( + hidden, + cu_seqlens=self.cu_seqlens, + position_embeddings=position_embeddings, + ) + return self.tower.merger(hidden) + + def example_input(self, dtype: torch.dtype | None = None) -> torch.Tensor: + dtype = dtype or next(self.tower.parameters()).dtype + generator = torch.Generator().manual_seed(0) + return torch.randn(self.shape.patches, self.shape.patch_dim, generator=generator).to(dtype) + + +def build( + model_dir: str | Path, + *, + image_size: int = DEFAULT_IMAGE_SIZE, + dtype: torch.dtype = torch.float32, +) -> FixedImageVisionTower: + shape = tower_shape(model_dir, image_size) + # ``.eval()`` on the wrapper too, not just the tower it holds: the exporter + # reads ``module.training`` on the top-level module and warns -- and any + # future dropout in a block would key off its own flag anyway. + return FixedImageVisionTower(load_tower(model_dir, dtype=dtype), shape).eval() + + +def export_onnx( + module: FixedImageVisionTower, + path: str | Path, + *, + opset: int = 21, + external_data: bool = True, +) -> dict[str, Any]: + """Export to ONNX and make the result loadable, which is a second step. + + ``sanitize_onnx`` is not optional here. torch's dynamo exporter emits + ``Reshape`` with ``allowzero=1``, and onnxruntime's DirectML provider + rejects that at **session initialisation** with a bare ``E_INVALIDARG`` -- + whereupon onnxruntime does not raise, it prints a fallback notice and hands + back a working CPU session. Measured here: the unsanitized 27B vision tower + placed 0 of 1171 nodes on the iGPU and computed correct features on the CPU; + the same graph with 58 ``allowzero`` attributes cleared placed 85 of 113. + The function was written for Qualcomm's converter and turns out to fix a + generic dynamo wart. + + opset 21 by default because A16W8 needs it: 16-bit QDQ is only expressible + without contrib ops from 21 on. + """ + import torch + + from vllm_omni.edge.qnn_export import sanitize_onnx + + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + example = module.example_input() + torch.onnx.export( + module, + (example,), + str(path), + input_names=["patches"], + output_names=["embeds"], + opset_version=opset, + dynamo=True, + optimize=True, + external_data=external_data, + ) + stats = sanitize_onnx(path) + return {"path": str(path), "opset": opset, "sanitize": stats, "shape": module.shape.to_dict()} + + +def _main(argv: list[str] | None = None) -> int: + import argparse + + import torch + + parser = argparse.ArgumentParser( + prog="python -m vllm_omni.edge.local.artifacts.vision_tower", + description="Export Qwen3.8-27B's vision tower as a fixed-size stage graph.", + ) + parser.add_argument("--model", required=True, help="checkpoint directory") + parser.add_argument("--out", required=True, help="destination .onnx") + parser.add_argument("--image-size", type=int, default=DEFAULT_IMAGE_SIZE) + parser.add_argument("--opset", type=int, default=21) + parser.add_argument("--dtype", default="float32", choices=["float32", "float16"]) + parser.add_argument("--check", action="store_true", + help="compare the fixed-shape wrapper against the real forward first") + args = parser.parse_args(argv) + + module = build(args.model, image_size=args.image_size, dtype=getattr(torch, args.dtype)) + print(f"shape: {module.shape.to_dict()}") + + if args.check: + grid = torch.tensor( + [[module.shape.grid_t, module.shape.grid_h, module.shape.grid_w]], dtype=torch.long + ) + example = module.example_input() + with torch.no_grad(): + wrapped = module(example) + reference = module.tower(example, grid).pooler_output + identical = torch.equal(wrapped, reference) + delta = (wrapped - reference).abs().max().item() + print(f"wrapper vs Qwen3_5VisionModel.forward: identical={identical} max_abs={delta:.3e}") + if not identical and delta > 1e-4: + print("the fixed-shape wrapper does not reproduce the model; refusing to export") + return 1 + + result = export_onnx(module, args.out, opset=args.opset) + print(f"exported {result['path']} (opset {result['opset']}), sanitize={result['sanitize']}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(_main()) diff --git a/vllm_omni/edge/local/capabilities.py b/vllm_omni/edge/local/capabilities.py new file mode 100644 index 000000000..23861a44d --- /dev/null +++ b/vllm_omni/edge/local/capabilities.py @@ -0,0 +1,458 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""What this *process* can actually execute, which is not what the machine has. + +:mod:`vllm_omni.edge.hardware_probe` answers "what is in the box": it +enumerates every CPU, iGPU and NPU and marks which ones a process could open. +That is the right question at install time and the wrong one at load time, +because a vLLM wheel is built for exactly one platform. The CUDA wheel in +``.venvs/omni-cuda`` cannot run a CPU stage no matter how many cores the probe +finds, and the CPU wheel in ``.venvs/omni-cpu`` cannot reach the 5090 sitting +next to it. So a device is *runnable* here only when the probe finds it **and** +``vllm.platforms.current_platform`` is the platform that drives it. + +The second thing a device decides is which weight formats it can execute, and +that gate is sharper than "has enough memory". Measured on this laptop +(2026-09-15): loading ``models/Spark-X2.5-1.7B-int8`` -- a compressed-tensors +int8 W8A8 checkpoint -- onto the RTX 5090 Laptop GPU initializes cleanly and +then dies in the first forward with + + dispatch_scaled_mm, csrc/.../c3x/scaled_mm_helper.hpp:34, + Int8 not supported on SM120. Use FP8 quantization instead, or run on + older arch (SM < 100). + +Twenty seconds and 2 GiB of weight loading to find out. The whole point of +:mod:`vllm_omni.edge.local.plan` is that this is knowable from the checkpoint's +``quantization_config`` and the device's compute capability before anything is +loaded, so the engine refuses with that sentence instead of crashing with it. + +Evidence markers follow the proposal's scale: ``D`` detected, ``B`` the backend +executed something here, ``E`` estimated. Nothing in this module claims ``C`` +or ``P``; it describes devices, not models. +""" + +from __future__ import annotations + +import os +from dataclasses import asdict, dataclass, field +from typing import Any + +from vllm_omni.edge.hardware_probe import ( + ACCEL_GPU_DISCRETE, + ACCEL_GPU_INTEGRATED, + ACCEL_NPU, + HardwareProfile, +) + +GiB = 2**30 + +# Weight formats this engine knows how to name. These are the spellings that +# appear in a checkpoint's ``quantization_config.quant_method`` (plus +# ``dense`` for an unquantized checkpoint), not vLLM's ``--quantization`` +# argument, which is a different vocabulary with overlapping words. +FORMAT_DENSE = "dense" +"""bf16/fp16/fp32 weights, no quantization_config.""" +FORMAT_INT8 = "compressed-tensors:int8" +FORMAT_FP8 = "compressed-tensors:fp8" +FORMAT_GPTQ = "gptq" +FORMAT_AWQ = "awq" +FORMAT_NVFP4 = "modelopt_fp4" + +# The external-stage devices do not execute checkpoints at all: they execute an +# exported graph, and the format that gates them is the graph's, not the +# checkpoint's. Kept in the same vocabulary so one ``weight_formats`` set can +# answer "can this device run this artifact" for every device on the row. +FORMAT_ONNX_FP32 = "onnx:fp32" +FORMAT_ONNX_FP16 = "onnx:fp16" +FORMAT_PT_FP32 = "pt:fp32" +FORMAT_PT_FP16 = "pt:fp16" +"""A TorchScript or ``torch.export`` artifact. The 890M has two routes and they +do not eat the same thing: onnxruntime-directml on Windows takes ONNX, while +torch-directml inside WSL takes a torch export -- which is also what +``vllm_omni.edge.decoder_export`` already produces. The device's format set +therefore depends on which route resolved, rather than promising both.""" + +FORMAT_ONNX_A16W8 = "onnx:a16w8" +"""16-bit activations, 8-bit weights. The **only** door on XDNA2: the overlays +shipped in the EP are named ``8x4_psu_model_a16w8_qdq.xclbin`` and friends, and +A8W8 -- which is what ``quantize_static`` produces by default -- is rejected as +completely as fp32, with no error and every node quietly on the CPU. See +:mod:`vllm_omni.edge.npu_ryzenai`.""" + +MASK_ENV = "VLLM_OMNI_LOCAL_MASK" +"""Comma-separated device kinds to hide, e.g. ``gpu_discrete``. The +capability-masked start of the M0 acceptance run uses this; it exists so the +"no NVIDIA" deployment class can be exercised without physically removing the +card, and it is honest about being a mask -- ``masked=True`` is reported, the +device is never silently dropped.""" + + +@dataclass(frozen=True) +class DeviceCapability: + """One device, as this process can use it.""" + + device_id: str + """Stable within a run: ``cuda:0``, ``cpu``. Never a bare index -- the + proposal's rule that 890M and the 5090 must not both be "GPU 0".""" + kind: str + """``gpu_discrete`` | ``cpu`` | ``gpu_integrated`` | ``npu``, from + :mod:`~vllm_omni.edge.hardware_probe`.""" + vendor: str + name: str + backend: str | None + """The runtime that would execute a stage here, e.g. ``vllm:cuda``, + ``vllm:cpu``. ``None`` when no backend in this process drives it.""" + runnable: bool + """Probe found it *and* a backend in this process can drive it.""" + memory_pool: str + """``vram`` (private) or ``host_ram`` (shared with the OS and every other + device on this row). Two devices sharing ``host_ram`` must never be given + two independent budgets.""" + memory_bytes: int + compute_capability: tuple[int, ...] | None + weight_formats: frozenset[str] + """Formats this device can *execute*. Not "can load": see the module + docstring.""" + evidence: str + """``D``, ``B`` or ``E`` for the runnable claim.""" + reason: str = "" + """Why it is not runnable, or how it is limited. The actionable half.""" + masked: bool = False + extra: dict[str, Any] = field(default_factory=dict) + + def to_dict(self) -> dict[str, Any]: + d = asdict(self) + d["weight_formats"] = sorted(self.weight_formats) + d["compute_capability"] = list(self.compute_capability) if self.compute_capability else None + return d + + +def cuda_weight_formats(capability: tuple[int, ...] | None) -> frozenset[str]: + """Formats a CUDA device of this compute capability can execute. + + The int8 exclusion is the one that is not arithmetic: cutlass' ``c3x`` + ``scaled_mm`` has no int8 path for SM >= 100, and vLLM's dispatch says so + by raising. Blackwell consumer parts (SM 120) are inside that range, so an + int8 W8A8 checkpoint is a *load-time refusal* on this laptop even though + it is 2 GiB and the card has 24. B: observed on sm_120, 2026-09-15. + """ + formats = {FORMAT_DENSE, FORMAT_GPTQ, FORMAT_AWQ} + if capability is None: + return frozenset(formats) + major = capability[0] + if major >= 9: # Hopper and later carry the fp8 tensor cores + formats.add(FORMAT_FP8) + if major >= 10: # Blackwell: NVFP4, and no int8 scaled_mm + formats.add(FORMAT_NVFP4) + else: + formats.add(FORMAT_INT8) + return frozenset(formats) + + +def cpu_weight_formats(profile: HardwareProfile) -> frozenset[str]: + """Formats the vLLM CPU platform can execute on this CPU. + + GPTQ is listed for x86 only because that is where the 4-bit CPU kernels + live (``int4_scaled_mm_cpu`` with AMX, tinygemm W4A16 without); see + :mod:`vllm_omni.edge.weight_path`, which picks *between* them. E on + aarch64: no 4-bit CPU path has been measured there in this repo. + """ + formats = {FORMAT_DENSE, FORMAT_INT8} + if profile.arch == "x86_64": + formats.add(FORMAT_GPTQ) + return frozenset(formats) + + +def _current_platform_device_type() -> str | None: + """What ``vllm`` in *this* interpreter will actually execute on. + + Imported lazily and defensively: the probe has to work in a process with no + CUDA context (that is why ``vllm_omni.__getattr__`` defers ``Omni``), and a + platform import failure must degrade to "no backend", not to a traceback. + """ + try: + from vllm.platforms import current_platform + except Exception: # pragma: no cover - vllm always present in practice + return None + device_type = getattr(current_platform, "device_type", None) + return str(device_type) if device_type else None + + +def parse_mask(mask: str | None = None) -> frozenset[str]: + """Device kinds to hide, from an explicit argument or :data:`MASK_ENV`.""" + raw = mask if mask is not None else os.environ.get(MASK_ENV, "") + return frozenset(part.strip() for part in raw.split(",") if part.strip()) + + +def enumerate_devices( + profile: HardwareProfile, + *, + mask: frozenset[str] | str | None = None, + platform_device_type: str | None = None, +) -> list[DeviceCapability]: + """Every device the probe found, annotated with what this process can do. + + ``platform_device_type`` is injectable so the CUDA-visible and + capability-masked plans of the M0 acceptance run can both be produced from + one snapshot, and so tests do not need a GPU. + """ + masked_kinds = mask if isinstance(mask, frozenset) else parse_mask(mask) + platform = platform_device_type if platform_device_type is not None else _current_platform_device_type() + devices: list[DeviceCapability] = [] + + for accel in profile.accelerators: + kind = str(accel.get("kind", "")) + if kind == ACCEL_GPU_DISCRETE and accel.get("vendor") == "nvidia": + devices.append(_cuda_device(profile, accel, masked_kinds, platform)) + elif kind in (ACCEL_GPU_INTEGRATED, ACCEL_NPU) and accel.get("vendor") == "amd": + devices.append(_amd_external_device(profile, accel, kind, masked_kinds)) + elif kind in (ACCEL_GPU_INTEGRATED, ACCEL_NPU): + devices.append(_unsupported_accelerator(accel, kind, masked_kinds)) + + devices.append(_cpu_device(profile, masked_kinds, platform)) + return devices + + +def _cuda_device( + profile: HardwareProfile, + accel: dict[str, Any], + masked: frozenset[str], + platform: str | None, +) -> DeviceCapability: + capability = tuple(profile.gpu_capability) if profile.gpu_capability else None + is_masked = ACCEL_GPU_DISCRETE in masked + probe_usable = bool(accel.get("usable")) + platform_ok = platform == "cuda" + runnable = probe_usable and platform_ok and not is_masked + + if is_masked: + reason = f"masked out by {MASK_ENV}/--mask ({ACCEL_GPU_DISCRETE}); the no-NVIDIA deployment class" + elif not probe_usable: + reason = str(accel.get("reason") or "the probe could not open this device from this process") + elif not platform_ok: + reason = ( + f"vllm.platforms.current_platform is {platform!r}, not 'cuda': this " + "interpreter's vLLM wheel does not drive the card. Use the CUDA venv " + "(.venvs/omni-cuda)." + ) + else: + reason = "" + + return DeviceCapability( + device_id="cuda:0", + kind=ACCEL_GPU_DISCRETE, + vendor="nvidia", + name=profile.gpu_name or str(accel.get("name", "")), + backend="vllm:cuda" if runnable else None, + runnable=runnable, + memory_pool="vram", + memory_bytes=int(profile.gpu_mem_bytes or accel.get("mem_bytes") or 0), + compute_capability=capability, + weight_formats=cuda_weight_formats(capability), + evidence="B" if runnable else "D", + reason=reason, + masked=is_masked, + extra={"sm_count": profile.gpu_sm_count, "source": accel.get("source")}, + ) + + +def _cpu_device(profile: HardwareProfile, masked: frozenset[str], platform: str | None) -> DeviceCapability: + is_masked = "cpu" in masked + platform_ok = platform == "cpu" + runnable = platform_ok and not is_masked + if is_masked: + reason = f"masked out by {MASK_ENV}/--mask (cpu)" + elif not platform_ok: + reason = ( + f"vllm.platforms.current_platform is {platform!r}, not 'cpu': this " + "interpreter's vLLM wheel has no CPU executor. The CPU build lives " + "in a separate venv (.venvs/omni-cpu); a CPU plan from here is a " + "plan for that interpreter, not something this one can run." + ) + else: + reason = "" + return DeviceCapability( + device_id="cpu", + kind="cpu", + vendor="amd" if "amd" in profile.cpu_model.lower() else "cpu", + name=profile.cpu_model, + backend="vllm:cpu" if runnable else None, + runnable=runnable, + memory_pool="host_ram", + memory_bytes=int(profile.ram_total_bytes), + compute_capability=None, + weight_formats=cpu_weight_formats(profile), + evidence="B" if runnable else "D", + reason=reason, + masked=is_masked, + extra={ + "n_physical": profile.n_physical, + "n_logical": profile.n_logical, + "ram_available_bytes": profile.ram_available_bytes, + "bf16": profile.supports_bf16, + }, + ) + + +def _amd_external_device( + profile: HardwareProfile, + accel: dict[str, Any], + kind: str, + masked: frozenset[str], +) -> DeviceCapability: + """The Radeon 890M or the XDNA2 NPU, reached through an external worker. + + Neither runs under vLLM, and neither can run in *this* interpreter: + + * the NPU is an MCDM device. WSL2's GPU-PV forwards WDDM display adapters + only, so there is no ``/dev/accel`` inside WSL, and VitisAI is a Windows + DLL besides. + * the 890M is reachable from WSL through DirectML, but ``torch-directml`` + pins torch 2.4.1 while Omni runs 2.13. + + So ``runnable`` here means "a worker route resolves", not "this process can + open the device" -- which is why the route is recorded in ``extra`` and the + reason names the missing interpreter when there is one. Being explicit + about that is the point: "the NPU is unusable" and "the NPU's venv is not + installed" call for different work. + + Both draw from ``host_ram``. They must never be given budgets of their own: + an iGPU allocation and a CPU stage's weights come out of the same pool that + the Windows host and the WSL quota are already splitting. + """ + from vllm_omni.edge.local.external import launch as _launch + + is_npu = kind == ACCEL_NPU + # Same ``kind:vendor`` spelling ``_unsupported_accelerator`` uses, so a + # device does not change identity the day it gains a backend -- these ids + # are in the M0 records already. + device_id = f"{kind}:{accel.get('vendor', 'amd')}" + # Ordered by preference. The iGPU has two routes and they are not equal -- + # DirectML-under-WSL measured *slower than the CPU it shares memory with* + # (1.23 TFLOP/s fp32 against 1.54, 5-7.4 GB/s against 20.7) because the + # D3D12 translation layer costs more than the iGPU wins, so the native + # route is tried first and M3.1 measures both. + candidates = ( + (_launch.ROUTE_VITISAI,) if is_npu else (_launch.ROUTE_DML, _launch.ROUTE_TORCH_DML) + ) + routes = [_launch.resolve(name) for name in candidates] + chosen = next((r for r in routes if r.available), None) + + is_masked = kind in masked + runnable = chosen is not None and not is_masked + + if is_masked: + reason = f"masked out by {MASK_ENV}/--mask ({kind})" + elif chosen is None: + reason = "; ".join(f"{r.name}: {r.reason}" for r in routes) + else: + reason = str(accel.get("reason") or "") + + if is_npu: + # A16W8 and nothing else. Measured, not a policy: A8W8 -- the default + # shape of ``quantize_static`` -- and fp32 are both rejected by the + # overlays with no error at all, leaving every node on the CPU while + # the session still reports the EP. See :mod:`vllm_omni.edge.npu_ryzenai`. + formats = frozenset({FORMAT_ONNX_A16W8}) + elif chosen is not None and chosen.name == _launch.ROUTE_TORCH_DML: + formats = frozenset({FORMAT_PT_FP16, FORMAT_PT_FP32}) + elif chosen is not None: + formats = frozenset({FORMAT_ONNX_FP16, FORMAT_ONNX_FP32}) + else: + # No route resolved, so nothing is executable here. Advertising the + # union of what *some* route would take would make the format gate pass + # for a device that cannot run anything. + formats = frozenset() + + backend = None + if runnable and chosen is not None: + backend = "torch:dml" if chosen.name == _launch.ROUTE_TORCH_DML else f"ort:{chosen.ep}" + + return DeviceCapability( + device_id=device_id, + kind=kind, + vendor=str(accel.get("vendor", "amd")), + name=str(accel.get("name", "")), + backend=backend, + runnable=runnable, + memory_pool="host_ram", + # Zero, and deliberately. Section 6 forbids giving the iGPU and the NPU + # budgets of their own: all three devices allocate out of the one pool + # the CPU row already carries, so repeating the number here would let a + # caller add three copies of the same RAM together. The external stage + # planner takes its capacity from the host row for exactly this reason. + memory_bytes=0, + compute_capability=None, + weight_formats=formats, + evidence="B" if runnable else "D", + reason=reason, + masked=is_masked, + extra={ + "source": accel.get("source"), + "probe_route": accel.get("route"), + "host_pool_bytes": int(profile.ram_total_bytes), + "host_pool_available_bytes": int(profile.ram_available_bytes), + "worker_route": chosen.name if chosen else None, + "worker_interpreter": chosen.interpreter if chosen else None, + "routes": [r.to_dict() for r in routes], + "executes": "exported graphs, not checkpoints", + # Dedicated VRAM, for the record. It is not a budget: 512 MiB is a + # carve-out of the same system RAM, and DirectML grows past it into + # the shared pool. + "dedicated_vram_bytes": int(accel.get("mem_bytes") or 0), + }, + ) + + +def _unsupported_accelerator(accel: dict[str, Any], kind: str, masked: frozenset[str]) -> DeviceCapability: + """An iGPU or NPU: found, and not a stage backend in this engine yet. + + Reported rather than dropped, because "the machine has an NPU" and "a stage + can run on it" are different claims and the plan has to be able to say + which one it is relying on. M3 is where this row grows a backend. + """ + return DeviceCapability( + device_id=f"{kind}:{accel.get('vendor', '?')}", + kind=kind, + vendor=str(accel.get("vendor", "")), + name=str(accel.get("name", "")), + backend=None, + runnable=False, + memory_pool="host_ram", + memory_bytes=0, + compute_capability=None, + weight_formats=frozenset(), + evidence="D", + reason=( + str(accel.get("reason") or "") + or "detected, but this engine has no stage backend for it yet (roadmap M3)" + ), + masked=kind in masked, + extra={"source": accel.get("source"), "route": accel.get("route")}, + ) + + +def runs_exported_graphs(device: DeviceCapability) -> bool: + """True for a device that executes an exported graph rather than a checkpoint. + + The distinction matters to :func:`~vllm_omni.edge.local.plan.plan_text_session`: + the 890M and the NPU are runnable, and a whole autoregressive text session + still cannot be placed on either. They take a stage, from + :func:`~vllm_omni.edge.local.external.stage.plan_external_stage`. + """ + # Keyed on the device, not on today's format set. When no worker route is + # installed the format set is empty -- correctly, since nothing here can + # run -- and a predicate reading that set would then answer "no", which + # would quietly drop the device out of the planner's explanation exactly + # when the caller most needs to be told why it was not used. + return device.kind in (ACCEL_GPU_INTEGRATED, ACCEL_NPU) + + +def describe(devices: list[DeviceCapability]) -> str: + lines = [] + for d in devices: + mark = "runnable" if d.runnable else ("masked" if d.masked else "not here") + mem = f"{d.memory_bytes / GiB:.1f}GiB {d.memory_pool}" if d.memory_bytes else d.memory_pool + tail = f"\n {d.reason}" if d.reason else "" + lines.append(f" [{mark:^8}] {d.device_id:<16} {d.name} ({mem}) backend={d.backend}{tail}") + return "\n".join(lines) if lines else " (none)" diff --git a/vllm_omni/edge/local/cli.py b/vllm_omni/edge/local/cli.py new file mode 100644 index 000000000..e39094683 --- /dev/null +++ b/vllm_omni/edge/local/cli.py @@ -0,0 +1,548 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""``python -m vllm_omni.edge.local`` -- plan, run, and the M0 acceptance run. + +``accept`` is the one that matters: it is the roadmap's M0 acceptance criteria +executed as a single command, writing one JSON file that carries the versions, +the plan, the twelve requests, the cancellation check and the memory record. +What it does *not* do is decide whether the numbers are good. It records the +conditions with them so somebody can. +""" + +from __future__ import annotations + +import argparse +import asyncio +import json +import os +import sys +import time +from pathlib import Path +from typing import Any + +from vllm_omni.edge.hardware_probe import load_profile +from vllm_omni.edge.local.capabilities import describe, enumerate_devices, parse_mask +from vllm_omni.edge.local.engine import LocalTextEngine, apply_runtime_env +from vllm_omni.edge.local.manifest import runtime_versions +from vllm_omni.edge.local.plan import ExecutionPlan, plan_text_session +from vllm_omni.edge.local.parity import ( + PromptParity, + compare_ids, + reference_token_ids, + summarize, +) +from vllm_omni.edge.local.prompts import acceptance_prompts + + +def _write(path: str | None, payload: dict[str, Any]) -> None: + if not path: + return + out = Path(path) + out.parent.mkdir(parents=True, exist_ok=True) + out.write_text(json.dumps(payload, indent=1, ensure_ascii=False, default=str), encoding="utf-8") + print(f"\nwrote {out}") + + +def cmd_devices(args: argparse.Namespace) -> int: + profile = load_profile() + devices = enumerate_devices(profile, mask=parse_mask(args.mask)) + print(describe(devices)) + _write(args.json, {"devices": [d.to_dict() for d in devices], "runtime": runtime_versions().to_dict()}) + return 0 + + +def _plan_from_args(args: argparse.Namespace) -> ExecutionPlan: + return plan_text_session( + args.model, + max_model_len=args.max_model_len, + max_num_seqs=args.max_num_seqs, + max_num_batched_tokens=args.max_num_batched_tokens, + enforce_eager=not args.cuda_graphs, + mask=parse_mask(args.mask), + digest_weights=getattr(args, "digest_weights", False), + ) + + +def cmd_plan(args: argparse.Namespace) -> int: + plan = _plan_from_args(args) + print(plan.summary()) + _write(args.json, plan.to_dict()) + # A refusal is a successful plan run. It gets exit code 2 so a script can + # tell "refused with reasons" from "the planner itself broke" (1). + return 0 if plan.admitted else 2 + + +def cmd_run(args: argparse.Namespace) -> int: + plan = _plan_from_args(args) + if not plan.admitted: + print(plan.summary(), file=sys.stderr) + return 2 + print(plan.summary()) + + async def go() -> dict[str, Any]: + async with LocalTextEngine(plan) as engine: + session = engine.open_session() + rid, stream = await engine.submit( + session, args.prompt, max_tokens=args.max_tokens, temperature=args.temperature + ) + async for event in stream: + if event.kind == "token": + print(event.payload["text"], end="", flush=True) + elif event.kind == "error": + print(f"\n[error] {event.error}", file=sys.stderr) + print() + return { + "plan": plan.to_dict(), + "placement": engine.report_placement(), + "usage": engine.report_usage(), + } + + payload = asyncio.run(go()) + _write(args.json, payload) + return 0 + + +async def _acceptance(plan: ExecutionPlan, args: argparse.Namespace) -> dict[str, Any]: + prompts = acceptance_prompts() + results: list[dict[str, Any]] = [] + async with LocalTextEngine(plan) as engine: + print(f"loaded in {engine.load_seconds:.2f} s") + + for name, text in prompts: + session = engine.open_session() + rid, stream = await engine.submit( + session, text, + max_tokens=args.max_tokens, + temperature=0.0, + # The acceptance criterion is "at least 128 tokens", and a + # short greedy answer that hits EOS would not meet it. The + # length is forced so the criterion means the same thing for + # every prompt; this is a measurement setting, not a sampling + # recommendation. + ignore_eos=True, + ) + async for _ in stream: + pass + record = engine.records[rid] + engine.close_session(session.session_id) + results.append({"prompt": name, **record.to_dict(), "stream": stream.stats()}) + print( + f" {name:<12} in={record.prompt_tokens} out={record.output_tokens} " + f"ttft={record.ttft_s:.3f}s decode={record.decode_tok_per_s:.2f} tok/s" + if record.ttft_s and record.decode_tok_per_s + else f" {name:<12} out={record.output_tokens} (timing unavailable)" + ) + + # Cancellation: interrupt a long request part-way and prove that + # nothing from the retired epoch reaches the consumer afterwards. + cancel_session = engine.open_session() + rid, stream = await engine.submit( + cancel_session, prompts[0][1], max_tokens=4096, ignore_eos=True + ) + seen = 0 + async for event in stream: + if event.kind == "token": + seen += 1 + if seen >= 8: + break + cancel_report = await engine.cancel(rid) + cancel_report["events_before_cancel"] = seen + leaked = [] + while True: + try: + event = await asyncio.wait_for(stream.get(), timeout=1.0) + except (asyncio.TimeoutError, Exception): + break + if event is None: + break + leaked.append(event.to_dict()) + cancel_report["events_after_cancel"] = leaked + cancel_report["clean"] = not leaked + print( + f" cancel: aborted={cancel_report['aborted_in_backend']} " + f"after {seen} events, leaked={len(leaked)}" + ) + + payload = { + "started_unix": time.time(), + "argv": sys.argv, + "acceptance": { + "prompts": len(prompts), + "min_tokens_required": args.max_tokens, + "all_reached_min": all(r["output_tokens"] >= args.max_tokens for r in results), + "cancel_clean": cancel_report["clean"], + }, + "plan": plan.to_dict(), + "placement": engine.report_placement(), + "usage": engine.report_usage(), + "requests": results, + "cancel": cancel_report, + } + return payload + + +def cmd_accept(args: argparse.Namespace) -> int: + """The M0 acceptance run: both startups, twelve prompts, cancel, memory.""" + profile = load_profile() + + # 1. The capability-masked startup. Planned first and on purpose: the + # criterion is that the no-NVIDIA class produces a compatible plan *or* + # an explicit refusal, and a refusal here is a pass, not a failure. + masked_plan = plan_text_session( + args.model, + profile=profile, + max_model_len=args.max_model_len, + max_num_seqs=args.max_num_seqs, + max_num_batched_tokens=args.max_num_batched_tokens, + enforce_eager=not args.cuda_graphs, + mask=frozenset({"gpu_discrete"}), + ) + print("=" * 72) + print("capability-masked startup (--mask gpu_discrete): the no-NVIDIA class") + print("=" * 72) + print(masked_plan.summary()) + + # 2. The device-visible startup, which is the one that runs. + plan = _plan_from_args(args) + print() + print("=" * 72) + print("device-visible startup") + print("=" * 72) + print(plan.summary()) + print() + + payload: dict[str, Any] = { + "masked_plan": masked_plan.to_dict(), + "masked_outcome": "admitted" if masked_plan.admitted else "refused_explicitly", + } + if not plan.admitted: + payload["plan"] = plan.to_dict() + payload["outcome"] = "refused_explicitly" + _write(args.json, payload) + return 2 + + payload.update(asyncio.run(_acceptance(plan, args))) + payload["outcome"] = "completed" + ok = payload["acceptance"]["all_reached_min"] and payload["acceptance"]["cancel_clean"] + print(f"\nacceptance: {'PASS' if ok else 'INCOMPLETE'} " + f"({payload['acceptance']})") + _write(args.json, payload) + return 0 if ok else 3 + + +async def _engine_token_ids(plan: ExecutionPlan, prompts, max_tokens: int) -> dict[str, Any]: + out: dict[str, Any] = {} + async with LocalTextEngine(plan) as engine: + for name, text in prompts: + session = engine.open_session() + request_id, stream = await engine.submit( + session, text, max_tokens=max_tokens, temperature=0.0, ignore_eos=True + ) + async for _ in stream: + pass + record = engine.records[request_id] + engine.close_session(session.session_id) + out[name] = { + "token_ids": list(record.output_token_ids), + "text": record.text, + "prompt_tokens": record.prompt_tokens, + } + return out + + +def cmd_parity(args: argparse.Namespace) -> int: + """Greedy token-id parity: the Omni stage path against plain vLLM. + + Exit 0 when every prompt matches, 3 when any diverges (the run succeeded; + the models disagreed), 2 when the plan refused. + """ + # Both sides get the same environment, applied before either engine starts. + # Not cosmetic: the reference is a plain vllm.LLM in this same process and + # loads *first*, so anything the engine would have set at its own start is + # set too late for it. Two of the three entries here are load-time fatal on + # this host (see RUNTIME_ENV), and a comparison between a process that has + # them and one that does not measures the configuration, not the path. + env = apply_runtime_env() + + plan = _plan_from_args(args) + if not plan.admitted: + print(plan.summary(), file=sys.stderr) + return 2 + print(plan.summary()) + prompts = acceptance_prompts() + + print(f"\nreference: plain vllm.LLM, {len(prompts)} prompts x {args.max_tokens} greedy tokens") + reference = reference_token_ids( + plan.manifest.model_dir, prompts, + engine_kwargs=plan.engine_kwargs, max_tokens=args.max_tokens, + ) + print("engine: AsyncOmni -> Omni StageRuntime -> EngineCore") + engine = asyncio.run(_engine_token_ids(plan, prompts, args.max_tokens)) + + rows: list[PromptParity] = [] + for name, _ in prompts: + ref = reference[name]["token_ids"] + eng = engine[name]["token_ids"] + exact, first, prefix = compare_ids(ref, eng) + lo = max(0, (first or 0) - 2) + rows.append(PromptParity( + name=name, + prompt_tokens=reference[name]["prompt_tokens"], + reference_tokens=len(ref), + engine_tokens=len(eng), + exact=exact, + first_divergence=first, + agreement_prefix=prefix, + reference_head=[] if exact else ref[lo:lo + 8], + engine_head=[] if exact else eng[lo:lo + 8], + )) + mark = "exact" if exact else f"diverges at {first} (of {len(ref)})" + print(f" {name:<12} ref={len(ref)} omni={len(eng)} {mark}") + + stats = summarize(rows) + print(f"\nparity: {stats['exact']}/{stats['prompts']} exact") + if stats["min_agreement_prefix"] is not None: + print(f" shortest agreement prefix: {stats['min_agreement_prefix']} tokens") + _write(args.json, { + "summary": stats, + "env": env, + "max_tokens": args.max_tokens, + "plan": plan.to_dict(), + "prompts": [r.to_dict() for r in rows], + "reference": reference, + "engine": engine, + }) + return 0 if stats["exact"] == stats["prompts"] else 3 + + +def _synthetic_inputs(graph: str, spec: tuple[dict[str, Any], ...], fmt: str) -> dict[str, Any]: + """Example inputs for the profiled placement run. + + Placement can only be read off a graph that has actually executed, so the + probe needs something to feed it. Shapes come from the graph itself; a + dynamic dimension is an error rather than a guess, because guessing one + wrong produces a session that fails for a reason unrelated to placement. + """ + import numpy as np + + if not spec and not fmt.startswith("onnx:"): + raise SystemExit( + f"a {fmt} artifact declares no input shapes, so they cannot be " + "synthesized. Pass --inputs with an .npz holding the real tensors; " + "placement is verified from an actual run, and a run needs real inputs." + ) + if not spec: + import onnx + + model = onnx.load(graph, load_external_data=False) + spec = tuple( + { + "name": i.name, + "shape": [d.dim_value if d.HasField("dim_value") else None + for d in i.type.tensor_type.shape.dim], + "elem_type": i.type.tensor_type.elem_type, + } + for i in model.graph.input + ) + + dtypes = {1: np.float32, 6: np.int32, 7: np.int64, 10: np.float16, 11: np.float64} + inputs: dict[str, Any] = {} + for item in spec: + shape = list(item.get("shape") or []) + if any(d is None or (isinstance(d, str)) or d == 0 for d in shape): + raise SystemExit( + f"input {item['name']!r} has a dynamic shape {shape}; pass --inputs " + "with an .npz holding real tensors instead of asking for synthetic ones" + ) + dtype = dtypes.get(int(item.get("elem_type", 1)), np.float32) + if np.issubdtype(dtype, np.integer): + inputs[item["name"]] = np.zeros(shape, dtype=dtype) + else: + inputs[item["name"]] = np.random.RandomState(0).randn(*shape).astype(dtype) + return inputs + + +def cmd_external(args: argparse.Namespace) -> int: + """Plan an exported graph onto the iGPU or the NPU, and prove the placement.""" + import numpy as np + + from vllm_omni.edge.local.external.stage import ( + ExternalStage, + PlacementRefused, + plan_external_stage, + ) + from vllm_omni.edge.local.manifest import build_graph_artifact + + profile = load_profile() + devices = enumerate_devices(profile, mask=parse_mask(args.mask)) + artifact = build_graph_artifact( + args.graph, + fmt=args.format, + opset=args.opset, + source_model=args.source_model or "(unrecorded)", + component=args.component, + exporter=args.exporter or "(unrecorded)", + ) + plan = plan_external_stage( + artifact, devices, + prefer=args.prefer, require=args.device, + min_fraction_on_target=args.min_placement, + ) + print(plan.summary()) + record: dict[str, Any] = { + "argv": sys.argv, "runtime": runtime_versions().to_dict(), + "devices": [d.to_dict() for d in devices], "plan": plan.to_dict(), + } + if not plan.admitted: + _write(args.json, record) + return 2 + + if args.inputs: + loaded = np.load(args.inputs) + example = {k: loaded[k] for k in loaded.files} + else: + example = _synthetic_inputs(artifact.path, artifact.input_spec, artifact.fmt) + + stage = ExternalStage(plan) + try: + report = stage.open(example, profile_dir=args.profile_dir) + print() + print(plan.summary()) + timings = [] + for _ in range(args.runs): + _, timing = stage.run(example) + timings.append(timing) + if timings: + worker = sorted(t.worker_s for t in timings) + trip = sorted(t.round_trip_s for t in timings) + transport = sorted(t.transport_s for t in timings) + mid = len(timings) // 2 + print( + f"\n {args.runs} runs: device {worker[mid] * 1e3:.3f} ms median, " + f"round trip {trip[mid] * 1e3:.3f} ms, " + f"transport {transport[mid] * 1e3:.3f} ms" + ) + print( + " transport is the cost of leaving the process; a stage only earns " + "its place here when the gain beats it" + ) + record["timings"] = [t.to_dict() for t in timings] + record["stats"] = stage.stats() + record["report"] = report.to_dict() + record["plan"] = plan.to_dict() + except PlacementRefused as exc: + print(f"\nREFUSED [{exc.refusal.code}] {exc.refusal.message}") + print(f" remedy: {exc.refusal.remedy}") + record["plan"] = plan.to_dict() + _write(args.json, record) + return 2 + finally: + stage.close() + + _write(args.json, record) + return 0 + + +def build_parser() -> argparse.ArgumentParser: + p = argparse.ArgumentParser( + prog="python -m vllm_omni.edge.local", + description="M0 local text mode: plan, run and audit a single-stage text session.", + ) + sub = p.add_subparsers(dest="cmd", required=True) + + def add_plan_args(sp: argparse.ArgumentParser) -> None: + sp.add_argument("--model", required=True, help="checkpoint directory") + sp.add_argument("--max-model-len", type=int, default=4096) + sp.add_argument("--max-num-seqs", type=int, default=1) + sp.add_argument("--max-num-batched-tokens", type=int, default=None) + sp.add_argument("--cuda-graphs", action="store_true", + help="enable graph capture (M0 runs eager; this also adds a 2 GiB reserve)") + sp.add_argument("--mask", default=None, + help="comma-separated device kinds to hide, e.g. gpu_discrete") + sp.add_argument("--json", default=None, help="write the record here") + + sp = sub.add_parser("devices", help="what this process can actually run on") + sp.add_argument("--mask", default=None) + sp.add_argument("--json", default=None) + sp.set_defaults(func=cmd_devices) + + sp = sub.add_parser("plan", help="plan a session; exit 2 means an explicit refusal") + add_plan_args(sp) + sp.add_argument("--digest-weights", action="store_true", + help="sha256 the weight files (slow; use when publishing a number)") + sp.set_defaults(func=cmd_plan) + + sp = sub.add_parser("run", help="stream one prompt") + add_plan_args(sp) + sp.add_argument("--prompt", required=True) + sp.add_argument("--max-tokens", type=int, default=128) + sp.add_argument("--temperature", type=float, default=0.0) + sp.set_defaults(func=cmd_run) + + sp = sub.add_parser( + "parity", + help="greedy token-id parity of the Omni path against plain vllm.LLM; exit 3 on divergence", + ) + add_plan_args(sp) + sp.add_argument("--max-tokens", type=int, default=128) + sp.set_defaults(func=cmd_parity) + + sp = sub.add_parser( + "external", + help="place an exported graph on the iGPU or the NPU; exit 2 is an explicit refusal", + ) + sp.add_argument("--graph", required=True, help="the exported .onnx") + sp.add_argument("--format", default="onnx:a16w8", + choices=["onnx:a16w8", "onnx:fp16", "onnx:fp32", "pt:fp16", "pt:fp32"]) + sp.add_argument("--opset", type=int, default=21) + sp.add_argument("--component", default="graph", help="which part of the model this is") + sp.add_argument("--source-model", default=None) + sp.add_argument("--exporter", default=None) + sp.add_argument("--device", default=None, + help="require this device id; without it a refusal falls through to the other") + sp.add_argument("--prefer", default=None, help="try this device id first, then fall through") + sp.add_argument("--min-placement", type=float, default=0.5, + help="fraction of the graph the target EP must take for the placement to count") + sp.add_argument("--inputs", default=None, help=".npz of real inputs; otherwise synthesized") + sp.add_argument("--runs", type=int, default=10) + sp.add_argument("--profile-dir", default=None, help="where ORT writes its placement profile") + sp.add_argument("--mask", default=None) + sp.add_argument("--json", default=None) + sp.set_defaults(func=cmd_external) + + sp = sub.add_parser("accept", help="the M0 acceptance run") + add_plan_args(sp) + sp.add_argument("--max-tokens", type=int, default=128, + help="tokens per prompt; the criterion is at least 128") + sp.set_defaults(func=cmd_accept) + + return p + + +def _use_selector_event_loop_on_windows() -> None: + """zmq.asyncio needs a selector loop; Windows defaults to Proactor. + + [edge-infer W1] Python 3.8+ on Windows defaults to + ``WindowsProactorEventLoopPolicy``, which does not implement the + ``add_reader`` family. ``zmq.asyncio`` -- which the Omni orchestrator and + every stage client use -- requires it, and the failure is a hard + ``RuntimeError`` at the first async socket, not a degradation. + + This is a no-op everywhere else, and it is set here rather than at import + so that importing the CLI does not reach into a host application's loop + policy. A synchronous ``vllm.LLM`` run never hits this, which is why it + surfaces only once the async engine is involved. + """ + if sys.platform != "win32": + return + policy = getattr(asyncio, "WindowsSelectorEventLoopPolicy", None) + if policy is None: # pragma: no cover - non-Windows + return + if not isinstance(asyncio.get_event_loop_policy(), policy): + asyncio.set_event_loop_policy(policy()) + + +def main(argv: list[str] | None = None) -> int: + _use_selector_event_loop_on_windows() + args = build_parser().parse_args(argv) + return int(args.func(args)) diff --git a/vllm_omni/edge/local/engine.py b/vllm_omni/edge/local/engine.py new file mode 100644 index 000000000..636fb7faf --- /dev/null +++ b/vllm_omni/edge/local/engine.py @@ -0,0 +1,814 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""The M0 local text engine: run an admitted plan, and record what happened. + +This is thin on purpose. It owns the plan, the session state references, the +bounded output stream, the cancellation epoch and the measurement; it owns no +kernels, no scheduler and no cache. vLLM keeps doing all of that inside the +Omni stage, which is the whole point of evolving the existing runtime instead +of writing a second one. + +**On measuring memory here, honestly.** Three counters are sampled and they are +three different things, which the architecture record is emphatic about not +adding together: + +* ``nvml_device_used`` is the whole card, every process on it. Under WSL2, + ``nvmlDeviceGetComputeRunningProcesses`` returns an empty list, so this + engine *cannot* attribute GPU bytes to its own PIDs on this host. The delta + across a load is still evidence -- 7.8 GiB appearing while the host RSS does + not move says where the weights went -- but it is a delta on a shared + counter, and it is labelled that way. +* ``tree_rss`` / ``tree_pss`` are this process and its stage subprocesses. + PSS includes a share of file-backed pages, so on an mmap'd safetensors load + it counts pages that are page cache, not anonymous memory. +* ``host_available`` is the system's, and moves for reasons that are not us. + +Sampling is at a fixed interval, so a peak shorter than the interval is missed. +That lower bound is reported with the numbers rather than assumed away. + +**On cancellation.** ``cancel`` does three things in order: retire the stream's +epoch so nothing already in flight can be delivered, abort the request in the +backend so it stops producing, and release the session's reservation in the +ledger. Doing them in that order is what makes "取消后旧数据不能重新进入输出流" +true rather than likely. +""" + +from __future__ import annotations + +import asyncio +import contextlib +import os +import threading +import time +from dataclasses import asdict, dataclass, field +from typing import Any + +from vllm.logger import init_logger + +from vllm_omni.edge.local.plan import ExecutionPlan +from vllm_omni.edge.local.session import ( + BoundedEventStream, + ChunkEvent, + StateHandle, + new_session_id, +) + +logger = init_logger(__name__) + +GiB = 2**30 +SAMPLE_INTERVAL_S = 0.25 +"""Memory sampling period. Anything shorter-lived than this is not seen; the +reported peaks are lower bounds and say so.""" + +RUNTIME_ENV: dict[str, str] = { + # FlashInfer's sampler JIT-compiles on first use and needs nvcc, which is + # not installed here -- nvidia-smi reporting a CUDA version is not a + # toolkit. Without this the *sampler* takes down the engine at the first + # token with "Could not find nvcc". + "VLLM_USE_FLASHINFER_SAMPLER": "0", + # Under WSL2 vLLM's own GPU model runner allocates a UvaBuffer for its + # input batch and refuses to start with "UVA is not available" unless + # pinned memory is enabled. Omni's AR model runner does not take that path, + # so this matters most when something compares the two. + "VLLM_WSL2_ENABLE_PIN_MEMORY": "1", + "TOKENIZERS_PARALLELISM": "false", +} +"""The environment the measured runs used, in one place. + +It lives here rather than inline in ``start()`` because more than one entry +point has to apply it: a comparison that configures the reference differently +from the engine is measuring the configuration, not the engine. Applied with +``setdefault``, so a caller who set one of these keeps their value -- and the +run record reports what was actually in force.""" + + +def apply_runtime_env() -> dict[str, str | None]: + """Set :data:`RUNTIME_ENV` without overriding the caller, and report it.""" + for key, value in RUNTIME_ENV.items(): + os.environ.setdefault(key, value) + return {key: os.environ.get(key) for key in RUNTIME_ENV} + + +CANCEL_DRAIN_TIMEOUT_S = 10.0 +"""How long ``cancel`` waits for the producer task to finish unwinding before +reporting ``producer_stopped: False`` and returning anyway. The consumer is +released before this wait starts, so exceeding this delays cleanup, not the +caller.""" + + +# ---- measurement ----------------------------------------------------------- + + +def _host_kind() -> str: + """``windows``, ``wsl2`` or ``linux`` -- the three hosts this code has run on.""" + import sys + + if sys.platform == "win32": + return "windows" + try: + with open("/proc/version", encoding="utf-8", errors="replace") as f: + if "microsoft" in f.read().lower(): + return "wsl2" + except OSError: + pass + return "linux" + + +def gpu_attribution_for_tree() -> dict[str, Any]: + """Try to attribute GPU bytes to this process tree, and say why when it fails. + + [edge-infer W4] The earlier version of this report hardcoded a WSL2 + explanation for an empty ``nvmlDeviceGetComputeRunningProcesses`` list. + The W1 acceptance run on *native Windows* then carried that WSL2 sentence + verbatim -- the conclusion (placement unverified) was right, the reason + offered to the operator was wrong. So this now measures instead of + assuming: it asks NVML, intersects with the process tree, and only when + that comes back empty does it name the host-specific mechanism (WSL2's + GPU-PV and Windows' WDDM driver model both hide per-process usage from + NVML; a bare Linux host normally reports it). + """ + host = _host_kind() + try: + import psutil + import pynvml + + proc = psutil.Process() + tree = {proc.pid} | {c.pid for c in proc.children(recursive=True)} + pynvml.nvmlInit() + try: + handle = pynvml.nvmlDeviceGetHandleByIndex(0) + running = pynvml.nvmlDeviceGetComputeRunningProcesses(handle) + finally: + with contextlib.suppress(Exception): + pynvml.nvmlShutdown() + except Exception as exc: + return { + "per_process_gpu_attribution": None, + "per_process_gpu_attribution_reason": ( + f"NVML per-process query unavailable on this {host} host: {exc}. " + "The load delta is a whole-card measurement." + ), + } + ours = [ + {"pid": int(r.pid), "used_gpu_memory": getattr(r, "usedGpuMemory", None)} + for r in running + if int(r.pid) in tree + ] + if ours: + return { + "per_process_gpu_attribution": ours, + "per_process_gpu_attribution_reason": None, + } + mechanism = { + "wsl2": "under WSL2 the GPU-PV driver reports no per-process compute list", + "windows": "under the Windows WDDM driver model NVML reports no per-process " + "compute list (the same NVML call that is empty under WSL2's GPU-PV)", + "linux": "no compute process from this tree appears in NVML's list", + }[host] + return { + "per_process_gpu_attribution": None, + "per_process_gpu_attribution_reason": ( + f"nvmlDeviceGetComputeRunningProcesses returned {len(running)} process(es), " + f"none from this engine's tree of {len(tree)} PID(s): {mechanism}, so GPU " + "bytes cannot be attributed to this engine's PIDs on this host. The load " + "delta is a whole-card measurement." + ), + } + + + +@dataclass +class MemorySample: + unix: float + nvml_device_used: int | None + nvml_device_total: int | None + tree_rss: int + tree_pss: int + host_available: int + host_total: int + + +class MemorySampler: + """Background sampler over the process tree, the host and the GPU. + + A thread rather than a task: the engine's event loop is the thing being + measured, and a sampler that stops whenever the loop is busy would miss + exactly the moments worth sampling. + """ + + def __init__(self, interval_s: float = SAMPLE_INTERVAL_S) -> None: + self.interval_s = interval_s + self.samples: list[MemorySample] = [] + self._stop = threading.Event() + self._thread: threading.Thread | None = None + self._nvml = None + self._handle = None + + def _init_nvml(self) -> None: + try: + import pynvml + + pynvml.nvmlInit() + self._nvml = pynvml + self._handle = pynvml.nvmlDeviceGetHandleByIndex(0) + except Exception as exc: # no GPU, or no NVML in this venv + logger.debug("NVML unavailable, GPU memory will not be sampled: %s", exc) + self._nvml = None + + def sample(self) -> MemorySample: + import psutil + + proc = psutil.Process() + rss = pss = 0 + for p in [proc] + proc.children(recursive=True): + try: + info = p.memory_full_info() + rss += info.rss + pss += getattr(info, "pss", 0) or 0 + except (psutil.Error, OSError): + continue + vm = psutil.virtual_memory() + used = total = None + if self._nvml is not None: + try: + mem = self._nvml.nvmlDeviceGetMemoryInfo(self._handle) + used, total = int(mem.used), int(mem.total) + except Exception: + used = total = None + return MemorySample( + unix=time.time(), nvml_device_used=used, nvml_device_total=total, + tree_rss=rss, tree_pss=pss, + host_available=vm.available, host_total=vm.total, + ) + + def start(self) -> None: + self._init_nvml() + + def loop() -> None: + while not self._stop.is_set(): + try: + self.samples.append(self.sample()) + except Exception as exc: # never take the engine down + logger.debug("memory sample failed: %s", exc) + self._stop.wait(self.interval_s) + + self._thread = threading.Thread(target=loop, name="local-mem-sampler", daemon=True) + self._thread.start() + + def stop(self) -> None: + self._stop.set() + if self._thread is not None: + self._thread.join(timeout=2.0) + if self._nvml is not None: + with contextlib.suppress(Exception): + self._nvml.nvmlShutdown() + + def peaks(self, since: float = 0.0) -> dict[str, Any]: + window = [s for s in self.samples if s.unix >= since] + if not window: + return {"samples": 0, "interval_s": self.interval_s} + gpu = [s.nvml_device_used for s in window if s.nvml_device_used is not None] + return { + "samples": len(window), + "interval_s": self.interval_s, + "note": ( + "peaks are lower bounds: anything shorter than interval_s is missed. " + "nvml_device_used is the whole card including other processes; " + "tree_pss includes a share of file-backed pages. These three " + "counters are different accountings and must not be summed." + ), + "nvml_device_used_peak": max(gpu) if gpu else None, + "nvml_device_used_min": min(gpu) if gpu else None, + "tree_rss_peak": max(s.tree_rss for s in window), + "tree_pss_peak": max(s.tree_pss for s in window), + "host_available_min": min(s.host_available for s in window), + } + + +# ---- per-request record ---------------------------------------------------- + + +@dataclass +class RequestRecord: + """Timing for one request, in the units the validation plan asks for.""" + + request_id: str + session_id: str + epoch: int + prompt_chars: int + prompt_tokens: int | None = None + output_tokens: int = 0 + submitted_unix: float = 0.0 + first_token_unix: float | None = None + last_token_unix: float | None = None + finished: bool = False + cancelled: bool = False + error: str | None = None + token_unix: list[float] = field(default_factory=list) + text: str = "" + output_token_ids: list[int] = field(default_factory=list) + """The sampled ids, kept so a run can be compared against a reference + decode token-for-token. Text is not enough for that: two different id + sequences can detokenize to the same string.""" + + @property + def ttft_s(self) -> float | None: + """Submission to first token. Includes admission and scheduling, which + is what a caller experiences; it is not an isolated prefill kernel + time and is never reported as one.""" + if self.first_token_unix is None: + return None + return self.first_token_unix - self.submitted_unix + + @property + def decode_tok_per_s(self) -> float | None: + """Tokens after the first, over the time between first and last. + + ``None`` below two tokens rather than a number computed from one + interval. + """ + if self.first_token_unix is None or self.last_token_unix is None: + return None + if self.output_tokens < 2: + return None + span = self.last_token_unix - self.first_token_unix + return (self.output_tokens - 1) / span if span > 0 else None + + def to_dict(self) -> dict[str, Any]: + d = asdict(self) + d["ttft_s"] = self.ttft_s + d["decode_tok_per_s"] = self.decode_tok_per_s + return d + + +# ---- the engine ------------------------------------------------------------ + + +class PlanNotAdmitted(RuntimeError): + """Raised when someone tries to run a plan that refused.""" + + +class LocalTextEngine: + """One admitted plan, one Omni stage, streamed text with a cancel epoch.""" + + def __init__(self, plan: ExecutionPlan, *, extra_engine_kwargs: dict[str, Any] | None = None) -> None: + if not plan.admitted: + raise PlanNotAdmitted(plan.summary()) + self.plan = plan + self.extra_engine_kwargs = dict(extra_engine_kwargs or {}) + self.sampler = MemorySampler() + self.records: dict[str, RequestRecord] = {} + self._omni: Any = None + self._sessions: dict[str, StateHandle] = {} + self._streams: dict[str, BoundedEventStream] = {} + self._pumps: dict[str, asyncio.Task] = {} + self._inflight: dict[str, str] = {} # request_id -> session_id + self.load_seconds: float | None = None + self.load_memory: dict[str, Any] = {} + self._closed = False + + # -- lifecycle ---------------------------------------------------------- + async def start(self) -> None: + """Load the model according to the plan, and price the load.""" + from vllm_omni import AsyncOmni + + apply_runtime_env() + self.sampler.start() + before = self.sampler.sample() + started = time.perf_counter() + kwargs = {**self.plan.engine_kwargs, **self.extra_engine_kwargs} + logger.info( + "[local] loading %s on %s with %s", + self.plan.manifest.model_dir, self.plan.selected.device_id, kwargs, + ) + self._omni = AsyncOmni(model=self.plan.manifest.model_dir, **kwargs) + self.load_seconds = time.perf_counter() - started + after = self.sampler.sample() + + self.load_memory = { + "before": asdict(before), + "after": asdict(after), + "nvml_device_used_delta": ( + (after.nvml_device_used - before.nvml_device_used) + if (after.nvml_device_used is not None and before.nvml_device_used is not None) + else None + ), + "tree_rss_delta": after.tree_rss - before.tree_rss, + "peaks_during_load": self.sampler.peaks(since=before.unix), + } + self._record_actuals() + + def _record_actuals(self) -> None: + """Turn the load into a bound on the reservations, without overclaiming. + + The NVML delta across the load is **not** "the weights". It is the + weights plus the CUDA context, the allocator's arenas and whatever else + the backend set up, measured on a counter shared with every other + process on the card. So it is recorded as an *upper bound* on the + weights reservation rather than as its actual value: a ratio above 1.0 + here means the budget was under the bound, not that the weights were + bigger than their files. + """ + pool = self.plan.selected.memory_pool if self.plan.selected else "vram" + if pool == "vram": + delta = self.load_memory.get("nvml_device_used_delta") + unit = "on the device (whole-card delta)" + else: + delta = self.load_memory.get("tree_rss_delta") + unit = "of process-tree RSS" + if not delta or delta <= 0: + return + for reservation in self.plan.reservations: + if reservation.purpose == "weights" and reservation.pool == pool: + reservation.actual_upper_bound_bytes = int(delta) + reservation.note += ( + f"; the load moved +{delta / GiB:.2f} GiB {unit}, which bounds " + "weights + backend context together and is not an attribution " + "to the weights alone" + ) + + def measured_peak(self) -> dict[str, Any]: + """Budgeted peak against the measured one, on the selected pool. + + The measured number subtracts a pre-load baseline, because the counter + is shared: on this host the card already held 2.75 GiB before the + engine started, and reporting that as the engine's would be the same + mistake the architecture record calls out about summing WSL ``used``, + Windows ``used`` and PSS. + """ + if self.plan.selected is None: + return {} + pool = self.plan.selected.memory_pool + peaks = self.sampler.peaks() + before = self.load_memory.get("before") or {} + if pool == "vram": + baseline = before.get("nvml_device_used") + peak = peaks.get("nvml_device_used_peak") + counter = "nvml_device_used" + else: + baseline = before.get("tree_rss") + peak = peaks.get("tree_rss_peak") + counter = "tree_rss" + attributable = (peak - baseline) if (peak is not None and baseline is not None) else None + budget = self.plan.peak_bytes + return { + "pool": pool, + "counter": counter, + "baseline_bytes": baseline, + "peak_bytes": peak, + "attributable_peak_bytes": attributable, + "budgeted_peak_bytes": budget, + "budget_over_measured": (budget / attributable) if attributable else None, + "admitted_without_oom": True, + "caveats": ( + "sampled every %.2f s, so the peak is a lower bound; the counter is " + "shared with other processes and the baseline subtraction only " + "removes what was there before the load, not what arrived during it. " + "A budget above the measured peak is the intended outcome -- this " + "budget is a safe upper bound for admission, not a prediction." + % self.sampler.interval_s + ), + } + + async def close(self) -> None: + if self._closed: + return + self._closed = True + for task in list(self._pumps.values()): + task.cancel() + for task in list(self._pumps.values()): + with contextlib.suppress(asyncio.CancelledError, Exception): + await task + self._pumps.clear() + for stream in list(self._streams.values()): + with contextlib.suppress(Exception): + await stream.close() + self._streams.clear() + if self._omni is not None: + for name in ("shutdown", "close"): + closer = getattr(self._omni, name, None) + if callable(closer): + with contextlib.suppress(Exception): + result = closer() + if asyncio.iscoroutine(result): + await result + break + self.sampler.stop() + + async def __aenter__(self) -> LocalTextEngine: + await self.start() + return self + + async def __aexit__(self, *exc: Any) -> None: + await self.close() + + # -- sessions ----------------------------------------------------------- + def open_session(self) -> StateHandle: + handle = StateHandle( + session_id=new_session_id(), + backend=self.plan.backend or "unknown", + artifact_id=self.plan.manifest.artifact_id, + ) + self._sessions[handle.session_id] = handle + return handle + + def session(self, session_id: str) -> StateHandle: + return self._sessions[session_id] + + def close_session(self, session_id: str) -> None: + self._sessions.pop(session_id, None) + + # -- generation --------------------------------------------------------- + async def submit( + self, + handle: StateHandle, + prompt: str, + *, + max_tokens: int = 128, + temperature: float = 0.0, + ignore_eos: bool = False, + request_id: str | None = None, + max_chunks: int = 64, + max_bytes: int = 1 << 20, + ) -> tuple[str, BoundedEventStream]: + """Start one request; return its id and its bounded event stream. + + Returning the stream rather than yielding from here keeps cancellation + possible from outside the consumer: ``cancel`` can retire the epoch + while the caller is still awaiting an event. + """ + if self._omni is None: + raise RuntimeError("engine not started") + live = self._sessions.get(handle.session_id) + if live is None or not live.accepts(handle): + raise RuntimeError( + f"state handle for session {handle.session_id} is stale " + f"(epoch {handle.epoch}, live {live.epoch if live else 'closed'}); " + "open a new session rather than reusing retired state" + ) + from vllm import SamplingParams + + request_id = request_id or f"{handle.session_id}-{len(self.records)}" + stream = BoundedEventStream(max_chunks=max_chunks, max_bytes=max_bytes) + if handle.epoch: + # A session that has been cancelled before starts its new + # stream already fenced against everything older. + stream.retire_epoch(handle.epoch - 1) + record = RequestRecord( + request_id=request_id, + session_id=handle.session_id, + epoch=handle.epoch, + prompt_chars=len(prompt), + submitted_unix=time.time(), + ) + self.records[request_id] = record + self._streams[request_id] = stream + self._inflight[request_id] = handle.session_id + + params = SamplingParams( + temperature=temperature, max_tokens=max_tokens, ignore_eos=ignore_eos + ) + self._pumps[request_id] = asyncio.create_task( + self._pump(request_id, handle, prompt, params, stream, record), + name=f"local-pump-{request_id}", + ) + return request_id, stream + + async def _pump( + self, + request_id: str, + handle: StateHandle, + prompt: str, + params: Any, + stream: BoundedEventStream, + record: RequestRecord, + ) -> None: + """Adapt Omni's cumulative outputs into incremental ChunkEvents. + + Omni yields the output-so-far, like vLLM. Re-emitting that verbatim + would make a dropped chunk invisible to the consumer, so only the delta + goes on the wire and ``seq`` is what detects loss. + """ + seq = 0 + emitted_text = "" + try: + async for out in self._omni.generate( + prompt, request_id=request_id, sampling_params_list=[params] + ): + completion = _first_completion(out) + if completion is None: + continue + text = completion.get("text", "") + token_ids = completion.get("token_ids", ()) + if record.prompt_tokens is None: + record.prompt_tokens = completion.get("prompt_tokens") + delta = text[len(emitted_text):] if text.startswith(emitted_text) else text + n_tokens = len(token_ids) + if n_tokens <= record.output_tokens and not delta: + continue + now = time.time() + if record.first_token_unix is None: + record.first_token_unix = now + record.last_token_unix = now + record.output_tokens = max(record.output_tokens, n_tokens) + record.token_unix.append(now) + emitted_text = text if text.startswith(emitted_text) else emitted_text + delta + if n_tokens >= len(record.output_token_ids): + record.output_token_ids = list(token_ids) + seq += 1 + await stream.put(ChunkEvent( + request_id=request_id, stage_id=0, seq=seq, epoch=handle.epoch, + kind="token", payload={"text": delta, "n_tokens": n_tokens}, + started_unix=record.submitted_unix, emitted_unix=now, + input_watermark=record.prompt_tokens or 0, + )) + record.text = emitted_text + record.finished = True + await stream.put(ChunkEvent( + request_id=request_id, stage_id=0, seq=seq + 1, epoch=handle.epoch, + kind="done", payload={"output_tokens": record.output_tokens}, + started_unix=record.submitted_unix, emitted_unix=time.time(), final=True, + )) + except asyncio.CancelledError: + record.cancelled = True + record.text = emitted_text + raise + except Exception as exc: + record.error = str(exc) + record.text = emitted_text + with contextlib.suppress(Exception): + await stream.put(ChunkEvent( + request_id=request_id, stage_id=0, seq=seq + 1, epoch=handle.epoch, + kind="error", payload=None, started_unix=record.submitted_unix, + emitted_unix=time.time(), final=True, error=str(exc), + )) + finally: + self._inflight.pop(request_id, None) + with contextlib.suppress(Exception): + await stream.close() + + async def cancel(self, request_id: str) -> dict[str, Any]: + """Retire the epoch, stop the backend, release the reservation. + + The order matters and is the contract: by the time this returns, no + event produced before the cancel can still reach the consumer. + """ + record = self.records.get(request_id) + stream = self._streams.get(request_id) + session_id = self._inflight.get(request_id) or (record.session_id if record else None) + before = self.sampler.sample() + + if stream is not None: + stream.retire_epoch(record.epoch if record else 0) + + aborted = False + if self._omni is not None: + for name in ("abort", "abort_request", "abort_requests_async"): + fn = getattr(self._omni, name, None) + if callable(fn): + with contextlib.suppress(Exception): + result = fn(request_id) + if asyncio.iscoroutine(result): + await result + aborted = True + break + + # Close the stream *before* waiting on the producer. The consumer is + # blocked in ``get()`` and the sentinel is what releases it; making + # that depend on the backend's unwind would mean a backend that hangs + # hangs the caller too, which is the opposite of what cancel is for. + if stream is not None: + with contextlib.suppress(Exception): + await stream.close() + + task = self._pumps.pop(request_id, None) + pump_stopped = True + if task is not None: + task.cancel() + # Bounded: tearing down an async generator runs the backend's own + # cleanup, and that cleanup can wait on a request this method has + # already aborted. Waiting forever for it would turn a cancel into + # a deadlock, so the wait has a deadline and the outcome is + # reported rather than assumed. + done, pending = await asyncio.wait({task}, timeout=CANCEL_DRAIN_TIMEOUT_S) + pump_stopped = not pending + if done: + with contextlib.suppress(asyncio.CancelledError, Exception): + task.result() + + if record is not None: + record.cancelled = True + if session_id and session_id in self._sessions: + # A cancelled session's state is retired, not reused: the next + # request on it must start from a fresh epoch. + self._sessions[session_id] = self._sessions[session_id].next_epoch() + + # Give the backend a moment to free blocks, then look. + await asyncio.sleep(SAMPLE_INTERVAL_S * 2) + after = self.sampler.sample() + return { + "request_id": request_id, + "aborted_in_backend": aborted, + "producer_stopped": pump_stopped, + "epoch_retired": stream.epoch if stream is not None else None, + "dropped_stale_events": stream.dropped_stale if stream is not None else 0, + "output_tokens_at_cancel": record.output_tokens if record else 0, + "session_epoch_now": self._sessions[session_id].epoch if session_id in self._sessions else None, + "inflight_after": len(self._inflight), + "nvml_device_used_before": before.nvml_device_used, + "nvml_device_used_after": after.nvml_device_used, + "note": ( + "a cancelled request's KV blocks return to vLLM's pool, which is " + "pre-allocated; the device counter is therefore expected NOT to " + "fall. What must fall is the in-flight count, and no retired-epoch " + "event may be delivered." + ), + } + + # -- reporting ---------------------------------------------------------- + def report_placement(self) -> dict[str, Any]: + """Where the stage is actually configured and running. + + Two claims of different strength: the *configured* device, read back + from the stage's own ``VllmConfig`` rather than from the plan that + requested it, and the *observed* device memory delta across the load. + Neither is a per-process attribution, because NVML cannot give one + here; that limitation is in the output, not only in the docstring. + """ + configured: dict[str, Any] = {} + try: + configs = getattr(self._omni.engine, "stage_vllm_configs", None) or {} + cfg = configs[0] if isinstance(configs, (list, tuple)) else configs.get(0) + device_config = getattr(cfg, "device_config", None) + model_config = getattr(cfg, "model_config", None) + cache_config = getattr(cfg, "cache_config", None) + configured = { + "device": str(getattr(device_config, "device", "")), + "device_type": str(getattr(device_config, "device_type", "")), + "dtype": str(getattr(model_config, "dtype", "")), + "quantization": getattr(model_config, "quantization", None), + "max_model_len": getattr(model_config, "max_model_len", None), + "enforce_eager": getattr(model_config, "enforce_eager", None), + "kv_cache_memory_bytes": getattr(cache_config, "kv_cache_memory_bytes", None), + "num_gpu_blocks": getattr(cache_config, "num_gpu_blocks", None), + "block_size": getattr(cache_config, "block_size", None), + } + except Exception as exc: + configured = {"error": f"stage config unavailable: {exc}"} + + return { + "planned_device": self.plan.selected.device_id if self.plan.selected else None, + "planned_backend": self.plan.backend, + "configured": configured, + "observed": { + "nvml_device_used_delta_at_load": self.load_memory.get("nvml_device_used_delta"), + "tree_rss_delta_at_load": self.load_memory.get("tree_rss_delta"), + **gpu_attribution_for_tree(), + }, + } + + def report_usage(self) -> dict[str, Any]: + """Budgeted vs measured, per reservation, plus the sampled peaks.""" + return { + "load_seconds": self.load_seconds, + "reservations": [ + { + **r.to_dict(), + "upper_bound_over_budget": ( + r.actual_upper_bound_bytes / r.bytes + if (r.actual_upper_bound_bytes and r.bytes) + else None + ), + } + for r in self.plan.reservations + ], + "budgeted_peak_bytes": self.plan.peak_bytes, + "measured_peak": self.measured_peak(), + "load_memory": self.load_memory, + "peaks": self.sampler.peaks(), + "requests": [r.to_dict() for r in self.records.values()], + "streams": {rid: s.stats() for rid, s in self._streams.items()}, + } + + +def _first_completion(out: Any) -> dict[str, Any] | None: + """Pull ``(text, token_ids, prompt_tokens)`` out of an Omni output. + + Kept in one function because the shape is the thing most likely to move + between vLLM/Omni versions, and a version bump should break here, loudly, + rather than silently start reporting zero tokens. + """ + outputs = getattr(out, "outputs", None) + if not outputs: + return None + first = outputs[0] + text = getattr(first, "text", None) + token_ids = getattr(first, "token_ids", None) + if text is None and token_ids is None: + return None + prompt_ids = getattr(out, "prompt_token_ids", None) + return { + "text": text or "", + "token_ids": tuple(token_ids or ()), + "prompt_tokens": len(prompt_ids) if prompt_ids is not None else None, + } diff --git a/vllm_omni/edge/local/external/__init__.py b/vllm_omni/edge/local/external/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/vllm_omni/edge/local/external/client.py b/vllm_omni/edge/local/external/client.py new file mode 100644 index 000000000..148c615c8 --- /dev/null +++ b/vllm_omni/edge/local/external/client.py @@ -0,0 +1,403 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""The engine's half of an out-of-process stage. + +Binds a listener, starts the worker, waits for it to connect *back*, and then +speaks :mod:`~vllm_omni.edge.local.external.protocol` over that socket. The +connect-back direction is what keeps this deployable: WSL reaching into Windows +needs an inbound firewall rule, while Windows reaching into WSL does not. + +Two things here are load-bearing rather than defensive. + +**The placement gate.** :meth:`ExternalWorker.load` returns a +:class:`LoadReport` carrying the per-provider node counts ORT wrote while the +graph actually ran. ``fraction_on_target`` is the number the planner refuses +on. It has to be measured this way because every cheaper check is wrong: a +VitisAI session lists the EP in ``get_providers()`` whether or not it took any +nodes, and when it takes none the outputs are bit-identical to the CPU's -- so +neither the provider list nor the numbers can tell a real NPU run from a silent +fallback. Only the node assignment can. + +**The worker's own memory.** The engine's sampler walks a psutil process tree, +which stops at the WSL boundary; a Windows worker's resident set is invisible +to it. Since both AMD devices allocate out of the *same host RAM* the CPU +stage is budgeted against, an unreported worker is a hole in the admission +ledger, not a rounding error. So the worker reports its own RSS and +:meth:`stats` is how that reaches the ledger. +""" + +from __future__ import annotations + +import secrets +import socket +import subprocess +import sys +import time +from dataclasses import asdict, dataclass, field +from pathlib import Path +from typing import Any + +import numpy as np + +from vllm_omni.edge.local.external import launch as _launch +from vllm_omni.edge.local.external import protocol as _proto +from vllm_omni.edge.local.external.protocol import ProtocolError, WorkerError + +DEFAULT_START_TIMEOUT_S = 120.0 +"""Generous on purpose: the first VitisAI session on a cold machine spends +over a second in ``InferenceSession`` alone, and a Windows interpreter starting +off a ``\\\\wsl.localhost`` share is slower still.""" + +DEFAULT_CALL_TIMEOUT_S = 300.0 + +DEFAULT_CLOSE_TIMEOUT_S = 10.0 +"""Bounded, for the same reason M0 bounds its cancellation drain: a worker that +will not exit must not hold the engine's shutdown open.""" + + +@dataclass(frozen=True) +class LoadReport: + """What the worker did when it opened the graph, and where it put it.""" + + ep: str + graph_path: str + session_providers: tuple[str, ...] = () + available_providers: tuple[str, ...] = () + node_counts: dict[str, int] = field(default_factory=dict) + total_nodes: int = 0 + target_nodes: int = 0 + fraction_on_target: float | None = None + """``None`` means *unverified*, which is not the same as zero and must not + be rounded into it: no run happened, so ORT wrote no assignments.""" + session_create_s: float = 0.0 + warmup_s: float = 0.0 + rss_bytes: int = 0 + inputs: tuple[dict[str, Any], ...] = () + outputs: tuple[dict[str, Any], ...] = () + profile_path: str = "" + ep_directory: str = "" + requested_provider_missing: str = "" + fallback_note: str = "" + onnxruntime: str = "" + note: str = "" + + @classmethod + def from_body(cls, body: dict[str, Any], *, ep: str, graph_path: str) -> LoadReport: + return cls( + ep=ep, + graph_path=graph_path, + session_providers=tuple(body.get("session_providers") or ()), + available_providers=tuple(body.get("available_providers") or ()), + node_counts=dict(body.get("node_counts") or {}), + total_nodes=int(body.get("total_nodes") or 0), + target_nodes=int(body.get("target_nodes") or 0), + fraction_on_target=body.get("fraction_on_target"), + session_create_s=float(body.get("session_create_s") or 0.0), + warmup_s=float(body.get("warmup_s") or 0.0), + rss_bytes=int(body.get("rss_bytes") or 0), + inputs=tuple(body.get("inputs") or ()), + outputs=tuple(body.get("outputs") or ()), + profile_path=str(body.get("profile_path") or ""), + ep_directory=str(body.get("ep_directory") or ""), + requested_provider_missing=str(body.get("requested_provider_missing") or ""), + fallback_note=str(body.get("fallback_note") or ""), + onnxruntime=str(body.get("onnxruntime") or ""), + note=str(body.get("placement_unknown_reason") or ""), + ) + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + def summary(self) -> str: + if self.fraction_on_target is None: + placement = f"placement unverified ({self.note or 'no profiled run'})" + else: + split = ", ".join(f"{k}={v}" for k, v in sorted(self.node_counts.items())) + placement = f"{self.target_nodes}/{self.total_nodes} nodes on {self.ep} ({split})" + return ( + f"{Path(self.graph_path).name} on {self.ep}: {placement}; " + f"session {self.session_create_s:.3f}s, warmup {self.warmup_s:.3f}s, " + f"worker rss {self.rss_bytes / 2**20:.0f} MiB" + ) + + +@dataclass(frozen=True) +class RunTiming: + """Where one forward's wall time went. + + Split because the composition rule needs it split: a stage only earns its + place on a second device when the end-to-end gain beats the cost of getting + there, and ``transport_s`` is that cost measured rather than assumed. + """ + + worker_s: float + """Inside ``session.run`` on the far side.""" + round_trip_s: float + """Wall time here, from send to the last byte of the reply.""" + + @property + def transport_s(self) -> float: + """Serialisation, both socket traversals, and the worker's dispatch.""" + return max(0.0, self.round_trip_s - self.worker_s) + + def to_dict(self) -> dict[str, float]: + return { + "worker_s": self.worker_s, + "round_trip_s": self.round_trip_s, + "transport_s": self.transport_s, + } + + +class WorkerStartError(RuntimeError): + """The worker never connected back. Carries whatever it printed.""" + + +def _default_gateway() -> str | None: + """The far side of the default route, i.e. the Windows host under WSL2 NAT.""" + try: + lines = Path("/proc/net/route").read_text().splitlines()[1:] + except OSError: + return None + for line in lines: + fields = line.split() + if len(fields) > 2 and fields[1] == "00000000": + packed = int(fields[2], 16) + return ".".join(str((packed >> shift) & 0xFF) for shift in (0, 8, 16, 24)) + return None + + +def _bind_address(*, is_windows: bool) -> str: + """The address the worker should dial. + + A WSL worker gets loopback, which nothing outside this machine can reach. A + native-Windows worker cannot use loopback -- under WSL2's NAT the two live + on different network stacks -- so it needs this distro's address on the + host-only network (``172.21.141.90`` here), which is the source address the + routing table picks for the *default gateway*. + + The gateway specifically, not any address that happens to be handy. The + obvious candidate is the nameserver in ``/etc/resolv.conf``, and it is + wrong: WSL puts ``10.255.255.254`` on a loopback alias, so a probe aimed at + it reports ``10.255.255.254`` as the source, the listener binds an address + only this distro can reach, and the Windows worker gets ECONNREFUSED after + a silent start. That is exactly the failure this comment exists to prevent + a future edit from reintroducing. + """ + if not is_windows or _launch.is_wsl() is False: + return "127.0.0.1" + # 1.1.1.1 is a fallback route target, not a destination: UDP ``connect`` + # sends no packets, it only resolves which local address would be used. + for target in (_default_gateway(), "1.1.1.1"): + if not target: + continue + probe = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) + try: + probe.connect((target, 53)) + address = str(probe.getsockname()[0]) + except OSError: + continue + finally: + probe.close() + if not address.startswith("127."): + return address + return "127.0.0.1" + + +class ExternalWorker: + """A stage running in another interpreter, possibly on another OS.""" + + def __init__( + self, + route: _launch.Route, + *, + start_timeout_s: float = DEFAULT_START_TIMEOUT_S, + call_timeout_s: float = DEFAULT_CALL_TIMEOUT_S, + ) -> None: + if not route.available: + raise WorkerStartError(f"route {route.name} is not available: {route.reason}") + self.route = route + self.start_timeout_s = start_timeout_s + self.call_timeout_s = call_timeout_s + self._process: subprocess.Popen[bytes] | None = None + self._sock: socket.socket | None = None + self._listener: socket.socket | None = None + self.hello: dict[str, Any] = {} + self.report: LoadReport | None = None + + # -- lifecycle ---------------------------------------------------------- + + def start(self) -> dict[str, Any]: + """Launch the worker and complete the handshake. Returns its ``hello``.""" + token = secrets.token_hex(16) + host = _bind_address(is_windows=self.route.is_windows) + listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + listener.bind((host, 0)) + listener.listen(1) + listener.settimeout(self.start_timeout_s) + self._listener = listener + port = listener.getsockname()[1] + + argv = [ + str(self.route.interpreter), + _launch.to_worker_path(self.route.worker, is_windows=self.route.is_windows), + "--host", host, + "--port", str(port), + "--token", token, + ] + self._process = subprocess.Popen(argv, stdout=subprocess.PIPE, stderr=subprocess.PIPE) + + try: + conn, _ = listener.accept() + except socket.timeout as exc: + raise WorkerStartError( + f"worker did not connect back within {self.start_timeout_s:.0f}s: " + f"{self._drain_process()}" + ) from exc + finally: + listener.close() + self._listener = None + + conn.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) + conn.settimeout(self.call_timeout_s) + self._sock = conn + + op, body, _ = _proto.recv_message(conn) + if op != _proto.OP_HELLO: + self.close() + raise WorkerStartError(f"first message was {op!r}, expected {_proto.OP_HELLO!r}") + if body.get("token") != token: + # Somebody else reached the listener first. Refuse rather than + # speak to it: the whole point of the token. + self.close() + raise WorkerStartError("connect-back presented the wrong token; refusing the peer") + self.hello = body + return body + + def _drain_process(self) -> str: + process = self._process + if process is None: + return "worker was never started" + if process.poll() is None: + process.kill() + try: + out, err = process.communicate(timeout=5) + except subprocess.TimeoutExpired: # pragma: no cover - kill just ran + return f"worker pid {process.pid} did not exit" + text = (err or b"").decode("utf-8", "replace").strip() + if not text: + text = (out or b"").decode("utf-8", "replace").strip() + return f"exit={process.returncode} output={text[:2000]!r}" if text else f"exit={process.returncode}" + + def close(self, timeout_s: float = DEFAULT_CLOSE_TIMEOUT_S) -> None: + """Ask the worker to exit; kill it if it will not. Safe to call twice.""" + sock, self._sock = self._sock, None + if sock is not None: + try: + sock.settimeout(min(timeout_s, 5.0)) + _proto.send_message(sock, _proto.OP_CLOSE, {}) + _proto.recv_message(sock) + except (OSError, ProtocolError, WorkerError): + pass + finally: + sock.close() + if self._listener is not None: + self._listener.close() + self._listener = None + process, self._process = self._process, None + if process is not None and process.poll() is None: + try: + process.wait(timeout=timeout_s) + except subprocess.TimeoutExpired: + process.kill() + process.wait(timeout=5) + if process is not None: + for stream in (process.stdout, process.stderr): + if stream is not None: + stream.close() + + def __enter__(self) -> ExternalWorker: + self.start() + return self + + def __exit__(self, *exc: object) -> None: + self.close() + + # -- calls -------------------------------------------------------------- + + def _call( + self, + op: str, + body: dict[str, Any] | None = None, + tensors: dict[str, np.ndarray] | None = None, + ) -> tuple[dict[str, Any], dict[str, np.ndarray]]: + sock = self._sock + if sock is None: + raise WorkerStartError("worker is not running; call start() first") + try: + _proto.send_message(sock, op, body, tensors) + reply_op, reply_body, reply_tensors = _proto.recv_message(sock) + except (OSError, ProtocolError) as exc: + # The socket is the only liveness signal we have; a dead worker + # shows up as a broken pipe or a short read, and the useful part of + # the diagnosis is on its stderr. + detail = self._drain_process() + self._sock = None + raise WorkerStartError(f"worker connection failed ({exc}); {detail}") from exc + if reply_op != _proto.OP_OK: + raise ProtocolError(f"unexpected reply {reply_op!r} to {op!r}") + return reply_body, reply_tensors + + def load( + self, + graph_path: str | Path, + *, + example_inputs: dict[str, np.ndarray] | None = None, + profile: bool = True, + profile_prefix: str | Path | None = None, + ep_dir: str | Path | None = None, + intra_op_num_threads: int | None = None, + ) -> LoadReport: + """Open a graph and measure where its nodes actually ran. + + ``example_inputs`` are not optional in practice: without a run ORT + writes no node assignments, so the report comes back with + ``fraction_on_target is None`` and the planner treats the placement as + unverified rather than assuming it is fine. + """ + body: dict[str, Any] = { + "graph_path": _launch.to_worker_path(graph_path, is_windows=self.route.is_windows), + "ep": self.route.ep, + "profile": bool(profile), + } + if profile_prefix is not None: + body["profile_prefix"] = _launch.to_worker_path( + profile_prefix, is_windows=self.route.is_windows + ) + if ep_dir is not None: + body["ep_dir"] = str(ep_dir) + if intra_op_num_threads is not None: + body["intra_op_num_threads"] = int(intra_op_num_threads) + + reply, _ = self._call(_proto.OP_LOAD, body, example_inputs or {}) + self.report = LoadReport.from_body(reply, ep=self.route.ep, graph_path=str(graph_path)) + return self.report + + def run(self, inputs: dict[str, np.ndarray]) -> tuple[dict[str, np.ndarray], RunTiming]: + """One forward, with the transport cost separated from the compute.""" + started = time.perf_counter() + reply, outputs = self._call(_proto.OP_RUN, {}, inputs) + round_trip = time.perf_counter() - started + return outputs, RunTiming(worker_s=float(reply.get("run_s", 0.0)), round_trip_s=round_trip) + + def stats(self) -> dict[str, Any]: + """Counters and the worker's own resident set, for the ledger.""" + reply, _ = self._call(_proto.OP_STATS, {}) + return reply + + +def start_worker(route_name: str, **kwargs: Any) -> ExternalWorker: + """Resolve a route by name and start it.""" + worker = ExternalWorker(_launch.resolve(route_name), **kwargs) + worker.start() + return worker diff --git a/vllm_omni/edge/local/external/launch.py b/vllm_omni/edge/local/external/launch.py new file mode 100644 index 000000000..001a1a7b8 --- /dev/null +++ b/vllm_omni/edge/local/external/launch.py @@ -0,0 +1,288 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Which interpreter runs a stage worker, and why one is missing. + +Four routes reach the two AMD devices on this laptop, and no two of them can +share an interpreter: + +============== ========================== ==================================== +route interpreter why it is separate +============== ========================== ==================================== +``ort-vitisai`` native Windows, plain ORT the NPU is an MCDM device; WSL's + >= 1.25 GPU-PV forwards WDDM adapters only, + so ``/dev/accel`` does not exist +``ort-dml`` native Windows, ``onnxruntime-directml`` is pinned + onnxruntime-directml at 1.24.4 and cannot load the + VitisAI EP, which needs >= 1.25 +``torch-dml`` WSL, ``.venvs/dml`` ``torch-directml`` pins torch 2.4.1; + Omni runs 2.13 +``ort-cpu`` anything with onnxruntime the reference arm, and what the + tests use on a machine with no AMD + hardware +============== ========================== ==================================== + +Every lookup returns a :class:`Route` either way. A missing interpreter is a +first-class answer with a ``reason`` a person can act on, not an exception and +not silence: "the NPU is unusable" and "the NPU venv is not installed" lead to +different work, and the planner has to be able to tell the caller which it hit. +""" + +from __future__ import annotations + +import os +import shutil +import subprocess +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Any + +ROUTE_VITISAI = "ort-vitisai" +ROUTE_DML = "ort-dml" +ROUTE_TORCH_DML = "torch-dml" +ROUTE_CPU = "ort-cpu" +ROUTE_WIN_CUDA = "win-cuda" +"""[W4] Native-Windows vLLM on the NVIDIA GPU. + +The first route whose far side is **torch**, not onnxruntime, and the first +that exists to escape WSL rather than to reach a device WSL cannot see. It is +declared here so a Windows engine is a row in this table instead of a parallel +universe -- but see :data:`WIN_CUDA_PLACEMENT_EVIDENCE`: it cannot currently +pass the placement gate, and that is deliberate, not an oversight. +""" + +ENV_PREFIX = "VLLM_OMNI_EXTERNAL_PYTHON_" +"""``VLLM_OMNI_EXTERNAL_PYTHON_ORT_VITISAI=...`` overrides one route's +interpreter. Route names are upper-cased with ``-`` turned into ``_``.""" + +WORKER_ENV_PREFIX = "VLLM_OMNI_EXTERNAL_WORKER_" +"""``VLLM_OMNI_EXTERNAL_WORKER_ORT_CPU=...`` overrides one route's worker +script. The tests point this at a fake that speaks the protocol and never +imports onnxruntime, so the client, the framing and the refusal paths are +covered on a machine with no AMD hardware and no ORT at all.""" + +# Where the routes live on this machine. Defaults, not requirements: the env +# override above is the supported way to point at a different install, and the +# reason string names this list when nothing is found. +_DEFAULT_CANDIDATES: dict[str, tuple[str, ...]] = { + ROUTE_VITISAI: ( + r"/mnt/c/Users/zhout/npu-ep/Scripts/python.exe", + r"/mnt/c/Users/zhout/ryzenai/Scripts/python.exe", + ), + ROUTE_DML: ( + r"/mnt/c/Users/zhout/dml-ep/Scripts/python.exe", + r"/mnt/c/Users/zhout/directml/Scripts/python.exe", + ), + ROUTE_TORCH_DML: (".venvs/dml/bin/python",), + ROUTE_CPU: (".venvs/omni-cuda/bin/python", ".venvs/omni-cpu/bin/python"), + ROUTE_WIN_CUDA: ( + r"/mnt/c/Users/zhout/w0/w0venv/Scripts/python.exe", + r"/mnt/c/Users/zhout/vllm-win/Scripts/python.exe", + ), +} + +_WORKER_FOR_ROUTE = { + ROUTE_VITISAI: "worker_ort.py", + ROUTE_DML: "worker_ort.py", + ROUTE_CPU: "worker_ort.py", + ROUTE_TORCH_DML: "worker_dml.py", + ROUTE_WIN_CUDA: "worker_vllm.py", +} + +_EP_FOR_ROUTE = { + ROUTE_VITISAI: "vitisai", + ROUTE_DML: "dml", + ROUTE_CPU: "cpu", + ROUTE_TORCH_DML: "dml", + ROUTE_WIN_CUDA: "cuda", +} + + +WIN_CUDA_PLACEMENT_EVIDENCE: str | None = None +"""What a ``win-cuda`` worker can prove about where its work ran. Currently: nothing. + +Every other route in this table earns its placement from a *record the runtime +writes while the graph runs*: onnxruntime reports per-node provider assignments, +and ``torch-directml`` reports the device of the output tensor. M3 made that a +hard admission condition after a VitisAI session listed its EP, took **zero** +nodes, and returned output bit-identical to the CPU's -- so neither the provider +list nor the numbers can tell a real device run from a silent fallback. + +**vLLM has no analogue.** It does not emit a per-operation device map, and the +cheap substitutes are exactly the ones M3 proved worthless: ``current_platform`` +is a *selection*, not an observation; allocated VRAM says weights are resident, +not that the decode step ran on them; and plausible output is what a silent CPU +fallback also produces. + +``None`` therefore means *unverified*, which +:func:`vllm_omni.edge.local.external.stage` already refuses with +``REFUSE_EP_PLACEMENT`` -- the same answer the NPU gets when it has not been +measured. Refusing our own engine on the rule we wrote for a vendor's is the +point: the rule is about evidence, not about whose code it is. + +To make this route admissible, give it a *measured* device attribution -- a +worker that runs a marked op under a profiler and reports where it landed -- and +set this to that fraction. Do not set it to 1.0 because the code looks right. +""" + + +@dataclass(frozen=True) +class Route: + """One way to reach a device, resolved or not.""" + + name: str + interpreter: str | None + worker: str + """Absolute path to the worker script, in this filesystem's spelling.""" + ep: str + """The execution provider the worker should ask for.""" + is_windows: bool + """Whether paths handed to this interpreter need Windows spelling.""" + available: bool + reason: str = "" + extra: dict[str, Any] | None = None + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +def repo_root() -> Path: + """The checkout this package was imported from, when there is one. + + Used only to resolve the relative interpreter candidates above. An + installed wheel has no checkout around it, in which case the relative + routes simply do not resolve and say so. + """ + override = os.environ.get("VLLM_OMNI_REPO_ROOT") + if override: + return Path(override) + here = Path(__file__).resolve() + # ``.venvs`` first, and in a separate pass: ``vllm-omni`` is a nested + # checkout with its own ``.git``, so a single loop that accepts either + # marker stops one level too deep and every relative candidate misses. + for marker in (".venvs", ".git"): + for parent in here.parents: + if (parent / marker).is_dir(): + return parent + return here.parents[-1] + + +def is_wsl() -> bool: + """True when this process can launch native-Windows executables. + + The marker is interop, not the kernel string: a WSL2 kernel without + ``/proc/sys/fs/binfmt_misc/WSLInterop`` cannot exec a ``.exe``, and that is + the capability the caller actually needs. + """ + if os.name == "nt": + return False + return Path("/proc/sys/fs/binfmt_misc/WSLInterop").exists() or Path( + "/proc/sys/fs/binfmt_misc/WSLInterop-late" + ).exists() + + +def to_worker_path(path: str | Path, *, is_windows: bool) -> str: + """Spell a path the way the worker's interpreter will understand it. + + A native-Windows worker running a script that lives in WSL sees it as a + ``\\\\wsl.localhost\\...`` UNC path. ``wslpath -w`` is the translation, and + it is the reason a Windows process can load graphs and write profiles + straight into the Linux tree with no copying. + """ + text = str(path) + if not is_windows or os.name == "nt": + return text + if not shutil.which("wslpath"): + return text + try: + out = subprocess.run( + ["wslpath", "-w", text], capture_output=True, text=True, timeout=20, check=True + ) + except (subprocess.SubprocessError, OSError): + return text + return out.stdout.strip() or text + + +def _resolve_interpreter(route: str, root: Path) -> tuple[str | None, str]: + env_key = f"{ENV_PREFIX}{route.upper().replace('-', '_')}" + override = os.environ.get(env_key) + if override: + if Path(override).exists(): + return override, "" + return None, f"{env_key}={override!r} does not exist" + + tried: list[str] = [] + for candidate in _DEFAULT_CANDIDATES.get(route, ()): + path = Path(candidate) if Path(candidate).is_absolute() else root / candidate + tried.append(str(path)) + if path.exists(): + return str(path), "" + return None, f"no interpreter found; tried {tried} and {env_key} is unset" + + +def resolve(route: str, *, root: Path | None = None) -> Route: + """Find the interpreter for one route, or say why it is not there.""" + if route not in _WORKER_FOR_ROUTE: + raise ValueError(f"unknown route {route!r}; expected one of {sorted(_WORKER_FOR_ROUTE)}") + root = root or repo_root() + is_windows = route in (ROUTE_VITISAI, ROUTE_DML, ROUTE_WIN_CUDA) + worker_key = f"{WORKER_ENV_PREFIX}{route.upper().replace('-', '_')}" + worker = os.environ.get(worker_key) or str( + Path(__file__).resolve().parent / _WORKER_FOR_ROUTE[route] + ) + + if is_windows and not is_wsl() and os.name != "nt": + return Route( + name=route, + interpreter=None, + worker=worker, + ep=_EP_FOR_ROUTE[route], + is_windows=True, + available=False, + reason=( + "this route needs a native-Windows interpreter, and this " + "process is neither on Windows nor in a WSL distro with " + "interop enabled, so it cannot start one" + ), + ) + + interpreter, reason = _resolve_interpreter(route, root) + # An interpreter without a worker script is not a route. Checked here + # rather than at start(), because ``available`` is what the device + # enumeration turns into ``runnable`` -- reporting a device as runnable and + # then failing to launch it is the same silent-substitution mistake this + # module exists to avoid, one layer down. + if interpreter is not None and not Path(worker).is_file(): + return Route( + name=route, + interpreter=interpreter, + worker=worker, + ep=_EP_FOR_ROUTE[route], + is_windows=is_windows, + available=False, + reason=f"worker script not found: {worker}", + ) + return Route( + name=route, + interpreter=interpreter, + worker=worker, + ep=_EP_FOR_ROUTE[route], + is_windows=is_windows, + available=interpreter is not None, + reason=reason, + ) + + +def resolve_all(*, root: Path | None = None) -> list[Route]: + """Every route, resolved or not, in a stable order.""" + root = root or repo_root() + return [resolve(name, root=root) for name in (ROUTE_VITISAI, ROUTE_DML, ROUTE_TORCH_DML, ROUTE_CPU)] + + +def describe(routes: list[Route]) -> str: + lines = [] + for route in routes: + mark = "available" if route.available else "missing" + lines.append(f" [{mark:^9}] {route.name:<12} ep={route.ep:<8} {route.interpreter or ''}") + if route.reason: + lines.append(f" {route.reason}") + return "\n".join(lines) if lines else " (none)" diff --git a/vllm_omni/edge/local/external/protocol.py b/vllm_omni/edge/local/external/protocol.py new file mode 100644 index 000000000..ea890c809 --- /dev/null +++ b/vllm_omni/edge/local/external/protocol.py @@ -0,0 +1,185 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""The wire between the engine and a stage that runs in another process. + +Deliberately tiny, and deliberately **stdlib + numpy only**. The worker on the +other end of this socket may be a native-Windows interpreter that has no torch, +no vllm and no vllm_omni on its path -- ``C:\\Users\\zhout\\npu-ep`` is exactly +that: Python 3.12, onnxruntime, onnx, numpy, nothing else. Anything imported +here has to be installable there, so this module imports neither. + +**Why a socket at all.** Two devices on this laptop cannot be reached from the +engine's own interpreter, and for different reasons: + +* the **XDNA2 NPU** is an MCDM device. WSL2's GPU-PV forwards WDDM display + adapters only, so there is no ``/dev/accel`` inside WSL and no amount of + installing changes that. VitisAI is a Windows DLL. +* the **Radeon 890M** is reachable from WSL through DirectML, but + ``torch-directml`` pins torch 2.4.1 and Omni runs 2.13. + +So out-of-process is forced, not preferred. One protocol covers both, plus the +native-Windows DirectML route, which needs its own venv anyway because +``onnxruntime-directml`` is pinned at 1.24.4 and cannot load the VitisAI EP +(that needs ORT >= 1.25) -- the two AMD devices cannot share an interpreter. + +**Framing.** ``[4-byte big-endian header length][JSON header][raw tensor bytes]``. +Tensors are described in the header (name, dtype, shape, nbytes) and their +buffers are concatenated after it in the order listed, so a receive is two +reads and one ``frombuffer`` per tensor -- no per-element work and no pickle. +Measured on this machine, WSL <-> native Windows: 11 us half round trip, +330 MB/s. That is ample for a stage boundary (a 2 s Code2Wav chunk is ~192 KB, +about 0.6 ms) and nowhere near enough for per-layer traffic inside a decode +loop, which is the arithmetic behind "coarse stages only". + +``pickle`` is not used anywhere in here on purpose: the header is JSON, the +payload is raw buffers, so a hostile or merely broken peer cannot execute code +by replying. +""" + +from __future__ import annotations + +import json +import socket +import struct +from typing import Any + +import numpy as np + +# Ops the client sends. +OP_LOAD = "load" +"""Open a session on a graph. Replies with the placement report.""" +OP_RUN = "run" +"""One forward. Inputs in, outputs out.""" +OP_STATS = "stats" +"""Counters and RSS. The worker's memory is invisible to the engine's sampler +when the worker is a Windows process, so it has to report its own.""" +OP_CLOSE = "close" + +# Ops the worker sends. +OP_HELLO = "hello" +"""First message after connect-back: interpreter, versions, available EPs.""" +OP_OK = "ok" +OP_ERR = "err" + +MAX_HEADER_BYTES = 1 << 20 +"""A header is metadata; a megabyte of it is a bug or an attack, not a graph.""" + +MAX_PAYLOAD_BYTES = 2 << 30 +"""Refuse rather than allocate on a corrupt length. Two gibibytes is far above +any legitimate stage payload here -- the largest is a vision tower's patch +tensor, single-digit MB.""" + + +class ProtocolError(RuntimeError): + """The peer sent something this protocol cannot represent.""" + + +class WorkerError(RuntimeError): + """The peer ran and failed. Carries the remote traceback when there is one.""" + + def __init__(self, message: str, *, remote_traceback: str = "", code: str = "") -> None: + # The remote traceback goes in the message, not only in an attribute: + # a bare "AssertionError:" surfacing here with the actual stack one + # attribute away is a debugging loop nobody should have to discover. + detail = f"{message}\n--- worker traceback ---\n{remote_traceback}" if remote_traceback else message + super().__init__(detail) + self.message = message + self.remote_traceback = remote_traceback + self.code = code + + +def _recv_exactly(sock: socket.socket, n: int) -> bytes: + """Read exactly ``n`` bytes or raise. + + ``recv`` is allowed to return short reads on a stream socket and does, once + payloads get past a page or two; a partial read here would be silently + misparsed as the next tensor. + """ + if n == 0: + return b"" + chunks: list[bytes] = [] + remaining = n + while remaining: + chunk = sock.recv(min(remaining, 1 << 20)) + if not chunk: + raise ProtocolError(f"peer closed with {remaining} of {n} bytes outstanding") + chunks.append(chunk) + remaining -= len(chunk) + return b"".join(chunks) if len(chunks) > 1 else chunks[0] + + +def send_message( + sock: socket.socket, + op: str, + body: dict[str, Any] | None = None, + tensors: dict[str, np.ndarray] | None = None, +) -> None: + """Send one message. ``tensors`` are sent as raw buffers after the header.""" + tensors = tensors or {} + described: list[dict[str, Any]] = [] + buffers: list[memoryview] = [] + for name, array in tensors.items(): + array = np.asarray(array) + # Record the shape *before* making it contiguous: ``ascontiguousarray`` + # promotes a 0-d array to shape (1,), which would silently turn a + # scalar graph input into a length-1 vector on the far side. + shape = list(array.shape) + contiguous = np.ascontiguousarray(array) + described.append( + { + "name": name, + # ``str(dtype)`` loses endianness for multibyte types; ``.str`` + # keeps the '<'/'>' so a big-endian peer would still decode. + "dtype": contiguous.dtype.str, + "shape": shape, + "nbytes": int(contiguous.nbytes), + } + ) + # ``reshape(-1)`` rather than a direct cast: ``memoryview.cast`` refuses + # a 0-dimensional source. + buffers.append(memoryview(contiguous.reshape(-1)).cast("B")) + + header = json.dumps({"op": op, "body": body or {}, "tensors": described}).encode("utf-8") + if len(header) > MAX_HEADER_BYTES: + raise ProtocolError(f"header is {len(header)} bytes, over the {MAX_HEADER_BYTES} limit") + sock.sendall(struct.pack(">I", len(header))) + sock.sendall(header) + for buffer in buffers: + sock.sendall(buffer) + + +def recv_message(sock: socket.socket) -> tuple[str, dict[str, Any], dict[str, np.ndarray]]: + """Receive one message as ``(op, body, tensors)``. + + Raises :class:`WorkerError` on an ``err`` reply so callers do not each have + to remember to check; every other op is returned as-is. + """ + (header_len,) = struct.unpack(">I", _recv_exactly(sock, 4)) + if header_len > MAX_HEADER_BYTES: + raise ProtocolError(f"peer announced a {header_len}-byte header, over the limit") + try: + header = json.loads(_recv_exactly(sock, header_len).decode("utf-8")) + except (UnicodeDecodeError, json.JSONDecodeError) as exc: + raise ProtocolError(f"header is not JSON: {exc}") from exc + + total = sum(int(t["nbytes"]) for t in header.get("tensors", ())) + if total > MAX_PAYLOAD_BYTES: + raise ProtocolError(f"peer announced a {total}-byte payload, over the limit") + + tensors: dict[str, np.ndarray] = {} + for spec in header.get("tensors", ()): + raw = _recv_exactly(sock, int(spec["nbytes"])) + array = np.frombuffer(raw, dtype=np.dtype(spec["dtype"])) + # ``frombuffer`` is read-only and aliases ``raw``; a consumer that + # writes in place (ORT's IOBinding does) needs its own buffer. + tensors[spec["name"]] = array.reshape(spec["shape"]).copy() + + op = str(header.get("op", "")) + body = header.get("body") or {} + if op == OP_ERR: + raise WorkerError( + str(body.get("message", "worker failed with no message")), + remote_traceback=str(body.get("traceback", "")), + code=str(body.get("code", "")), + ) + return op, body, tensors diff --git a/vllm_omni/edge/local/external/stage.py b/vllm_omni/edge/local/external/stage.py new file mode 100644 index 000000000..11bc082b6 --- /dev/null +++ b/vllm_omni/edge/local/external/stage.py @@ -0,0 +1,522 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Planning and opening a stage that runs on the 890M or the NPU. + +Same contract as :mod:`vllm_omni.edge.local.plan`, one device further out: pick +a device, budget it, and either place the work or refuse with a code and a +remedy. Three things differ, and each is the reason for a refusal code that a +checkpoint-on-vLLM plan never needs. + +**These devices execute graphs, not checkpoints.** The gate is the exported +artifact's format against the device's, so "no export exists yet" is its own +answer (:data:`~vllm_omni.edge.local.plan.REFUSE_NO_ARTIFACT`) rather than a +vague unsupported -- the work it implies is an export, not a port. + +**Placement has to be measured after the fact.** A VitisAI session lists the EP +in ``get_providers()`` whether or not it took a single node, and a graph it +declined runs entirely on the CPU and returns bit-identical numbers. So neither +the provider list nor the outputs can distinguish a real NPU run from a silent +fallback; only ORT's per-node assignments can, and those exist only after +something has run. :meth:`ExternalStage.open` therefore loads, runs once under +the profiler, and refuses on the node split. An *unverified* placement is +refused too: ``None`` is not rounded to "fine". + +**The memory is not on a private pool.** Both devices allocate out of the same +host RAM the CPU stage is budgeted against, and the worker is a different +process -- on Windows, a different OS -- so the engine's psutil sampler cannot +see it at all. The worker reports its own resident set and it lands in the same +ledger; an unreported worker is a hole in the budget, not a rounding error. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +import numpy as np + +from vllm_omni.edge.local.capabilities import MASK_ENV, DeviceCapability +from vllm_omni.edge.local.external import launch as _launch +from vllm_omni.edge.local.external.client import ExternalWorker, LoadReport, RunTiming +from vllm_omni.edge.local.manifest import GraphArtifact +from vllm_omni.edge.local.plan import ( + EXTERNAL_RESERVE_BYTES, + MIN_FRACTION_ON_TARGET, + REFUSE_CAPACITY, + REFUSE_EP_PLACEMENT, + REFUSE_NO_ARTIFACT, + REFUSE_NO_DEVICE, + REFUSE_ROUTE, + MemoryReservation, + Refusal, +) + +MiB = 2**20 + +WORKER_BASE_BYTES: dict[str, int] = { + # Measured on this laptop, 2026-09-15: a bare Windows 3.12 interpreter is + # 28.5 MiB resident, and the same process holding one VitisAI session on a + # 2 MiB graph is 270 MiB. Nearly all of that is the execution provider, not + # the graph, so it is charged per worker rather than scaled by artifact + # size. E, with a measured basis; ``open()`` replaces it with the worker's + # own number as soon as there is one. + _launch.ROUTE_VITISAI: 320 * MiB, + # Measured 369 MiB holding the 27B vision tower, against an earlier 256 MiB + # guess -- the budget has to bound the measurement, so it was raised rather + # than the overrun explained away. + _launch.ROUTE_DML: 512 * MiB, + _launch.ROUTE_CPU: 128 * MiB, + # torch-directml drags a whole torch build in beside the D3D12 runtime. + _launch.ROUTE_TORCH_DML: 768 * MiB, +} + +WORKER_WORKSPACE_BYTES = 256 * MiB +"""Activations and the EP's scratch. E: not separable from the worker's RSS by +any counter available here, so it is budgeted and then subsumed by the measured +total rather than checked on its own.""" + + +class PlacementRefused(RuntimeError): + """The stage could not be placed. Carries the refusal, not just a message.""" + + def __init__(self, refusal: Refusal, report: LoadReport | None = None) -> None: + super().__init__(f"[{refusal.code}] {refusal.message}") + self.refusal = refusal + self.report = report + + +@dataclass +class StagePlan: + """Where one exported graph would run, and what it would cost.""" + + artifact: GraphArtifact + device: DeviceCapability | None + route: _launch.Route | None + reservations: list[MemoryReservation] = field(default_factory=list) + refusals: list[Refusal] = field(default_factory=list) + min_fraction_on_target: float = MIN_FRACTION_ON_TARGET + report: LoadReport | None = None + + @property + def admitted(self) -> bool: + return self.device is not None and self.route is not None + + @property + def budget_bytes(self) -> int: + return sum(r.bytes for r in self.reservations) + + def to_dict(self) -> dict[str, Any]: + return { + "artifact": self.artifact.to_dict(), + "device": self.device.to_dict() if self.device else None, + "route": self.route.to_dict() if self.route else None, + "reservations": [r.to_dict() for r in self.reservations], + "refusals": [r.to_dict() for r in self.refusals], + "min_fraction_on_target": self.min_fraction_on_target, + "budget_bytes": self.budget_bytes, + "admitted": self.admitted, + "report": self.report.to_dict() if self.report else None, + } + + def summary(self) -> str: + head = f"{self.artifact.component} ({self.artifact.fmt}, {self.artifact.bytes / MiB:.1f} MiB)" + if not self.admitted: + lines = [f"{head}: REFUSED"] + lines += [f" [{r.code}] {r.device_id}: {r.message}\n remedy: {r.remedy}" for r in self.refusals] + return "\n".join(lines) + assert self.device is not None and self.route is not None + lines = [ + f"{head} -> {self.device.device_id} via {self.route.name}", + f" budget {self.budget_bytes / MiB:.0f} MiB on {self.device.memory_pool} " + f"(shared with every other device on this pool)", + ] + for reservation in self.reservations: + actual = "" + if reservation.actual_upper_bound_bytes is not None: + measured = reservation.actual_upper_bound_bytes + # A budget that does not bound its measurement is a failed + # budget. Saying so here is the whole discipline: the number is + # supposed to be a provable ceiling, not a guess that is usually + # close. + verdict = "measured <=" if measured <= reservation.bytes else "OVER, measured" + actual = f", {verdict} {measured / MiB:.0f} MiB" + lines.append( + f" {reservation.purpose:<18} {reservation.bytes / MiB:>7.0f} MiB " + f"[{reservation.evidence}] {reservation.lifetime}{actual}" + ) + if self.report is not None: + lines.append(f" {self.report.summary()}") + for refusal in self.refusals: + lines.append(f" also considered: [{refusal.code}] {refusal.device_id}: {refusal.message}") + return "\n".join(lines) + + +def _reservations_for(artifact: GraphArtifact, route: _launch.Route, pool: str) -> list[MemoryReservation]: + return [ + MemoryReservation( + pool=pool, + purpose="graph weights", + bytes=artifact.bytes, + lifetime="session", + reclaimable="no", + evidence="D", + note=f"the exported file, {artifact.fmt}, sha {artifact.sha256[:12]}", + ), + MemoryReservation( + pool=pool, + purpose="worker runtime", + bytes=WORKER_BASE_BYTES.get(route.name, 256 * MiB), + lifetime="session", + reclaimable="on_release", + evidence="E", + note=( + "interpreter plus execution provider in the worker process; " + "measured 270 MiB for a VitisAI session on a 2 MiB graph and " + "369 MiB for a DirectML session on the 27B vision tower. " + "NOTE: on an integrated GPU the device allocations come out of " + "this same system RAM and do not appear in the worker's working " + "set, so the measured figure is a lower bound for this pool, " + "not a total" + ), + ), + MemoryReservation( + pool=pool, + purpose="worker workspace", + bytes=WORKER_WORKSPACE_BYTES, + lifetime="request", + reclaimable="on_release", + evidence="E", + note="activations and EP scratch; not separable by any counter here", + ), + MemoryReservation( + pool=pool, + purpose="external reserve", + bytes=EXTERNAL_RESERVE_BYTES, + lifetime="session", + reclaimable="no", + evidence="E", + note="host RAM this engine does not control; the Windows side included", + ), + ] + + +def host_pool_capacity(devices: list[DeviceCapability]) -> int: + """Bytes available on the shared host pool, taken from the row that owns it. + + The iGPU and the NPU report ``memory_bytes == 0`` on purpose: all three + devices allocate out of the CPU's RAM, and carrying the number on each row + would let a caller add three copies of the same memory together. So the + capacity a stage is checked against is the host row's, which is also what + the CPU stage is being budgeted against at the same time. + """ + for device in devices: + if device.memory_pool == "host_ram" and device.memory_bytes: + available = int((device.extra or {}).get("ram_available_bytes") or 0) + return available or device.memory_bytes + for device in devices: + hinted = int((device.extra or {}).get("host_pool_available_bytes") or 0) + if hinted: + return hinted + return 0 + + +def plan_external_stage( + artifact: GraphArtifact, + devices: list[DeviceCapability], + *, + prefer: str | None = None, + require: str | None = None, + min_fraction_on_target: float = MIN_FRACTION_ON_TARGET, + pool_capacity_bytes: int | None = None, +) -> StagePlan: + """Choose a device for one exported graph, or refuse with reasons. + + ``prefer`` names a ``device_id`` to try first and falls through to the + others if it is refused. ``require`` names the only one that may be used, + so a refusal is the answer rather than a quieter substitution -- which is + what you want when the question is "does the NPU run this", not "does + anything run this". Answering the first question with the iGPU would be the + silent-fallback mistake this module exists to prevent, one level up. + + Every device that was rejected leaves a refusal behind even when a later + one succeeds, so the record says what was considered, not only what won. + """ + plan = StagePlan(artifact=artifact, device=None, route=None, + min_fraction_on_target=min_fraction_on_target) + + candidates = [d for d in devices if d.kind in ("gpu_integrated", "npu")] + if require: + candidates = [d for d in candidates if d.device_id == require] + if not candidates: + plan.refusals.append( + Refusal( + REFUSE_NO_DEVICE, + require, + f"{require!r} was required and is not among the enumerated " + f"external devices ({[d.device_id for d in devices]})", + "run `python -m vllm_omni.edge.local devices` to see what this " + "process can reach, and why anything missing is missing", + ) + ) + return plan + elif prefer: + candidates.sort(key=lambda d: d.device_id != prefer) + if not candidates: + plan.refusals.append( + Refusal( + REFUSE_NO_DEVICE, + "-", + "no integrated GPU or NPU was enumerated on this machine", + "run `python -m vllm_omni.edge.local devices`; if the hardware is " + "present but absent here, the probe could not open it and its " + "reason field says why", + ) + ) + return plan + + for device in candidates: + if not device.runnable: + plan.refusals.append( + Refusal( + REFUSE_ROUTE, + device.device_id, + device.reason or "no worker route resolves from this process", + _route_remedy(device), + ) + ) + continue + if artifact.fmt not in device.weight_formats: + plan.refusals.append( + Refusal( + REFUSE_NO_ARTIFACT, + device.device_id, + f"{device.device_id} executes {sorted(device.weight_formats)}, " + f"and this graph is {artifact.fmt}", + _format_remedy(artifact, device), + ) + ) + continue + + route_name = (device.extra or {}).get("worker_route") + route = _launch.resolve(str(route_name)) if route_name else None + if route is None or not route.available: + plan.refusals.append( + Refusal( + REFUSE_ROUTE, + device.device_id, + f"the route recorded for this device ({route_name!r}) no longer resolves", + "re-run device enumeration; an interpreter was removed between " + "enumeration and planning", + ) + ) + continue + + reservations = _reservations_for(artifact, route, device.memory_pool) + capacity = ( + pool_capacity_bytes + if pool_capacity_bytes is not None + else host_pool_capacity(devices) + ) + budget = sum(r.bytes for r in reservations) + if capacity and budget > capacity: + plan.refusals.append( + Refusal( + REFUSE_CAPACITY, + device.device_id, + f"the stage budgets {budget / MiB:.0f} MiB and the shared host " + f"pool has {capacity / MiB:.0f} MiB available; this is the same " + "pool the CPU stage is budgeted against, not a pool of its own", + "free host memory, or keep this stage on the device that " + "already holds the model so the graph is not resident twice", + ) + ) + continue + + plan.device = device + plan.route = route + plan.reservations = reservations + break + + return plan + + +def _route_remedy(device: DeviceCapability) -> str: + if device.masked: + # A masked device is unreachable *by request*. Explaining the MCDM + # forwarding rules here would send someone after a problem they do not + # have; the fix is one flag. + return ( + f"drop {device.kind!r} from --mask / {MASK_ENV} to re-enable it. The " + "device itself was not tested -- masking is how the no-accelerator " + "deployment class is exercised without removing hardware" + ) + routes = (device.extra or {}).get("routes") or [] + missing = [r for r in routes if not r.get("available")] + if device.kind == "npu": + return ( + "the NPU is not reachable from inside WSL at all -- it is an MCDM " + "device and GPU-PV forwards display adapters only. It needs a " + "native-Windows interpreter with onnxruntime >= 1.25 (not " + "onnxruntime-directml, which is pinned at 1.24.4 and cannot load the " + "VitisAI EP). Point VLLM_OMNI_EXTERNAL_PYTHON_ORT_VITISAI at one." + + (f" Tried: {[r['reason'] for r in missing]}" if missing else "") + ) + return ( + "the iGPU needs either a native-Windows venv with onnxruntime-directml " + "(VLLM_OMNI_EXTERNAL_PYTHON_ORT_DML) or a WSL venv with torch-directml, " + "which pins torch 2.4.1 and so cannot share Omni's environment " + "(VLLM_OMNI_EXTERNAL_PYTHON_TORCH_DML)." + + (f" Tried: {[r['reason'] for r in missing]}" if missing else "") + ) + + +def _format_remedy(artifact: GraphArtifact, device: DeviceCapability) -> str: + if device.kind == "npu": + return ( + f"export {artifact.component} as onnx:a16w8 at opset >= 21 -- 16-bit " + "activations with 8-bit weights, via npu_ryzenai.quantization_kwargs(). " + "A8W8, which is what quantize_static does by default, and fp32 are " + "both declined by the XDNA2 overlays with no error and every node " + "left on the CPU." + ) + return ( + f"export {artifact.component} as onnx:fp16 or onnx:fp32; DirectML has no " + "path for this graph's precision" + ) + + +class ExternalStage: + """A planned stage, opened on its worker, with the placement gate enforced.""" + + def __init__(self, plan: StagePlan) -> None: + if not plan.admitted: + raise PlacementRefused( + plan.refusals[0] + if plan.refusals + else Refusal(REFUSE_NO_DEVICE, "-", "no device was selected", "see the plan") + ) + self.plan = plan + assert plan.route is not None + self._worker = ExternalWorker(plan.route) + self._open = False + + def open( + self, + example_inputs: dict[str, np.ndarray], + *, + profile_dir: str | Path | None = None, + ) -> LoadReport: + """Start the worker, load the graph, and verify where it actually ran. + + ``example_inputs`` are required, not optional: ORT writes node + assignments only for nodes that have executed, so without a run there + is no placement to check and the gate would be passing on faith. + """ + if not example_inputs: + raise ValueError( + "example_inputs are required: placement is verified from a " + "profiled run, and without one the node assignments do not exist" + ) + self._worker.start() + self._open = True + prefix = Path(profile_dir) / f"{self.plan.artifact.component}_" if profile_dir else None + report = self._worker.load( + self.plan.artifact.path, example_inputs=example_inputs, profile_prefix=prefix + ) + self.plan.report = report + + refusal = verify_placement(self.plan, report) + if refusal is not None: + self.plan.refusals.append(refusal) + self.close() + raise PlacementRefused(refusal, report) + + # Close the ledger loop the way M0 does: the budget was an upper bound, + # and this is the measurement it is checked against. + for reservation in self.plan.reservations: + if reservation.purpose == "worker runtime" and report.rss_bytes: + reservation.actual_upper_bound_bytes = report.rss_bytes + return report + + def run(self, inputs: dict[str, np.ndarray]) -> tuple[dict[str, np.ndarray], RunTiming]: + if not self._open: + raise RuntimeError("stage is not open; call open() first") + return self._worker.run(inputs) + + def stats(self) -> dict[str, Any]: + return self._worker.stats() + + def close(self) -> None: + self._open = False + self._worker.close() + + def __enter__(self) -> ExternalStage: + return self + + def __exit__(self, *exc: object) -> None: + self.close() + + +def verify_placement(plan: StagePlan, report: LoadReport) -> Refusal | None: + """The gate. ``None`` means the device really ran the graph. + + Two failures, deliberately not merged: a graph the EP *declined* and a + graph whose placement was never *measured* are different situations with + different remedies, and collapsing the second into the first would let an + unverified run be reported as a verified one. + """ + device_id = plan.device.device_id if plan.device else "-" + if report.requested_provider_missing: + # The most precise form of the failure, so it is reported before the + # node count -- "the provider never initialised" is actionable, while + # "0 of 1171 nodes" leaves the reader to guess whether the graph was + # declined op by op or the EP was never there at all. + return Refusal( + REFUSE_EP_PLACEMENT, + device_id, + report.fallback_note or f"{report.requested_provider_missing} did not initialise", + _placement_remedy(plan, report), + ) + if report.fraction_on_target is None: + return Refusal( + REFUSE_EP_PLACEMENT, + device_id, + "placement was not measured: no profiled run produced node " + f"assignments ({report.note or 'no run'})", + "call open() with example inputs so the graph executes once under " + "the ORT profiler; an unverified placement is refused rather than " + "assumed, because a session that ran entirely on the CPU looks " + "identical from every cheaper angle", + ) + if report.fraction_on_target < plan.min_fraction_on_target: + split = ", ".join(f"{k}={v}" for k, v in sorted(report.node_counts.items())) + return Refusal( + REFUSE_EP_PLACEMENT, + device_id, + f"{report.ep} took {report.target_nodes} of {report.total_nodes} nodes " + f"({report.fraction_on_target:.0%}, below the {plan.min_fraction_on_target:.0%} " + f"required); the split was {split}", + _placement_remedy(plan, report), + ) + return None + + +def _placement_remedy(plan: StagePlan, report: LoadReport) -> str: + if plan.device is not None and plan.device.kind == "npu": + return ( + "the graph is not in a shape XDNA2 partitions. Check, in this order: " + "the quantization is A16W8 (QUInt16 activations, QInt8 weights) at " + "opset >= 21; the EP loaded is onnxruntime_vitisai_ep.dll and not the " + "RyzenAI *Light* provider beside it, which registers cleanly and " + "claims zero nodes from every graph; and that the graph is not mostly " + "requantization, since the partitioner leaves dequantize nodes on the " + "CPU by its own account. Running it anyway would be a CPU fallback " + "reported as NPU execution." + ) + return ( + "DirectML declined most of this graph. Check for ops with no DML kernel " + "(integer Gather in an RVQ codebook lookup is the one this project has " + "hit) and consider exporting the variant that takes the embedding as an " + "input instead." + ) diff --git a/vllm_omni/edge/local/external/worker_dml.py b/vllm_omni/edge/local/external/worker_dml.py new file mode 100644 index 000000000..ceb9e05ab --- /dev/null +++ b/vllm_omni/edge/local/external/worker_dml.py @@ -0,0 +1,290 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""A torch-DirectML stage worker for the Radeon 890M. + +Speaks the same protocol as :mod:`worker_ort`, in its own interpreter for the +usual reason: ``torch-directml`` pins torch 2.4.1 and Omni runs 2.13, so the +two cannot share an environment. Here that interpreter is a WSL venv +(``.venvs/dml``) rather than a Windows one -- the 890M is a WDDM display +adapter, which GPU-PV *does* forward, so DirectML reaches it over ``/dev/dxg`` +without leaving Linux. The NPU is the one that cannot be reached this way. + +**It consumes torch exports, not ONNX.** That is not a second way of doing the +same thing: ``vllm_omni.edge.decoder_export`` already emits ``.pt2`` via +``torch.export``, and the alternative iGPU route (onnxruntime-directml on +Windows, served by ``worker_ort``) needs its own venv because that package is +pinned at 1.24.4 and cannot load the VitisAI EP the NPU needs. Two routes to +one device, two artifact formats; the plan records which one ran. + +**What placement evidence means here.** There is no per-node execution-provider +report, because there is no partitioner -- torch dispatches to the device the +tensors are on. So the evidence is the device the outputs actually came back +from: a module that silently fell back to the CPU returns CPU tensors, and +``fraction_on_target`` is 0. It is a coarser signal than ORT's node counts and +it catches the failure that matters, which is the same one: work reported on an +accelerator that ran on the CPU. +""" + +from __future__ import annotations + +import argparse +import importlib.util +import os +import socket +import sys +import time +import traceback +from pathlib import Path +from typing import Any + +import numpy as np + +_HERE = Path(__file__).resolve().parent +spec = importlib.util.spec_from_file_location("_edge_external_protocol", _HERE / "protocol.py") +proto = importlib.util.module_from_spec(spec) +sys.modules["_edge_external_protocol"] = proto +spec.loader.exec_module(proto) + +DEVICE_TYPE = "privateuseone" +"""What torch calls a DirectML tensor's device. Checking for this string is how +a CPU fallback is caught.""" + + +def _rss_bytes() -> int: + try: + for line in Path("/proc/self/status").read_text().splitlines(): + if line.startswith("VmRSS:"): + return int(line.split()[1]) * 1024 + except OSError: + pass + return 0 + + +def _peak_rss_bytes() -> int: + try: + for line in Path("/proc/self/status").read_text().splitlines(): + if line.startswith("VmHWM:"): + return int(line.split()[1]) * 1024 + except OSError: + pass + return 0 + + +class _Session: + def __init__(self) -> None: + self.module: Any = None + self.device: Any = None + self.dtype: Any = None + self.input_order: list[str] = [] + self.output_names: list[str] = [] + self.runs = 0 + self.total_run_s = 0.0 + self.report: dict[str, Any] = {} + + +def _select_device(index: int | None) -> tuple[Any, str]: + """The 890M, by name rather than by index. + + ``torch_directml`` enumerates every DirectML adapter, and on this laptop + that includes the RTX 5090 at index 0 -- taking index 0 would quietly + benchmark the discrete GPU and call it the iGPU. + """ + import torch_directml + + count = torch_directml.device_count() + if index is not None: + return torch_directml.device(index), torch_directml.device_name(index) + for i in range(count): + name = torch_directml.device_name(i) + if "890M" in name or "Radeon" in name: + return torch_directml.device(i), name + raise RuntimeError( + "no integrated Radeon adapter among " + f"{[torch_directml.device_name(i) for i in range(count)]}; pass device_index " + "explicitly if the intended adapter is named differently" + ) + + +def _load(state: _Session, body: dict[str, Any], tensors: dict[str, np.ndarray]) -> dict[str, Any]: + import torch + + graph = str(body["graph_path"]) + device, device_name = _select_device(body.get("device_index")) + state.device = device + # No cast unless one is asked for. The artifact's precision is part of its + # identity -- it is in the manifest and in the format string -- so quietly + # casting an fp16 export to fp32 because that is the friendlier default + # would make the record describe a graph that did not run. + requested = body.get("dtype") + state.dtype = getattr(torch, str(requested)) if requested else None + + started = time.perf_counter() + if not graph.endswith(".pt2"): + # Measured on this laptop, torch-directml 0.2.5 / torch 2.4.1: a + # TorchScript module cannot be put on a DirectML device by either + # route. ``torch.jit.load(map_location=dml)`` raises + # "Cannot access storage of OpaqueTensorImpl", and loading to CPU then + # ``.to(device)`` trips ``assert isinstance(param, Parameter)`` inside + # ``Module._apply``, because a RecursiveScriptModule's parameters are + # not ``Parameter`` instances. Refused with the reason rather than + # left to surface as a bare AssertionError from inside torch. + raise ValueError( + f"{Path(graph).name}: torch-directml cannot load a TorchScript " + "module onto the iGPU (both map_location and .to(device) fail " + "inside torch). Export with torch.export to a .pt2 instead -- " + "which is what vllm_omni.edge.decoder_export already produces." + ) + module = torch.export.load(graph).module() + if state.dtype is not None: + module = module.to(state.dtype) + # No ``.eval()``: an exported program's module raises + # "Calling eval() is not supported yet" on torch 2.4.1, and it is already + # captured in inference mode, so there is nothing to switch off. + module = module.to(device) + state.module = module + create_s = time.perf_counter() - started + + state.input_order = list(body.get("input_order") or tensors.keys()) + state.output_names = list(body.get("output_names") or []) + + info: dict[str, Any] = { + "ep": "dml", + "device_name": device_name, + "torch": torch.__version__, + "session_create_s": create_s, + "session_providers": ["torch-directml"], + "available_providers": ["torch-directml"], + "inputs": [{"name": n, "shape": list(tensors[n].shape)} for n in state.input_order + if n in tensors], + } + + if tensors: + started = time.perf_counter() + outputs = _forward(state, tensors) + info["warmup_s"] = time.perf_counter() - started + on_device = sum(1 for t in outputs if getattr(t, "device", None) is not None + and t.device.type == DEVICE_TYPE) + total = len(outputs) + info["node_counts"] = {"dml": on_device, "cpu": total - on_device} + info["total_nodes"] = total + info["target_nodes"] = on_device + info["fraction_on_target"] = (on_device / total) if total else 0.0 + info["outputs"] = [{"name": f"out_{i}", "shape": list(t.shape)} + for i, t in enumerate(outputs)] + if not state.output_names: + state.output_names = [f"out_{i}" for i in range(total)] + else: + info["node_counts"] = {} + info["fraction_on_target"] = None + info["placement_unknown_reason"] = ( + "no example inputs were sent, so nothing ran and no output device " + "was observed; placement is unverified, not verified-empty" + ) + + info["rss_bytes"] = _rss_bytes() + state.report = info + return info + + +def _forward(state: _Session, tensors: dict[str, np.ndarray]) -> list[Any]: + import torch + + args = [] + for name in state.input_order: + if name not in tensors: + continue + tensor = torch.from_numpy(np.ascontiguousarray(tensors[name])) + if state.dtype is not None and tensor.is_floating_point(): + tensor = tensor.to(state.dtype) + args.append(tensor.to(state.device)) + with torch.no_grad(): + result = state.module(*args) + if isinstance(result, torch.Tensor): + return [result] + if isinstance(result, dict): + state.output_names = list(result.keys()) + return list(result.values()) + return list(result) + + +def _run(state: _Session, tensors: dict[str, np.ndarray]) -> tuple[dict[str, Any], dict[str, np.ndarray]]: + if state.module is None: + raise RuntimeError("run before load") + started = time.perf_counter() + outputs = _forward(state, tensors) + # ``.cpu()`` is the synchronisation point: DirectML queues work, so timing + # without the copy back would measure submission, not execution. + materialised = [t.detach().to("cpu").numpy() for t in outputs] + elapsed = time.perf_counter() - started + state.runs += 1 + state.total_run_s += elapsed + names = state.output_names or [f"out_{i}" for i in range(len(materialised))] + return {"run_s": elapsed, "runs": state.runs}, dict(zip(names, materialised)) + + +def serve(host: str, port: int, token: str) -> int: + sock = socket.create_connection((host, port), timeout=60.0) + sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) + state = _Session() + try: + import torch + import torch_directml + + versions = { + "torch": torch.__version__, + "torch_directml": getattr(torch_directml, "__version__", "unknown"), + "available_providers": [ + torch_directml.device_name(i) for i in range(torch_directml.device_count()) + ], + } + except Exception as exc: # pragma: no cover - reported, not raised + versions = {"torch": f"unavailable: {exc}", "available_providers": []} + + proto.send_message( + sock, proto.OP_HELLO, + {"token": token, "executable": sys.executable, "python": sys.version, + "platform": sys.platform, "pid": os.getpid(), "onnxruntime": "n/a (torch-directml)", + "numpy": np.__version__, "rss_bytes": _rss_bytes(), **versions}, + ) + + while True: + try: + op, body, tensors = proto.recv_message(sock) + except proto.ProtocolError: + return 0 + try: + if op == proto.OP_LOAD: + proto.send_message(sock, proto.OP_OK, _load(state, body, tensors)) + elif op == proto.OP_RUN: + reply, outputs = _run(state, tensors) + proto.send_message(sock, proto.OP_OK, reply, outputs) + elif op == proto.OP_STATS: + proto.send_message(sock, proto.OP_OK, { + "runs": state.runs, "total_run_s": state.total_run_s, + "rss_bytes": _rss_bytes(), "peak_rss_bytes": _peak_rss_bytes(), + "report": state.report, + }) + elif op == proto.OP_CLOSE: + proto.send_message(sock, proto.OP_OK, {}) + return 0 + else: + raise ValueError(f"unknown op {op!r}") + except Exception as exc: + proto.send_message(sock, proto.OP_ERR, { + "message": f"{type(exc).__name__}: {exc}", + "code": type(exc).__name__, + "traceback": traceback.format_exc(), + }) + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--host", required=True) + parser.add_argument("--port", type=int, required=True) + parser.add_argument("--token", default="") + args = parser.parse_args(argv) + return serve(args.host, args.port, args.token) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/vllm_omni/edge/local/external/worker_ort.py b/vllm_omni/edge/local/external/worker_ort.py new file mode 100644 index 000000000..3bbf6a3d0 --- /dev/null +++ b/vllm_omni/edge/local/external/worker_ort.py @@ -0,0 +1,370 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""An ONNX Runtime stage worker, run by an interpreter the engine cannot use. + +Started by :mod:`vllm_omni.edge.local.external.client`, which binds a listener +first and passes its address in; this process connects *back*. That direction +is deliberate: a Windows process reaching into WSL needs no firewall rule, +while WSL reaching into Windows does. + +Run with **stdlib + numpy + onnxruntime only**. It is executed by +``C:\\Users\\zhout\\npu-ep\\Scripts\\python.exe`` -- Python 3.12.10, +onnxruntime 1.30.0, onnx 1.22.0, numpy 2.5.3 -- which has no torch, no vllm and +no vllm_omni. The one thing it does borrow is +:mod:`vllm_omni.edge.npu_ryzenai`, loaded **by file path** rather than +imported: that module is plain stdlib (it only imports ``onnxruntime`` inside +its functions), so loading it as a standalone file reuses the EP recipe instead +of copying it, without dragging in the package ``__init__``. + +**What "it ran on the NPU" means here.** A VitisAI session that lists the EP in +``get_providers()`` may still be executing every node on the CPU, and when it +does its outputs are bit-identical to the CPU's -- which is precisely what a +*correct* NPU run does not look like. So identical numbers cannot confirm +placement and small differences cannot confirm it either. The only evidence is +the per-node provider assignment, which ORT emits into its profile, so +:func:`_load` always does one profiled warm-up run and reports the counts. The +engine's planner refuses on that number; this process does not decide. +""" + +from __future__ import annotations + +import argparse +import importlib.util +import json +import os +import socket +import sys +import time +import traceback +from pathlib import Path +from typing import Any + +import numpy as np + +_HERE = Path(__file__).resolve().parent +_PROTOCOL = _HERE / "protocol.py" +_NPU_RYZENAI = _HERE.parents[1] / "npu_ryzenai.py" + + +def _load_sibling(path: Path, name: str) -> Any: + """Import a single .py file without importing its package. + + ``vllm_omni.edge.npu_ryzenai`` cannot be imported normally here: the + package ``__init__`` chain reaches vllm, which is not installed on the + Windows side. The module itself has no such dependency. + """ + spec = importlib.util.spec_from_file_location(name, path) + if spec is None or spec.loader is None: # pragma: no cover - path is ours + raise ImportError(f"cannot load {path}") + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + return module + + +proto = _load_sibling(_PROTOCOL, "_edge_external_protocol") + + +def _memory_bytes() -> tuple[int, int]: + """``(resident, peak resident)`` for this process, on either OS. + + The engine's :class:`~vllm_omni.edge.local.engine.MemorySampler` walks a + psutil process tree, and that walk stops at the WSL boundary -- a Windows + worker's memory is simply invisible to it. Since both AMD devices allocate + out of the *same host RAM* the CPU stage is budgeted against, an unreported + worker is a hole in the admission ledger rather than a rounding error. So + the worker reports its own. + + The Windows path needs explicit ``argtypes``/``restype``: ``GetCurrentProcess`` + returns a pseudo-handle that ctypes truncates to a 32-bit int without them, + and ``GetProcessMemoryInfo`` then fails and returns zero -- which reads as + "the worker used no memory" instead of "the query failed". + """ + if sys.platform == "win32": + import ctypes + from ctypes import wintypes + + class _Counters(ctypes.Structure): + _fields_ = [ + ("cb", wintypes.DWORD), + ("PageFaultCount", wintypes.DWORD), + ("PeakWorkingSetSize", ctypes.c_size_t), + ("WorkingSetSize", ctypes.c_size_t), + ("QuotaPeakPagedPoolUsage", ctypes.c_size_t), + ("QuotaPagedPoolUsage", ctypes.c_size_t), + ("QuotaPeakNonPagedPoolUsage", ctypes.c_size_t), + ("QuotaNonPagedPoolUsage", ctypes.c_size_t), + ("PagefileUsage", ctypes.c_size_t), + ("PeakPagefileUsage", ctypes.c_size_t), + ] + + kernel32 = ctypes.WinDLL("kernel32", use_last_error=True) + kernel32.GetCurrentProcess.restype = wintypes.HANDLE + kernel32.GetCurrentProcess.argtypes = [] + # Since Windows 7 the psapi entry points are forwarders onto kernel32's + # K32* symbols; prefer those and keep psapi as the fallback. + query = getattr(kernel32, "K32GetProcessMemoryInfo", None) + if query is None: + query = ctypes.WinDLL("psapi", use_last_error=True).GetProcessMemoryInfo + query.restype = wintypes.BOOL + query.argtypes = [wintypes.HANDLE, ctypes.POINTER(_Counters), wintypes.DWORD] + + counters = _Counters() + counters.cb = ctypes.sizeof(_Counters) + if not query(kernel32.GetCurrentProcess(), ctypes.byref(counters), counters.cb): + return 0, 0 + return int(counters.WorkingSetSize), int(counters.PeakWorkingSetSize) + + resident = peak = 0 + try: + for line in Path("/proc/self/status").read_text().splitlines(): + if line.startswith("VmRSS:"): + resident = int(line.split()[1]) * 1024 + elif line.startswith("VmHWM:"): + peak = int(line.split()[1]) * 1024 + except OSError: + pass + return resident, peak + + +def _rss_bytes() -> int: + return _memory_bytes()[0] + + +class _Session: + """One loaded graph, plus the placement evidence for it.""" + + def __init__(self) -> None: + self.session: Any = None + self.report: dict[str, Any] = {} + self.runs = 0 + self.total_run_s = 0.0 + self.output_names: list[str] = [] + + +def _provider_counts(profile_path: str) -> dict[str, int]: + """Per-EP node counts out of an ORT profile. + + Same rule as :func:`vllm_omni.edge.npu_ryzenai.nodes_on_npu`, generalised + to every provider so the report can show the split rather than one number: + a graph that is 3 nodes on the NPU and 6 on the CPU is a very different + result from 9 and 0, and only the split says which one happened. + """ + try: + events = json.loads(Path(profile_path).read_text()) + except (OSError, json.JSONDecodeError): + return {} + counts: dict[str, int] = {} + for event in events: + if event.get("cat") != "Node": + continue + provider = (event.get("args") or {}).get("provider") + if provider: + counts[str(provider)] = counts.get(str(provider), 0) + 1 + return counts + + +def _make_session(body: dict[str, Any]) -> tuple[Any, dict[str, Any]]: + """Build the ORT session for the requested execution provider.""" + import onnxruntime as ort + + ep = str(body.get("ep", "cpu")) + graph = str(body["graph_path"]) + info: dict[str, Any] = {"ep": ep, "onnxruntime": ort.__version__} + + options = ort.SessionOptions() + if body.get("profile"): + options.enable_profiling = True + if body.get("profile_prefix"): + options.profile_file_prefix = str(body["profile_prefix"]) + if body.get("intra_op_num_threads"): + options.intra_op_num_threads = int(body["intra_op_num_threads"]) + + if ep == "vitisai": + # The recipe, not a copy of it: the EP directory has to go on the DLL + # search path before registration or the load fails with an unhelpful + # "no dependency", and the package directory denies listing so it is + # found through Get-AppxPackage. + npu = _load_sibling(_NPU_RYZENAI, "_edge_npu_ryzenai") + ep_dir = body.get("ep_dir") + name = npu.register(Path(ep_dir) if ep_dir else None) + merged = npu.session_options(name) + merged.enable_profiling = options.enable_profiling + if body.get("profile_prefix"): + merged.profile_file_prefix = str(body["profile_prefix"]) + options = merged + info["ep_directory"] = str(body.get("ep_dir") or npu.find_ep_directory()) + providers = [name] + elif ep == "dml": + providers = ["DmlExecutionProvider"] + elif ep == "cpu": + providers = ["CPUExecutionProvider"] + else: + raise ValueError(f"unknown execution provider {ep!r}; expected vitisai, dml or cpu") + + available = list(ort.get_available_providers()) + info["available_providers"] = available + if ep == "dml" and "DmlExecutionProvider" not in available: + raise RuntimeError( + "DmlExecutionProvider is not in this interpreter's onnxruntime. " + "The DirectML build is a separate package (onnxruntime-directml) " + "and cannot share a venv with the plain onnxruntime that VitisAI " + "needs, because that one is pinned at 1.24.4 and VitisAI wants " + ">= 1.25. Two venvs, one per device." + ) + + started = time.perf_counter() + session = ort.InferenceSession(graph, sess_options=options, providers=providers) + info["session_create_s"] = time.perf_counter() - started + info["session_providers"] = list(session.get_providers()) + + # ORT's Python binding does not fail when a requested provider cannot + # initialise: it prints "Falling back to ['CPUExecutionProvider'] and + # retrying" and hands back a working CPU session. Observed here with + # DirectML raising E_INVALIDARG on a Reshape carrying allowzero=1 -- the + # session then ran the whole vision tower on the CPU and returned correct + # features. The node-assignment gate downstream catches that, but only + # after a profiled run; naming it at the source turns a 0%-placement number + # into the actual reason. + if providers and providers[0] not in info["session_providers"]: + info["requested_provider_missing"] = providers[0] + info["fallback_note"] = ( + f"{providers[0]} failed to initialise and onnxruntime silently fell " + f"back to {info['session_providers']}. Anything this session computes " + "runs there, not on the requested device." + ) + return session, info + + +def _load(state: _Session, body: dict[str, Any], tensors: dict[str, np.ndarray]) -> dict[str, Any]: + session, info = _make_session(body) + state.session = session + state.output_names = [o.name for o in session.get_outputs()] + + info["inputs"] = [ + {"name": i.name, "type": i.type, "shape": list(i.shape)} for i in session.get_inputs() + ] + info["outputs"] = [ + {"name": o.name, "type": o.type, "shape": list(o.shape)} for o in session.get_outputs() + ] + + # The placement evidence. One profiled run, because ORT only writes node + # assignments once nodes have actually executed -- a session that has never + # run has no placement to report, only an intention. + if body.get("profile") and tensors: + started = time.perf_counter() + session.run(state.output_names, {k: v for k, v in tensors.items()}) + info["warmup_s"] = time.perf_counter() - started + profile_path = session.end_profiling() + info["profile_path"] = str(profile_path) + counts = _provider_counts(str(profile_path)) + info["node_counts"] = counts + total = sum(counts.values()) + target = str(body.get("ep", "cpu")) + target_key = {"vitisai": "vitisai", "dml": "DmlExecutionProvider", "cpu": "CPUExecutionProvider"}[target] + info["total_nodes"] = total + info["target_nodes"] = counts.get(target_key, 0) + info["fraction_on_target"] = (counts.get(target_key, 0) / total) if total else 0.0 + elif body.get("profile"): + info["node_counts"] = {} + info["fraction_on_target"] = None + info["placement_unknown_reason"] = ( + "no example inputs were sent, so no run happened and ORT wrote no " + "node assignments; placement is unverified, not verified-empty" + ) + + info["rss_bytes"] = _rss_bytes() + state.report = info + return info + + +def _run(state: _Session, tensors: dict[str, np.ndarray]) -> tuple[dict[str, Any], dict[str, np.ndarray]]: + if state.session is None: + raise RuntimeError("run before load") + started = time.perf_counter() + outputs = state.session.run(state.output_names, dict(tensors)) + elapsed = time.perf_counter() - started + state.runs += 1 + state.total_run_s += elapsed + return {"run_s": elapsed, "runs": state.runs}, dict(zip(state.output_names, outputs)) + + +def serve(host: str, port: int, token: str) -> int: + sock = socket.create_connection((host, port), timeout=60.0) + sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) + sock.settimeout(None) + state = _Session() + try: + import onnxruntime as ort + + ort_version, providers = ort.__version__, list(ort.get_available_providers()) + except Exception as exc: # pragma: no cover - reported, not raised + ort_version, providers = f"unavailable: {exc}", [] + + proto.send_message( + sock, + proto.OP_HELLO, + { + "token": token, + "executable": sys.executable, + "python": sys.version, + "platform": sys.platform, + "pid": os.getpid(), + "onnxruntime": ort_version, + "numpy": np.__version__, + "available_providers": providers, + "rss_bytes": _rss_bytes(), + }, + ) + + while True: + try: + op, body, tensors = proto.recv_message(sock) + except proto.ProtocolError: + return 0 # the engine went away; that is a normal shutdown here + try: + if op == proto.OP_LOAD: + proto.send_message(sock, proto.OP_OK, _load(state, body, tensors)) + elif op == proto.OP_RUN: + reply, outputs = _run(state, tensors) + proto.send_message(sock, proto.OP_OK, reply, outputs) + elif op == proto.OP_STATS: + proto.send_message( + sock, + proto.OP_OK, + { + "runs": state.runs, + "total_run_s": state.total_run_s, + "rss_bytes": _memory_bytes()[0], + "peak_rss_bytes": _memory_bytes()[1], + "report": state.report, + }, + ) + elif op == proto.OP_CLOSE: + proto.send_message(sock, proto.OP_OK, {}) + return 0 + else: + raise ValueError(f"unknown op {op!r}") + except Exception as exc: + proto.send_message( + sock, + proto.OP_ERR, + { + "message": f"{type(exc).__name__}: {exc}", + "code": type(exc).__name__, + "traceback": traceback.format_exc(), + }, + ) + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--host", required=True) + parser.add_argument("--port", type=int, required=True) + parser.add_argument("--token", default="") + args = parser.parse_args(argv) + return serve(args.host, args.port, args.token) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/vllm_omni/edge/local/manifest.py b/vllm_omni/edge/local/manifest.py new file mode 100644 index 000000000..8ff7c7b5e --- /dev/null +++ b/vllm_omni/edge/local/manifest.py @@ -0,0 +1,405 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""What was run, precisely enough to run it again. + +A performance number in this repo is worthless without the row that produced +it, and the audit behind the 2026-09-15 proposal found three separate ways the +row goes wrong: + +* **A git SHA is not what executed.** ``vllm-omni`` is checked out *and* + installed, and the installed wheel's ``interfaces.py``, UVA offloader and + Qwen model files hashed differently from the checkout at the recorded SHA. + So this records both, plus which file the import actually resolved to. +* **A model name is not a checkpoint.** ``Qwen3.8-27B`` names three artifacts + (official FP8, NVIDIA NVFP4, community GGUF) that behave differently; Spark + names nineteen directories under ``models/``. The manifest is keyed on the + directory contents, not the name. +* **A checkpoint is not its weights alone.** The tokenizer and the chat + template decide what the model is even asked, and they change independently + of the weights. + +Weights are identified by name, size and count rather than by content hash: +digesting 7.7 GB on every plan costs ~20 s and answers a question nobody asked +at plan time. ``digest_weights=True`` does the full sha256 when the question is +actually being asked -- publishing a number, or chasing a checkpoint that +changed underneath a result. +""" + +from __future__ import annotations + +import hashlib +import importlib +import json +import os +import platform +import subprocess +import sys +from dataclasses import asdict, dataclass, field +from pathlib import Path +from typing import Any + +# Files that decide what the model is asked, as opposed to how it answers. +# Hashed in full on every manifest: together they are a few hundred KB. +_TOKENIZER_FILES = ( + "tokenizer.json", + "tokenizer_config.json", + "special_tokens_map.json", + "chat_template.jinja", + "vocab.json", + "merges.txt", + "generation_config.json", +) + + +def sha256_file(path: Path, chunk: int = 1 << 20) -> str: + h = hashlib.sha256() + with path.open("rb") as fh: + while block := fh.read(chunk): + h.update(block) + return h.hexdigest() + + +def _sha256_text(text: str) -> str: + return hashlib.sha256(text.encode("utf-8")).hexdigest() + + +@dataclass(frozen=True) +class RuntimeVersions: + """The interpreter, the wheels, and the checkouts -- all three.""" + + python: str + platform: str + torch: str | None + vllm: str | None + vllm_omni: str | None + vllm_omni_file: str + """Where ``import vllm_omni`` actually landed. An editable install and a + site-packages copy are different code with the same version string.""" + vllm_file: str + checkouts: dict[str, str] = field(default_factory=dict) + """``{repo_path: git describe}`` for the source trees, when they are git + repositories. Dirty trees are marked; this repo's engine work happens on a + dirty tree by design and a clean-looking SHA would be a lie.""" + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +def _module_version(name: str) -> tuple[str | None, str]: + try: + mod = importlib.import_module(name) + except Exception: + return None, "" + return getattr(mod, "__version__", None), getattr(mod, "__file__", "") or "" + + +def _git_describe(repo: Path) -> str | None: + if not (repo / ".git").exists(): + return None + try: + sha = subprocess.run( + ["git", "-C", str(repo), "rev-parse", "HEAD"], + capture_output=True, text=True, timeout=10, check=True, + ).stdout.strip() + dirty = subprocess.run( + ["git", "-C", str(repo), "status", "--porcelain"], + capture_output=True, text=True, timeout=20, check=True, + ).stdout.strip() + except Exception: + return None + return f"{sha}{'-dirty' if dirty else ''}" + + +def runtime_versions() -> RuntimeVersions: + torch_v, _ = _module_version("torch") + vllm_v, vllm_file = _module_version("vllm") + omni_v, omni_file = _module_version("vllm_omni") + + checkouts: dict[str, str] = {} + for file in (omni_file, vllm_file): + if not file: + continue + # walk up out of the package to the repo root + for parent in Path(file).resolve().parents: + described = _git_describe(parent) + if described is not None: + checkouts[str(parent)] = described + break + + return RuntimeVersions( + python=sys.version.split()[0], + platform=platform.platform(), + torch=torch_v, + vllm=vllm_v, + vllm_omni=omni_v, + vllm_omni_file=omni_file, + vllm_file=vllm_file, + checkouts=checkouts, + ) + + +@dataclass(frozen=True) +class ArtifactManifest: + """One checkpoint, identified well enough to be told from its siblings.""" + + model_dir: str + model_type: str + architectures: tuple[str, ...] + dtype: str | None + weight_format: str + """One of :mod:`vllm_omni.edge.local.capabilities`'s ``FORMAT_*``.""" + quantization: dict[str, Any] | None + """The checkpoint's own ``quantization_config``, verbatim. Kept whole + because ``int-quantized`` vs ``float-quantized`` and the ``ignore`` list + are what decide whether a device can run it.""" + config_sha256: str + tokenizer_sha256: str + """One digest over every tokenizer/template file present, in a fixed order, + each entry salted with its filename so a rename is not invisible.""" + tokenizer_files: tuple[str, ...] + weight_files: tuple[str, ...] + weight_bytes: int + weight_sha256: str | None + """Only when ``digest_weights=True``; ``None`` otherwise, never a stand-in.""" + revision: str | None + """From a sibling HF cache ``refs``/snapshot path when the directory came + from one. ``None`` for a locally produced checkpoint, which is the honest + answer -- most of ``models/`` was quantized here.""" + runtime: RuntimeVersions + notes: list[str] = field(default_factory=list) + + @property + def artifact_id(self) -> str: + """A short stable key for a performance profile. + + Deliberately built from config + tokenizer + weight *sizes*, so two + quantizations of one model never collide, and re-running the same plan + on the same directory reproduces the key without reading 7.7 GB. + """ + seed = f"{self.config_sha256}:{self.tokenizer_sha256}:{self.weight_bytes}:{len(self.weight_files)}" + return _sha256_text(seed)[:16] + + def to_dict(self) -> dict[str, Any]: + d = asdict(self) + d["architectures"] = list(self.architectures) + d["tokenizer_files"] = list(self.tokenizer_files) + d["weight_files"] = list(self.weight_files) + d["artifact_id"] = self.artifact_id + return d + + +def _detect_weight_format(config: dict[str, Any]) -> str: + from vllm_omni.edge.local.capabilities import ( + FORMAT_AWQ, + FORMAT_DENSE, + FORMAT_FP8, + FORMAT_GPTQ, + FORMAT_INT8, + FORMAT_NVFP4, + ) + + quant = config.get("quantization_config") + if not quant: + return FORMAT_DENSE + method = str(quant.get("quant_method", "")).lower() + if method == "compressed-tensors": + # ``format`` is the discriminator: ``int-quantized`` is the int8 W8A8 + # build that SM >= 100 cannot execute, ``float-quantized`` is fp8. + fmt = str(quant.get("format", "")).lower() + if "float" in fmt or "fp8" in fmt: + return FORMAT_FP8 + return FORMAT_INT8 + if method == "gptq": + return FORMAT_GPTQ + if method == "awq": + return FORMAT_AWQ + if method in ("modelopt", "modelopt_fp4", "nvfp4"): + return FORMAT_NVFP4 + return method or FORMAT_DENSE + + +def _hf_revision(model_dir: Path) -> str | None: + """Recover the revision when the directory is inside an HF cache snapshot. + + Layout is ``.../models--org--name/snapshots//``. A directory that + is not one returns ``None`` rather than a guess. + """ + parts = model_dir.resolve().parts + if "snapshots" in parts: + idx = parts.index("snapshots") + if idx + 1 < len(parts): + return parts[idx + 1] + return None + + +def build_manifest(model_dir: str | os.PathLike[str], *, digest_weights: bool = False) -> ArtifactManifest: + """Read a checkpoint directory into an :class:`ArtifactManifest`.""" + root = Path(model_dir) + config_path = root / "config.json" + if not config_path.is_file(): + raise FileNotFoundError(f"{root} has no config.json; not a checkpoint directory") + config_text = config_path.read_text(encoding="utf-8") + config = json.loads(config_text) + + tokenizer_present = [name for name in _TOKENIZER_FILES if (root / name).is_file()] + tok_hash = hashlib.sha256() + for name in tokenizer_present: + tok_hash.update(name.encode("utf-8")) + tok_hash.update(sha256_file(root / name).encode("ascii")) + + weight_paths = sorted(root.glob("*.safetensors")) + weight_bytes = sum(p.stat().st_size for p in weight_paths) + weight_digest: str | None = None + if digest_weights and weight_paths: + h = hashlib.sha256() + for p in weight_paths: + h.update(p.name.encode("utf-8")) + h.update(sha256_file(p).encode("ascii")) + weight_digest = h.hexdigest() + + notes: list[str] = [] + if not weight_paths: + notes.append("no *.safetensors in the directory; weight_bytes is 0 and no size budget can be formed") + if weight_digest is None and weight_paths: + notes.append("weights identified by name+size, not content; pass digest_weights=True to hash them") + + return ArtifactManifest( + model_dir=str(root.resolve()), + model_type=str(config.get("model_type", "")), + architectures=tuple(config.get("architectures") or ()), + dtype=config.get("dtype") or config.get("torch_dtype"), + weight_format=_detect_weight_format(config), + quantization=config.get("quantization_config"), + config_sha256=_sha256_text(config_text), + tokenizer_sha256=tok_hash.hexdigest(), + tokenizer_files=tuple(tokenizer_present), + weight_files=tuple(p.name for p in weight_paths), + weight_bytes=weight_bytes, + weight_sha256=weight_digest, + revision=_hf_revision(root), + runtime=runtime_versions(), + notes=notes, + ) + + +# ---- exported graphs ------------------------------------------------------- + + +@dataclass(frozen=True) +class GraphArtifact: + """An exported graph, with everything needed to trust and to reproduce it. + + :class:`ArtifactManifest` describes a checkpoint; this describes what an + exporter made *from* one. The extra fields are not bookkeeping. AGENTS.md + requires a quantized or compiled artifact to carry its version, layout, + calibration and numeric record, because the same model exported twice is + two different things -- and this project has the scars: a Qwen3-TTS vocoder + quantized to int8 measured -4.3 dB with correlation 0.00 while its fp16 + export was faithful at 39.2 dB. A file called ``code2wav.onnx`` says + nothing about which of those two it is. + + ``parity`` is therefore part of the artifact's identity, not a test result + filed next to it: an export with no numeric record is not usable, and a + failure constrains *this artifact* rather than the model or the precision. + """ + + path: str + sha256: str + bytes: int + fmt: str + """``onnx:fp32`` / ``onnx:fp16`` / ``onnx:a16w8`` -- matched against a + device's ``weight_formats``.""" + opset: int + """A16W8 needs >= 21, where 16-bit QDQ is expressible without contrib ops.""" + source_model: str + """The checkpoint this came from.""" + source_revision: str | None + component: str + """Which part of the model: ``code2wav``, ``vision_tower``, ``spark_sliding``.""" + exporter: str + """Dotted path of the module that produced it, so the export is repeatable.""" + input_spec: tuple[dict[str, Any], ...] = () + output_spec: tuple[dict[str, Any], ...] = () + calibration: dict[str, Any] | None = None + """Which data the quantizer saw. ``None`` for an unquantized export; a + quantized one without this is not reproducible.""" + parity: dict[str, Any] | None = None + """The numeric record against the reference implementation.""" + notes: str = "" + + @property + def artifact_id(self) -> str: + return _sha256_text(f"{self.sha256}:{self.fmt}:{self.opset}:{self.component}")[:16] + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +def external_data_files(graph: Path) -> list[Path]: + """Weight files an ONNX graph keeps beside itself. + + Read without materialising the tensors -- the point is to size them, and + loading a 1.8 GB checkpoint to ask how big it is defeats the purpose. + """ + try: + import onnx + except ImportError: # pragma: no cover - onnx is a dependency in practice + sidecar = graph.with_name(graph.name + ".data") + return [sidecar] if sidecar.is_file() else [] + + try: + model = onnx.load(str(graph), load_external_data=False) + except Exception: + return [] + found: dict[str, Path] = {} + for init in model.graph.initializer: + if init.data_location != onnx.TensorProto.EXTERNAL: + continue + for entry in init.external_data: + if entry.key == "location": + candidate = graph.parent / entry.value + if candidate.is_file(): + found[entry.value] = candidate + return sorted(found.values()) + + +def build_graph_artifact( + path: str | os.PathLike[str], + *, + fmt: str, + opset: int, + source_model: str, + component: str, + exporter: str, + source_revision: str | None = None, + input_spec: tuple[dict[str, Any], ...] = (), + output_spec: tuple[dict[str, Any], ...] = (), + calibration: dict[str, Any] | None = None, + parity: dict[str, Any] | None = None, + notes: str = "", +) -> GraphArtifact: + """Digest an exported graph into a :class:`GraphArtifact`.""" + graph = Path(path) + if not graph.is_file(): + raise FileNotFoundError(f"exported graph not found: {graph}") + sidecars = external_data_files(graph) + return GraphArtifact( + path=str(graph), + sha256=sha256_file(graph), + # The graph *and* its weights. An ONNX file with external data is a few + # megabytes of topology pointing at gigabytes of tensors beside it; + # charging only the .onnx would tell the admission ledger a 1.8 GB + # vision tower costs 2 MiB. + bytes=graph.stat().st_size + sum(f.stat().st_size for f in sidecars), + fmt=fmt, + opset=int(opset), + source_model=str(source_model), + source_revision=source_revision, + component=component, + exporter=exporter, + input_spec=input_spec, + output_spec=output_spec, + calibration=calibration, + parity=parity, + notes=notes, + ) diff --git a/vllm_omni/edge/local/parity.py b/vllm_omni/edge/local/parity.py new file mode 100644 index 000000000..c80956e8a --- /dev/null +++ b/vllm_omni/edge/local/parity.py @@ -0,0 +1,147 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Does the Omni path decode the same tokens as plain vLLM? + +The validation plan makes reference consistency a *hard* initial acceptance +condition, and it is the one claim the rest of M0 cannot make for itself. +Everything else the local text mode reports -- the plan, the budget, the +cancellation -- is about the engine's own behaviour. This is about whether the +engine changed the model's answer. + +The comparison is **token ids, not text**. Two different id sequences can +detokenize to the same string, and a text comparison would hide exactly the +kind of drift worth catching. Both sides run greedy (`temperature=0`) with +`ignore_eos` so the sequences are the same length and a divergence is a +divergence rather than an early stop. + +Both sides also run on the same checkpoint, the same device and the same +request limits, so the only difference is the path: ``vllm.LLM`` straight into +vLLM, against ``AsyncOmni`` → Omni StageRuntime → the stage's EngineCore. A +mismatch here localises to the Omni stage path, which is precisely what M0 +added and what nothing else tests. + +**What a mismatch would and would not mean.** Greedy decoding is not bitwise +stable across batch shapes -- a different prefill chunking or batch composition +reorders reductions and can flip an argmax at a near-tie. So the first divergent +position is reported, not just a yes/no: a run that agrees for 120 of 128 tokens +and parts on one near-tie is a different finding from one that parts at token 3. + +The two engines are loaded **sequentially in one process**, reference first, +with the reference released before the second starts. Two 7.7 GB models +resident at once would not fit the budget this milestone is about respecting. +""" + +from __future__ import annotations + +import contextlib +import gc +from dataclasses import asdict, dataclass, field +from typing import Any + +from vllm.logger import init_logger + +logger = init_logger(__name__) + + +@dataclass +class PromptParity: + """One prompt, both paths.""" + + name: str + prompt_tokens: int | None + reference_tokens: int + engine_tokens: int + exact: bool + first_divergence: int | None + """Index of the first differing id, or ``None`` when they agree. Reported + even when the lengths differ, so a truncation is distinguishable from a + disagreement.""" + agreement_prefix: int + """How many leading ids matched. ``== reference_tokens`` iff exact.""" + reference_head: list[int] = field(default_factory=list) + engine_head: list[int] = field(default_factory=list) + """The eight ids around the first divergence, for a look at what changed.""" + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +def compare_ids(reference: list[int], engine: list[int]) -> tuple[bool, int | None, int]: + """``(exact, first_divergence, agreement_prefix)``.""" + prefix = 0 + for a, b in zip(reference, engine): + if a != b: + return False, prefix, prefix + prefix += 1 + if len(reference) != len(engine): + return False, prefix, prefix + return True, None, prefix + + +def summarize(rows: list[PromptParity]) -> dict[str, Any]: + exact = sum(1 for r in rows if r.exact) + divergences = [r.agreement_prefix for r in rows if not r.exact] + return { + "prompts": len(rows), + "exact": exact, + "exact_fraction": (exact / len(rows)) if rows else 0.0, + "min_agreement_prefix": min(divergences) if divergences else None, + "mean_agreement_prefix": (sum(divergences) / len(divergences)) if divergences else None, + "note": ( + "Greedy decoding is not bitwise stable across batch shapes: a " + "different prefill chunking reorders reductions and can flip an " + "argmax at a near-tie. agreement_prefix is reported so a late " + "single-token split is distinguishable from an early structural one." + ), + } + + +def reference_token_ids( + model_dir: str, + prompts: list[tuple[str, str]], + *, + engine_kwargs: dict[str, Any], + max_tokens: int, +) -> dict[str, dict[str, Any]]: + """Greedy decode through plain ``vllm.LLM``, then release the engine. + + ``engine_kwargs`` comes from the same :class:`~vllm_omni.edge.local.plan. + ExecutionPlan` the Omni side will use, so the two differ in path and not in + configuration. + """ + import vllm_omni # noqa: F401 - registers Spark's arch and config + from vllm import LLM, SamplingParams + + from vllm_omni.edge.local.engine import apply_runtime_env + + # Applied here, not left to the caller: this function *starts an engine*, + # and two of the three entries are load-time fatal on this host. A caller + # that forgot would not get a subtly different comparison, it would get + # "UVA is not available" -- which is how this line came to exist. + apply_runtime_env() + + kwargs = dict(engine_kwargs) + # ``LLM`` does not take Omni's spelling of these; everything else is shared. + kwargs.pop("kv_cache_memory_bytes", None) + llm = LLM(model=model_dir, trust_remote_code=False, disable_log_stats=True, **kwargs) + try: + params = SamplingParams(temperature=0.0, max_tokens=max_tokens, ignore_eos=True) + outputs = llm.generate([text for _, text in prompts], params, use_tqdm=False) + return { + name: { + "token_ids": list(out.outputs[0].token_ids), + "text": out.outputs[0].text, + "prompt_tokens": len(out.prompt_token_ids or ()), + } + for (name, _), out in zip(prompts, outputs) + } + finally: + # Release before the Omni engine loads: the budget this milestone is + # about does not have room for two copies of the weights. + del llm + gc.collect() + with contextlib.suppress(Exception): + import torch + + if torch.cuda.is_available(): + torch.cuda.empty_cache() diff --git a/vllm_omni/edge/local/plan.py b/vllm_omni/edge/local/plan.py new file mode 100644 index 000000000..0904df7d2 --- /dev/null +++ b/vllm_omni/edge/local/plan.py @@ -0,0 +1,679 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Decide where a text session runs, and be able to say why -- before loading. + +This is the M0 deliverable the proposal calls "一个可审计的 Omni 本地文本模式": +a plan that names the device, the backend, every byte it intends to reserve and +the evidence class of each claim, and that **refuses explicitly** when no +combination works. Refusing is a first-class outcome here, not an exception on +the way to one. + +Two refusals are worth naming because they are the reason this module exists +rather than a ``try: load()``: + +*The format refusal.* ``models/Spark-X2.5-1.7B-int8`` is 1.9 GB and the card +has 24 GiB, so every capacity check passes. It still cannot run: cutlass has no +int8 ``scaled_mm`` for SM >= 100 and vLLM's dispatch raises in the first +forward. Without this module you find out after ~40 s of engine startup, from a +stack trace in a subprocess. With it you find out in 30 ms, from a sentence +that names the remedy. + +*The capacity refusal.* ``Spark-X2.5-4B`` is 7.66 GiB of weights, and the +budget it needs is not 7.66 GiB -- it is weights plus the KV cache for the +context you asked for plus activations plus the loader's transient plus a +margin. Sizing that from ``free VRAM`` at load time is what the proposal's +section 6.2 calls out: the number to check is the sum over the session's +lifetime, formed before anything allocates. + +**What is measured and what is arithmetic.** The weight bytes are read off the +files (D). The KV/state bytes come from +:func:`vllm_omni.edge.kv_budget.kv_budget_for`, which reads the real +``layer_types`` rather than charging every layer full length (E, exact +arithmetic on a documented layout). Activations, workspace and the load +transient are **estimates** (E) from the constants below, and they are +deliberately generous: this module's job is to be a safe upper bound that +refuses before an OOM, not to predict the peak. + +The estimates are checked rather than believed: +:meth:`vllm_omni.edge.local.engine.LocalTextEngine.measured_peak` reports the +budgeted peak against the sampled one after every run, so a budget that is +wildly off is visible in the record instead of surviving as a constant nobody +re-examined. On Spark-X2.5-4B at a 4096 context this budget is 12.51 GiB +against a measured 8.93 GiB attributable peak -- 1.40x, i.e. safe and loose, +which is the intended direction. +""" + +from __future__ import annotations + +import math +from dataclasses import asdict, dataclass, field +from typing import Any + +from vllm_omni.edge.hardware_probe import HardwareProfile, load_profile +from vllm_omni.edge.kv_budget import KVBudget, kv_budget_for, text_config +from vllm_omni.edge.local.capabilities import ( + DeviceCapability, + enumerate_devices, + parse_mask, + runs_exported_graphs, +) +from vllm_omni.edge.local.manifest import ArtifactManifest, build_manifest + +GiB = 2**30 +MiB = 2**20 + +# ---- estimate constants ---------------------------------------------------- +# Kept as module constants rather than env vars, per project policy, and named +# so a plan's arithmetic can be read off the output. + +ACTIVATION_BYTES_PER_TOKEN_PER_HIDDEN = 2 +"""bf16 activations. A prefill chunk's peak is dominated by a handful of +``[tokens, hidden]`` tensors live at once inside a layer.""" + +ACTIVATION_LIVE_TENSORS = 12 +"""How many ``[tokens, hidden]``-sized tensors to charge for concurrently. +Generous: a gated MLP with a 4x intermediate is the widest point, and 12 covers +it with room for the attention workspace. E.""" + +LOGITS_BYTES_PER_TOKEN_PER_VOCAB = 4 +"""The logits tensor is fp32 in vLLM's sampler. At Spark's 131072 vocab this is +512 KiB per sampled row, which is small per request and stops being small if a +plan ever batches prefill logits.""" + +BACKEND_WORKSPACE_BYTES = 512 * MiB +"""Attention backend scratch, NCCL/driver context, allocator slack. E.""" + +GRAPH_CAPTURE_BYTES = 2 * GiB +"""Charged only when CUDA graphs are enabled. The same conservative constant +:func:`vllm_omni.engine.stage_admission.graph_reserve_bytes` uses, kept equal on +purpose so a local plan and a multi-stage admission do not disagree about the +same pool. M0 runs eager, so this is 0 in practice.""" + +LOAD_TRANSIENT_FRACTION = 0.10 +"""Peak above the resident weights while loading, as a fraction of them. + +Safetensors is mmap'd and copied shard by shard, so the transient tracks one +shard rather than the whole checkpoint. This fraction is the floor; +:func:`load_transient_bytes` takes the larger of it and one average shard, +which is what covers a single-file checkpoint. E.""" + +SAFETY_MARGIN_BYTES = 1 * GiB +"""Fragmentation and allocator slack. Matches +``stage_admission._DEFAULT_SAFETY_MARGIN_BYTES``.""" + +EXTERNAL_RESERVE_BYTES = 1 * GiB +"""Memory this engine does not control: the compositor on a VRAM pool, every +other process on a host_ram pool. Matches +``stage_admission._DEFAULT_EXTERNAL_RESERVE_BYTES``.""" + + +# ---- refusals -------------------------------------------------------------- + +REFUSE_NO_DEVICE = "no_runnable_device" +REFUSE_FORMAT = "weight_format_unsupported" +REFUSE_CAPACITY = "budget_exceeds_device" +REFUSE_CONFIG = "config_unreadable" + +# External-stage refusals (M3). A stage placed on the 890M or the NPU can fail +# for three reasons a checkpoint-on-vLLM plan never has to name. +REFUSE_NO_ARTIFACT = "no_exported_graph" +"""The device runs exported graphs, and no export exists for this stage at +this precision. Distinct from "unsupported": the work is an export, not a port.""" +REFUSE_EP_PLACEMENT = "execution_provider_declined_graph" +"""The session opened, ran, and left the graph on the CPU. This is the refusal +that stops a CPU fallback being reported as NPU execution -- and it cannot be +replaced by a cheaper check, because a VitisAI session lists the EP in +``get_providers()`` whether or not it claimed a single node, and when it claims +none the outputs are bit-identical to the CPU's.""" +REFUSE_NUMERICS = "numeric_gate_failed" +"""The graph ran on the device and disagreed with the reference beyond +tolerance. Constrains *this artifact*, never the model or the precision in +general -- see AGENTS.md on quantization failures.""" +REFUSE_ROUTE = "device_unreachable_from_this_process" +"""The device exists and works, and no interpreter here can drive it: the NPU +from inside WSL, or a route whose venv is not installed.""" + +MIN_FRACTION_ON_TARGET = 0.5 +"""How much of a graph the target EP must actually take for the placement to +count. Not tuned -- chosen as "most of it", because the failure this guards is +bimodal: XDNA2 either partitions a graph or declines it wholesale, and the +observed decline is 0 nodes, not a poor split. Callers override it per stage; +what must not happen is accepting an unverified placement by default.""" + + +@dataclass(frozen=True) +class Refusal: + """A named reason a device was not chosen, with the action that fixes it.""" + + code: str + device_id: str + message: str + remedy: str + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +@dataclass +class MemoryReservation: + """One claim on one pool, budgeted before the load and checked after it. + + ``actual_upper_bound_bytes`` starts ``None`` and stays ``None`` for + anything the runtime does not bound. Most reservations never get one: this + host has no counter that isolates a KV pool from an activation buffer, and + a reservation whose measurement is unknown is reported as unknown rather + than backfilled with its own estimate. The plan-level check that does + happen every run is ``LocalTextEngine.measured_peak``. + """ + + pool: str + """``vram`` or ``host_ram``. Two reservations in the same pool add up; + reservations in different pools do not.""" + purpose: str + bytes: int + lifetime: str + """``load`` (freed once weights are resident), ``session`` (held until the + engine closes) or ``request`` (held per in-flight request).""" + reclaimable: str + """``no``, ``on_evict`` (a rebuildable cache) or ``on_release``.""" + evidence: str + """``D`` read from the artifact, ``E`` estimated.""" + note: str = "" + actual_upper_bound_bytes: int | None = None + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +@dataclass +class ExecutionPlan: + """Everything needed to start the session, and to explain the choice.""" + + manifest: ArtifactManifest + devices: list[DeviceCapability] + selected: DeviceCapability | None + backend: str | None + stage: dict[str, Any] + engine_kwargs: dict[str, Any] + reservations: list[MemoryReservation] + kv_budget: KVBudget | None + refusals: list[Refusal] + fallbacks: list[str] + admitted: bool + request_limits: dict[str, int] + notes: list[str] = field(default_factory=list) + + # -- accounting --------------------------------------------------------- + def reserved_bytes(self, pool: str, *, lifetimes: tuple[str, ...] | None = None) -> int: + return sum( + r.bytes + for r in self.reservations + if r.pool == pool and (lifetimes is None or r.lifetime in lifetimes) + ) + + @property + def peak_bytes(self) -> int: + """The high-water mark on the selected device's pool. + + Load-lifetime reservations and the KV cache do not both peak: vLLM sizes + the cache *after* the weights are resident and the transient is gone. So + the peak is the larger of (load path) and (steady state), not the sum -- + which is the arithmetic ``free VRAM at load time`` gets wrong in the + other direction. + """ + if self.selected is None: + return 0 + pool = self.selected.memory_pool + load = self.reserved_bytes(pool, lifetimes=("load", "session_weights")) + steady = self.reserved_bytes(pool) + return max(load, steady) + + def to_dict(self) -> dict[str, Any]: + return { + "admitted": self.admitted, + "backend": self.backend, + "selected_device": self.selected.to_dict() if self.selected else None, + "devices": [d.to_dict() for d in self.devices], + "stage": self.stage, + "engine_kwargs": self.engine_kwargs, + "reservations": [r.to_dict() for r in self.reservations], + "reserved_total_bytes": self.peak_bytes, + "kv_budget": _kv_to_dict(self.kv_budget), + "refusals": [r.to_dict() for r in self.refusals], + "fallbacks": self.fallbacks, + "request_limits": self.request_limits, + "manifest": self.manifest.to_dict(), + "notes": self.notes, + } + + def summary(self) -> str: + if not self.admitted: + lines = ["REFUSED: no device in this process can run this checkpoint as asked."] + for r in self.refusals: + lines.append(f" [{r.code}] {r.device_id}: {r.message}") + lines.append(f" remedy: {r.remedy}") + return "\n".join(lines) + assert self.selected is not None + pool = self.selected.memory_pool + capacity = _device_capacity(self.selected) + capacity_note = "available" if pool == "host_ram" else "total" + lines = [ + f"ADMITTED on {self.selected.device_id} ({self.selected.name}) via {self.backend}", + f" artifact {self.manifest.artifact_id} {self.manifest.weight_format} " + f"{self.manifest.weight_bytes / GiB:.2f}GiB", + f" budget {self.peak_bytes / GiB:.2f}GiB of {capacity / GiB:.2f}GiB {pool} ({capacity_note})", + ] + for r in self.reservations: + lines.append( + f" {r.purpose:<18} {r.bytes / MiB:9.0f} MiB {r.pool:<9} " + f"{r.lifetime:<15} {r.evidence} {r.note}" + ) + if self.kv_budget is not None: + lines.append(f" kv {self.kv_budget.summary()}") + for r in self.refusals: + lines.append(f" not used: [{r.code}] {r.device_id}: {r.message}") + return "\n".join(lines) + + +def _kv_to_dict(kv: KVBudget | None) -> dict[str, Any] | None: + if kv is None: + return None + d = asdict(kv) + d["summary"] = kv.summary() + d["saving_vs_flat"] = kv.saving_vs_flat + return d + + +# ---- estimates ------------------------------------------------------------- + + +def activation_bytes(hf_config: Any, *, max_num_batched_tokens: int, max_num_seqs: int) -> int: + """Peak activation working set for one prefill chunk. E.""" + cfg = text_config(hf_config) + hidden = int(getattr(cfg, "hidden_size", 0) or 0) + vocab = int(getattr(cfg, "vocab_size", 0) or 0) + hidden_bytes = ( + max_num_batched_tokens + * hidden + * ACTIVATION_BYTES_PER_TOKEN_PER_HIDDEN + * ACTIVATION_LIVE_TENSORS + ) + # Logits are charged per *sequence*, not per batched token: vLLM gathers + # only the last position of each sequence before sampling. + logits_bytes = max_num_seqs * vocab * LOGITS_BYTES_PER_TOKEN_PER_VOCAB + return hidden_bytes + logits_bytes + + +def load_transient_bytes(manifest: ArtifactManifest) -> int: + """Headroom above the resident weights during the load. E. + + Charged as the larger of a flat fraction of the checkpoint and one average + shard, because a single-file checkpoint (Spark-1.7B: one 1.9 GB + ``model.safetensors``) has a transient the fraction alone understates. + """ + n_shards = max(len(manifest.weight_files), 1) + per_shard = manifest.weight_bytes // n_shards + return int(max(manifest.weight_bytes * LOAD_TRANSIENT_FRACTION, per_shard)) + + +# ---- planning -------------------------------------------------------------- + + +def _load_hf_config(manifest: ArtifactManifest) -> Any: + from transformers import AutoConfig + + # Spark's config class is vendored precisely so this does not need + # ``trust_remote_code``; registration happens on the omni import path. + from vllm_omni.engine.arg_utils import _register_omni_hf_configs + + _register_omni_hf_configs() + return AutoConfig.from_pretrained(manifest.model_dir, trust_remote_code=False) + + +def _device_capacity(device: DeviceCapability) -> int: + """Bytes this engine may plan against on the device's pool. + + On ``host_ram`` that is *available*, not total: the OS and everything else + already hold the difference, and the proposal is explicit that Windows RAM, + the WSL quota and iGPU/NPU shared memory are separate constraints that must + not be added together. + """ + if device.memory_pool == "host_ram": + available = int(device.extra.get("ram_available_bytes") or 0) + return available or device.memory_bytes + return device.memory_bytes + + +def _candidate_order(devices: list[DeviceCapability]) -> list[DeviceCapability]: + """Runnable devices, best first. + + "Best" here is only "private memory pool before shared one", which on this + class of machine means the discrete GPU before the CPU. It is deliberately + not a performance ranking: the proposal's rule is that placement is decided + by profiles from measured full chains, and this module has none. + """ + runnable = [d for d in devices if d.runnable and not runs_exported_graphs(d)] + return sorted(runnable, key=lambda d: (d.memory_pool != "vram", -d.memory_bytes)) + + +def plan_text_session( + model_dir: str, + *, + profile: HardwareProfile | None = None, + max_model_len: int = 4096, + max_num_seqs: int = 1, + max_num_batched_tokens: int | None = None, + enforce_eager: bool = True, + kv_dtype: str = "auto", + mask: str | frozenset[str] | None = None, + platform_device_type: str | None = None, + digest_weights: bool = False, +) -> ExecutionPlan: + """Plan one single-stage text session, or refuse with reasons. + + The caller gets a plan object either way: ``admitted`` says which, and + ``refusals`` is never empty when it is ``False``. + """ + manifest = build_manifest(model_dir, digest_weights=digest_weights) + profile = profile if profile is not None else load_profile() + devices = enumerate_devices( + profile, + mask=mask if isinstance(mask, frozenset) else parse_mask(mask), + platform_device_type=platform_device_type, + ) + batched_tokens = max_num_batched_tokens or min(2048, max_model_len) + + refusals: list[Refusal] = [] + notes: list[str] = [] + + try: + hf_config = _load_hf_config(manifest) + except Exception as exc: + return ExecutionPlan( + manifest=manifest, devices=devices, selected=None, backend=None, + stage={}, engine_kwargs={}, reservations=[], kv_budget=None, + refusals=[Refusal( + REFUSE_CONFIG, "-", + f"the checkpoint config could not be read as a registered architecture: {exc}", + "check that the architecture is in vllm_omni.model_executor.models.registry " + "and its config class in engine.arg_utils._register_omni_hf_configs", + )], + fallbacks=[], admitted=False, + request_limits={}, notes=notes, + ) + + # KV geometry is a property of the model, not the device, so it is computed + # once and reused for every candidate. + effective_kv_dtype = manifest.dtype or "bfloat16" if kv_dtype == "auto" else kv_dtype + kv = kv_budget_for( + hf_config, + max_model_len=max_model_len, + max_num_seqs=max_num_seqs, + kv_dtype=effective_kv_dtype, + ) + + for device in devices: + if not device.runnable: + refusals.append(Refusal( + REFUSE_NO_DEVICE, device.device_id, + device.reason or "not runnable from this process", + "see the reason; masked devices are re-enabled by dropping them from --mask", + )) + elif runs_exported_graphs(device): + # Runnable, and still not a candidate for *this* plan. Said out + # loud rather than left out: once these devices gained a backend + # they stopped appearing in the refusal list at all, which reads as + # "not considered" when the truth is "cannot host a whole session". + refusals.append(Refusal( + REFUSE_NO_ARTIFACT, device.device_id, + f"{device.device_id} executes exported graphs " + f"({sorted(device.weight_formats)}), not checkpoints, so a whole " + "autoregressive text session cannot be placed on it", + "place a single stage there instead, via " + "vllm_omni.edge.local.external.stage.plan_external_stage; the " + "decode loop stays on the CPU or the GPU because both of these " + "devices read the same system memory it does", + )) + + for device in _candidate_order(devices): + if manifest.weight_format not in device.weight_formats: + refusals.append(Refusal( + REFUSE_FORMAT, device.device_id, + f"the checkpoint is {manifest.weight_format} and {device.device_id} " + f"({device.name}, sm_{''.join(map(str, device.compute_capability or ())) or '-'}) " + f"executes only {sorted(device.weight_formats)}", + _format_remedy(manifest, device), + )) + continue + + capacity = _device_capacity(device) + reservations = _reservations_for( + manifest, hf_config, kv, device, + max_num_seqs=max_num_seqs, + max_num_batched_tokens=batched_tokens, + enforce_eager=enforce_eager, + ) + plan = ExecutionPlan( + manifest=manifest, devices=devices, selected=device, + backend=device.backend, stage=_stage_spec(manifest), + engine_kwargs={}, reservations=reservations, kv_budget=kv, + refusals=list(refusals), fallbacks=[], admitted=False, + request_limits={ + "max_model_len": max_model_len, + "max_num_seqs": max_num_seqs, + "max_num_batched_tokens": batched_tokens, + }, + notes=notes, + ) + required = plan.peak_bytes + if required > capacity: + refusals.append(Refusal( + REFUSE_CAPACITY, device.device_id, + f"the session needs {required / GiB:.2f} GiB on {device.memory_pool} and " + f"{capacity / GiB:.2f} GiB is available " + f"({manifest.weight_bytes / GiB:.2f} GiB weights + " + f"{kv.bytes_total / GiB:.2f} GiB KV at {max_model_len} tokens x {max_num_seqs})", + _capacity_remedy(kv, max_model_len, required - capacity), + )) + continue + + plan.refusals = list(refusals) + plan.engine_kwargs = _engine_kwargs( + manifest, device, kv, + max_model_len=max_model_len, + max_num_seqs=max_num_seqs, + max_num_batched_tokens=batched_tokens, + enforce_eager=enforce_eager, + capacity=capacity, + required=required, + ) + plan.fallbacks = _fallbacks(manifest, devices, device) + plan.admitted = True + plan.notes = notes + [ + "activations, workspace and the load transient are estimates (E); " + "engine.py records the measured peak against them after the load", + ] + return plan + + return ExecutionPlan( + manifest=manifest, devices=devices, selected=None, backend=None, + stage=_stage_spec(manifest), engine_kwargs={}, reservations=[], + kv_budget=kv, refusals=refusals, fallbacks=[], admitted=False, + request_limits={ + "max_model_len": max_model_len, + "max_num_seqs": max_num_seqs, + "max_num_batched_tokens": batched_tokens, + }, + notes=notes, + ) + + +def _stage_spec(manifest: ArtifactManifest) -> dict[str, Any]: + """The single AR text stage, in the shape ``StagePipelineConfig`` declares. + + M0 is one model and one stage on purpose. The fields are spelled out rather + than implied so that adding a second stage later is a change to this dict, + not a change to the engine. + """ + return { + "stage_id": 0, + "model_stage": manifest.model_type, + "execution_type": "LLM_AR", + "input_sources": [], + "final_output": True, + "final_output_type": "text", + "owns_tokenizer": True, + "state": "kv_cache", + "streaming_granularity": "token", + } + + +def _reservations_for( + manifest: ArtifactManifest, + hf_config: Any, + kv: KVBudget, + device: DeviceCapability, + *, + max_num_seqs: int, + max_num_batched_tokens: int, + enforce_eager: bool, +) -> list[MemoryReservation]: + pool = device.memory_pool + acts = activation_bytes( + hf_config, + max_num_batched_tokens=max_num_batched_tokens, + max_num_seqs=max_num_seqs, + ) + reservations = [ + MemoryReservation( + pool, "weights", manifest.weight_bytes, "session_weights", "no", "D", + f"{len(manifest.weight_files)} safetensors shard(s), on-disk size", + ), + MemoryReservation( + pool, "kv_cache", kv.bytes_total, "session", "on_evict", "E", + f"flat budget at {max_num_seqs}x context; " + f"{kv.bytes_total_hybrid / MiB:.0f} MiB if the backend accounts " + f"{kv.sliding_layers} sliding layers separately", + ), + MemoryReservation( + pool, "activations", acts, "request", "on_release", "E", + f"{max_num_batched_tokens}-token prefill chunk x {ACTIVATION_LIVE_TENSORS} live tensors", + ), + MemoryReservation( + pool, "backend_workspace", BACKEND_WORKSPACE_BYTES, "session", "no", "E", + "attention scratch, driver/allocator context", + ), + MemoryReservation( + pool, "load_transient", load_transient_bytes(manifest), "load", "on_release", "E", + "freed once the weights are resident; peaks with the largest shard", + ), + MemoryReservation( + pool, "external_reserve", EXTERNAL_RESERVE_BYTES, "session", "no", "E", + "memory this engine does not own on a shared pool", + ), + MemoryReservation( + pool, "safety_margin", SAFETY_MARGIN_BYTES, "session", "no", "E", + "fragmentation and allocator slack", + ), + ] + if not enforce_eager: + reservations.append(MemoryReservation( + pool, "graph_capture", GRAPH_CAPTURE_BYTES, "session", "no", "E", + "CUDA graph pool; zero under enforce_eager", + )) + return reservations + + +def _engine_kwargs( + manifest: ArtifactManifest, + device: DeviceCapability, + kv: KVBudget, + *, + max_model_len: int, + max_num_seqs: int, + max_num_batched_tokens: int, + enforce_eager: bool, + capacity: int, + required: int, +) -> dict[str, Any]: + """Turn the plan into flags the backend actually has. + + The KV pool is set as an **absolute** byte count, not left to + ``gpu_memory_utilization``. The proposal's section 6.2 records the failure + mode this avoids: a 4096-token single-sequence run whose logs still showed + a far larger cache allocated, because the fraction, not the request limits, + decided the pool size. + """ + kwargs: dict[str, Any] = { + "max_model_len": max_model_len, + "max_num_seqs": max_num_seqs, + "max_num_batched_tokens": max_num_batched_tokens, + "enforce_eager": enforce_eager, + "enable_prefix_caching": False, + # Derived from this plan's own budget, not left at the default. The + # flag is named for the GPU and is **not** GPU-only: on the CPU backend + # vLLM reads the same field as the fraction of *host RAM* to reserve. + # Leaving it alone there is not "no opinion", it is asking for 92% of + # the machine -- which is how a 6.55 GiB CPU plan turned into + # "desired CPU memory utilization (0.92, 28.43 GiB)" and refused to + # start. Setting it on both pools is what makes the plan's number the + # number the backend uses. + "gpu_memory_utilization": min(0.95, math.ceil(required / capacity * 100) / 100), + # Absolute KV pool, so the cache is the size the request limits imply + # rather than whatever is left over inside the fraction above. + "kv_cache_memory_bytes": int(kv.bytes_total), + } + if manifest.dtype: + kwargs["dtype"] = manifest.dtype + return kwargs + + +def _fallbacks( + manifest: ArtifactManifest, + devices: list[DeviceCapability], + selected: DeviceCapability, +) -> list[str]: + """Pre-approved alternatives, named so a fallback is never a surprise. + + The proposal forbids silently swapping the model, the precision or the + context. So the list is descriptive: these are the moves an operator may + make, not moves the engine will make on its own. + """ + out: list[str] = [] + for d in devices: + if d.device_id == selected.device_id or not d.runnable: + continue + if manifest.weight_format in d.weight_formats: + out.append(f"{d.device_id} runs this format; re-plan with --mask {selected.kind}") + out.append("shorter --max-model-len reduces the KV reservation linearly") + return out + + +def _format_remedy(manifest: ArtifactManifest, device: DeviceCapability) -> str: + from vllm_omni.edge.local.capabilities import FORMAT_FP8, FORMAT_INT8 + + if manifest.weight_format == FORMAT_INT8 and FORMAT_FP8 in device.weight_formats: + return ( + "this device has fp8 tensor cores but no int8 scaled_mm path " + "(cutlass c3x has none for SM >= 100). Use an fp8 or a dense build " + "here, or run the int8 build on the CPU backend (.venvs/omni-cpu)." + ) + return ( + f"build or obtain a {sorted(device.weight_formats)} artifact for this device, " + "or plan for a device that executes this format" + ) + + +def _capacity_remedy(kv: KVBudget, max_model_len: int, over_by: int) -> str: + per_token = max(kv.bytes_per_token_flat, 1) + shed_tokens = math.ceil(over_by / per_token) + if shed_tokens < max_model_len: + return ( + f"dropping --max-model-len to about {max_model_len - shed_tokens} frees " + f"{over_by / MiB:.0f} MiB of KV. The engine will not do this for you: " + "silently shortening a user's context is exactly what the design forbids." + ) + return ( + "the weights alone do not leave room for a usable context on this device. " + "Use a smaller or more compressed build, or a device with a larger pool." + ) diff --git a/vllm_omni/edge/local/prompts.py b/vllm_omni/edge/local/prompts.py new file mode 100644 index 000000000..c9097bb8c --- /dev/null +++ b/vllm_omni/edge/local/prompts.py @@ -0,0 +1,90 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""The fixed prompt set the M0 acceptance run uses, and why it is this one. + +The roadmap asks for "同一组 12 个提示至少生成 128 token" -- one fixed set of 12 +prompts, at least 128 tokens each. The set is not new: it is the same twelve +this repo already measures Spark's quantized builds against +(``analysis/experiments/spark_edge/quality_int4.py``), which is the point. A +new prompt set would make M0's numbers incomparable with the W4A16/W4A8/Q4_K_M +fidelity work that produced "2/12 and 1/12 greedy", and comparability is the +only reason to fix a set at all. + +It is three chat prompts and nine plain completions: + +* ``chat_short`` (28 tokens), ``chat_tools`` (137) and ``chat_long`` (2465) + are stored **already rendered** through Spark's chat template, copied from + the reference captures in ``analysis/experiments/spark_edge/ref*.json``. They + are verbatim so that the tokens fed to the model here are byte-identical to + the ones the parity and fidelity runs fed it -- a re-render through a newer + template would silently change the comparison. +* the nine ``plain_*`` are raw completions. They exercise the path with no + template at all, which is also the path a base-model deployment takes. + +The spread of lengths is deliberate: ``chat_long`` at 2465 tokens crosses +Spark's 512-token sliding window several times, so a run over this set touches +window rollover rather than only short prefixes. + +:func:`render_chat` is here as the adapter seam -- turning messages into a +prompt is a model-adapter job, not an engine job -- but the acceptance set does +not use it, for the reason above. +""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + +_HERE = Path(__file__).parent +_CHAT_FILE = _HERE / "acceptance_chat_prompts.json" + +# Verbatim from analysis/experiments/spark_edge/quality_int4.py::PLAIN. +PLAIN_PROMPTS: tuple[str, ...] = ( + "The capital of Australia is", + "In 1969, the Apollo 11 mission", + 'def fibonacci(n):\n """Return the nth Fibonacci number."""\n', + "The three laws of thermodynamics state that", + "To make a proper espresso you need", + "The difference between TCP and UDP is", + "Photosynthesis converts", + "A binary search tree is a data structure that", + "The Treaty of Westphalia was signed in", +) + + +def acceptance_prompts() -> list[tuple[str, str]]: + """The twelve ``(name, prompt)`` pairs, chat first then plain. + + Raises if the chat captures are missing rather than falling back to nine + prompts: a run that quietly measured a different set would be worse than + one that failed. + """ + chat: dict[str, str] = json.loads(_CHAT_FILE.read_text(encoding="utf-8")) + ordered = [(name, chat[name]) for name in ("chat_short", "chat_tools", "chat_long")] + ordered += [(f"plain_{i}", text) for i, text in enumerate(PLAIN_PROMPTS)] + if len(ordered) != 12: + raise RuntimeError(f"the acceptance set must be 12 prompts, got {len(ordered)}") + return ordered + + +def render_chat(model_dir: str, messages: list[dict[str, Any]], *, add_generation_prompt: bool = True) -> str: + """Apply the checkpoint's own chat template. The adapter seam. + + Transformers 5 returns a list from ``apply_chat_template`` under some + argument combinations and a string under others; the architecture record + has a failed run from exactly that (``fp8_offload10``). Rendering to text + explicitly, and asserting it, is cheaper than discovering it again. + """ + from transformers import AutoTokenizer + + tok = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=False) + rendered = tok.apply_chat_template( + messages, tokenize=False, add_generation_prompt=add_generation_prompt + ) + if not isinstance(rendered, str): + raise TypeError( + f"apply_chat_template returned {type(rendered).__name__}, not str; " + "pass tokenize=False and render before tokenizing" + ) + return rendered diff --git a/vllm_omni/edge/local/session.py b/vllm_omni/edge/local/session.py new file mode 100644 index 000000000..5d563e913 --- /dev/null +++ b/vllm_omni/edge/local/session.py @@ -0,0 +1,249 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""The three data boundaries the local engine needs, and nothing else. + +The proposal's section 4 is explicit that the first version needs exactly three +contracts: **a request/event stream, a tensor-or-buffer reference, and an +opaque state handle**. This module is the first and the third. The second does +not appear in M0 because a single text stage never hands a buffer to anybody -- +it is the thing that arrives in M2, when a talker hands codec frames to a +vocoder, and inventing it now would be inventing it wrong. + +Two properties here are correctness, not bookkeeping: + +**An epoch fences cancellation.** Cancelling a request is not "stop calling +next()". The backend has work in flight, and tokens from that work arrive after +the cancel. Every event carries the epoch it was produced under, and +:class:`BoundedEventStream` drops events from a retired epoch instead of +delivering them. Without that, an interrupted turn's tail lands in the next +turn's output -- the "重复输出和跨请求污染" the proposal names. + +**A queue bounds chunks and bytes, both.** A consumer that stops reading must +slow the producer down, not grow the heap. Bounding only the chunk count is not +enough once a chunk can be a second of PCM rather than one token, so the bound +is on both from the start even though M0's chunks are small. + +``StateHandle`` is deliberately opaque. The engine holds a *reference* to the +session's state; vLLM holds the actual KV blocks, the position counter and the +RNG. The handle records enough to know when that state is no longer valid -- +a different artifact, a different layout version, a retired epoch -- and +nothing that would let a caller reach inside it. +""" + +from __future__ import annotations + +import asyncio +import time +import uuid +from dataclasses import asdict, dataclass, field +from typing import Any + +STATE_LAYOUT_VERSION = 1 +"""Bumped when the meaning of a handle's state changes. A handle from a +different version is not migrated, it is rejected: the proposal's item 10 says +to treat state as non-migratable until a specific backend pair has been +measured, and a version that silently matches would hide that.""" + + +@dataclass(frozen=True) +class StateHandle: + """A reference to session state the backend owns. + + The engine never reads the state. It reads this, to decide whether a + request may continue a session or must start a new one. + """ + + session_id: str + backend: str + artifact_id: str + """Ties the state to the checkpoint that produced it. The same shapes from + a different quantization are not the same state.""" + layout_version: int = STATE_LAYOUT_VERSION + epoch: int = 0 + """Incremented on every cancel. Events from a lower epoch are stale.""" + replayable: bool = True + """A text session can be rebuilt by re-submitting its committed prompt, so + recovery after a backend crash is possible at an output boundary.""" + migratable: bool = False + """Never true in M0. Moving KV between backends needs a measured layout + conversion for the specific pair; assuming it is the failure the proposal + warns about.""" + created_unix: float = field(default_factory=time.time) + + def next_epoch(self) -> StateHandle: + from dataclasses import replace + + return replace(self, epoch=self.epoch + 1) + + def accepts(self, other: StateHandle) -> bool: + """Whether ``other`` may continue this session's state.""" + return ( + other.session_id == self.session_id + and other.artifact_id == self.artifact_id + and other.layout_version == self.layout_version + and other.epoch == self.epoch + ) + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +@dataclass(frozen=True) +class ChunkEvent: + """One unit of output, with everything needed to order and retire it.""" + + request_id: str + stage_id: int + seq: int + """Monotonic within (request, epoch). Gaps mean loss; repeats mean a bug.""" + epoch: int + kind: str + """``token`` | ``done`` | ``error`` | ``cancelled``.""" + payload: Any + """For ``token``: the incremental text and token ids. Never the whole + output so far -- a cumulative payload makes a dropped chunk invisible.""" + started_unix: float + emitted_unix: float + final: bool = False + error: str | None = None + input_watermark: int = 0 + """How much of the input the producer had consumed when this was emitted. + One number in M0 (prompt tokens); it is what a duplex turn will align on.""" + release_token: str = "" + """Returned by the consumer to give credit back. Empty when the event + carries no resource the producer is holding on its behalf.""" + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +class StreamClosed(RuntimeError): + """Raised on a stream that has been closed while a consumer waited.""" + + +class BoundedEventStream: + """An async queue with a credit bound on chunks *and* bytes. + + Producers ``await put(...)``, which blocks while the consumer is behind: + that is the backpressure. Consumers ``async for`` over it and the credit is + returned as each event leaves the queue. + """ + + def __init__(self, *, max_chunks: int = 64, max_bytes: int = 4 << 20) -> None: + if max_chunks < 1 or max_bytes < 1: + raise ValueError("bounds must be positive") + self.max_chunks = max_chunks + self.max_bytes = max_bytes + self._queue: asyncio.Queue[ChunkEvent | None] = asyncio.Queue() + self._bytes = 0 + self._chunks = 0 + self._credit = asyncio.Condition() + self._closed = False + self._epoch = 0 + self.dropped_stale = 0 + """Events discarded because their epoch had been retired. Reported + rather than silently swallowed: a non-zero count on a run with no + cancellation is a bug worth seeing.""" + self.high_water_chunks = 0 + self.high_water_bytes = 0 + + # -- producer side ------------------------------------------------------ + @staticmethod + def _sizeof(event: ChunkEvent) -> int: + payload = event.payload + if isinstance(payload, (bytes, bytearray, memoryview)): + return len(payload) + if isinstance(payload, str): + return len(payload.encode("utf-8")) + if isinstance(payload, dict): + return sum( + len(str(k)) + (len(v.encode("utf-8")) if isinstance(v, str) else 8 * (len(v) if isinstance(v, (list, tuple)) else 1)) + for k, v in payload.items() + ) + return 64 + + async def put(self, event: ChunkEvent) -> None: + """Enqueue, waiting for credit. Stale-epoch events are dropped here.""" + if self._closed: + raise StreamClosed("stream is closed") + if event.epoch < self._epoch: + self.dropped_stale += 1 + return + size = self._sizeof(event) + async with self._credit: + await self._credit.wait_for( + lambda: self._closed + or (self._chunks < self.max_chunks and self._bytes + size <= self.max_bytes) + # A single event larger than the whole byte bound must still go + # through once the queue is empty, or it deadlocks forever. + or (self._chunks == 0 and size > self.max_bytes) + ) + if self._closed: + raise StreamClosed("stream closed while producing") + self._chunks += 1 + self._bytes += size + self.high_water_chunks = max(self.high_water_chunks, self._chunks) + self.high_water_bytes = max(self.high_water_bytes, self._bytes) + await self._queue.put(event) + + async def close(self) -> None: + """End the stream. Idempotent; wakes every waiter.""" + if self._closed: + return + self._closed = True + async with self._credit: + self._credit.notify_all() + await self._queue.put(None) + + def retire_epoch(self, epoch: int) -> int: + """Retire everything at or below ``epoch``; return what was discarded. + + Called by cancel. Queued events from the retired epoch are dropped on + the way out rather than scanned here, so this stays O(1) and cannot + race a producer mid-``put``. + """ + self._epoch = max(self._epoch, epoch + 1) + return self.dropped_stale + + @property + def epoch(self) -> int: + return self._epoch + + # -- consumer side ------------------------------------------------------ + async def get(self) -> ChunkEvent | None: + """Next live event, or ``None`` at end of stream.""" + while True: + event = await self._queue.get() + size = 0 if event is None else self._sizeof(event) + if event is not None: + async with self._credit: + self._chunks -= 1 + self._bytes -= size + self._credit.notify_all() + if event.epoch < self._epoch: + self.dropped_stale += 1 + continue + return event + + def __aiter__(self) -> BoundedEventStream: + return self + + async def __anext__(self) -> ChunkEvent: + event = await self.get() + if event is None: + raise StopAsyncIteration + return event + + def stats(self) -> dict[str, int]: + return { + "max_chunks": self.max_chunks, + "max_bytes": self.max_bytes, + "high_water_chunks": self.high_water_chunks, + "high_water_bytes": self.high_water_bytes, + "dropped_stale": self.dropped_stale, + "epoch": self._epoch, + } + + +def new_session_id() -> str: + return uuid.uuid4().hex[:16] diff --git a/vllm_omni/edge/npu_ryzenai.py b/vllm_omni/edge/npu_ryzenai.py new file mode 100644 index 000000000..aeb2b799d --- /dev/null +++ b/vllm_omni/edge/npu_ryzenai.py @@ -0,0 +1,184 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Running graphs on a Ryzen AI (XDNA2) NPU, from the Windows side. + +The NPU is an **MCDM** device. WSL2's GPU-PV forwards WDDM display adapters +only, so there is no ``/dev/accel`` node and no amount of installing inside +WSL changes that -- this module is meant to be imported by a *Windows* Python +process that a WSL process talks to, not by the engine. + +Three things had to be right at once, and getting any one of them wrong looks +identical from the outside: the session builds, runs, returns correct numbers, +and every node silently executes on the CPU. + +**1. The right EP.** The appx package ships two. +``onnxruntime_providers_ryzenai.dll`` registers cleanly, enumerates the NPU as +``OrtHardwareDeviceType.NPU``, opens a session -- and claims **zero nodes** from +every graph tried (fp32, INT8 QDQ symmetric and asymmetric, per-channel, +power-of-two scales, MatMul and Conv, and via ahead-of-time +``ModelCompiler``, which produced zero ``EPContext`` nodes). It is the *Light* +EP. The general one is ``onnxruntime_vitisai_ep.dll``. + +**2. The DLL search path.** ``onnxruntime_vitisai_ep.dll`` fails to load with a +bare ``register_execution_provider_library`` -- ORT reports ``Error loading +... But no dependency``, which is as unhelpful as it sounds. It needs its own +directory on the search path first (``os.add_dll_directory``), because it pulls +in ``dyn_dispatch_core.dll`` and ``vaiml.dll`` from beside it. + +**3. A16W8, not A8W8.** The overlays are named in the DLL: + + 4x2_psf_model_a8w8_qdq.xclbin + 8x4_psu_model_a16w8_qdq.xclbin <- most of them look like this + 2x4x4_bmm_model_a16w16.xclbin + +XDNA2 wants **16-bit activations with 8-bit weights**. Quantizing A8W8 -- the +default shape of ``quantize_static`` -- is rejected as completely as fp32, with +no error. Once A16W8 is used the partitioner starts talking, and names its +target: ``AMD_AIE2P_4x8_CMC_Overlay``. + +Measured once all three are right, against the same INT8 graph on the CPU +(Ryzen AI 9 HX 370, 12 cores): + + GEMM CPU NPU speedup NPU + 256 x 1024 x 1024 1.47 ms 0.56 ms 2.64x 0.97 TOPS + 512 x 2048 x 2048 9.64 ms 2.14 ms 4.51x 2.01 TOPS + 1024 x 2048 x 2048 19.89 ms 3.76 ms 5.29x 2.29 TOPS + +with a max relative error of 3.2e-05 against the CPU. It scales the right way: +the bigger the GEMM, the better, because a small one cannot amortise the +dispatch. Dequantize nodes stay on the CPU -- the partitioner says so +explicitly -- so a graph that is mostly requantization will not benefit. + +**What this is worth for LLM serving.** Not decode: 2.29 TOPS against the +dGPU's 48.2 TFLOP/s, and it reads the same 20.7 GB/s system memory the CPU +does, so it cannot help a bandwidth-bound step (see +``docs/edge/laptop-heterogeneous.md``). It is worth having for the things that +are compute-bound and off the per-token path -- a vision tower, a draft model +for speculative decoding -- where 5x the CPU on large GEMMs is real. +""" + +from __future__ import annotations + +import os +from pathlib import Path +from typing import Any + +APPX_GLOB = "WindowsWorkload.EP.AMD.VitisAI.Framework*" +WINDOWSAPPS = Path(r"C:\Program Files\WindowsApps") + +EP_GENERAL = "onnxruntime_vitisai_ep.dll" +"""The EP that partitions general graphs. Not the one whose name looks right.""" +EP_LIGHT = "onnxruntime_providers_ryzenai.dll" +"""Enumerates the NPU and claims nothing. Kept named so nobody re-discovers it.""" + +PARTITION_TARGET = "AMD_AIE2P_4x8_CMC_Overlay" +"""What the partitioner calls XDNA2 when it accepts a subgraph. Seeing this +string in the log is the only positive confirmation that the NPU took work.""" + + +def find_ep_directory(root: Path = WINDOWSAPPS) -> Path: + """The ``ExecutionProvider`` folder inside the appx package. + + ``WindowsApps`` denies directory listing to a normal process, so a plain + ``glob`` returns nothing even though the files are readable by full path. + PowerShell's ``Get-AppxPackage`` is the reliable way to get the path; + globbing is tried first because it works when the ACL permits it. + """ + for base in sorted(root.glob(APPX_GLOB)): + candidate = base / "ExecutionProvider" + if (candidate / EP_GENERAL).exists(): + return candidate + import subprocess + + out = subprocess.run( + ["powershell.exe", "-NoProfile", "-Command", + f"(Get-AppxPackage -Name '{APPX_GLOB}').InstallLocation"], + capture_output=True, text=True, check=False, + ).stdout.strip() + if out: + candidate = Path(out) / "ExecutionProvider" + if (candidate / EP_GENERAL).exists(): + return candidate + raise FileNotFoundError( + "Ryzen AI execution provider not found. It ships as the appx package " + f"{APPX_GLOB!r}; on a machine without it, AMD's Ryzen AI Software " + "installs the same EP." + ) + + +def register(ep_dir: Path | None = None, name: str = "vitisai") -> str: + """Put the EP on the DLL search path and register it. Returns the EP name. + + Order matters: ``add_dll_directory`` before ``register_execution_provider_library``, + or the load fails with "no dependency". + """ + import onnxruntime as ort + + if tuple(int(p) for p in ort.__version__.split(".")[:2]) < (1, 25): + raise RuntimeError( + f"onnxruntime {ort.__version__} is too old: the EP requests ORT API " + "25, so 1.25+ is required. Note that onnxruntime-directml is pinned " + "at 1.24.4 and cannot load it -- use plain onnxruntime." + ) + ep_dir = Path(ep_dir) if ep_dir else find_ep_directory() + os.environ["PATH"] = f"{ep_dir}{os.pathsep}{os.environ.get('PATH', '')}" + os.add_dll_directory(str(ep_dir)) + ort.register_execution_provider_library(name, str(ep_dir / EP_GENERAL)) + return name + + +def npu_devices(name: str = "vitisai") -> list[Any]: + import onnxruntime as ort + + return [ + d for d in ort.get_ep_devices() + if d.ep_name == name and str(d.device.type).endswith("NPU") + ] + + +def session_options(name: str = "vitisai") -> Any: + """SessionOptions bound to the NPU, or a clear error if it is not there.""" + import onnxruntime as ort + + devices = npu_devices(name) + if not devices: + raise RuntimeError( + "no NPU device from the VitisAI EP. Check the driver is present " + "(Get-PnpDevice -Class ComputeAccelerator) and that register() ran." + ) + options = ort.SessionOptions() + options.add_provider_for_devices(devices, {}) + return options + + +def quantization_kwargs() -> dict[str, Any]: + """A16W8 -- what the overlays are built for. + + Pass to ``onnxruntime.quantization.quantize_static`` alongside + ``quant_format=QuantFormat.QDQ``. The model must be opset 21 or later, so + that 16-bit QDQ is expressible without contrib ops. + """ + from onnxruntime.quantization import QuantType + + return { + "activation_type": QuantType.QUInt16, + "weight_type": QuantType.QInt8, + "per_channel": False, + } + + +def nodes_on_npu(profile_path: str | Path, name: str = "vitisai") -> int: + """How many nodes the EP actually took, from an ORT profile. + + The only trustworthy check. A session that reports the EP in + ``get_providers()`` may still be running every node on the CPU, and the + outputs will be bit-identical to the CPU when it does -- which is exactly + what a *correct* NPU run does not look like. + """ + import json + + events = json.loads(Path(profile_path).read_text()) + return sum( + 1 for e in events + if e.get("cat") == "Node" and (e.get("args") or {}).get("provider") == name + ) diff --git a/vllm_omni/edge/placement.py b/vllm_omni/edge/placement.py new file mode 100644 index 000000000..51ecf44de --- /dev/null +++ b/vllm_omni/edge/placement.py @@ -0,0 +1,346 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Where each parameter lives when the checkpoint is bigger than the VRAM. + +:mod:`vllm_omni.edge.weight_path` picks a *format* from the device. This picks +a *placement*, which is a different question and only arises when the two +numbers do not fit: Qwen3.8-27B-FP8 is 30.89 GB of weights against a 24 GiB +card. vLLM will do this for you -- ``cpu_offload_gb`` maps host memory into the +GPU's address space and the kernels read it over PCIe -- but it offloads +"non-selectively until the memory limit is reached", and that is the expensive +way to do it. + +**The cost of getting it wrong, measured.** On this laptop (RTX 5090 Laptop +24 GiB, PCIe 4.0 x8, 13.9 GB/s pinned H2D, 472 GB/s VRAM) Qwen3.8-27B-FP8 runs +at **1.05 tok/s** with the smallest bulk offload that fits (10 GB; 9 GB fails +to leave room for the cache). Decode reads every weight once per token, so a +byte on the host side of a 13.9 GB/s link costs 34x what the same byte costs in +472 GB/s VRAM. At 10 GB offloaded that is 719 ms of a ~950 ms step: **the +offload is 75% of the decode time and the other 18.8 GB of weights are 4%.** + +**Which bytes those are is a choice.** Not every parameter is read once per +token, and the ones that are not are free to evict. Qwen3.8-27B-FP8's +``outside.safetensors`` and ``mtp.safetensors`` hold 6.5 GB, of which: + + lm_head 2.543 GB HOT -- one 248320 x 5120 GEMM per token + embed_tokens 2.543 GB COLD -- a gather; one row, ~10 KB per token + visual.* 0.922 GB COLD -- unless the prompt has an image + mtp.* 0.478 GB COLD -- unless speculative decoding is on + +``embed_tokens`` and ``lm_head`` are the same shape and the same dtype, and one +of them is the single best parameter in the model to offload while the other is +among the worst. A size-ordered, name-blind policy cannot tell them apart; +``cpu_offload_params`` can, because it matches name segments. + +So the plan here is: close the deficit out of the cold set first, and only +spend the hot set when the cold set runs out. On this checkpoint the cold set +is **3.94 GB**, which is most of the 8.4 GB deficit, and the arithmetic says +the remainder costs ~324 ms/token instead of 719. + +**Two limits, both found the hard way (2026-09-15).** + +*The cold set is not actually offloadable in vLLM today.* The plan below sorts +parameters into hot and cold and emits ``cpu_offload_params`` for the cold ones. +That flag exists and matches name segments as documented -- but the offloader is +only ever **offered decoder layers**: ``wrap_modules`` is called from inside +``make_layers`` (``vllm/model_executor/models/utils.py:812``) and nowhere else. +``embed_tokens``, ``lm_head``, the vision tower and the MTP head are built +outside it and can never be reached, whatever the flag says. Measured: asking +for ``{embed_tokens, visual, mtp}`` on a 21.92 GB checkpoint with a 4.31 GB cold +set moved **~0.5 GiB**, and the run failed identically to one with no offload at +all. The classification here is still right about which bytes are cheap to move; +the *mechanism* is missing, and closing it means an upstream change. Until then +treat ``cold_offload`` plans as advisory. + +*Fewer bytes is not always faster.* This module prices weights by traffic, which +assumes decode is bandwidth-bound. That held for Qwen3.8-27B-FP8 to within 4.3% +(see the test). It does **not** hold for the NVFP4 build of the same model: +21.95 GB of NVFP4 measured **0.67 tok/s** against 30.87 GB of FP8 at +**1.05 tok/s** -- slower while 30% smaller, and 4.7x off what the traffic model +predicts. Enabling torch.compile and CUDA graphs moved it to 0.72, so it is not +an eager-mode artifact: that path is compute-bound, not bandwidth-bound, and a +model of bytes cannot see it. Never recommend a format change on this module's +estimate alone -- measure it. + +Estimates here are arithmetic, not measurements: bytes over a bandwidth. They +are labelled ``est_`` and the measured points they were checked against are in +the table above. +""" + +from __future__ import annotations + +import json +import struct +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +GiB = 2**30 + +# Name segments that are read once per decode step, and segments that are not. +# Matching is on dot-separated segments, the same rule vLLM's +# ``cpu_offload_params`` uses, so what is planned here is what it will do. +COLD_SEGMENTS: tuple[str, ...] = ( + "embed_tokens", # a gather: one row per token, not the whole table + "visual", # vision tower: only on a prompt that carries an image + "vision_tower", + "mtp", # multi-token-prediction head: only under speculation + "audio_tower", +) +"""Parameters a text decode step does not read. Offloading these is close to +free -- ``embed_tokens`` still costs one row's worth of PCIe per token, which +at 5120 x 2 bytes is 10 KiB against a 13.9 GB/s link, i.e. under a microsecond.""" + +HOT_ALWAYS: tuple[str, ...] = ("lm_head",) +"""Named so the planner can never classify them cold by accident: ``lm_head`` +is the same shape and dtype as ``embed_tokens`` and the opposite decision.""" + + +@dataclass(frozen=True) +class ParamGroup: + name: str + bytes: int + cold: bool + + +@dataclass(frozen=True) +class Placement: + strategy: str + """``resident`` | ``cold_offload`` | ``hot_offload`` | ``infeasible``.""" + device_bytes: int + host_bytes: int + host_read_per_token: int + """Bytes crossing the link every decode step. This, not ``host_bytes``, is + what costs time -- offloading a parameter nothing reads costs nothing.""" + deficit_bytes: int + """How far the checkpoint overran the device budget before any offload.""" + est_link_s_per_token: float + est_vram_s_per_token: float + engine_kwargs: dict[str, Any] + rationale: str + notes: list[str] = field(default_factory=list) + + @property + def est_decode_tok_s(self) -> float: + """Weight-traffic bound only: no attention, no kernel launch, no + sampling. An upper bound, and on the measured point it was 1.24x the + observed rate.""" + total = self.est_link_s_per_token + self.est_vram_s_per_token + return 1.0 / total if total > 0 else float("inf") + + def summary(self) -> str: + return ( + f"{self.strategy}: {self.device_bytes / 1e9:.1f} GB on device, " + f"{self.host_bytes / 1e9:.1f} GB on host of which " + f"{self.host_read_per_token / 1e9:.1f} GB is read per token; " + f"est {self.est_decode_tok_s:.1f} tok/s on weight traffic alone" + ) + + +def is_cold(name: str) -> bool: + """Segment match, so ``model.visual.blocks.0`` is cold and ``visualise`` + is not. ``lm_head`` is never cold whatever else it matches.""" + segments = set(name.split(".")) + if segments & set(HOT_ALWAYS): + return False + return bool(segments & set(COLD_SEGMENTS)) + + +def _safetensors_header(path: Path) -> dict[str, Any]: + with path.open("rb") as f: + n = struct.unpack(" str: + """The unit the planner is allowed to move: a prefix, but never coarser + than one decoder layer. + + A plain 3-segment prefix puts all 64 of Qwen3.8's decoder layers in one + 24.4 GB group, and a planner that can only take groups whole then has a + choice between offloading no layers and offloading every layer. Keeping the + layer index makes each layer its own 0.38 GB group, which is also the + granularity vLLM's prefetch offloader works at. + """ + parts = name.split(".") + for i, part in enumerate(parts): + if part.isdigit() and i > 0 and parts[i - 1].endswith("layers"): + return ".".join(parts[: i + 1]) + return ".".join(parts[:top_segments]) + + +def checkpoint_param_groups(model_dir: str | Path, top_segments: int = 3) -> list[ParamGroup]: + """Bytes per parameter group, read from the safetensors headers. + + Reads headers only -- a few kilobytes per shard -- so this is cheap enough + to run at deploy-planning time on a 31 GB checkpoint. Walks every + ``*.safetensors``, which is how Qwen3.8's per-layer shards are laid out. + """ + root = Path(model_dir) + shards = sorted(root.glob("*.safetensors")) + if not shards: + raise FileNotFoundError(f"no *.safetensors under {root}") + + totals: dict[str, int] = {} + cold: dict[str, bool] = {} + for shard in shards: + for key, meta in _safetensors_header(shard).items(): + if key == "__metadata__": + continue + offs = meta.get("data_offsets") + if not offs: + continue + prefix = group_key(key, top_segments) + totals[prefix] = totals.get(prefix, 0) + (offs[1] - offs[0]) + # A group is cold only if every tensor in it is. + cold[prefix] = cold.get(prefix, True) and is_cold(key) + return sorted( + (ParamGroup(k, v, cold[k]) for k, v in totals.items()), + key=lambda g: -g.bytes, + ) + + +def _engine_kwargs( + offloaded: list[ParamGroup], hot: list[ParamGroup], freed: int +) -> dict[str, Any]: + """Translate a plan into flags vLLM actually has. + + The cold set is expressible exactly: ``cpu_offload_params`` matches name + segments, and every cold group is named by one. The hot set is not -- it is + "these 12 specific layers", and the only knob for that is the prefetch + offloader's ``offload_group_size`` / ``offload_num_in_group``, which takes + the last N of every group of M. So the hot half is emitted as the closest + (M, N) whose count matches, and the caller is told it is an approximation. + """ + segments = sorted( + { + seg + for g in offloaded + if g.cold + for seg in g.name.split(".") + if seg in COLD_SEGMENTS + } + ) + kwargs: dict[str, Any] = { + "cpu_offload_gb": round(freed / GiB + 0.5, 1), + "cpu_offload_params": segments, + } + if hot: + n_layers = sum(1 for g in hot if group_key(g.name) == g.name and + any(p.isdigit() for p in g.name.split("."))) + n_layers = n_layers or len(hot) + # Offload 1 layer in every ceil(total/n) -- the coarsest pattern that + # reaches the count. The caller re-plans if the fit is not exact. + kwargs["offload_group_size"] = max(2, round(64 / max(n_layers, 1))) + kwargs["offload_num_in_group"] = 1 + return kwargs + + +def plan_placement( + groups: list[ParamGroup], + *, + device_bytes_available: int, + link_bytes_per_s: float, + vram_bytes_per_s: float, + embedding_row_bytes: int = 0, +) -> Placement: + """Decide what goes to the host, cheapest bytes first. + + ``device_bytes_available`` is what is left for *weights* after the KV + cache, the recurrent state, activations and the driver's own reservation -- + :func:`vllm_omni.edge.kv_budget.kv_budget_for` sizes the first two. + """ + weights = sum(g.bytes for g in groups) + deficit = weights - device_bytes_available + + if deficit <= 0: + return Placement( + strategy="resident", + device_bytes=weights, host_bytes=0, host_read_per_token=0, + deficit_bytes=0, + est_link_s_per_token=0.0, + est_vram_s_per_token=weights / vram_bytes_per_s, + engine_kwargs={}, + rationale=( + "The checkpoint fits. Every byte is read at VRAM speed and " + "nothing crosses the link, which is the only placement that " + "does not pay 34x for some of its weights." + ), + ) + + # Cheapest first: parameters a decode step never reads. + offloaded: list[ParamGroup] = [] + freed = 0 + for g in sorted((g for g in groups if g.cold), key=lambda g: -g.bytes): + if freed >= deficit: + break + offloaded.append(g) + freed += g.bytes + + cold_only = freed >= deficit + hot_offloaded: list[ParamGroup] = [] + if not cold_only: + # Still short. Spend hot parameters, largest first -- they all cost the + # same per byte, so the only thing that matters is hitting the target + # in as few groups as possible. + for g in sorted((g for g in groups if not g.cold), key=lambda g: -g.bytes): + if freed >= deficit: + break + hot_offloaded.append(g) + offloaded.append(g) + freed += g.bytes + + if freed < deficit: + return Placement( + strategy="infeasible", + device_bytes=device_bytes_available, host_bytes=freed, + host_read_per_token=sum(g.bytes for g in hot_offloaded), + deficit_bytes=deficit - freed, + est_link_s_per_token=float("inf"), est_vram_s_per_token=0.0, + engine_kwargs={}, + rationale=( + f"{(deficit - freed) / 1e9:.1f} GB short even with every " + "parameter on the host. A smaller checkpoint is the only " + "remaining lever." + ), + ) + + hot_bytes = sum(g.bytes for g in hot_offloaded) + # A cold embedding still pays one row per token; everything else cold is 0. + read_per_token = hot_bytes + ( + embedding_row_bytes + if any("embed_tokens" in g.name for g in offloaded) + else 0 + ) + device_bytes = weights - freed + notes = [ + f"cold offload: {', '.join(g.name for g in offloaded if g.cold)}" + or "cold offload: none", + ] + if hot_offloaded: + notes.append( + "hot offload (read every token): " + + ", ".join(g.name for g in hot_offloaded) + ) + return Placement( + strategy="cold_offload" if cold_only else "hot_offload", + device_bytes=device_bytes, + host_bytes=freed, + host_read_per_token=read_per_token, + deficit_bytes=deficit, + est_link_s_per_token=read_per_token / link_bytes_per_s, + est_vram_s_per_token=device_bytes / vram_bytes_per_s, + engine_kwargs=_engine_kwargs(offloaded, hot_offloaded, freed), + rationale=( + f"{deficit / 1e9:.1f} GB over budget. " + + ( + "Closed entirely out of parameters a decode step does not " + "read, so the link is idle between steps." + if cold_only + else f"The cold set gave {(freed - hot_bytes) / 1e9:.1f} GB; " + f"the remaining {hot_bytes / 1e9:.1f} GB is read every token " + f"and is what the decode rate is bound by." + ) + ), + notes=notes, + ) diff --git a/vllm_omni/edge/qnn_export.py b/vllm_omni/edge/qnn_export.py index 4e1b7fc4a..a11f91bf8 100644 --- a/vllm_omni/edge/qnn_export.py +++ b/vllm_omni/edge/qnn_export.py @@ -76,7 +76,10 @@ __all__ = [ "DecodeStepConfig", "DecodeStep", "config_from_qwen3_tts", + "config_from_minicpmo", + "load_decoder_weights", "load_qwen3_tts_weights", + "load_minicpmo_tts_weights", "example_inputs", "input_names", "output_names", @@ -106,6 +109,8 @@ class DecodeStepConfig: vocab_size: int rms_norm_eps: float = 1e-6 rope_theta: float = 1e6 + qk_norm: bool = True + """Qwen3-style per-head RMSNorm on q and k. Llama-style decoders (MiniCPM-o's TTS head) do not have it.""" def __post_init__(self) -> None: if self.num_attention_heads % self.num_key_value_heads: @@ -147,8 +152,11 @@ class _Layer(nn.Module): cfg = self.cfg kv, d, g = cfg.num_key_value_heads, cfg.head_dim, cfg.group_size h = _rms(x, self.input_layernorm, cfg.rms_norm_eps) - q = _rms(self.q_proj(h).view(1, cfg.num_attention_heads, 1, d), self.q_norm, cfg.rms_norm_eps) - k = _rms(self.k_proj(h).view(1, kv, 1, d), self.k_norm, cfg.rms_norm_eps) + q = self.q_proj(h).view(1, cfg.num_attention_heads, 1, d) + k = self.k_proj(h).view(1, kv, 1, d) + if cfg.qk_norm: + q = _rms(q, self.q_norm, cfg.rms_norm_eps) + k = _rms(k, self.k_norm, cfg.rms_norm_eps) v = self.v_proj(h).view(1, kv, 1, d) q, k = _rope(q, cos, sin), _rope(k, cos, sin) keys = torch.cat([k_cache, k], dim=2) @@ -206,6 +214,75 @@ def example_inputs(cfg: DecodeStepConfig, cache_len: int, position: int | None = return (x, cos, sin, mask, *caches) +def config_from_minicpmo(hf_config: dict[str, Any]) -> DecodeStepConfig: + """Build a :class:`DecodeStepConfig` for MiniCPM-o 4.5's TTS head from its ``config.json`` dict. + + A Llama-style decoder (no q/k norm), single codebook (``num_vq == 1``), speech tokens at 25 Hz. + """ + cfg = hf_config["tts_config"] if "tts_config" in hf_config else hf_config + heads = cfg["num_attention_heads"] + return DecodeStepConfig( + hidden_size=cfg["hidden_size"], + num_hidden_layers=cfg["num_hidden_layers"], + num_attention_heads=heads, + num_key_value_heads=cfg.get("num_key_value_heads", heads), + head_dim=cfg.get("head_dim", cfg["hidden_size"] // heads), + intermediate_size=cfg["intermediate_size"], + vocab_size=cfg["num_audio_tokens"], + rms_norm_eps=cfg.get("rms_norm_eps", 1e-5), + rope_theta=cfg.get("rope_theta", 10000.0), + qk_norm=False, + ) + + +def load_decoder_weights( + module: DecodeStep, + safetensors_path: str | Path, + prefix: str, + head_key: str, + *, + layer_prefix: str = "model.layers.", + norm_key: str = "model.norm.weight", + head_is_weight_normed: bool = False, +) -> int: + """Copy decoder weights into ``module`` from one safetensors file. + + ``prefix`` is prepended to every key (``"talker."`` for Qwen3-TTS, ``"tts."`` for MiniCPM-o). ``head_key`` names + the output projection. ``head_is_weight_normed`` reconstructs a ``torch.nn.utils.parametrizations.weight_norm`` + head from its stored ``original0`` (magnitude) and ``original1`` (direction) tensors. + """ + from safetensors import safe_open + + loaded = 0 + with safe_open(str(safetensors_path), "pt") as st: + get = lambda key: st.get_tensor(prefix + key).float() # noqa: E731 + for i, layer in enumerate(module.layers): + p = f"{layer_prefix}{i}." + for name, param in (("input_layernorm", layer.input_layernorm), ("post_attention_layernorm", layer.post_attention_layernorm)): + param.data.copy_(get(p + name + ".weight")) + loaded += 1 + if module.cfg.qk_norm: + for name, param in (("q_norm", layer.q_norm), ("k_norm", layer.k_norm)): + param.data.copy_(get(p + "self_attn." + name + ".weight")) + loaded += 1 + for name in ("q_proj", "k_proj", "v_proj", "o_proj"): + getattr(layer, name).weight.data.copy_(get(p + "self_attn." + name + ".weight")) + loaded += 1 + for name in ("gate_proj", "up_proj", "down_proj"): + getattr(layer, name).weight.data.copy_(get(p + "mlp." + name + ".weight")) + loaded += 1 + module.norm.data.copy_(get(norm_key)) + loaded += 1 + if head_is_weight_normed: + g = get(head_key + ".parametrizations.weight.original0") + v = get(head_key + ".parametrizations.weight.original1") + module.lm_head.weight.data.copy_(g * v / v.norm(dim=tuple(range(1, v.dim())), keepdim=True)) + else: + module.lm_head.weight.data.copy_(get(head_key)) + loaded += 1 + return loaded + + def config_from_qwen3_tts(hf_config: dict[str, Any], component: str) -> DecodeStepConfig: """Build a :class:`DecodeStepConfig` from a Qwen3-TTS ``config.json`` dict. ``component``: talker | predictor.""" talker = hf_config["talker_config"] if "talker_config" in hf_config else hf_config @@ -225,31 +302,17 @@ def config_from_qwen3_tts(hf_config: dict[str, Any], component: str) -> DecodeSt def load_qwen3_tts_weights(module: DecodeStep, safetensors_path: str | Path, component: str, lm_head_index: int = 0) -> int: """Copy real weights into ``module``. ``component``: ``talker`` (codec head) or ``predictor`` (``lm_head.``).""" - from safetensors import safe_open + return load_decoder_weights( + module, + safetensors_path, + prefix="talker." if component == "talker" else "talker.code_predictor.", + head_key="codec_head.weight" if component == "talker" else f"lm_head.{lm_head_index}.weight", + ) - prefix = "talker." if component == "talker" else "talker.code_predictor." - head_key = "codec_head.weight" if component == "talker" else f"lm_head.{lm_head_index}.weight" - loaded = 0 - with safe_open(str(safetensors_path), "pt") as st: - get = lambda key: st.get_tensor(prefix + key).float() # noqa: E731 - for i, layer in enumerate(module.layers): - p = f"model.layers.{i}." - for name, param in (("input_layernorm", layer.input_layernorm), ("post_attention_layernorm", layer.post_attention_layernorm)): - param.data.copy_(get(p + name + ".weight")) - loaded += 1 - for name, param in (("q_norm", layer.q_norm), ("k_norm", layer.k_norm)): - param.data.copy_(get(p + "self_attn." + name + ".weight")) - loaded += 1 - for name in ("q_proj", "k_proj", "v_proj", "o_proj"): - getattr(layer, name).weight.data.copy_(get(p + "self_attn." + name + ".weight")) - loaded += 1 - for name in ("gate_proj", "up_proj", "down_proj"): - getattr(layer, name).weight.data.copy_(get(p + "mlp." + name + ".weight")) - loaded += 1 - module.norm.data.copy_(get("model.norm.weight")) - module.lm_head.weight.data.copy_(get(head_key)) - loaded += 2 - return loaded + +def load_minicpmo_tts_weights(module: DecodeStep, safetensors_path: str | Path) -> int: + """Copy MiniCPM-o 4.5 TTS-head weights (``tts.*``) into ``module``; its code head is weight-normalised.""" + return load_decoder_weights(module, safetensors_path, prefix="tts.", head_key="head_code.0", head_is_weight_normed=True) def sanitize_onnx(path: str | Path) -> dict[str, int]: @@ -261,6 +324,15 @@ def sanitize_onnx(path: str | Path) -> dict[str, int]: """ import onnx + path = Path(path) + # Whether the file on disk keeps its weights beside it. ``onnx.load`` + # materialises them either way, so this has to be read before the load or + # the information is gone -- and saving a 1.8 GB tower back inline would + # both orphan the ``.data`` file and run at the 2 GB protobuf ceiling. + external = any( + init.data_location == onnx.TensorProto.EXTERNAL + for init in onnx.load(str(path), load_external_data=False).graph.initializer + ) model = onnx.load(str(path)) graph = model.graph produced = {out: node for node in graph.node for out in node.output} @@ -309,7 +381,23 @@ def sanitize_onnx(path: str | Path) -> dict[str, int]: for out in graph.output: if out.name in renamed: out.name = renamed[out.name] - onnx.save(model, str(path)) + if external: + # Rewrite the pair. ``convert_model_to_external_data`` only sets the + # location metadata; ``onnx.save`` is what writes the blob, and the old + # one has to go first or stale tensors survive under the same name. + data_file = path.name + ".data" + stale = path.with_name(data_file) + if stale.exists(): + stale.unlink() + onnx.save( + model, + str(path), + save_as_external_data=True, + all_tensors_to_one_file=True, + location=data_file, + ) + else: + onnx.save(model, str(path)) return stats diff --git a/vllm_omni/engine/stage_engine_startup.py b/vllm_omni/engine/stage_engine_startup.py index 71ba17c7c..153a20567 100644 --- a/vllm_omni/engine/stage_engine_startup.py +++ b/vllm_omni/engine/stage_engine_startup.py @@ -1169,7 +1169,22 @@ def launch_stage_replica( from vllm_omni.engine.stage_engine_core_proc_manager import StageEngineCoreProcManager - addresses = get_engine_zmq_addresses(vllm_config) + # [edge-infer W1] Ask for concrete ports, not "tcp://host:0" placeholders. + # + # vLLM mints ":0" by default and expects the *consumer* to bind, recover the + # kernel-assigned port via zmq.LAST_ENDPOINT and write it back before the + # engine handshake -- a recovery that needs + # client_addresses["actual_address_pipe"]. This colocated path cannot do it: + # it starts the child *first*, handing it the raw placeholder, and + # StageRuntime._client_addresses_from_zmq never passes that pipe. The child + # then "connects" to port 0, nothing is ever transmitted, and the client + # waits out VLLM_ENGINE_READY_TIMEOUT_S for a ready message that cannot come. + # + # No-op on Linux, where client_local_only is true and _addr() takes the + # ipc:// branch regardless of this flag -- which is exactly why the bug is + # invisible until the engine runs somewhere without ipc:// (measured: native + # Windows mints ":0" here, WSL mints real ipc:// paths). + addresses = get_engine_zmq_addresses(vllm_config, defer_api_server_ports=False) handshake_address = get_open_zmq_ipc_path() engines_to_handshake = [CoreEngine(index=0, local=True)] with scoped_spawn_device_env(stage_visible_devices, spawn_device_lock, stage_id=stage_id, replica_id=replica_id): diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index f2bc42ab1..176cfc679 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -12,7 +12,12 @@ out of StageEngineCoreClient into reusable functions. from __future__ import annotations import copy -import fcntl +# [edge-infer W1] fcntl does not exist on Windows; this module is in the +# import closure of AsyncOmni, so the bare import makes the engine +# unimportable before anything can be configured. The locking logic is +# sound and torch-free -- only the primitive changes. The shim keeps the +# BlockingIOError contract this file's stale-lock path depends on. +from vllm_omni._filelock_compat import flock_exclusive_nb, funlock import importlib import json import multiprocessing as mp @@ -1750,7 +1755,7 @@ def acquire_device_locks( try: lock_fd, lock_writable = open_device_lock_file(lock_file) try: - fcntl.flock(lock_fd, fcntl.LOCK_EX | fcntl.LOCK_NB) + flock_exclusive_nb(lock_fd) record_lock_holder_pid(lock_fd, lock_writable) lock_acquired = True lock_fds.append(lock_fd) @@ -1803,7 +1808,7 @@ def release_device_locks(lock_fds: list[int]) -> None: """Release file locks acquired by acquire_device_locks.""" for lock_fd in lock_fds: try: - fcntl.flock(lock_fd, fcntl.LOCK_UN) + funlock(lock_fd) os.close(lock_fd) logger.debug("Released initialization lock (fd=%s)", lock_fd) except (OSError, ValueError): diff --git a/vllm_omni/model_executor/models/spark2_5/pipeline.py b/vllm_omni/model_executor/models/spark2_5/pipeline.py new file mode 100644 index 000000000..ef52adc27 --- /dev/null +++ b/vllm_omni/model_executor/models/spark2_5/pipeline.py @@ -0,0 +1,43 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Spark-X2.5 as an Omni pipeline: one autoregressive text stage. + +Spark's model definition has been in ``vllm_omni.model_executor.models`` for a +while, but only as a vLLM architecture -- there was no pipeline, so +``Omni(model=)`` had nothing to resolve and the model was +reachable only through plain ``vllm.LLM``. That is the gap M0 closes: the text +mode has to run *through* the Omni stage runtime, or none of the stage +lifecycle, admission or cancellation work applies to it. + +There is exactly one stage, and that is the design, not a simplification to be +undone later. The proposal's first decision is that an autoregressive session +stays on one backend by default: prefill and decode share the KV cache, and a +per-token boundary between devices buys nothing and costs a synchronization. +Spark's tokenizer, chat template and sampling are model-adapter concerns that +run in-process around this stage; they are not stages. +""" + +from vllm_omni.config.stage_config import ( + PipelineConfig, + StageExecutionType, + StagePipelineConfig, +) + +SPARK2_5_PIPELINE = PipelineConfig( + model_type="spark2_5", + model_arch="Spark2_5ForCausalLM", + hf_architectures=("Spark2_5ForCausalLM",), + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="spark2_5", + execution_type=StageExecutionType.LLM_AR, + input_sources=(), + final_output=True, + final_output_type="text", + owns_tokenizer=True, + engine_output_type="text", + model_arch="Spark2_5ForCausalLM", + ), + ), +) diff --git a/vllm_omni/model_executor/models/spark2_5/spark2_5.py b/vllm_omni/model_executor/models/spark2_5/spark2_5.py index 2ecda4cec..f19de7bdc 100644 --- a/vllm_omni/model_executor/models/spark2_5/spark2_5.py +++ b/vllm_omni/model_executor/models/spark2_5/spark2_5.py @@ -494,9 +494,12 @@ class Spark2_5ForCausalLM(nn.Module, SupportsLoRA, SupportsPP): inputs_embeds: torch.Tensor | None = None, **kwargs, ) -> torch.Tensor | IntermediateTensors: - return self.model( - input_ids, positions, intermediate_tensors, inputs_embeds, **kwargs - ) + # Omni's AR model runner hands every stage model ``sampling_metadata``, + # ``logits_index`` and ``sampler``; the omni-native models take them in + # ``**kwargs``. Spark samples through vLLM's own sampler, and the + # decoder stack reads none of them, so they stop here rather than + # widening ``Spark2_5Model.forward`` to a signature it ignores. + return self.model(input_ids, positions, intermediate_tensors, inputs_embeds) def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None: return self.logits_processor(self.lm_head, hidden_states) diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index 3d689d11a..5b65d2061 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -1985,6 +1985,44 @@ class OmniGPUModelRunner(GPUModelRunner): update_dict = {out_key[0]: {out_key[1]: code_predictor_codes[idx : idx + 1]}} self._update_intermediate_buffer(req_id, update_dict) + def _accepted_forward_kwargs(self) -> frozenset[str] | None: + """Names the stage model's ``forward`` will accept, or ``None`` for any. + + Omni-native stage models take ``**kwargs`` and pick out what they need, + so the AR runner hands every stage ``sampling_metadata``, + ``logits_index`` and ``sampler`` unconditionally. A model that vLLM owns + does not: ``LlamaForCausalLM.forward`` is + ``(input_ids, positions, intermediate_tensors, inputs_embeds)`` and + nothing else, so passing the omni extras is an immediate ``TypeError`` + on the first forward -- which is what "the model is in the registry" not + meaning "the model runs under Omni" looks like in practice (proposal + 2.4, item 12). + + Dropping them is the correct behaviour rather than a workaround: those + arguments exist for stages that sample inside the model, and a + vLLM-native causal LM samples through vLLM's own sampler afterwards. + Computed once per runner -- ``inspect.signature`` on every decode step + would be absurd. + """ + cached = getattr(self, "_omni_forward_kwargs_cache", "unset") + if cached != "unset": + return cached + + import inspect + + accepted: frozenset[str] | None + try: + parameters = inspect.signature(type(self.model).forward).parameters + except (TypeError, ValueError): + accepted = None # not introspectable; assume it copes, as before + else: + if any(p.kind is inspect.Parameter.VAR_KEYWORD for p in parameters.values()): + accepted = None + else: + accepted = frozenset(parameters) - {"self"} + self._omni_forward_kwargs_cache = accepted + return accepted + def _model_forward( self, input_ids: torch.Tensor | None = None, @@ -1995,6 +2033,20 @@ class OmniGPUModelRunner(GPUModelRunner): ): """Inject omni-specific kwargs into forward and cache model output""" model_kwargs_extra = self._build_model_kwargs_extra() + accepted = self._accepted_forward_kwargs() + if accepted is not None: + dropped = sorted((set(model_kwargs) | set(model_kwargs_extra)) - accepted) + if dropped and not getattr(self, "_omni_logged_dropped_kwargs", False): + self._omni_logged_dropped_kwargs = True + logger.info( + "[Omni] %s.forward does not accept %s; passing the rest. This is " + "the expected path for a vLLM-native model used as an omni AR " + "stage -- it samples through vLLM's sampler, not its own.", + type(self.model).__name__, + dropped, + ) + model_kwargs = {k: v for k, v in model_kwargs.items() if k in accepted} + model_kwargs_extra = {k: v for k, v in model_kwargs_extra.items() if k in accepted} update_decode_metadata = getattr(self.model, "update_decode_step_metadata", None) if getattr(self.model, "supports_omni_decode_step_metadata", False) and callable(update_decode_metadata): update_decode_metadata(