[edge-infer] WIP snapshot: uncommitted M0/W1/W4/W5 edge work as of 2026-09-18
Verbatim copy of the edge/omni-edge working tree (tracked diffs + untracked edge/local, tests/edge/local, docs/edge, pipeline files). Taken so that vLLM-Omni v0.29.0rc1 can be merged on a branch without touching the user's working tree. logs_edge/ and benchmark .npy outputs excluded. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
This commit is contained in:
tzhouamandClaude Fable 5.1 committed
1 parent
3834de3105
commit
4e3e8a1ebc
62 files changed
+11043
-56
No files matched your search
@@ -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
|
||||
|
||||
@@ -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)
|
||||
@@ -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 <phone-ip>: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.
|
||||
@@ -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 <ckpt> # exit 2 == explicit refusal
|
||||
python -m vllm_omni.edge.local run --model <ckpt> --prompt "..."
|
||||
python -m vllm_omni.edge.local accept --model <ckpt> --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=<Spark checkpoint>)` 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 <ckpt> --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).
|
||||
@@ -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",)
|
||||
Whitespace-only changes.
@@ -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()
|
||||
Whitespace-only changes.
+96
@@ -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
|
||||
+110
@@ -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())
|
||||
+103
@@ -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
|
||||
+94
@@ -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()
|
||||
+209
@@ -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
|
||||
+64
@@ -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)
|
||||
+106
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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"]
|
||||
@@ -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)
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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("<Q", len(blob)) + blob + b"\0" * off)
|
||||
|
||||
|
||||
def test_param_groups_are_read_from_headers_without_loading_weights(tmp_path):
|
||||
_write_shard(tmp_path / "a.safetensors", [
|
||||
("model.language_model.embed_tokens.weight", 4096),
|
||||
("lm_head.weight", 4096),
|
||||
("model.language_model.layers.0.mlp.down_proj.weight", 2048),
|
||||
("model.language_model.layers.1.mlp.down_proj.weight", 2048),
|
||||
])
|
||||
groups = {g.name: g for g in checkpoint_param_groups(tmp_path)}
|
||||
assert groups["model.language_model.embed_tokens"].cold
|
||||
assert not groups["lm_head.weight"].cold
|
||||
assert set(groups) >= {"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
|
||||
@@ -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"]
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
Executable
+69
@@ -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"]
|
||||
@@ -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,
|
||||
|
||||
@@ -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/<name>/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=<checkpoint>)`` 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``).
|
||||
"""
|
||||
@@ -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()
|
||||
|
||||
@@ -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=<json>`` 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())
|
||||
+178
-16
@@ -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,
|
||||
)
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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())
|
||||
Whitespace-only changes.
@@ -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()
|
||||
@@ -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())
|
||||
@@ -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)"
|
||||
@@ -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))
|
||||
@@ -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,
|
||||
}
|
||||
Whitespace-only changes.
+403
@@ -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
|
||||
+288
@@ -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)"
|
||||
+185
@@ -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
|
||||
+522
@@ -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."
|
||||
)
|
||||
+290
@@ -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())
|
||||
+370
@@ -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())
|
||||
@@ -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/<revision>/``. 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,
|
||||
)
|
||||
@@ -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()
|
||||
@@ -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."
|
||||
)
|
||||
@@ -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
|
||||
@@ -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]
|
||||
@@ -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
|
||||
)
|
||||
@@ -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("<Q", f.read(8))[0]
|
||||
return json.loads(f.read(n))
|
||||
|
||||
|
||||
def group_key(name: str, top_segments: int = 3) -> 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,
|
||||
)
|
||||
+115
-27
@@ -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.<i>``)."""
|
||||
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
|
||||
|
||||
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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=<a Spark checkpoint>)`` 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",
|
||||
),
|
||||
),
|
||||
)
|
||||
@@ -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)
|
||||
|
||||
@@ -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(
|
||||
|
||||
Reference in new issue
Block a user