diff --git a/.gitignore b/.gitignore index c3b7432..baad4ad 100644 --- a/.gitignore +++ b/.gitignore @@ -13,6 +13,9 @@ dist/ # Trained adapter artifacts (never committed; see README) artifacts/ +# Downloaded checkpoints (README quick start puts them here) +models/ + # macOS .DS_Store diff --git a/README.md b/README.md index 8d174e0..994bbad 100644 --- a/README.md +++ b/README.md @@ -216,7 +216,7 @@ edge0-35b, 2.8 for edge0-8b** (MMLU-Pro is even above the base). Max 100: ## Benchmark -Measured with `examples/bench.py` (3.3k-token prompt prefill → 10 sampled +Measured with `examples/bench.py` (3.3k-token prompt prefill → 10 greedy warmup steps → 200 timed sampled decode tokens, 2 runs per tier): | Tier | Decode speed | Prefill throughput (cold / warm)* | Peak active memory** | Test machine | diff --git a/README_zh.md b/README_zh.md index 89aef98..f05f5ec 100644 --- a/README_zh.md +++ b/README_zh.md @@ -198,7 +198,7 @@ engine.close() # 释放 mmap / 专家缓存 ## 性能实测 -`examples/bench.py` 实测(3.3k token prompt prefill → 10 步采样 warmup → +`examples/bench.py` 实测(3.3k token prompt prefill → 10 步贪心 warmup → 200 token 计时段,每档 2 轮): | 档位 | 解码速度 | Prefill 吞吐(冷/热)* | 峰值 active 内存** | 测试机器 | diff --git a/docs/nvidia.md b/docs/nvidia.md index b86994e..a5d6fcf 100644 --- a/docs/nvidia.md +++ b/docs/nvidia.md @@ -1,5 +1,492 @@ # NVIDIA / CUDA support — status +## Jetson Orin Nano bring-up + +**Status (16 September 2026): edge0-8b runs on a physical Orin Nano 8 GB, +and the bring-up plan is complete through its final validation.** +Decode **0.77–0.82 tok/s**, prefill 42 tok/s at a 3.3k-token prompt, +inside a byte budget that rejects what will not fit before the model +loads, with token choices matching the reference implementation on this +device. The claim covers exactly the board, build, checkpoint, profile +and workloads in the [tested operating +envelope](#tested-operating-envelope-task-7-final-validation): single +request, greedy sampling, ≤ 4096 declared context, ≤ 256 validated +generation, 25W mode. The route there is the +[Task 3 baseline](#task-3-orin-baseline-summary) (0.553 tok/s +unoptimized), the [byte budget](#memory-budgeted-profile-task-4), the +[transfer knobs](#bounded-expert-reuse-and-predicted-prefetch-task-5-opt-in) +and the [quantized paths](#task-6-orin-acceptance) — **+42% end to end** +over the baseline. + +Before attempting inference on an Orin Nano, capture a machine-readable +capability report with the read-only probe: + +```bash +EDGE0_BACKEND=cuda python scripts/jetson_probe.py \ + --model-dir "$EDGE0_8B_MODEL" \ + --output artifacts/orin-probe.json +``` + +The command exits non-zero and lists actionable blockers when it cannot confirm +a Jetson Orin, CUDA-enabled PyTorch, and NVMe-backed model storage. It reports +facts only; passing the probe does not imply that inference fits in 8 GB. + +The staged implementation and measurement gates are documented in +[`docs/plans/jetson-orin-nano.md`](plans/jetson-orin-nano.md) (change +justifications in [`jetson-orin-nano-review.md`](plans/jetson-orin-nano-review.md)). +Start with the 8B tier and retain the existing correctness path while +optimizing memory and I/O. The benchmark's machine-readable report, which the +Orin baseline uses, is described in [Benchmark reporting](#benchmark-reporting) +below. + +### Task 3 Orin baseline summary + +Target: Jetson Orin Nano Developer Kit (Super), 8 GB (7.85 GB observed), +NVMe/ext4, L4T R39.2 (kernel 6.8.12-tegra), 25W power mode (observed, not +changed), Python 3.12.3. Torch: the official **2.14.0+cu130 aarch64 +wheel** — the same build the GB10 work used. It warns that its SASS list +excludes `sm_87`; empirically the `sm_80` binaries run on this Orin +(CC 8.x binary compatibility), and the smoke suite below is the evidence +gate that decides this, not the warning or the docs. Environment: +edge0 installed `pip install -e . --no-deps` plus +numpy/safetensors/tokenizers/psutil/transformers (5.17.0) — no MLX +package on the device; the 8B torch import path needs none. + +Evidence chain (raw reports in the untracked evidence set, sanitized +summary here): + +* **Probe** (`scripts/jetson_probe.py --model-dir …`): `ready: true`, + no blockers; checkpoint manifest sha256 `a83dee78…` (48 files, + 4.60 GB). +* **CUDA acceptance** — `pytest tests/test_torch_cuda_smoke.py + --require-cuda -q`: 6/6 pass on the device (matmul vs float64 + reference, packed-uint32 int4 `gather_qmm` vs a hand-computed + dequantization, RMSNorm, unsigned-word view semantics, core ops). + The suite FAILS rather than skips when CUDA is missing (verified + with `CUDA_VISIBLE_DEVICES=""`). +* **32-token greedy correctness** (`scripts/orin_reference_check.py`, + CUDA generation vs torch-CPU teacher-forced logits; torch-CPU is the + path checked against MLX layer-by-layer on the GB10): **32/32 token + choices identical, zero near-ties consumed** (smallest top-1→top-2 + margin 0.067, median 4.6). The pre-registered vocab-wide logit + bound (max |Δ| ≤ 1.0) was **exceeded: max |Δ| = 2.06** — reported as + the FAIL it is, not re-run with a looser bound. Post-hoc analysis + (recorded in the evidence set): median |Δ| 0.087, p99 0.52; every + excursion above 1.0 sits on deep-tail tokens (worst: reference rank + 429); at the reference's top-8 tokens of every step the difference + is ≤ 0.22. This is accumulated bfloat16 reduction-order noise, not + a kernel defect; a top-k-scoped logit criterion should be registered + in review before the next comparison run, and MLX-generated fixtures + remain the intended reference once available. +* **Short baseline** (`BENCH_TEMP=0 BENCH_SEED=0 examples/bench.py + edge0-8b --ntok 32 --warmup 0`, probe linked, tegrastats at 1 s + alongside, three independent process launches, two internal runs + each, `EDGE0_TORCH_WEIGHT_CACHE=0 EDGE0_PREWARM=0`): + + | Launch | Prefill (37 tok) | Decode (32 tok) | CUDA alloc peak | Peak RSS | + |---|---|---|---|---| + | 1 | 15.9 / 11.7 s | 0.53 / 0.55 tok/s | 0.92 / 1.00 GiB | 5.31 GiB | + | 2 | 13.4 / 11.4 s | 0.55 / 0.55 tok/s | 0.92 / 1.00 GiB | 5.69 GiB | + | 3 | 14.9 / 13.2 s | 0.56 / 0.58 tok/s | 0.92 / 1.00 GiB | 5.76 GiB | + + Six runs, descriptive statistics: decode mean **0.55 tok/s** (min + 0.53, max 0.58); engine build ~11.4 s. GPU utilization pegged at + 99% during decode (tegrastats); no OOM, no swap activity; RSS and + CUDA peaks are separate views (unified memory — RSS includes + touched mmapped expert pages) and are not added. A desktop session + (~3.5 GB) was resident throughout — headroom, not a guarantee. + +This is the **unoptimized reference implementation**: `gather_qmm` +dequantizes every gathered expert on every call, exactly the cost +Tasks 4–6 of the plan exist to remove (the GB10 measured 2.5× from the +dequantized-weight cache alone; it is off here pending Task 4's byte +budget). Do not read 0.55 tok/s as the platform's capability. + +### Memory-budgeted profile (Task 4) + +`EDGE0_MEMORY_BUDGET=auto` (or an explicit byte count) makes the 8B +engine resolve an explicit byte budget immediately before it builds the +streaming caches: observed available RAM (tightened by any address-space +rlimit), minus an OS-growth reserve, an allocator allowance, KV at the +DECLARED context (`EDGE0_BUDGET_CONTEXT`, default 1024 tokens), the +whole-layer prefill transient and in-flight expert builds — every +payload priced from the checkpoint's safetensors header, not from slot +counts. What remains bounds the shared LRU and the prefetch buffer in +bytes; the budget only ever LOWERS the tested profile, and it enforces +two bounds the unbudgeted path leaves open (`PrefetchBuffer.max_cap` +against `prefetch_all()`'s growth, `LayerOptions.max_inflight` against +the speculative build queue). + +Impossible profiles fail BEFORE inference with the itemized arithmetic: +a context whose KV cannot fit, a cache allowance below the working set +(a calculated zero is rejected — `SharedExpertCache(0)` / +`PrefetchBuffer(0)` mean *unbounded*, never "off"), or +`EDGE0_TORCH_WEIGHT_CACHE=1` whose measured ~4.1 GB does not fit the +post-cache headroom. The resolved policy (or the rejection) lands in +the bench report's `caches.memory_budget` group. On the Orin Nano +8 GB baseline workload, `auto` resolves ≈ 4.2 GB usable, keeps the +tested 64-slot / 48-cap profile with ≈ 1.9 GB headroom, and measures +identically to the unbudgeted baseline; a declared 4096-token context +is rejected (KV alone ≈ 4.5 GB), as is the weight cache. Unset, the +engine behaves exactly as before Task 4. + +### Bounded expert reuse and predicted prefetch (Task 5, opt-in) + +The measured Orin profile found the tested 64-slot LRU never hits on +this tier (working set 184 bundles/step) — and also that decode is +compute-bound, so transfers barely show end-to-end until the Task 6 +kernel work lands. Two opt-in knobs, both scheduling-only (token +output is bit-identical; wrong or late predictions fall back to demand +loads, never a staged zero row): + +* `EDGE0_CACHE_SLOTS=512` — request a working-set-sized shared LRU + (the budget may lower it). 85% hit rate on the baseline workload. +* `EDGE0_PREDICT_PREFETCH=1` — the prerouter's per-step predictions + are prefetched for non-staged layers while the current forward runs + (98.8% of predicted builds consumed, zero waited on, zero wasted). + +Together on the baseline workload: decode 0.568–0.623 tok/s (mean +0.594) vs 0.53–0.58 (mean 0.552) — **+7.7%**, for ~0.6 GiB more +resident, inside the resolved budget. Off by default until the +post-Task-6 re-measurement; the pinned-buffer CUDA-stream pipeline is +deferred on the same evidence (see the plan's Task 5 status). + +### Quantized paths (Task 6, developed on a workstation GPU) + +The Orin profile (plan Task 6 scoping) put 37% of a decode step in +`quant.gather_qmm` and 26% in the dense `QuantizedLinear.forward`, both +dequantize-then-matmul. Three paths were developed and parity-tested +on an RTX PRO 6000 (torch 2.14.0+cu130, the Orin's torch version) per +[`plans/rtx6000-task6-handoff.md`](plans/rtx6000-task6-handoff.md). All are +guarded (every unsupported case falls through to the reference +implementation unchanged), priced for the Task 4 budget and +recorded in the bench report's `caches.quant_paths` group. Which one the +Orin adopts, and any tok/s, is decided on the Orin — see +[the acceptance](#task-6-orin-acceptance) below, which made +`EDGE0_QMM_BATCHED` the **default** on this backend and left the two +dense knobs opt-in. + +| Knob | What it does | Exact? | Extra memory | +|---|---|---|---| +| `EDGE0_QMM_BATCHED=1` | `gather_qmm` dequantizes the distinct experts of a call ONCE and runs ONE batched matmul over a padded per-expert slab instead of a Python loop of dequantize+matmul pairs (decode: 8 pairs -> 1 + 1). Affine 4-bit only; 2/8-bit and any call whose priced transient exceeds `EDGE0_QMM_BATCHED_MAX_BYTES` (default 256 MiB) take the reference loop. | yes: same float32 arithmetic, parity 1e-4 to a float64 hand reference and 1e-5 to the loop | per-call transient only, priced by `quant.batched_transient_bytes`; the cap is deducted once by the budget | +| `EDGE0_TORCH_WEIGHT_CACHE_BYTES=` | keeps the dequantized bf16 weight of the dense `QuantizedLinear` modules that fit `n` bytes, admitted first-fit in checkpoint order with each module's own build transient counted (`nn.WeightCachePolicy`); the engine shrinks `n` to the budget's post-cache headroom; the admitted module list is in the report | the dequantized weight is bit-identical to the on-the-fly path (same per-chunk expression, asserted on CUDA); the cached forward then issues one GEMM where the reference issues one per 4096-row chunk, so for the two modules wider than a chunk (lm_head, layer-0 dense gate/up) cuBLAS may split differently — identical on this GPU and on the CPU, and the reference check decides on the device | the admitted bytes plus one build transient at a time: the fill is chunked, so it costs `nn.fill_transient_bytes` (edge0-8b's widest module, lm_head 157184 x 1536: 460.5 MiB resident, 556.9 MiB measured peak) rather than the 3702 MiB a whole-weight dequantization peaks at. edge0-8b's cacheable dense linears total 1.40 GB in bf16: lm_head 0.46 GiB, attention and shared experts the rest. Routers are not quantized in this checkpoint and `word_embeddings` is a `QuantizedEmbedding` that dequantizes only the rows it looks up, so neither is a candidate | +| `EDGE0_INT4PACK=1` | dense `QuantizedLinear` through torch's built-in `_weight_int4pack_mm` after a one-time repack of the MLX layout (uint32 LSB-first codes -> uint8 high-nibble-first pairs, `zero = bias + 8 * scale`), lazily on the first eligible call; CUDA + bf16 activations + 4-bit + group 32/64/128/256 + `in_features % 128 == 0` + `out_features % 8 == 0` only, and `int4pack.probe` must execute the kernel once on the device first (the Orin's wheel warns `sm_87` is outside its SASS list, so execution evidence gates, not the warning); a runtime failure disables the path for that module permanently | **no**: `zero` is rounded to bf16 and the kernel dequantizes in bf16; measured (identity-vector probe, n=512, k=1536): per-weight error median 3.7e-4, p99 3.7e-3, max 1.1e-2; 68% of weights equal the reference's bf16-rounded weights, the rest within ~1-2 bf16 ulps; outputs within 3.5e-3 of the output scale. Pre-registered test bounds: 1e-2 of output scale, 2^-6 (max|w| + max|zero|) per weight | duplicates the int4 payload it covers (`int4pack_extra_bytes`) next to the original buffers; for all of edge0-8b's dense linears that is 0.35 GB | + +Precedence inside `QuantizedLinear.forward`: full cache, then the capped +cache, then the int4 kernel, then the chunked on-the-fly dequantization — +exact paths first. + +**Both dense knobs need bfloat16 activations, and most calls are not.** +Traced on the real edge0-8b model over one prefill plus one decode step, +400 of 470 dense `QuantizedLinear` calls arrive as float32 and only 70 as +bfloat16. The cache therefore fills 35 of the 235 modules it admits +(0.174 GB of the 1.402 GB reserved) and the int4 kernel repacks exactly +those same 35. The reservation stays conservative on purpose — it never +under-reserves — and the report states `filled_bytes` next to +`admitted_bytes` so residency is never overstated. The routed path +(`EDGE0_QMM_BATCHED`) has no such gate and engages on every call. The batched gather touches only the routed-expert +path; the dense knobs touch only the dense path; nothing changes token +choice by construction except the int4 kernel, whose acceptability the +32-token reference check decides (`scripts/orin_reference_check.py` now +takes `--test-env KEY=VALUE` / `--reference-env` so the reference path and an +opt-in path can be compared on the SAME device, isolating the path from +device noise). + +Tests (no MLX, hand-computed references, tolerances written before any Orin +run): `tests/test_quant_paths.py` (every broadcast shape of +`tests/test_cuda_backend.py`, 2/8-bit fall-through, index boundaries, +negative scales, code extremes, bf16 dtypes, transient cap fall-back, the +allocator-measured transient against the price, kernel-launch counts), +`tests/test_weight_cache_policy.py` (admission, caps, re-finalize, bit +identity, loader registration), `tests/test_int4pack.py` (exact repack round +trip, the zero mapping, every guard, the registered kernel bounds, the LoRA +wrapper, exact-cache precedence), `tests/test_task6_budget.py` (pricing). +The CUDA cases skip without a device and FAIL under `--require-cuda`. + +Custom fused kernel (design C): not written. Toolchain recorded for the next +session: WSL Ubuntu nvcc 13.0.88 + g++ 13.3 compile +`-gencode arch=compute_87,code=sm_87 -gencode arch=compute_120,code=sm_120` +fat binaries that run on this GPU; Windows has VS 2022 BuildTools + CUDA +12.8 as an alternative. + +### Task 6 Orin acceptance + +Run on the target (Orin Nano 8 GB, L4T R39.2, CUDA 13.2, torch +2.14.0+cu130, 25W, budget active) on 16 September 2026. Same workload +and instrumentation as the Task 3 baseline: 37-token prompt, 32 greedy +tokens, `BENCH_TEMP=0 BENCH_SEED=0`, three independent launches x two +internal runs, tegrastats alongside. + +**Environment and correctness.** Full suite 341 passed on-device; the +CUDA suites pass under `--require-cuda` (96 tests). `int4pack.probe` +**executes on sm_87** — the kernel the wheel's SASS list does not +advertise runs here, answered by execution as the design intended. + +Same-device 32-token reference checks (`--test-env`, so the only +variable is the knob; the registered bound is Task 3's max |Δlogit| ≤ 1.0): + +| Path | Token choices | max abs logit diff | Verdict | +|---|---|---|---| +| `EDGE0_QMM_BATCHED=1` | 32/32 identical | **0.777** | **PASS** | +| `EDGE0_TORCH_WEIGHT_CACHE_BYTES=1.5e9` | 32/32 identical | 1.58 | FAIL of that bound, recorded | +| `EDGE0_INT4PACK=1` | 32/32 identical | 2.17 | FAIL of that bound, recorded | + +**Decode (6 runs each, mean ± sd):** + +| Configuration | decode tok/s | mean | vs baseline | CUDA peak / RSS | +|---|---|---|---|---| +| baseline (Task 3) | 0.530–0.580 | 0.553 | — | 1.00 / 5.3–5.8 GiB | +| batched gather | 0.623–0.699 | **0.668** | **+20.7%** | 1.00 / 5.6–5.7 GiB | +| batched + Task 5 knobs | 0.743–0.815 | **0.785** | **+41.8%** | 1.58 / 5.8–5.9 GiB | +| capped weight cache (screen) | 0.517 | 0.517 | +0.8% | 1.16 GiB | +| int4 kernel (screen) | 0.524–0.540 | 0.532 | +3.7% | 1.07 GiB | + +**Decision.** `EDGE0_QMM_BATCHED` is now **on by default** on this +backend (`EDGE0_QMM_BATCHED=0` selects the reference loop): it is the +only path that both passed the registered correctness bound on the +target and delivered a measured benefit there, it engages on every +gather call, its transient is bounded and the budget deducts the cap. +The two dense knobs stay **opt-in**: they move decode by ~1–4% because +of the dtype gate below, and both exceed the registered bound. + +**The Task 5 re-measurement Gate D asked for.** With compute reduced, +transfers matter more, exactly as predicted: the Task 5 knobs were worth ++7.7% before Task 6 and are worth **+17.5% on top of the batched gather** +now (0.668 → 0.785), with expert load wall falling 72.5 s → 17.0 s over +the same six runs and 30,159 cache hits against 5,169 loads. They remain +opt-in because they are board-specific tuning that needs the budget +active (`EDGE0_MEMORY_BUDGET=auto EDGE0_CACHE_SLOTS=512 +EDGE0_PREDICT_PREFETCH=1` is the recommended Orin profile, and is what +the 0.785 figure uses). + +**Two corrections to the workstation's report, found only by running here:** + +1. *The capped weight cache is not bit-identical on this GPU.* The + workstation measured max |Δlogit| exactly 0; the Orin measures 1.58. + Isolated: the cached dequantized weight IS bit-identical to the + chunked path's (verified directly), but feeding those same weights + through one GEMM versus the reference's per-4096-row GEMMs differs by + 0.03125 on a 30.6 output scale — one bf16 ulp of accumulation-order + difference in cuBLAS's kernel choice for the two shapes. The + exactness claim is a property of the workstation's cuBLAS heuristics, + not of the code; on sm_87 the path is exact in the weights and + near-exact in the outputs. +2. *The float32 activation gate is an MLA/RoPE promotion, not an + accident.* Confirmed here (35 bf16 vs 200 float32 dense calls per + decode step) and traced to its origin: layer 3 is the first + `BailingMLA`, whose RoPE path computes in float32 and concatenates + `q_nope.to(q_pe.dtype)` with the float32 `q_pe` — deliberate MLX + promotion parity, commented as such in the vendored model. From that + layer on, the residual stream stays float32, so only layers 0–2 (KDA) + present bf16 activations to their dense linears. That is the 35 + modules, and it is a numerics-parity decision, not a kernel + limitation: reaching the rest of the dense 26% means deciding to + diverge from MLX's promotion in the MLA path, which belongs in review + with its own reference check, not in a kernel increment. + +### Tested operating envelope (Task 7 final validation) + +Everything below was measured on the target on 16 September 2026. It +describes **this** board, build, checkpoint and profile on the workloads +listed; it is descriptive evidence from small samples, not a performance +guarantee, and nothing here is claimed for other Orin models, power +modes or context lengths. + +**Configuration under test.** Jetson Orin Nano Developer Kit (Super), +8 GB (7.85 GB observed), NVMe/ext4, L4T R39.2 (kernel 6.8.12-tegra), +CUDA 13.2 driver, 25W power mode (observed, unchanged), Python 3.12.3, +torch 2.14.0+cu130, no MLX installed. Checkpoint +`Edge0/Edge0-8B-A1B-preview`, manifest sha256 `a83dee78…`, with its +bundled LoRA and prerouter adapters. Profile: the shipped default +(batched expert gather on) plus `EDGE0_MEMORY_BUDGET=auto +EDGE0_CACHE_SLOTS=512 EDGE0_PREDICT_PREFETCH=1`. A desktop session +(~3.5 GB) was resident throughout — this is the realistic condition on +this board, not a cleared one. + +| Dimension | Tested envelope | +|---|---| +| Decode, short prompt | 0.744–0.786 tok/s, mean **0.772**, sd 0.016 (n=6: three launches x two runs) | +| Decode, sustained | mean **0.817** tok/s, sd 0.028, min 0.745, max 0.853 (30 requests / 960 tokens / 25 min, one process) | +| Prefill | 13.6–15.6 s at 37 tokens; **42.4–42.8 tok/s** at 3292 tokens | +| Generation length | validated to **256 tokens** in one request (0.807 / 0.820 tok/s over two runs — no decay with generation length) | +| Prompt / context | validated at **3292 prompt tokens**; the budget admits up to a **4096-token declared context** on this board, and rejects beyond that before loading | +| Memory | CUDA allocator peak 1.58 GiB short-prompt, 3.14 GiB at 3292 tokens; process RSS 5.45–5.82 GiB | +| Thermal | junction 54.5–61.8 °C over 25 minutes, **no throttling**; GPU busy in 749 of 774 samples | +| Correctness | 32/32 token choices identical to the reference loop, max abs logit diff **0.777** against the registered bound of 1.0 (PASS); all 30 sustained requests produced byte-identical tokens | + +**KV growth was measured, not inherited.** `kv_bytes_per_token` for +this tier was 1,100,000 — the README's Apple/MLX figure. Measured here +across 128…2048-token prefills, the marginal cost is **561,320 B/token +(548 KiB)**, roughly half of that; growth is sublinear below the top of +the range because most layers of this hybrid are KDA, whose state is +fixed-size, so only the MLA layers' KV grows with context. The constant +is now 600,000 (≈7% margin over the measurement). The practical effect +is a larger honest envelope: the declared context the budget admits on +this board goes from ~2560 to 4096 tokens. Over-reserving was safe but +was rejecting contexts that fit. + +**Repeated requests are stable, and the RSS movement is not a leak.** +Across 30 sequential requests RSS moved between 5.45 and 5.70 GiB, +rising *and falling* (5.45 → 5.65 → 5.53 → 5.70 → 5.69): mmapped expert +pages being reclaimed and re-faulted under memory pressure, which is +what RSS counting mmapped files looks like, not monotonic growth. CUDA +allocator peaks were flat at 1.57 GiB and every request produced +identical tokens. + +**Swap, stated explicitly.** The sustained run *did* swap: +1,165 MB in +the first fifth of the run, then flat (+24, 0, +18, −2 MB across the +remaining four fifths), with no throughput degradation — in fact the +later quartiles were slightly faster (0.791 / 0.834 / 0.816 / 0.824). +That profile is a one-time eviction of the resident desktop's idle pages +to make room, not the model's working set thrashing. Per the plan's +rule a swapping run cannot establish a resident-memory target, so the +memory figures above are stated as observed peaks under these +conditions, and a swap-free claim is **not** made for a board with a +desktop session resident. + +**Known limitations of this envelope.** Single concurrent request only +(no batching, no concurrent sessions tested). Greedy sampling at +temperature 0 for every measurement. Contexts beyond 4096 declared +tokens are rejected by the budget, not degraded gracefully. Prefill at +short prompts is dominated by fixed overhead (2–3 tok/s at 37 tokens vs +42 tok/s at 3292) and was not optimized. edge0-35b remains out of scope +on this board. + +## Benchmark reporting + +`examples/bench.py` keeps its two-run protocol and its human-readable output +and can additionally write a machine-readable report: + +```bash +EDGE0_BACKEND=cuda EDGE0_TORCH_DEVICE=cuda \ + python examples/bench.py edge0-8b --ntok 32 --warmup 0 \ + --probe-json "$EDGE0_RUN_DIR/orin-probe.json" \ + --json-output "$EDGE0_RUN_DIR/orin-short-baseline.json" +``` + +The pure builder (`examples/benchmark_report.py`, standard library only, +imports neither torch nor MLX) turns measured numbers into the schema; the +benchmark-owned collectors (`examples/benchmark_measure.py`) talk to the +backend, the process and the machine behind lazy, injectable imports. + +### What is measured + +* **Protocol** `fixed-length-sample-and-step`, version 1. Per run: + `engine.reset()`, a timed chat-templated prefill, `--warmup` **greedy** + (argmax) steps that are not timed, then `--ntok` timed iterations of "sample + one token from the current logits, run one engine step". The final step's + logits are computed but never sampled and the loop does not stop at EOS. + None of the numbers is a time-to-first-token or an EOS-aware completion + latency; the report's `protocol` block says so. +* **Counts** per run: `prompt_tokens` (after chat templating), + `warmup_tokens`, `timed_decode_tokens`, + `total_generated_tokens = warmup + timed`, + `decode_start_context_tokens = prompt + warmup`; the requested counts are + kept separately (`requested_*`). `decode_tokens_per_second` divides the + timed tokens only. A zero measured duration gives `null` plus a reason, + never a zero rate; `--ntok 0` is refused. +* **Wall time** is `time.perf_counter()` and includes storage, transfer and + CPU sampling costs by design. On an actual CUDA device (`EDGE0_BACKEND=cuda` + *and* a resolved `cuda` device; `EDGE0_TORCH_DEVICE=cpu`/`mps` is not one) + the benchmark calls `torch.cuda.synchronize(device)` after the engine reset, + at the end of prefill, at the end of warmup and after the final timed step: + phase boundaries only, no per-token synchronization. On MLX nothing changes, + the engine's `core.eval(logits)` already materializes every step. Torch on + `mps` is not synchronized by the benchmark (the plan mandates CUDA only), so + its wall times may exclude the final step's in-flight device work; the + report says so in `runtime.measurement.synchronization`. +* **Memory** is integer bytes everywhere, each value with its `method` and + `scope`. Per run (`runs[i].memory`): `cuda_peak_allocated` from the torch + caching allocator, reset immediately before prefill and read after the + synchronized decode. It covers prefill + warmup + decode, includes whatever + was resident at the reset (recorded as `allocated_at_reset`), does not + recover an earlier load transient and does not see the CUDA context, + driver or non-torch memory. `cuda_peak_reserved` is recorded next to it + with its own scope: `reset_peak_memory_stats` does not drop the reserved + peak below the segments the allocator already holds, so its floor is every + segment cached at the reset (`reserved_at_reset`) and run 0 and run 1 are + expected to be close. The headline value (the human `peak_active` line and + the legacy `peak_gib`) is `cuda_peak_allocated`, the analogue of MLX's peak + active memory; `runtime.measurement.headline_peak` names it. On MLX the + same reset point feeds `mlx_peak_active` (`mlx.core.get_peak_memory`); for + torch on any non-CUDA device (cpu, mps) the per-run `backend_peak` is + `null` with the reason. Process-wide (`memory`): `process_peak_rss`, a + lifetime high-water mark (`ru_maxrss`; `peak_wset` on Windows) that + includes model load and every run, and the sampled peaks from a sampler + thread started before the model loads (`--rss-sample-interval`, default + 0.25 s, `0` disables; spikes shorter than the interval are missed): + `process_sampled_peak_rss`, plus on Linux (`/proc/self/status`) + `process_sampled_peak_rss_anon`, `process_sampled_peak_rss_file` and + `process_sampled_peak_swap`. Resident set size counts the resident + file-backed pages of the mmap'd checkpoint (whole-layer prefill touches + every expert of `model.safetensors`, reclaimable page cache) together with + anonymous memory, so it is not the process's anonymous footprint and not + comparable to the "peak anonymous memory" figure further down this page; + the `*_rss_anon` peak is. A non-zero swap peak marks a swapping run, which + cannot establish a resident-memory target. RSS and allocator peaks are + separate views of the same shared DRAM on Jetson; never add them. +* **Unavailable values** are `null` with a reason: every data group + (`identity`, `model`, `runtime`, `workload`, `caches`, `memory`, `probe`, + each entry of `runs[]`, each statistic block of `summary` and every nested + object that can hold a null) carries an `unavailable_reasons` map + (`{field: reason}`), and memory metric objects carry `unavailable_reason`; + `protocol` and the `summary` container hold no nullable fields. An invalid *required* measurement (negative or + non-finite timing, count or byte value) aborts the run instead of being + written. Serialization is strict JSON (`allow_nan=False`). + +### Report layout (`schema_version` 1, independent of the probe's schema) + +| Group | Contents | +|---|---| +| `protocol` | name, version and the labels above | +| `identity` | start/finish timestamps, git commit and tracked-file dirty flag of this checkout (untracked files do not count), `run_count`, `command.argv`, the requested `BENCH_*` / `EDGE0_*` / `MLX_CACHE_LIMIT_MB` / `LING_HIDDEN_CLIP` / `LING_PREWARM` / `PREROUTER_*` knobs as raw strings, snapshotted before the engine is built (the 8B engine sets a default `LING_HIDDEN_CLIP` the user did not request), the requested arguments | +| `model` | resolved local checkpoint path, tier, `model_type` / `architectures`, `config.json` SHA-256, the checkpoint manifest (`files`: every regular file's relative path and size, `algorithm`, and the `sha256` over that list, so the digest is reproducible from the retained list; weights are never content-hashed), SHA-256 of every non-weight file up to 32 MiB (config, chat template, `tokenizer.json`, vocab, merges), safetensors header metadata, tensor counts and header hashes, `tokenizer_files` (name, size, SHA-256), and the LoRA / prerouter adapters as the engine resolved them (`engine.cfg`, which may point at the repo's `artifacts/` fallback rather than the checkpoint directory: path, whether it is inside the checkpoint directory, size, SHA-256 up to 256 MiB, header metadata, r/alpha and prerouter settings) | +| `runtime` | backend and version, resolved device, the measurement adapter's description (synchronization policy, peak source, `headline_peak`), execution evidence (class and device of the logits the engine produced), Python / torch / CUDA or MLX versions, platform, Jetson power mode (`nvpmodel -q`, else the nvpmodel status file, else a reason) | +| `workload` | prompt source, SHA-256 of the prompt text, SHA-256 of the prompt token ids after chat templating (`prompt_token_ids_sha256`: two launches are the same workload only if prompt digest, token count and token-id digest all match), length, text and token count; requested counts; resolved seed, temperature, top-k, top-p, repetition penalty, prefill chunk and think flag; `model_env`, the model-read knobs resolved after the build exactly as the bailing_hybrid code parses them (`ling_hidden_clip`, `prerouter_feature_topk`, `prerouter_intra`), which change numerics and per-token work and must be equal on both sides of a comparison | +| `caches` | resolved `LayerOptions`, the actual shared-LRU and prefetch capacities read from the installed streaming layers, `weight_cache` (the torch backend's resolved `EDGE0_TORCH_WEIGHT_CACHE`), `prewarm` (the requested `EDGE0_PREWARM` / `LING_PREWARM` flag; only the 8B engine honors it), thread settings, the MLX cache limit, streaming statistics after the runs | +| `runs[]` | per-run counts, `prefill_seconds` / `prefill_tokens_per_second`, `decode_seconds` / `decode_tokens_per_second`, per-run memory metrics, execution evidence | +| `summary` | per metric: `n`, `n_valid`, the per-run `values` (nulls kept in run order), mean, min, max and sample standard deviation over the runs that have a value (null with a reason when none has one; the standard deviation needs two); descriptive only, two sequential runs in one process are not independent launches | +| `memory` | the process-lifetime RSS peak and the sampled RSS / anonymous / file-backed / swap peaks, with notes | +| `probe` | with `--probe-json`: the probe's own `schema_version`, the SHA-256 of the file as read, its `ready` flag, its timestamp and local path | + +### Safety rails + +* `--json-output` must be a new path. An existing destination, an invalid + `--probe-json` (missing, not JSON, no integer `schema_version`), + `--ntok <= 0` (also via `BENCH_NTOK`), `--warmup < 0`, a non-integer + `BENCH_SEED` or a non-numeric `BENCH_TEMP` are rejected **before** the model + loads, and none of that needs a backend: `bench.py` imports the framework + lazily, so `--help` and the argument checks work on a host without torch + or MLX. `BENCH_SEED=0` is a seed (both runs re-seed with it right after + prefill); at `BENCH_TEMP=0` the sampler is greedy, so the output does not + depend on the seed even though it is recorded. A tier name resolves to + `EDGE0__MODEL` by the same rule as the CLI. +* The report is written atomically (temporary file in the destination + directory, fsync, then linked or renamed into place without ever + overwriting an existing file). A failed engine build, benchmark, report + assembly or write exits non-zero, prints `[bench] FAILED during ` and + leaves no report at the destination; the engine and the sampler are closed + in `finally` either way. +* The human output is unchanged, except that `peak_active` prints `n/a` when + the backend peak is unavailable and measured CUDA numbers now include the + synchronized device work. +* Keep the JSON outside tracked source: it records local paths and the prompt + text. Share sanitized summaries and evidence digests. + +### What this does and does not show + +`tests/test_benchmark_report.py`, `tests/test_benchmark_measure.py` and +`tests/test_bench_cli.py` validate the instrumentation with a fake engine, a +fake array namespace and fake device APIs: dispatch, synchronization order, +count arithmetic, units, pre-load validation, failure paths and cleanup. They +run on a host with neither torch nor MLX. `tests/test_bench_backend.py` +repeats the core cases through the real `core` ops and sampler of whichever +backend is importable (on a host without MLX it selects the torch backend on +the CPU unless the environment already chose otherwise) and is skipped when +there is none. None of this is CUDA or Orin evidence. The Orin baseline +(Task 3 of the plan) needs the physical device, a validated environment and +the probe. + `edge0` does **not** run on NVIDIA through MLX, at any version tested: the first forward pass fails, differently at each version (the table below). It does run on NVIDIA through the torch backend in diff --git a/docs/plans/jetson-orin-nano-review.md b/docs/plans/jetson-orin-nano-review.md new file mode 100644 index 0000000..8aa4140 --- /dev/null +++ b/docs/plans/jetson-orin-nano-review.md @@ -0,0 +1,145 @@ +# Edge0 Orin plan review and change justifications + +Reviewed 14 September 2026. + +**Assessment:** retain the staged Torch/CUDA approach, but correct the execution gates and benchmark prerequisites before implementing Task 2. The reviewed benchmark cannot currently run through its first memory reset on the CUDA backend. The baseline smoke command tests MLX, and the proposed memory/prefetch policies omit behaviors that could defeat their correctness or bounds. + +The attached file was treated as a document to review, not authorization to execute its embedded Claude prompt. Its referenced full plan, relevant implementation and tests were inspected through a separate public-repository checkout. No Edge0 implementation was changed, no commit/push was made, and no Jetson inference was run. + +## Deliverables + +- Revised Claude handoff: delivered outside the repository (it names a local checkout path and a private evidence directory). +- [Revised full engineering plan](jetson-orin-nano.md): now the repository plan. +- This review: justification and evidence for every substantive change, identified as C01–C16 in the revised documents. + +## Source and verification scope + +The original `edge0-orin-claude.md` remains unchanged. SHA-256: + +```text +7E2EC53ED180A5BCF91800A61E435B2C18AD57A434AF90C05F59A3F00D8FA455 +``` + +The public [draft PR #1](https://github.com/KarlTaylorKnight/Edge0/pull/1) was inspected on the review date. The code and full plan were read at commit `dc1d9b3b78f51eb1a5f3705d63ed16c5b2cc9dfe`; links below pin that snapshot. The Mac checkout path, its working-tree state and historical test outputs in the attachment were not independently inspected. Upstream PR discussion is useful provenance, not physical Orin evidence. + +## Changes and reasons + +### C01 — Pin the reviewed source and qualify inherited claims + +**Change:** add the reviewed commit, distinguish the supplied Mac path from the inspected public checkout, and label prior Mac checks as historical reports. Retain the PR #19 base and separate it from issue #20's MLX experiment. + +**Why:** mutable branches and upstream discussions can change. Without a source revision, a later implementer cannot tell whether a finding still applies. The attachment's test claims and another GPU's performance are not fresh Orin results. The [PR](https://github.com/KarlTaylorKnight/Edge0/pull/1) establishes the public branch context; the [reviewed NVIDIA notes](https://github.com/KarlTaylorKnight/Edge0/blob/dc1d9b3b78f51eb1a5f3705d63ed16c5b2cc9dfe/docs/nvidia.md) describe the different experiments. + +### C02 — Make Task 2's stopping point consistent + +**Change:** finish portable reporting in Task 2; put physical benchmark execution in Task 3. Replace “baseline captured and committed” with retained evidence and a reviewed summary. Permit a documented failed baseline attempt to lead to memory-policy correction, while requiring a successful reference run before performance optimization. + +**Why:** the handoff says to stop after Task 2, but the [full plan](https://github.com/KarlTaylorKnight/Edge0/blob/dc1d9b3b78f51eb1a5f3705d63ed16c5b2cc9dfe/docs/plans/jetson-orin-nano.md) requires an Orin run inside Task 2 and again in Task 3. It also alternates between committing results and forbidding generated artifacts. These contradictions can block useful host work or encourage premature hardware claims. An OOM baseline needs a recovery path rather than a circular prerequisite. + +### C03 — Repair the benchmark's missing CUDA measurement interface + +**Change:** add a benchmark-owned backend-aware measurement adapter and fake-engine integration tests before serializing results. Keep optional runtime imports out of the pure builder. + +**Why:** the [benchmark, lines 95 and 120](https://github.com/KarlTaylorKnight/Edge0/blob/dc1d9b3b78f51eb1a5f3705d63ed16c5b2cc9dfe/examples/bench.py#L95) calls peak-memory methods absent from the [CUDA core](https://github.com/KarlTaylorKnight/Edge0/blob/dc1d9b3b78f51eb1a5f3705d63ed16c5b2cc9dfe/src/edge0/backends/cuda/core.py). JSON-builder tests would miss the immediate runtime failure. A small benchmark adapter fixes the consumer without expanding the task into a backend rewrite. + +### C04 — Define synchronized wall-clock timing + +**Change:** synchronize the actual CUDA device at prefill/decode boundaries, including the end of warmup and the final timed step. Keep storage, transfer waits and sampling inside the relevant wall-time interval. + +**Why:** CUDA work is asynchronous and [CUDA `core.eval()`](https://github.com/KarlTaylorKnight/Edge0/blob/dc1d9b3b78f51eb1a5f3705d63ed16c5b2cc9dfe/src/edge0/backends/cuda/core.py#L292) does not synchronize. Existing incidental synchronizations do not define a reliable measurement contract. Kernel event time alone also omits the host/I/O costs central to this application. [PyTorch documents synchronization requirements for accurate CUDA timing](https://docs.pytorch.org/docs/main/notes/cuda.html#asynchronous-execution). + +### C05 — Specify what the report actually measures + +**Change:** version the benchmark schema; record requested and resolved configuration, backend and actual device, per-run results, sampling settings, input/model identity and an actual probe digest/reference. Distinguish prompt, warmup, timed and total generated tokens. Preserve and label the existing fixed-length sample-and-step protocol. + +**Why:** the [current loop](https://github.com/KarlTaylorKnight/Edge0/blob/dc1d9b3b78f51eb1a5f3705d63ed16c5b2cc9dfe/examples/bench.py#L93) uses greedy warmup despite its comment, executes an additional forward for every emitted timed token, and ignores EOS. Reporting these numbers as user-visible completion latency would be misleading. Also, the [Torch backend can select CPU](https://github.com/KarlTaylorKnight/Edge0/blob/dc1d9b3b78f51eb1a5f3705d63ed16c5b2cc9dfe/src/edge0/backends/cuda/core.py#L38); a backend label is not device evidence. A probe version alone cannot identify which machine-state report accompanied a run. + +### C06 — Give memory metrics correct units, scope and limitations + +**Change:** use integer bytes and measurement methods; separate process-lifetime RSS peaks, sampled peaks, CUDA allocated/reserved peaks and system telemetry. State precisely when CUDA peak statistics reset and exclude earlier loading transients from that per-run claim. Use null/reasons when metrics are unavailable. + +**Why:** an endpoint sample is not a peak, and a process high-water mark cannot be reset by subtracting earlier values. PyTorch allocation statistics cover its allocator, not every device allocation. On Tegra, CPU and GPU allocations share physical DRAM; RSS plus CUDA peaks is not a valid total-footprint formula. The combination of scoped metrics and system observations supports an honest fit assessment. [PyTorch allocator documentation](https://docs.pytorch.org/docs/main/notes/cuda.html#memory-management), [NVIDIA Tegra memory architecture](https://docs.nvidia.com/cuda/archive/13.0.2/cuda-for-tegra-appnote/index.html#memory-management). + +### C07 — Test collection, CLI failure and cleanup as well as fields + +**Change:** extend TDD to dispatch, synchronization order, count arithmetic, unavailable/invalid metrics, strict JSON, atomic writing and engine cleanup. Require a new output path so stale success artifacts cannot be mistaken for a failed run's result. Keep mocked and device evidence separate. + +**Why:** a schema test cannot detect an incorrect timer, CPU fallback, partial file or leaked engine. The [current CLI](https://github.com/KarlTaylorKnight/Edge0/blob/dc1d9b3b78f51eb1a5f3705d63ed16c5b2cc9dfe/examples/bench.py#L134) closes its engine only after successful benchmarking. Python's JSON encoder permits NaN/infinity by default, so a nominal JSON result can fail strict downstream consumers unless validated. [Python JSON documentation](https://docs.python.org/3.10/library/json.html). These are focused behavioral tests for the new reporting feature, not an unrelated testing expansion. + +### C08 — Add a reproducible Jetson installation prerequisite + +**Change:** resolve the actual board/L4T/JetPack/Python/Torch/CUDA combination, required features and tokenizer imports before inference. Audit backend dependency separation and preserve Apple defaults. + +**Why:** “Python 3.10+, PyTorch for Jetson” does not define an installable stack. The reviewed [package metadata](https://github.com/KarlTaylorKnight/Edge0/blob/dc1d9b3b78f51eb1a5f3705d63ed16c5b2cc9dfe/pyproject.toml) mandates MLX while [CUDA tokenizer loading](https://github.com/KarlTaylorKnight/Edge0/blob/dc1d9b3b78f51eb1a5f3705d63ed16c5b2cc9dfe/src/edge0/backends/cuda/io.py#L141) imports Transformers, and CUDA modules require specific Torch APIs. NVIDIA ties its builds and prerequisites to particular JetPack versions. The plan therefore requires a tested combination rather than inventing a wheel/version from an unknown device. [Installation requirements](https://docs.nvidia.com/deeplearning/frameworks/install-pytorch-jetson-platform/index.html), [compatibility matrix](https://docs.nvidia.com/deeplearning/frameworks/install-pytorch-jetson-platform-release-notes/pytorch-jetson-rel.html). + +### C09 — Limit what a successful probe certifies + +**Change:** retain the completed probe but explicitly check Nano 8 GB identity, missing memory evidence, environment compatibility and checkpoint completeness separately. Describe Linux-specific test prerequisites. + +**Why:** [the probe](https://github.com/KarlTaylorKnight/Edge0/blob/dc1d9b3b78f51eb1a5f3705d63ed16c5b2cc9dfe/scripts/jetson_probe.py#L244) matches generic Jetson/Orin model strings, queries CUDA properties and checks storage. Its blockers do not validate the 8 GB Nano variant, model contents, enough RAM or successful kernel execution. This is a limitation of its current contract, not grounds to discard the probe. A passing probe must not silently advance the entire acceptance gate. + +### C10 — Replace the incorrect hardware acceptance test + +**Change:** add a Torch-only test that requires and exercises CUDA, plus identified reference fixtures for on-device correctness. Retain existing MLX comparisons on supported hosts and define numerical tolerances before the run. + +**Why:** [the nominated NVIDIA smoke test](https://github.com/KarlTaylorKnight/Edge0/blob/dc1d9b3b78f51eb1a5f3705d63ed16c5b2cc9dfe/tests/test_nvidia_backend_smoke.py#L1) imports MLX directly. Setting `EDGE0_BACKEND=cuda` does not turn it into a Torch test. Nor is [test_cuda_backend.py](https://github.com/KarlTaylorKnight/Edge0/blob/dc1d9b3b78f51eb1a5f3705d63ed16c5b2cc9dfe/tests/test_cuda_backend.py#L14) a substitute: it requires MLX and its helper creates CPU tensors. The original “tokens or logits” criterion also allows the metric to be chosen after a mismatch; separate predefined comparisons avoid that ambiguity. + +### C11 — Replace the guessed footprint threshold with enforceable budgets + +**Change:** derive an explicit byte budget from observed conditions and the declared workload, accounting for load/prefill/KV/transient, staging and inflight allocations. Reject calculated zero capacities unless an explicit disabled-cache mode is implemented. Inspect configuration and cache construction, not only CUDA namespace files. + +**Why:** a universal 6.5 GB RSS ceiling does not establish system headroom. More concretely, [SharedExpertCache and PrefetchBuffer](https://github.com/KarlTaylorKnight/Edge0/blob/dc1d9b3b78f51eb1a5f3705d63ed16c5b2cc9dfe/src/edge0/streaming/cache.py#L19) evict only for positive capacities: zero is unbounded. The [prefetch implementation](https://github.com/KarlTaylorKnight/Edge0/blob/dc1d9b3b78f51eb1a5f3705d63ed16c5b2cc9dfe/src/edge0/streaming/layer.py#L456) also has inflight work and a path that grows capacity. Ignoring those behaviors defeats a nominally bounded LRU policy. + +### C12 — Preserve the actual 8B routing profile + +**Change:** connect future prefetch/cache work to the existing exact path; retain required expert contributions on misses. Do not enable staged decoding merely to activate prefetch. + +**Why:** [the 8B configuration](https://github.com/KarlTaylorKnight/Edge0/blob/dc1d9b3b78f51eb1a5f3705d63ed16c5b2cc9dfe/src/edge0/models/edge0_8b/__init__.py#L48) selects `prod_k8`, whose [definition deliberately disables staging](https://github.com/KarlTaylorKnight/Edge0/blob/dc1d9b3b78f51eb1a5f3705d63ed16c5b2cc9dfe/src/edge0/streaming/options.py#L112). The [engine](https://github.com/KarlTaylorKnight/Edge0/blob/dc1d9b3b78f51eb1a5f3705d63ed16c5b2cc9dfe/src/edge0/engine/ling.py#L70) passes only staged layers to the slot stager, while existing staged misses can use a zero row. A scheduling optimization must not silently become a routing/quality change. + +### C13 — Complete the asynchronous ownership contract + +**Change:** require both copy-completion and consumer-completion ordering, slot generations, pinned-buffer lifetime, bounded queues, failure/reset/close handling and stress tests. Treat pinned staging as a measured candidate on shared DRAM. + +**Why:** a readiness event prevents reading an incomplete copy but does not prevent overwriting a slot still used by a previous kernel. Reusing a pinned source too early can also corrupt a copy. Persistent asynchronous resources need stronger lifecycle handling than the current [reset/close logic](https://github.com/KarlTaylorKnight/Edge0/blob/dc1d9b3b78f51eb1a5f3705d63ed16c5b2cc9dfe/src/edge0/streaming/layer.py#L1378). [PyTorch stream semantics](https://docs.pytorch.org/docs/main/notes/cuda.html#cuda-streams) support these ordering requirements; [NVIDIA's Tegra guide](https://docs.nvidia.com/cuda/archive/13.0.2/cuda-for-tegra-appnote/index.html#memory-management) explains why duplicated host/device buffers share the same physical capacity. + +### C14 — Make kernel work conditional and align Gate D with its scope + +**Change:** select kernels from the measured profile, add guarded dispatch/fallback and target toolchain validation, and distinguish expert gather from dense quantized-linear optimization. Limit the allocation-free promise to expert-transfer payload reuse. + +**Why:** an INT4-only implementation cannot replace all the [existing quantization cases](https://github.com/KarlTaylorKnight/Edge0/blob/dc1d9b3b78f51eb1a5f3705d63ed16c5b2cc9dfe/src/edge0/backends/cuda/quant.py), and tensor outputs/intermediates still allocate. A fused expert gather also leaves [dense QuantizedLinear dequantization](https://github.com/KarlTaylorKnight/Edge0/blob/dc1d9b3b78f51eb1a5f3705d63ed16c5b2cc9dfe/src/edge0/backends/cuda/nn.py#L93) untouched. The revised gate avoids promising a larger kernel/allocator rewrite than the staged plan supplies, while retaining the opportunity to address dense work when measured evidence supports it. + +### C15 — Make before/after measurements comparable + +**Change:** freeze the workload, collect at least three independent launches on both baseline and optimized versions, retain each internal run, distinguish warm/cold assumptions, and record sustained thermal behavior and failed attempts. A reduced context is a separate experiment. + +**Why:** the [benchmark has two internal runs](https://github.com/KarlTaylorKnight/Edge0/blob/dc1d9b3b78f51eb1a5f3705d63ed16c5b2cc9dfe/examples/bench.py#L93), and engine reset does not guarantee cold storage or all caches. Comparing a single cold baseline with a warm optimized mean can attribute caching or thermal differences to the code. Power/thermal/system telemetry supplies the context needed to interpret the result. [NVIDIA tegrastats documentation](https://docs.nvidia.com/jetson/archives/r36.5/DeveloperGuide/AT/JetsonLinuxDevelopmentTools/TegrastatsUtility.html). + +### C16 — Separate reproducible local evidence from shared artifacts + +**Change:** store exact local paths and raw reports in an untracked evidence set; put sanitized summaries and content references in documentation. Hash/reference the specific probe and checkpoint manifest outside timed regions. Permit small reviewed deterministic test fixtures as test assets. Keep the supplied source unchanged and deliver revised copies. + +**Why:** exact paths and command strings can reveal usernames/machine details, conflicting with the original no-identifiers/no-generated-artifacts rules when copied into Git. A digest provides identity but must reference retained inputs to support reproduction. Keeping both source and revised copies makes the review auditable. This clarifies the existing artifact restriction rather than introducing a publication requirement. + +## Validation performed + +The review checked the pinned source, test imports, benchmark calls, cache capacity semantics, profile selection and supporting NVIDIA/PyTorch documentation. The source checkout was kept free of implementation changes. The original attachment's SHA-256 was recorded for preservation verification. + +On Windows/Python 3.14, this command ran: + +```text +uv run --no-project --with pytest python -m pytest tests/test_jetson_probe.py -q +14 failed, 10 passed in 0.79s +``` + +Failures include POSIX sysfs fixture names such as `259:1`, Linux mount-path assumptions and unavailable `os.major` on Windows. They are actual failures on that host, not evidence of an Orin regression and not a reproduction of the historical Mac test result. + +The same unchanged tests then ran under Ubuntu/WSL with Python 3.12 and pytest 9.1.1, using isolated test packages in the review workspace: + +```text +PYTHONPATH=/work/pytest-linux-packages python3 -m pytest tests/test_jetson_probe.py -q +24 passed in 0.36s +``` + +The initial WSL interpreter had neither pytest nor working venv bootstrapping. Test packages were installed into a task-local target directory; no system Python packages or operating-system packages were changed. This is a Linux fixture test result, not an on-device readiness result. + +No full backend/parity suite or physical CUDA/Orin benchmark was run during the review. The current benchmark failure is established by its calls and the missing functions in the selected backend; a model was not downloaded to reproduce it. diff --git a/docs/plans/jetson-orin-nano.md b/docs/plans/jetson-orin-nano.md new file mode 100644 index 0000000..afefb23 --- /dev/null +++ b/docs/plans/jetson-orin-nano.md @@ -0,0 +1,427 @@ +# Jetson Orin Nano CUDA implementation plan — revised + +Review date: 14 September 2026. Source snapshot: `KarlTaylorKnight/Edge0` commit `dc1d9b3b78f51eb1a5f3705d63ed16c5b2cc9dfe`. Change IDs link this plan to the accompanying [justifications](jetson-orin-nano-review.md). + +## Goal and scope [C01, C02] + +Establish reproducible Edge0-8B inference on a physical **8 GB Jetson Orin Nano with NVMe**, then improve the measured bottleneck while retaining the current Torch reference implementation and MLX behavior. Keep the branch stacked on `cuda-backend-base`, derived from upstream PR #19. Do not substitute the different MLX CUDA experiment in issue #20. + +Task 1 is present in the reviewed branch. Task 2 is the next implementation increment and its stopping point. Tasks 3–7 remain later increments. No reviewed evidence establishes Orin inference support or an Orin performance number. + +Preserve the existing model engines and facade. Add narrow measurement/configuration/streaming hooks as required, with regression tests. Edge0-35B, a TensorRT rewrite and new model families remain outside the initial target. Do not project GB10 results onto Orin. + +Use RED → GREEN → REFACTOR for meaningful behaviors. Document changes and evidence after each increment. Retain raw measurements outside Git; a documented baseline means a reviewed summary plus retrievable evidence references, not committed raw artifacts. + +## Acceptance gates + +### Gate A — target and installation evidence [C08, C09] + +Record board/module identity, nominal memory variant, observed RAM, architecture, L4T/JetPack, Python ABI, installed Torch build, CUDA runtime, GPU properties, storage backing and checkpoint identity. Verify the target is the intended Nano 8 GB model; a generic Orin match also accepts other Orin models. + +The existing probe provides preliminary evidence. Its `ready` field does not certify the 8 GB SKU, enough free memory, a valid/complete checkpoint, package compatibility or a successful CUDA kernel. Missing required observations remain blockers. Resolve L4T to JetPack only from a supported mapping or package evidence; otherwise report it as unknown. + +Before inference, establish a reproducible environment compatible with the detected board/software stack. Record exact package versions and wheel URLs/digests or container digest. Verify actual imports and required APIs, including `torch.nn.RMSNorm`, the unsigned dtypes used by the backend, tokenizer dependencies and CUDA execution. Use NVIDIA's [installation requirements](https://docs.nvidia.com/deeplearning/frameworks/install-pytorch-jetson-platform/index.html) and [compatibility matrix](https://docs.nvidia.com/deeplearning/frameworks/install-pytorch-jetson-platform-release-notes/pytorch-jetson-rel.html); select a combination supported for this Orin, rather than the newest table entry automatically. + +The reviewed packaging unconditionally depends on MLX/MLX-LM and omits explicit Torch/Transformers dependencies. Audit the actual 8B import path and provide an isolated, tested Torch installation recipe. If dependency metadata must be separated, make that a small prerequisite change and preserve the default Apple installation path. Do not use the unrelated Mac wheel blocker as a diagnosis of Jetson compatibility. + +### Gate B — actual CUDA and model correctness [C10, C12] + +Require a Torch-only smoke test that executes representative backend operations on CUDA and synchronizes. The acceptance invocation must fail if CUDA is required but absent, rather than quietly skipping or falling back to CPU. Assert tensor/output device identity, finite output and comparison against a deterministic reference. + +Then run Edge0-8B with a pinned checkpoint, tokenizer/chat template, LoRA/prerouter identity, prompt IDs, routing profile and numerical settings. Capture at least 32 greedy tokens and compare teacher-forced logits and token choices with identified reference data. Define tolerances before examining the Orin result. Any allowed near-tie token divergence must be explained by the logit comparison; do not retroactively loosen tolerances to pass. + +Run live MLX/Torch comparison suites where MLX is supported. Provide small, reviewed deterministic fixtures with generating commit, input/model hashes, dtype and tolerance for use on Jetson. An optimized CUDA path must also agree with the current Torch correctness path on the same Orin and workload. Keep free-running output checks separate from teacher-forced numerical comparisons. + +### Gate C — bounded shared memory [C06, C11] + +Require no OOM across initialization, prefill, decode and repeated requests for the declared workload. Replace the unsupported universal “RSS below 6.5 GB” acceptance threshold with an explicit byte budget and observed system headroom. Record the selected reserve and why it is adequate for this device's services and workload; it is a tested policy, not a hardware guarantee. + +Start from observed available RAM and any tighter process/container limit. Deduct allocations still to be made, OS/application growth reserve, KV/context growth, prefill and dequantization transients, pinned staging, in-flight work, retained caches and allocator/driver allowance. Account for each allocation once at the chosen observation point. Compute tensor payload sizes from checkpoint metadata; slot counts alone are insufficient. + +Validate minimum available RAM, RAM/swap trends, process RSS and Torch allocator peaks during the run. RSS and CUDA peaks are separate views and must not be added as a total. Mark swap activity explicitly; a swapping run cannot establish a resident-memory target. Reject impossible profiles before inference when the calculation can identify them, and preserve evidence for unexpected OOMs. + +### Gate D — measured optimization [C13, C14] + +For expert transfers, reuse a bounded payload pool after initialization; do not require every model operation to become allocation-free. Preserve exact fallback for required experts that are late, absent or predicted incorrectly. Demonstrate correct asynchronous ownership and a measured benefit before enabling a new path by default. + +Treat elimination of dense-layer dequantization as a separate, conditional objective. An expert gather kernel does not meet it. Any dense-weight cache must fit the byte budget; whole-model expansion is not a default Orin setting. + +## Task 1 — retain the probe; record its limits [C09] + +Keep `scripts/jetson_probe.py`, its stable schema and focused tests. Do not redo the completed increment. Add narrowly scoped corrections only if required evidence is missing or defective, with schema changes documented rather than silently redefining `ready`. + +Treat its Linux storage/sysfs behavior as platform-specific. Use Linux for those tests, and distinguish fixtures exercised on a supported host from behavior on the actual Jetson. The report builder introduced next must remain independently testable without this platform dependency. + +## Task 2 — functioning, machine-readable benchmark reporting [C03–C07] + +**Scope:** reporting, backend-aware measurement, focused tests and documentation. No cache/kernel optimization and no mandatory physical-device run in this increment. + +**Files:** modify `examples/bench.py`; add a pure helper such as `examples/benchmark_report.py`; add `tests/test_benchmark_report.py` and focused integration tests as appropriate; update `docs/nvidia.md` and this plan. Keep optional imports and metric collection outside the pure builder. + +### Measurement behavior [C03, C04, C05] + +The current benchmark calls peak-memory methods absent from the CUDA core. Add a benchmark-owned adapter: retain MLX metric behavior where available, use Torch CUDA metrics for an actual CUDA device, and use explicit unavailable values elsewhere. `EDGE0_BACKEND=cuda` alone does not prove the selected device is CUDA. + +Measure wall time with a monotonic clock. On CUDA, synchronize the resolved device before each timed phase and before its endpoint. Finish prefill before recording its time; complete untimed warmup before starting decode; finish the final timed step before recording decode. Use phase boundaries, not new per-token synchronizations. Include storage, transfer and CPU sampling costs in wall time; kernel-only event timings are a separate diagnostic. [PyTorch timing and stream semantics](https://docs.pytorch.org/docs/main/notes/cuda.html#asynchronous-execution). + +Keep the current two-run protocol, defaults, sampling behavior and human-output format. The existing warmup is greedy; timed iterations perform sampling followed by an engine step, including a final forward whose logits are not sampled. The loop is fixed-length and does not stop at EOS. Label this protocol explicitly; do not call it TTFT or EOS-aware completion latency. Correct the warmup comment without changing it to sampled warmup. + +For each run, record prompt tokens after chat templating; actual warmup tokens; timed decode tokens; total generated tokens; and context length at timed decode start. For the current loop, `total_generated = warmup + timed_decode` and `decode_start_context = prompt + warmup`. Timed throughput uses only timed decode tokens. Keep requested counts separately. Reject invalid negative counts and require a positive timed-token count; do not silently produce a successful empty benchmark. + +### Report contract [C05, C06, C16] + +Use a benchmark schema version independent of the probe schema. Include: + +| Group | Required information | +|---|---| +| Run identity | Timestamp, benchmark protocol version, Git commit, dirty status, run count/index and command/configuration needed to reproduce | +| Model | Local resolved checkpoint path, model/revision identity, manifest digest where available, tokenizer/template and adapter/prerouter identity | +| Runtime | Backend name, resolved device, observed execution evidence, Python/Torch/CUDA versions, platform; observed power mode or an explicit unavailable reason | +| Workload | Prompt identity/digest and token count; counts defined above; resolved seed, temperature, top-k, top-p, repetition penalty, context/prefill settings | +| Caches | Resolved LRU/prefetch/staging settings, weight-cache and prewarm flags, and relevant thread settings; do not record only requested environment variables | +| Timings | Per-run prefill/decode seconds and their separate token/s values; all runs retained; summary calculation and variability defined | +| Memory | Integer byte values, method and scope for every metric, with explicit unavailable reasons | +| Probe link | Optional supplied probe file's schema version, content SHA-256 and local reference; linkage to that actual report, not only a schema number | + +Record process-lifetime peak RSS with a platform-specific method and unit normalization. If only periodic samples are available, call it a sampled peak and record the interval. A single endpoint RSS value is not a peak. Do not subtract lifetime high-water marks to invent a per-run peak. + +For CUDA, synchronize after engine reset, reset allocator peak statistics immediately before prefill, and read allocated/reserved peaks after synchronized decode. State that this covers prefill, warmup and decode, and includes allocations already resident at reset; it does not recover an earlier initialization transient. Start process/system observation before model load in physical baseline runs. Allocator peaks do not measure all driver or non-Torch memory. [Torch allocator scope](https://docs.pytorch.org/docs/main/notes/cuda.html#memory-management). + +Keep hashing, power queries, collector setup and report construction outside timed intervals. Preconfigured lightweight memory/system samplers may run throughout the benchmark; record their interval and use identical instrumentation for comparisons. The pure builder performs no collection. Prepare a checkpoint manifest outside measurement; do not hash large weights before each alleged cold-cache trial. A digest identifies an input but does not reproduce it: retain the referenced manifest/prompt in the private evidence set. + +Represent unavailable metrics as `null` with a reason. Distinguish unavailable metadata from a failed required measurement. Reject negative/nonfinite timings and memory; zero measured duration gives unavailable throughput, not a fabricated zero-rate success. Use strict JSON serialization (`allow_nan=False`). [Python JSON behavior](https://docs.python.org/3.10/library/json.html). + +### CLI and acceptance tests [C07] + +Add `--json-output PATH` and optional `--probe-json PATH`. Validate an explicitly requested probe before inference. Keep probe collection separate so generic benchmarks do not acquire Jetson prerequisites. Preserve human-output format, except measured numeric values may correctly change after synchronization. + +Require a new output path and reject an existing destination before inference. Write complete JSON atomically through a temporary file in the destination directory; report output errors with a nonzero exit status. Close the engine and collectors in `finally`, including generation and write failures. Failed inference must not leave a success report at the requested path; identify failures separately. + +Use TDD for schema/units, count arithmetic, requested versus resolved configuration, CPU/MLX/CUDA dispatch, phase synchronization order, missing metrics, malformed probes, nonfinite inputs, output failure and cleanup. Exercise the CLI and fake engine without weights. Test output-format compatibility without pinning real measured numbers. Host tests validate instrumentation; CUDA tests validate device behavior; neither substitutes for the Orin baseline. + +**Done:** applicable checks pass, exact results and host limitations are reported, documentation describes the protocol, and Task 3's next command is supplied. Stop here in the current handoff. + +**Status (14 September 2026):** implemented on this branch. `examples/benchmark_report.py` (pure builder, schema version 1), `examples/benchmark_measure.py` (backend-aware measurement adapter, process-memory sampler, git/power/checkpoint/adapter identity), `examples/bench.py` (`--json-output`, `--probe-json`, `--rss-sample-interval`, phase-boundary synchronization, pre-load validation, `finally` cleanup, lazy backend import), tests `tests/test_benchmark_report.py`, `tests/test_benchmark_measure.py`, `tests/test_bench_cli.py` (fakes only) and `tests/test_bench_backend.py` (real backend namespace, no weights), and the protocol/schema section in `docs/nvidia.md`. `pyproject.toml` gained `pythonpath = ["."]` so the plain `pytest` invocation collects the repo-root `scripts`/`examples` imports. Host checks only: no CUDA execution and no Orin run were performed; Task 3 starts with the commands below. + +## Task 3 — establish the unoptimized Orin baseline [C08–C10, C15] + +First complete Gate A's installation and model checks. Add a small Torch-only acceptance module, proposed as `tests/test_torch_cuda_smoke.py`, with a documented `--require-cuda` option or equivalent hard-fail contract. Add an on-device reference checker using the pinned fixtures described in Gate B. These are future deliverables, not commands that exist in the reviewed branch. + +The reviewed plan's `tests/test_nvidia_backend_smoke.py` is an MLX smoke test. The existing `test_cuda_backend.py` also imports MLX and uses CPU tensors in its comparison helper. Keep those suites, but do not use their skip/pass result as CUDA execution evidence. + +On the actual Orin, use an isolated validated environment and an untracked `$EDGE0_RUN_DIR`: + +```bash +export EDGE0_BACKEND=cuda +export EDGE0_TORCH_DEVICE=cuda +export EDGE0_TORCH_WEIGHT_CACHE=0 +export EDGE0_PREWARM=0 +python scripts/jetson_probe.py --model-dir "$EDGE0_8B_MODEL" \ + --output "$EDGE0_RUN_DIR/orin-probe.json" +``` + +After implementing the dedicated acceptance test, its intended invocation is: + +```bash +python -m pytest tests/test_torch_cuda_smoke.py --require-cuda -q +``` + +Run the separate 32-token greedy correctness check and retain its comparison evidence. After that passes and Task 2 reporting exists, a short benchmark is: + +```bash +BENCH_TEMP=0 BENCH_SEED=0 python examples/bench.py edge0-8b \ + --ntok 32 --warmup 0 \ + --probe-json "$EDGE0_RUN_DIR/orin-probe.json" \ + --json-output "$EDGE0_RUN_DIR/orin-short-baseline.json" +``` + +For performance comparisons, freeze a workload manifest and use at least three independent process launches before and after optimization. Keep both internal runs from each launch and distinguish their order/cache state. Use the same prompt, counts, seed, checkpoint, numerical/routing settings and power profile on both sides. A benchmark reset is not an OS or expert-cache flush. + +Capture `tegrastats` beside the run, from before loading until cleanup, with interval and timestamps. Record power mode, clock policy, cooling, storage/filesystem and thermal/throttling evidence. Observe existing settings before proposing changes. Do not assume a fresh process makes storage cold or drop system caches automatically. [NVIDIA tegrastats](https://docs.nvidia.com/jetson/archives/r36.5/DeveloperGuide/AT/JetsonLinuxDevelopmentTools/TegrastatsUtility.html). + +If inference fails, preserve its stage, configuration, available diagnostic evidence and exit status. OOM diagnostics may not reveal an exact failing allocation, especially for an OS kill; state unknown when necessary. Reducing context or generation produces a new workload, so retain both results. If default inference cannot fit, a documented failed attempt permits a bounded-memory correction in Task 4; establish a successful reference baseline before asynchronous or kernel optimization. + +**Status (14 September 2026, on the physical device):** executed on an +Orin Nano Developer Kit (Super) 8 GB, L4T R39.2, CUDA 13.2 driver, +25W mode, NVMe. Gate A: validated venv (Python 3.12.3, official torch +2.14.0+cu130 aarch64 wheel — the GB10 combination; edge0 installed +`--no-deps`, no MLX on device; `torch.nn.RMSNorm`, `torch.uint32`, +tokenizer imports and CUDA kernel execution verified). The wheel warns +`sm_87` is outside its SASS list; execution evidence, not the warning, +was used as the gate. Probe: `ready: true`, no blockers. +`tests/test_torch_cuda_smoke.py` + `tests/conftest.py` (`--require-cuda`) +added: 6/6 pass on device; hard-fail contract verified with +`CUDA_VISIBLE_DEVICES=""`. `scripts/orin_reference_check.py` added +(CUDA greedy vs torch-CPU teacher-forced logits, tolerances registered +in the script before the run): 32/32 token choices identical, no +near-ties consumed; the pre-registered vocab-wide logit bound (1.0) +was exceeded (max |Δ| 2.06, tail-token bfloat16 noise — ≤ 0.22 at every +step's top-8 tokens; analysis retained in the evidence set). Recorded +as a FAIL of that bound, not loosened retroactively; a top-k-scoped +criterion is proposed for review before the next comparison run, and +MLX fixtures remain the intended reference. Short baseline: three +independent launches × two runs (`--ntok 32 --warmup 0`, temp 0, seed 0, +weight cache and prewarm off, tegrastats alongside): decode 0.53–0.58 +tok/s (mean 0.55), prefill 11.3–15.9 s at 37 prompt tokens, CUDA +allocator peak ≤ 1.00 GiB, process peak RSS 5.3–5.8 GiB, GPU 99% +utilized, no OOM/swap. Raw JSON reports, npz logit captures and +tegrastats logs live in the untracked `$EDGE0_RUN_DIR` evidence set; +sanitized summary in `docs/nvidia.md`. Task 4's profiling starts from +this baseline; the dominant cost is per-call expert dequantization in +`gather_qmm` (GB10 already measured 2.5× from the weight cache that +Task 4's byte budget must first make safe to enable here). + +## Task 4 — memory-budgeted profile [C11] + +Implement the pure budget calculation first, then integrate at existing configuration/cache-construction points. Inspect `models/base.py`, the 8B configuration, `streaming/install.py`, `streaming/cache.py` and `streaming/layer.py`; the CUDA backend namespace alone cannot enforce the profile. + +Resolve byte limits for completed caches and in-flight builds, prefetch buffers, pinned memory, staged/assembled copies, full-layer prefill, transient dequantization and KV growth. Bound the producer queues as well as their completed outputs. Audit `prefetch_all()`, which can raise buffer capacity beyond its configured starting value. + +**Existing zero semantics:** `SharedExpertCache(0)` and `PrefetchBuffer(0)` disable eviction. Never pass a calculated zero allowance to them as if it disabled caching. Initially reject such impossible settings; if cache-off operation is needed, implement an explicit bounded/disabled state with regression tests rather than silently changing shared MLX semantics. + +Record resolved policy and rejected overrides. Keep full dequantized weight caching off unless the measured budget permits it. Test exhausted/tiny budgets, missing observations, overflow, invalid overrides, varying tensor sizes, inflight backpressure and repeated-request memory stability. + +**Status (15 September 2026, implemented and validated on the device):** +`edge0/streaming/budget.py` (pure: `Observation` / `ExpertFootprint` / +`WorkloadDecl` / `Reserves` → `resolve_budget()` → `ResolvedBudget` with +itemized deductions; `bundle_bytes_from_entries()` prices the per-expert +payload from the safetensors header, not slot counts). Enforcement: +`PrefetchBuffer` gained `max_cap` (a ceiling `set_cap` cannot exceed — +bounds `prefetch_all()`'s silent growth to `num_experts + 32`); +`LayerOptions.max_inflight` bounds QUEUED speculative prefetch builds +(demand loads are never dropped); both default to the historical +unbounded semantics when no budget is active. Integration: +`EDGE0_MEMORY_BUDGET=auto|` resolved in the 8B engine's +`load_installed` immediately before cache construction +(`EDGE0_BUDGET_CONTEXT` declares the context, default 1024 tokens; +`ModelConfig.kv_bytes_per_token` prices it — measured, 1.1 MB/token for +this tier). The resolved policy raises `BudgetError` pre-inference for +impossible profiles — a calculated starved cache is rejected, never +constructed as `SharedExpertCache(0)`/`PrefetchBuffer(0)` — and is +recorded in the engine (`engine.memory_budget`) and the bench report's +caches group. `EDGE0_TORCH_WEIGHT_CACHE=1` under a budget is vetoed +unless its measured 4.1 GB fits the post-cache headroom. 25 unit tests +in `tests/test_memory_budget.py`; the streaming package init became +lazy for backend-bound names so the pure modules test everywhere. +On-device validation (Orin Nano 8 GB): `auto` resolved usable ≈ 4.2 GB, +kept the tested 64/48 profile (real bundle ≈ 1.30 MiB/expert, full-layer +transient 162 MiB), benchmark equivalent to the unoptimized baseline +(0.52 tok/s, same peaks, RSS 5.55 GiB in the baseline range); a +4096-token declaration was rejected with itemized arithmetic (KV +4.5 GB > usable), and the weight cache was vetoed (4.1 GB > 2.2 GB +headroom). Not yet done: staged/assembled-copy pricing (the 8B +`prod_k8` profile has staging off — budgeting a staged profile is +future work with Task 5), and a long repeated-request soak (two +same-process runs were stable; a sustained test belongs to Task 7). + +## Task 5 — bounded asynchronous expert transfers [C12, C13] + +The 8B `prod_k8` profile deliberately disables staged decode. Do not enable `staged_k8` or replace routing to activate prefetch. Integrate the new transfer cache with the existing exact expert-consumption path. Predictions may change scheduling, not required expert identity or contribution. Missing required data must wait/load through the reference path rather than use the existing staged zero-row behavior. + +Start with bounded reusable pinned buffers and CUDA slots as a candidate design. Both consume the same physical DRAM pool on Jetson; measure their duplication and benefit. Keep alternative memory paths and GPUDirect Storage deferred until support and benefit are established. [NVIDIA Tegra memory architecture](https://docs.nvidia.com/cuda/archive/13.0.2/cuda-for-tegra-appnote/index.html#memory-management). + +Define slot ownership from free → filling → ready → in-use → reusable, with generation IDs. Publish readiness only after copy completion. Reuse a slot only after all GPU consumers finish; retain pinned source contents until its copy finishes. Bound queued work; handle stale completion events, failed reads, cancellation, reset, close and partial initialization. Use explicit stream dependencies and allocator lifetime handling as required. [PyTorch stream ownership](https://docs.pytorch.org/docs/main/notes/cuda.html#cuda-streams). + +Test state transitions with fake events, then real CUDA stress/parity cases: wrong/late predictions, eviction during consumption, repeated requests and shutdown. Define counters and scopes for cache/prediction hits, bytes requested/read/copied, I/O wait and consumer stalls. Distinguish logical file reads from physical NVMe traffic. Demonstrate overlap and end-to-end improvement on the baseline workload before making the path default. + +**Status (15 September 2026, measured on the device):** the measured +profile reshaped this task. First finding: the tested 64-slot shared +LRU is smaller than this tier's per-step working set (23 MoE layers x +K=8 = 184 bundles), so on the baseline it NEVER hits (`hits: 0`, +~184 rebuilds per token). Second finding: fixing that (512 slots, 65% +hit rate, loads 11,776 → 4,075 over the workload) moved end-to-end +decode by roughly nothing (0.49–0.57 vs 0.50–0.52 tok/s) — the +unoptimized decode is COMPUTE-bound (GPU pegged at 99%; per-call +dequantization). Per this task's own gate, the pinned-buffer / +CUDA-stream slot pipeline (free→filling→ready→in-use→reusable, +generation IDs) is therefore **deferred** — it would optimize a cost +that is currently invisible end-to-end — and is to be revisited after +Task 6 shrinks compute, when transfers become a meaningful fraction. + +What WAS implemented (opt-in, exact path untouched): +`EDGE0_CACHE_SLOTS=` overrides the REQUESTED LRU size (the budget +may still lower it; zero/negative rejected), and +`LayerOptions.predict_prefetch` / `EDGE0_PREDICT_PREFETCH=1` routes +each step's prerouter predictions into the existing bounded +`prefetch()` for non-staged layers (`prod_k8` keeps staged decode OFF; +`stage_experts` is untouched; wrong/late predictions fall back to +demand loads through `_get_bundles` — never a staged zero row). The +prefetch buffer is raised to one full predicted step (owners x K = 128) +and the per-layer in-flight bound defaults to K, both priced by the +Task 4 budget (in-flight transient = per-layer bound x producing +layers). Reuse-safety of consumed bundles rests on Python/torch +reference counting (consumers hold the tensors they use); the manual +slot-generation lifecycle belongs to the deferred pinned design. + +Measured (3 independent launches x 2 runs, identical workload and +instrumentation as the Task 3 baseline, budget active): decode +**0.568–0.623 tok/s, mean 0.594** vs baseline 0.53–0.58 mean 0.552 +(**+7.7%**); hit rate 85% (10,059 hits / 1,717 loads); 2,038 of 2,063 +predicted prefetches consumed (98.8%), 0 wasted, prefetch_wait 0.0 s — +the builds genuinely overlap the forward; load wall 26.5 s → ~6 s; +CUDA peak 1.45/1.58 GiB and RSS 5.48–5.82 GiB, both inside the +resolved budget (headroom ≈ 1.27 GB) with no OOM. Token choices are +bit-identical to the Task 3 reference sequence with both knobs on. +Kept **off by default** per Gate D: the gain is real but small while +compute dominates; defaulting is a profile change to make together +with the Task 6 re-measurement. Remaining for a later increment: +logical-vs-physical NVMe read accounting, CUDA stress cases for +eviction-during-consumption, and the deferred pinned-slot pipeline. + +## Task 6 — conditional quantized-kernel work [C14] + +**Scoping (15 September 2026, measured on the device):** per-op decode +shares — `gather_qmm` 37% (69 calls/step), dense `QuantizedLinear` 26% +(235 calls/step), everything else 37%. Both dequantization paths +qualify. Development is handed to a workstation GPU per +[`rtx6000-task6-handoff.md`](rtx6000-task6-handoff.md); decisions and +acceptance stay on the Orin. + +**Workstation increment (15 September 2026, RTX PRO 6000, torch +2.14.0+cu130 — parity and pricing only, no Orin decision):** three +opt-in, guarded paths, documented in `docs/nvidia.md` ("Opt-in quantized +paths"): `EDGE0_QMM_BATCHED=1` (exact batched expert gather, transient +priced at 140 MB for the decode shape and deducted by the budget), +`EDGE0_TORCH_WEIGHT_CACHE_BYTES=` (exact capped dequantized-weight +cache, first-fit, shrunk to the budget headroom; the full cache is now +priced at the measured bf16 total, 1.40 GB for this tier, not 4.1 GB), +and `EDGE0_INT4PACK=1` (torch's built-in int4 kernel: approximate at +bf16 rounding, bounds registered in `tests/test_int4pack.py` before any +Orin run). Design C (custom fused kernel) was scoped to a cross-compile +toolchain record only. Orin acceptance for each knob: the Task 3 +reference check with `--test-env ` against the reference path on +the same device, then the three-launch benchmark protocol. + +**Finding that reshapes this task's dense half:** traced on the real +model, 400 of 470 dense `QuantizedLinear` calls per step arrive with +float32 activations and only 70 with bfloat16. Both dense knobs require +bf16, so the weight cache fills 35 of the 235 modules it admits +(0.174 GB of 1.402 GB reserved) and the int4 kernel repacks exactly +those same 35. Neither moves decode measurably as a result. The dense +26% is therefore gated by activation dtype, not by the kernel; whether +those call sites can run bf16, and what that costs numerically, is the +question to settle on the device before more kernel effort. The batched +expert gather is unaffected: it covers the routed 37% and engages on +every call. + +Use the measured profile to choose expert gather, dense quantized linear work, or no kernel change. For an INT4 prototype, verify the actual checkpoint layout: packed words, group size, signed/negative scale behavior, bias handling, accumulation and output dtype. Pin the CUDA extension/CUTLASS/compiler combination supported by the detected Orin stack and verify runtime loading. + +Dispatch only supported dtype/shape/stride/transpose/group-size/device combinations to the new kernel. Keep the current implementation for unsupported cases, including 2/8-bit layouts. Test all existing broadcast and gather shapes, valid index boundaries, quantization extremes and checkpoint dtypes. Compare numerical results against both deterministic MLX fixtures and the current Torch path on Orin. Define tolerances in advance and measure scratch-memory peaks as well as speed. + +A fused expert gather does not remove dense `QuantizedLinear` dequantization. If dense work remains dominant, handle it as its own bounded, parity-tested change or document the residual cost. Gate D does not require eliminating all model allocations. + +**Status — ACCEPTED ON THE DEVICE (16 September 2026).** Orin Nano 8 GB, +same workload and instrumentation as the Task 3 baseline, budget active, +three launches x two runs per configuration; sanitized tables in +`docs/nvidia.md` (Task 6 Orin acceptance), raw reports in the evidence +set. Suite 341 passed on-device, 96 under `--require-cuda`. + +*Adopted:* the batched expert gather is now the **default** on the CUDA +backend (`EDGE0_QMM_BATCHED=0` selects the reference loop). It is the +only candidate that both PASSED the registered 32-token reference bound +on the target (32/32 token choices, max |Δlogit| 0.777 vs the bound of +1.0) and delivered a measured benefit there: decode 0.623–0.699, +mean **0.668 tok/s, +20.7%** over the 0.553 baseline, at unchanged CUDA +peak (1.00 GiB) and RSS. Its transient is bounded and the budget +deducts the cap (verified in the report: `kernel_transient` 268,435,456 +with no env var set). `int4pack.probe` executes on sm_87 — the kernel +runs on hardware the wheel's SASS list does not advertise, decided by +execution as designed. + +*Not adopted (kept opt-in):* the capped weight cache (+0.8%) and the +int4 kernel (+3.7%), both because the activation-dtype gate limits them +to 35 of 235 modules, and both exceed the registered bound (1.58 and +2.17) — recorded as that bound's failures, not loosened. + +*Gate D's Task 5 re-measurement:* with compute reduced the transfer +knobs are worth much more than before — +7.7% pre-Task-6, **+17.5% on +top of the batched gather** now (0.668 → 0.785 tok/s, expert load wall +72.5 s → 17.0 s over the same six runs). `EDGE0_MEMORY_BUDGET=auto +EDGE0_CACHE_SLOTS=512 EDGE0_PREDICT_PREFETCH=1` is the recommended Orin +profile and stays opt-in as board-specific tuning that requires the +budget. End to end the tier now runs at **0.785 tok/s, +42% over the +Task 3 baseline**, inside the resolved budget with no OOM. + +*Two workstation claims corrected here, both only visible on the target:* +the capped weight cache is NOT bit-identical on sm_87 — the cached +weights are, but one GEMM versus the reference's per-chunk GEMMs differs +by one bf16 ulp of accumulation order (isolated directly: same weights, +0.03125 on a 30.6 scale), so that exactness was a property of the +workstation's cuBLAS heuristics; and the float32 activation gate is +deliberate MLX promotion parity in the MLA/RoPE path (layer 3 is the +first `BailingMLA`; the residual stream is float32 from there on), so +reaching the remaining dense 26% is a numerics-parity decision for +review, not a kernel increment. + +*Open for a later increment:* design C (custom fused kernel) remains +unwritten with its toolchain recorded; the MLA float32 promotion +question; and the deferred pinned-slot transfer pipeline from Task 5, +which is now more attractive than it was at the compute-bound baseline. + +## Task 7 — final validation and evidence [C15, C16] + +Repeat the identical workload manifest using at least three independent launches. Retain individual prefill/decode results, explain first-versus-repeated-run state, and report sample count, mean, standard deviation and min/max separately. Treat the small sample as descriptive evidence rather than a broad performance guarantee. + +Repeat correctness, maximum declared context/generation, memory pressure and repeated-request checks, plus a declared sustained-run test long enough to expose thermal behavior. Capture failures and throttling. State the tested operating envelope: exact board/software/checkpoint/profile, context limits, memory reserve and remaining limitations. + +Only claim physical Orin inference once Gate B passes, and qualify the claim to its tested workload. Claim the bounded profile or optimization only after its own gates pass. Commit reviewed code/tests/docs and sanitized result summaries; keep raw benchmarks, telemetry, private paths, identifiers and checkpoints in the separate evidence set. Small reviewed deterministic test fixtures are test assets, not raw benchmark dumps. + +**Status — COMPLETE (16 September 2026).** The tested operating +envelope is stated in `docs/nvidia.md`; raw reports, telemetry logs and +the npz logit captures stay in the untracked evidence set. + +*Repeat measurement.* Final adopted profile, three independent launches +x two runs: decode 0.744–0.786 tok/s, **mean 0.772, sd 0.016** (n=6), +prefill 13.6–15.6 s, CUDA peak 1.58 GiB, RSS 5.57–5.69 GiB. Reported as +descriptive statistics of a small sample, not a guarantee. + +*Correctness repeat.* Same-device 32-token reference check of the final +profile against the reference loop: **PASS**, 32/32 token choices +identical, max |Δlogit| 0.777 against the registered bound of 1.0. All +30 requests of the sustained run produced byte-identical tokens. + +*Maximum declared context and generation.* KV growth was MEASURED here +rather than inherited: the marginal cost is 561,320 B/token (548 KiB), +about half the 1,100,000 taken from the README's Apple/MLX figure, so +`kv_bytes_per_token` for this tier is corrected to 600,000 (≈7% margin). +Growth is sublinear below the top of the range because most layers of +this hybrid are KDA with fixed-size state; only the MLA layers' KV grows +with context, which is why the top-of-range marginal slope is the right +constant to extrapolate from. With it the budget admits a **4096-token +declared context** on this board (it admitted ~2560 before), validated +by a real 3292-token prompt run: prefill 42.4–42.8 tok/s, decode +0.72–0.81 tok/s, CUDA peak 3.14 GiB, RSS 5.8 GiB, no OOM, 39 MB of +resolved headroom at a 3584-token declaration. Generation validated to +**256 tokens** in one request (0.807 / 0.820 tok/s — no decay with +length). + +*Memory pressure, repeated requests, sustained run.* 30 sequential +requests in one process over 25 minutes (960 tokens): decode mean +**0.817 tok/s**, sd 0.028, quartile means 0.791 / 0.834 / 0.816 / 0.824 +— no degradation over time. Junction temperature 54.5–61.8 °C with +**no throttling** (GPU busy in 749 of 774 samples). RSS moved between +5.45 and 5.70 GiB, rising and falling rather than growing monotonically: +mmapped expert pages reclaimed and re-faulted, not a leak; CUDA +allocator peaks flat at 1.57 GiB. + +*Swap, marked explicitly.* The sustained run DID swap: +1,165 MB during +the first fifth, then flat (+24, 0, +18, −2 MB), with no throughput +loss. That is a one-time eviction of the resident desktop session's +idle pages, not the model's working set thrashing — but per Gate C a +swapping run cannot establish a resident-memory target, so no swap-free +claim is made for this board with a desktop resident. + +*The claim, qualified.* edge0-8b runs on a physical Jetson Orin Nano +8 GB at **0.77–0.82 tok/s decode**, 42 tok/s prefill at 3.3k tokens, +within a byte budget that rejects what does not fit before loading, with +token choices matching the reference implementation on this device. The +claim covers exactly the board, build, checkpoint, profile and workloads +above: single request, greedy sampling, ≤ 4096 declared context, ≤ 256 +validated generation, 25W mode, desktop resident. + +*Remaining, none of it blocking:* design C (custom fused kernel, +toolchain recorded); the MLA float32 promotion question that gates the +dense 26%; the deferred pinned-slot transfer pipeline; concurrent / +batched requests, which were never in scope here; and short-prompt +prefill overhead (2–3 tok/s at 37 tokens against 42 tok/s at 3292), +which no increment targeted. diff --git a/docs/plans/rtx6000-task6-handoff.md b/docs/plans/rtx6000-task6-handoff.md new file mode 100644 index 0000000..0366f35 --- /dev/null +++ b/docs/plans/rtx6000-task6-handoff.md @@ -0,0 +1,241 @@ +# Task 6 handoff: quantized-kernel development on a workstation GPU + +Written 15 September 2026 on the target Orin, for kernel development on +a separate CUDA workstation (RTX 6000-class). Read +[`jetson-orin-nano.md`](jetson-orin-nano.md) Task 6 and Gate D first — +this document adds the measured profile that picks the work and the +division of labor between the two machines. Branch: +`feat/jetson-orin-nano-foundation` (Tasks 1–5 complete; Task 3–5 +statuses in the plan carry the Orin evidence). + +## Why this split + +Kernel iteration (nvcc/CUTLASS builds, parity sweeps) is slow and +memory-cramped on the 8 GB Orin. Correctness of an int4 kernel is +arch-portable and can be developed anywhere torch + CUDA runs; sm_87 +SASS cross-compiles from any x86 CUDA 12/13 toolchain +(`-gencode arch=compute_87,code=sm_87` — no Jetson needed to build). +What does NOT transfer is performance: an RTX 6000 (discrete GDDR, +different SM generation) has a different bottleneck shape than the +unified-LPDDR5 sm_87 board, exactly the way the plan forbids projecting +GB10 numbers. So: **develop and parity-test there, decide and accept +here.** + +## The measured Orin profile (what to optimize) + +Decode is compute-bound (GPU 99%; transfers measured near-irrelevant — +plan Task 5 status). Per-op shares of one decode step, synchronized +timers, unoptimized profile (`EDGE0_TORCH_WEIGHT_CACHE=0`; shares are +the signal, absolute times are sync-inflated): + +| Op | Share | Calls per step | What it is | +|---|---:|---:|---| +| `backends/cuda/quant.py::gather_qmm` | **37%** | 69 (23 layers × 3 projections) | routed experts: python loop over ≤8 experts, each dequantized in full to float32, then matmul | +| `backends/cuda/nn.py::QuantizedLinear.forward` | **26%** | 235 | dense attention / shared expert / router / lm_head: dequantize per call, then linear | +| everything else | 37% | — | attention, routing, sampling, expert loads, sync overhead | + +Both consumers read the same storage format: MLX affine int4, +group_size 64, codes packed 8-per-uint32 little-end-first along the +input axis, `w = code * scale + bias`, scales/biases bf16, **scales can +be negative and there is no zero-point symmetry** — see the layout +notes at the top of `backends/cuda/quant.py`, verified against real +`mx.quantize` output in `tests/test_cuda_backend.py`. + +## Candidate designs (in rising effort order) + +A. **Repack to torch's built-in int4 kernel.** `torch.ops.aten. + _weight_int4pack_mm` (the gpt-fast / torchao path) runs group-wise + int4×bf16 on CUDA. One-time repack at load (MLX layout → tensor-core + tiled layout; MLX's `scale/bias` maps to the kernel's + `scales_and_zeros`) would accelerate `QuantizedLinear` with no + custom extension. Check sm_87 support in the pinned torch build + FIRST (the Orin runs the official `2.14.0+cu130` aarch64 wheel); + verify the affine (bias, not zero-point) mapping is exact, not + approximate. +B. **Batch the expert loop.** `gather_qmm` currently dequantizes and + matmuls one distinct expert at a time in python. A batched variant + (dequantize the ≤8 gathered experts into one `[E, out, in]` + transient, one `bmm`) cuts launch count ~8× at the cost of a + bounded, budget-priceable transient. Pure torch, no extension — + cheap to try before any custom kernel. +C. **Custom fused kernel** (CUDA C++/CUTLASS extension): int4 dequant + in registers fused into the matmul, covering both the gathered + (`rhs_indices`) and dense cases. Highest ceiling, highest cost; pin + the extension/compiler combination and prove runtime loading on the + Orin stack per the plan. + +Recommendation: try A and B first — they are hours, not days, and their +Orin acceptance can run immediately. + +## Workstation setup + +```bash +git clone && cd Edge0 +git checkout feat/jetson-orin-nano-foundation +python3.12 -m venv .venv +.venv/bin/pip install "numpy>=1.24" "safetensors>=0.4" "tokenizers>=0.15" \ + "psutil>=5.9" "pytest>=8.0" transformers +.venv/bin/pip install torch --index-url https://download.pytorch.org/whl/cu130 +.venv/bin/pip install -e . --no-deps # keep MLX out +.venv/bin/python -m pytest -q # expect: all pass, MLX suites skip +.venv/bin/python -m pytest tests/test_torch_cuda_smoke.py --require-cuda -q +``` + +Real weights are optional for kernel work (the parity tests build +synthetic quantized tensors); for end-to-end runs fetch +`Edge0/Edge0-8B-A1B-preview` (~4.2 GB) per the README and export +`EDGE0_BACKEND=cuda EDGE0_8B_MODEL=...`. + +## Deliverables (what to bring back) + +1. The kernel/path behind a **guarded dispatch**: only supported + dtype/shape/stride/transpose/group-size/device combinations take the + new path; everything else — including 2-bit and 8-bit layouts — + falls through to the current reference implementation unchanged. + Off by default (env-gated), like every Task 5 knob. +2. **Parity tests in the smoke-suite style** (hand-computed references, + no MLX): every broadcast/gather shape already covered in + `tests/test_cuda_backend.py`, index boundaries, negative scales, + quantization extremes, all checkpoint dtypes. Tolerances DEFINED IN + THE TEST before any Orin run (the Task 3 reference check shows the + format for pre-registering them). +3. Build pinned and cross-arch: compile for the workstation's arch AND + `compute_87`; record the exact toolchain versions. For design A, + record the repack transform and its inverse test. +4. Scratch/transient memory of the new path measured and stated in + bytes, so the Task 4 budget can price it (`streaming/budget.py` + takes it as a deduction or per-call transient). + +## What stays on the Orin (do not claim from the workstation) + +* Which design wins, and any tok/s number: re-run the profile above and + the 3-launch benchmark protocol (`docs/nvidia.md`, Task 3 status) on + the Orin only. +* Runtime loading of any compiled extension on the L4T R39 / CUDA 13.2 + / torch 2.14.0+cu130 stack. +* The 32-token reference check (`scripts/orin_reference_check.py`) + with the new path enabled — token choices must stay explainable + under the pre-registered comparison, same as Task 3. +* The Task 5 re-measurement: once compute shrinks, transfers matter + more — `EDGE0_CACHE_SLOTS` / `EDGE0_PREDICT_PREFETCH` (+7.7% + measured pre-Task-6) and the deferred pinned-slot pipeline get + re-evaluated then. + +## Workstation results (15 September 2026) + +Workstation: RTX PRO 6000 Blackwell (sm_120, 96 GB), Windows 11; venv +Python 3.12 + **torch 2.14.0+cu130** (the Orin's torch version, x86 +build) for every GPU measurement; a torch-CPU venv (Python 3.11) and a +Python 3.9 venv for portability; WSL Ubuntu 24.04 for the Linux checks. +Nothing below is an Orin number. + +### Delivered (all opt-in, off by default, guarded, priced) + +| Design | Knob | Files | Status | +|---|---|---|---| +| B batched expert gather | `EDGE0_QMM_BATCHED=1`, `EDGE0_QMM_BATCHED_MAX_BYTES` | `backends/cuda/quant.py` | exact (float32, parity 1e-4 to a float64 hand reference, 1e-5 to the loop); 4-bit affine only, 2/8-bit and over-cap calls fall through; transient priced by `batched_transient_bytes` and asserted against allocator peaks in `tests/test_quant_paths.py`; kernel launches per call (reference loop vs batched) counted with `torch.profiler`, >= 3x fewer asserted | +| A' capped exact weight cache | `EDGE0_TORCH_WEIGHT_CACHE_BYTES=` | `backends/cuda/nn.py` (`WeightCachePolicy`), `backends/cuda/io.py`, `engine/ling.py` | first-fit in checkpoint order, each module admitted only if its resident bytes AND its own build transient fit; engine shrinks the cap to the budget headroom. The dequantized weight is bit-identical to the on-the-fly path (same per-chunk expression, asserted on CUDA); the cached forward issues one GEMM instead of one per 4096-row chunk, which only the two modules wider than a chunk (lm_head, layer-0 dense gate/up) can notice. **The fill is chunked**: whole-weight dequantization of lm_head (157184 x 1536) peaks at 3702 MiB for a 460.5 MiB result and would OOM the Orin's stated headroom; chunked it measures 556.9 MiB, and `fill_transient_bytes` prices that. The full cache shares the branch and inherits the fix; it is priced from the loader's measured total (235 cacheable linears, 1.402 GB bf16 — routers are not quantized here and `word_embeddings` is a lookup-only `QuantizedEmbedding`) plus one fill transient | +| A torch int4 kernel | `EDGE0_INT4PACK=1` | `backends/cuda/int4pack.py`, `backends/cuda/nn.py` | runs on torch 2.14.0+cu130; **approximate**: per-weight error median 3.7e-4 / p99 3.7e-3 / max 1.1e-2 (n=512, k=1536, MLX-like statistics), 68% of weights identical to the reference's bf16 weights, outputs within 3.5e-3 of the output scale; bounds registered in `tests/test_int4pack.py` (1e-2 of output scale; 2^-6 (max\|w\| + max\|zero\|) per weight). Constraints verified on this build: bf16 activations only, `in_features % 128 == 0` (every edge0-8b dense k qualifies: 256, 512, 1536, 2048, 4608), `out_features % 8 == 0` up to the 157184-row lm_head; the dequantization is exactly `(q - 8) * scale + zero` on an identity probe. **sm_87 is not answered from a table**: `int4pack.probe` executes the kernel once per module on the target device before committing, and any runtime failure disables the path for that module — so a wheel without an image for the Orin's arch falls through instead of raising. Extra resident bytes: the repacked payload, 0.351 GB for all dense linears; originals stay for the fall-through paths | +| C custom fused kernel | — | — | not written; toolchain proven: WSL `nvcc` 13.0.88 + g++ 13.3.0 compile `-gencode arch=compute_87,code=sm_87 -gencode arch=compute_120,code=sm_120` and the binary runs on this GPU; Windows alternative: VS 2022 BuildTools + CUDA 12.8 | + +Budget: `streaming/budget.py` gained `Reserves.kernel_transient_bytes` +(the batched gather's cap, deducted once when the path is on) and +`dense_cache_request`, which prices the dense cache as resident bytes +plus one build transient (`full` keeps its veto, `capped` shrinks to the +headroom, `off` prices nothing); `engine/ling.py` calls it and +re-finalizes the policy after the budget resolves. +Report: `examples/bench.py` records `caches.quant_paths` (batched gate +and cap, the policy's admitted/skipped modules with their bytes and fill +transients, int4pack repacked modules and extra bytes) and echoes the +four knobs. +Reference check: `scripts/orin_reference_check.py --test-env KEY=VALUE` +/ `--reference-env` run the two phases with different knobs on the same +device. The report now carries **host identity** (GPU name, compute +capability, torch build, `/proc/device-tree/model`, machine, hashed +hostname) and each phase's **resolved** knobs, not just the CLI +overrides — a workstation report can no longer be mistaken for an Orin +one, and a same-device run whose two phases resolve identical knobs +warns that it is comparing a configuration with itself. The schema is +`edge0-reference-check/2`. + +### Same-device reference checks on the workstation (real weights) + +`scripts/orin_reference_check.py models/edge0-8b --test-device cuda +--reference-device cuda --test-env ` — the reference path +teacher-forces the tokens the knob's path generated, on the same GPU, so +the only variable is the knob. 32 greedy tokens, edge0-8b, RTX PRO 6000: + +| Knob | Token choices | Divergences | max abs logit diff | Verdict against the registered bound (1.0) | +|---|---|---|---|---| +| `EDGE0_TORCH_WEIGHT_CACHE_BYTES=2000000000` | 32/32 identical | 0 | **0 exactly** | PASS — bit-identical through the whole model, the strongest form of the exactness claim | +| `EDGE0_QMM_BATCHED=1` | 32/32 identical | 0 | 1.01 | FAIL of that bound, reported as such | +| `EDGE0_INT4PACK=1` | 32/32 identical | 0 | 2.06 | FAIL of that bound, reported as such | + +All three keep every token choice with zero near-ties consumed. The two +that exceed the bound are **not** loosened: as in Task 3, the bound's +failure is recorded. For the batched gather the mechanism is identified — +at the real decode shapes it lands within one bfloat16 ulp of the +reference loop (`tests/test_quant_paths.py` asserts exactly that), and 24 +layers of one-ulp differences accumulate; note the Task 3 CUDA-vs-CPU +baseline measured 2.06 by the same mechanism, so 1.01 sits below the +noise the tier already exhibits. For the int4 kernel the difference is +expected by construction (bf16 zero-point and bf16 dequantization). The +Orin's own reference check decides acceptability; a top-k-scoped +criterion is the change to propose in review first, per the Task 3 status. + +**A caution about this machine's shape.** Decode here runs at roughly +1 tok/s against the Orin's 0.55 — under 2x, on a GPU with vastly more +compute. That is the handoff's point restated as data: the bottleneck is +launch and host overhead, not arithmetic, so a workstation speedup would +not transfer. No tok/s comparison between configurations is reported +here for that reason; only launch counts, bytes and parity. + +### The finding that reshapes the dense half of Task 6 + +**Most dense `QuantizedLinear` calls run with float32 activations, so +both dense paths engage on a small minority of them.** Traced on the real +model, one prefill plus one decode step: **400 of 470 calls arrive as +float32, only 70 as bfloat16.** Both dense knobs require bf16 — the cache +because it is priced at the checkpoint's scales itemsize, the int4 kernel +because `_weight_int4pack_mm` accepts nothing else. So: + +* the weight cache admits 235 modules / 1.402 GB but **actually fills 35 + modules / 0.174 GB**; the reservation is conservative and never + under-reserves, and the bench report now states `filled_bytes` beside + `admitted_bytes` so the two are never confused; +* the int4 kernel repacks **exactly those same 35 modules** — that is + where its `repacked_modules: 35` comes from, not from a shape guard; +* neither knob moves decode measurably on this GPU (0.95 and 0.96 against + a 0.92 baseline), which is what covering 70 of 470 calls predicts. + +This is the thing to settle before spending more effort on design A or C: +the dense half of the 26% is currently gated by **activation dtype, not by +the kernel**. Whether those float32 call sites can run bf16 (and what that +costs numerically) is a model-path question for the Orin session, and it +governs how much of that 26% any kernel can reach. The batched expert +gather is unaffected — it covers the routed 37% and engages on every call. + +### What the benchmarks showed (workstation only, 3 launches x 2 runs each) + +| Configuration | decode tok/s (6 runs) | CUDA alloc peak | Engaged | +|---|---|---|---| +| baseline | 0.90-0.95 (mean 0.92) | 1.02 GiB | — | +| `EDGE0_QMM_BATCHED=1` | 1.41-1.52 (mean 1.46) | 1.02 GiB | every gather call | +| `EDGE0_TORCH_WEIGHT_CACHE_BYTES=1.5e9` | 0.94-0.96 (mean 0.95) | 1.18 GiB | 35/235 modules filled | +| `EDGE0_INT4PACK=1` | 0.95-0.96 (mean 0.96) | 1.07 GiB | 35 modules repacked | +| batched + weight cache | 1.46-1.47 (mean 1.46) | 1.18 GiB | both | + +Every configuration ran to completion with the budget active, no OOM. +**Do not project the batched gather's 1.59x onto the Orin.** The +mechanism (fewer kernel launches) is the same one that dominates the +Orin's decode, but this is a different SM generation with different +launch costs, and the plan forbids projecting between machines — the +whole reason this workstation decodes at only ~1 tok/s against the +Orin's 0.55 despite far more compute. The number is here as evidence the +path works end to end, not as a prediction. + +### How the Orin accepts a knob + +1. `EDGE0_BACKEND=cuda EDGE0_TORCH_DEVICE=cuda python -m pytest tests/test_quant_paths.py tests/test_weight_cache_policy.py tests/test_int4pack.py tests/test_task6_budget.py --require-cuda -q` on the board (the CUDA cases run for real there). +2. `scripts/orin_reference_check.py edge0-8b --test-device cuda --reference-device cuda --test-env = --output "$EDGE0_RUN_DIR/ref-.json"` — the reference path vs the knob on the same device; the exact knobs must match 32/32 tokens with logits at float32-noise level, the int4 kernel is judged by the pre-registered bounds and the near-tie rule. +3. The three-launch benchmark protocol (Task 3 status) with `EDGE0_MEMORY_BUDGET=auto` so the transient / cache bytes are priced, comparing against the Task 5 baseline; then the Task 5 re-measurement. diff --git a/examples/bench.py b/examples/bench.py index 971c459..12faba3 100644 --- a/examples/bench.py +++ b/examples/bench.py @@ -2,17 +2,40 @@ """edge0 speed + memory benchmark (any tier). Runs the same measurement shape as the deployment's bench: prefill -> -sampled warmup steps -> timed sampled decode -> report tok/s and MLX -peak active memory. Works for both tiers; the tier is auto-detected -from the checkpoint (or forced via the model name). +greedy warmup steps (untimed) -> timed sampled decode -> report tok/s and +the backend's peak memory (MLX peak active memory; the torch CUDA +allocator's peak allocated bytes on a CUDA device; n/a for torch on +cpu/mps). Works for both tiers; the tier is auto-detected from the +checkpoint (or forced via the model name). + +What the numbers are (and are not). Two runs per process; each run +resets the engine, times the chat-templated prefill, runs ``--warmup`` +greedy steps untimed, then times ``--ntok`` iterations of "sample one +token, run one engine step". The loop is fixed-length and does not stop +at EOS; the final step's logits are computed but never sampled. This is +NOT time-to-first-token and NOT an EOS-aware completion latency. On a +CUDA device the resolved device is synchronized at those phase +boundaries only (after reset, end of prefill, end of warmup, after the +final timed step) so every wall time includes the device work it +launched; storage, transfer and CPU sampling costs are inside the wall +time by design. Sampling at temperature 0 is greedy (see +``edge0.sampling``), so ``BENCH_TEMP=0`` output does not depend on the +seed; the seed is still recorded. Usage: python examples/bench.py /path/to/model [--ntok 200] [--warmup 10] python examples/bench.py edge0-35b # via $EDGE0_35B_MODEL BENCH_LONG=1 python examples/bench.py edge0-35b # ~3k-token prefill + python examples/bench.py edge0-8b --json-output "$RUN/bench.json" \\ + --probe-json "$RUN/orin-probe.json" # machine-readable report + +``--json-output`` must name a NEW file (an existing one is refused before +the model loads); the report schema is described in docs/nvidia.md. +Keep the JSON outside tracked source: it records local paths and the +prompt text. Env (sampling knobs, defaults follow the tier's GenerationConfig): - BENCH_TEMP / BENCH_PROMPT / BENCH_SEED + BENCH_TEMP / BENCH_PROMPT / BENCH_SEED / BENCH_NTOK BENCH_LONG=1 use a ~3k-token synthetic prompt so the prefill timing is meaningful (short prompts only measure fixed overhead; set BENCH_PROMPT for your own long text) @@ -21,13 +44,26 @@ from __future__ import annotations import argparse +import dataclasses +import enum +import hashlib +import math import os import sys import time +import traceback +from datetime import datetime, timezone +from pathlib import Path +from typing import Any + +try: # imported as ``examples.bench`` (tests) + from examples import benchmark_measure as measure + from examples import benchmark_report as report +except ImportError: # run as a script: helpers sit next to this file + import benchmark_measure as measure # type: ignore[no-redef] + import benchmark_report as report # type: ignore[no-redef] -from edge0 import AutoEngine -from edge0.backends import core -from edge0.sampling import sample +REPO_ROOT = Path(__file__).resolve().parents[1] PROMPTS = { "edge0-35b": "什么是混合专家模型(MoE)?它和普通 Transformer 有什么区别?", @@ -55,12 +91,49 @@ "hidden state, SSD read latency hides behind compute entirely, which " "matters for large-model deployment on edge hardware.") +#: Environment knobs echoed into the report as REQUESTED (raw strings, +#: snapshotted before the engine is built: ``engine/ling.py`` sets a +#: default ``LING_HIDDEN_CLIP`` the user did not ask for). +ENV_KNOBS = ("BENCH_TEMP", "BENCH_PROMPT", "BENCH_SEED", "BENCH_NTOK", + "BENCH_LONG", "EDGE0_BACKEND", "EDGE0_TORCH_DEVICE", + "EDGE0_MEMORY_BUDGET", "EDGE0_BUDGET_CONTEXT", + "EDGE0_CACHE_SLOTS", "EDGE0_PREDICT_PREFETCH", + "EDGE0_QMM_BATCHED", "EDGE0_QMM_BATCHED_MAX_BYTES", + "EDGE0_TORCH_WEIGHT_CACHE_BYTES", "EDGE0_INT4PACK", + "EDGE0_TORCH_WEIGHT_CACHE", "EDGE0_PREWARM", "LING_PREWARM", + "MLX_CACHE_LIMIT_MB", "LING_HIDDEN_CLIP", + "PREROUTER_FEATURE_TOPK", "PREROUTER_INTRA") + +_TOKENIZER_FILE_PREFIXES = ("tokenizer", "chat_template", "special_tokens", + "vocab", "merges", "added_tokens") + + +def _framework(): + """The active backend, its array namespace and the sampler. + + Imported lazily: ``edge0.backends`` binds to ``EDGE0_BACKEND`` at + import time, and nothing in argument validation, ``--help`` or the + pre-load output/probe checks needs a backend. Tests replace this + with fakes. + """ + from edge0.backends import backend, core + from edge0.sampling import sample + return backend, core, sample + def _long_prompt() -> str: paras = [_PARA_ZH if i % 2 == 0 else _PARA_EN for i in range(40)] return "请仔细阅读以下材料并用一句话总结其核心思想:" + " ".join(paras) +def _prompt_source() -> str: + if os.environ.get("BENCH_PROMPT"): + return "BENCH_PROMPT" + if os.environ.get("BENCH_LONG") == "1": + return "BENCH_LONG" + return "tier default" + + def _prompt_for(engine) -> str: if os.environ.get("BENCH_PROMPT"): return os.environ["BENCH_PROMPT"] @@ -69,46 +142,115 @@ def _prompt_for(engine) -> str: return PROMPTS.get(engine.name, DEFAULT_PROMPT) -def run_bench(engine, ntok: int, warmup: int) -> dict: +def _encode_prompt(engine, prompt: str) -> list[int]: tok = engine._tok - prompt = _prompt_for(engine) if hasattr(engine, "encode_chat"): - ids = engine.encode_chat( - [{"role": "user", "content": prompt}], think=False) - else: - ids = list(tok(tok.apply_chat_template( - [{"role": "user", "content": prompt}], - tokenize=False, add_generation_prompt=True, - enable_thinking=False))["input_ids"]) + return [int(t) for t in engine.encode_chat( + [{"role": "user", "content": prompt}], think=False)] + return [int(t) for t in tok(tok.apply_chat_template( + [{"role": "user", "content": prompt}], + tokenize=False, add_generation_prompt=True, + enable_thinking=False))["input_ids"]] + +def _sampling_knobs(engine) -> dict: + """Resolved sampling settings: the tier's GenerationConfig, overridden + by BENCH_TEMP / BENCH_SEED (parsed once; ``BENCH_SEED=0`` is a seed).""" gen = getattr(engine.cfg, "gen", None) - temperature = float(os.environ.get( - "BENCH_TEMP", getattr(gen, "temperature", 0.7))) - top_k = getattr(gen, "top_k", 64) - top_p = getattr(gen, "top_p", 0.95) - rp = getattr(gen, "repetition_penalty", 1.0) - seed = os.environ.get("BENCH_SEED") + seed_env = os.environ.get("BENCH_SEED") + return { + "temperature": float(os.environ.get( + "BENCH_TEMP", getattr(gen, "temperature", 0.7))), + "top_k": getattr(gen, "top_k", 64), + "top_p": getattr(gen, "top_p", 0.95), + "repetition_penalty": getattr(gen, "repetition_penalty", 1.0), + "seed": int(seed_env) if seed_env is not None else None, + } + + +def _headline_bytes(peaks: dict, measurement): + """The one per-run peak the human output and the legacy ``peak_gib`` + show (allocated on CUDA, peak active on MLX); ``None`` when that + metric is unavailable.""" + name = getattr(measurement, "headline", None) + metric = peaks.get(name) if name else None + if metric is None: + for candidate in peaks.values(): + if candidate.get("bytes") is not None: + return candidate["bytes"] + return None + return metric.get("bytes") + + +def _fmt_gib(value) -> str: + return "n/a" if value is None else f"{value:.2f} GiB" + + +def run_bench(engine, ntok: int, warmup: int, *, measurement=None, + runs: int = 2) -> dict: + """Run the fixed-length sample-and-step benchmark. + + Returns the legacy dict (``runs`` with prefill_s/decode_s/ntok/ + peak_gib/tok_s, ``mean_tok_s``, ``peak_gib``) plus ``report_runs`` + (validated ``benchmark_report.run_result`` entries), ``workload``, + ``measurement`` and ``execution_evidence`` for the JSON report. + ``measurement`` defaults to the adapter for the active backend and + its resolved device (``benchmark_measure.make_measurement``). + """ + if not isinstance(ntok, int) or isinstance(ntok, bool) or ntok <= 0: + raise ValueError( + "ntok must be a positive integer: an empty timed window is not " + "a benchmark") + if not isinstance(warmup, int) or isinstance(warmup, bool) or warmup < 0: + raise ValueError("warmup must be an integer >= 0") + backend, core, sample = _framework() + if measurement is None: + measurement = measure.make_measurement(backend.name, core) + + prompt = _prompt_for(engine) + ids = _encode_prompt(engine, prompt) + knobs = _sampling_knobs(engine) + temperature = knobs["temperature"] + top_k, top_p, rp, seed = (knobs["top_k"], knobs["top_p"], + knobs["repetition_penalty"], knobs["seed"]) results = [] - for run in range(2): + report_runs = [] + evidence = None + for run in range(runs): engine.reset() - core.reset_peak_memory() + # Finish any device work the reset queued, then reset the peak + # statistics right before prefill: the per-run peak covers + # prefill + warmup + decode (plus whatever is resident now). + measurement.synchronize() + measurement.reset_peak() t0 = time.perf_counter() engine.prefill(ids) + # Prefill's final device work belongs to prefill wall time. + measurement.synchronize() t_prefill = time.perf_counter() - t0 logits = engine.next_logits() + if evidence is None: + evidence = measure.execution_evidence(logits) if seed is not None: core.random.seed(int(seed)) history = list(ids) - # warmup: sampled, untimed (first tokens are dominated by expert - # cold-starts; the timed window is the steady state) + # warmup: GREEDY (argmax), untimed -- the first tokens are + # dominated by expert cold-starts; the timed window below is the + # steady state. Only the timed loop samples. + warmup_done = 0 for _ in range(warmup): tid = int(core.argmax(logits, axis=-1).item()) history.append(tid) logits = engine.step(tid) + warmup_done += 1 + # Warmup must be complete before the decode clock starts. + measurement.synchronize() out: list[int] = [] t0 = time.perf_counter() + # timed: sample from the current logits, then one forward. The + # last forward's logits are never sampled; no EOS stop. for _ in range(ntok): tid = sample(logits, temperature=temperature, top_k=top_k, top_p=top_p, repetition_penalty=rp, @@ -116,48 +258,542 @@ def run_bench(engine, ntok: int, warmup: int) -> dict: history.append(tid) out.append(tid) logits = engine.step(tid) + # The final timed step's device work is inside decode wall time. + measurement.synchronize() t_decode = time.perf_counter() - t0 - peak_gib = core.get_peak_memory() / (1024 ** 3) - results.append(dict(prefill_s=t_prefill, decode_s=t_decode, - ntok=len(out), peak_gib=peak_gib, - tok_s=len(out) / t_decode if t_decode else 0.0)) + peaks = measurement.peak_metrics() # after the synchronized decode + report_runs.append(report.run_result( + index=run, prompt_tokens=len(ids), + requested_warmup_tokens=warmup, warmup_tokens=warmup_done, + requested_timed_tokens=ntok, timed_decode_tokens=len(out), + prefill_seconds=t_prefill, decode_seconds=t_decode, + memory=peaks, execution_evidence=evidence)) + peak_bytes = _headline_bytes(peaks, measurement) + peak_gib = (peak_bytes / (1024 ** 3)) if peak_bytes is not None else None + results.append({"prefill_s": t_prefill, "decode_s": t_decode, + "ntok": len(out), "peak_gib": peak_gib, + "tok_s": len(out) / t_decode if t_decode else 0.0}) pf_tps = (len(ids) / t_prefill) if t_prefill else 0.0 print(f"prompt={len(ids)} tok prefill={t_prefill:.2f}s " f"({pf_tps:.0f} tok/s) " f"decode={len(out)}/{t_decode:.2f}s " f"tok/s={results[-1]['tok_s']:.1f} " - f"peak_active={peak_gib:.2f} GiB", flush=True) + f"peak_active={_fmt_gib(peak_gib)}", flush=True) mean_ts = sum(r["tok_s"] for r in results) / len(results) - max_peak = max(r["peak_gib"] for r in results) - print(f"[bench] mean tok/s={mean_ts:.1f} peak_active<={max_peak:.2f} GiB", - flush=True) - return {"runs": results, "mean_tok_s": mean_ts, "peak_gib": max_peak} + max_peak = max((r["peak_gib"] for r in results + if r["peak_gib"] is not None), default=None) + print(f"[bench] mean tok/s={mean_ts:.1f} " + f"peak_active<={_fmt_gib(max_peak)}", flush=True) + workload = { + "prompt_source": _prompt_source(), + "prompt_sha256": hashlib.sha256(prompt.encode("utf-8")).hexdigest(), + "prompt_token_ids_sha256": hashlib.sha256( + ",".join(str(t) for t in ids).encode()).hexdigest(), + "prompt_chars": len(prompt), + "prompt_text": prompt, + "prompt_tokens": len(ids), + "requested_warmup_tokens": warmup, + "requested_timed_tokens": ntok, + **knobs, + } + return {"runs": results, "mean_tok_s": mean_ts, "peak_gib": max_peak, + "report_runs": report_runs, "workload": workload, + "measurement": measurement.describe(), + "execution_evidence": evidence} -def main() -> int: - ap = argparse.ArgumentParser(description=__doc__) - ap.add_argument("model", help="checkpoint dir or tier name") - ap.add_argument("--ntok", type=int, - default=int(os.environ.get("BENCH_NTOK", "200"))) - ap.add_argument("--warmup", type=int, default=10) - args = ap.parse_args() - # tier name -> checkpoint dir via $EDGE0__MODEL (cli parity) - model = args.model - name = None +# ---- report assembly (outside every timed interval) --------------------------- + + +def _jsonable(value) -> Any: + """Config objects -> JSON: enums by name, paths as strings, tuples as + lists, dataclasses as dicts, anything else by repr.""" + if isinstance(value, bool) or value is None or isinstance( + value, (str, int, float)): + return value + if isinstance(value, enum.Enum): + return value.name + if isinstance(value, dict): + return {str(k): _jsonable(v) for k, v in value.items()} + if isinstance(value, (list, tuple, set, frozenset)): + return [_jsonable(v) for v in value] + if isinstance(value, os.PathLike): + return os.fspath(value) + if dataclasses.is_dataclass(value) and not isinstance(value, type): + return _jsonable(dataclasses.asdict(value)) + return repr(value) + + +def _with_reasons(obj: dict, reason: str) -> dict: + """Give a plain mapping the schema's null convention: every ``None`` + value gets ``reason`` in ``unavailable_reasons``.""" + out = dict(obj) + reasons = dict(out.get("unavailable_reasons") or {}) + for key, value in out.items(): + if key != "unavailable_reasons" and value is None: + reasons.setdefault(key, reason) + out["unavailable_reasons"] = reasons + return out + + +def _env_snapshot() -> dict: + """The requested environment knobs, taken BEFORE the engine is built.""" + return {k: os.environ[k] for k in ENV_KNOBS if k in os.environ} + + +def _model_env() -> dict: + """Model-read knobs as the bailing_hybrid code parses them, resolved + AFTER the engine is built (``engine/ling.py`` fills in a default + ``LING_HIDDEN_CLIP``). They change numerics and per-token work, so + before/after comparisons must hold them equal.""" + out: dict[str, Any] = { + "ling_hidden_clip": None, + "prerouter_feature_topk": os.environ.get("PREROUTER_FEATURE_TOPK", + "teacher"), + "prerouter_intra": os.environ.get("PREROUTER_INTRA", "0") == "1", + "note": ("parsed exactly as the bailing_hybrid model code does; " + "ling_hidden_clip 0.0 means the hidden-state clip is off"), + "unavailable_reasons": {}, + } + raw = os.environ.get("LING_HIDDEN_CLIP", "0") + try: + out["ling_hidden_clip"] = float(raw) + except ValueError: + out["unavailable_reasons"]["ling_hidden_clip"] = ( + f"LING_HIDDEN_CLIP={raw!r} is not a number") + return out + + +def _model_group(engine, model_dir: str) -> dict: + root = Path(model_dir) + manifest = measure.checkpoint_manifest(root) + config = manifest["config"] + cfg = engine.cfg + prerouter = getattr(cfg, "prerouter", None) + adapters = { + "lora": measure.adapter_identity(getattr(cfg, "lora", None), root), + "prerouter": measure.adapter_identity( + getattr(prerouter, "weights_file", None), root), + "lora_settings": _with_reasons(_jsonable({ + "r": getattr(cfg, "lora_r", None), + "alpha": getattr(cfg, "lora_alpha", None)}), + "tier config has no LoRA settings"), + "prerouter_settings": _with_reasons(_jsonable({ + "start_layer": getattr(prerouter, "start_layer", None), + "hidden": getattr(prerouter, "hidden", None), + "dtype": getattr(prerouter, "dtype", None), + "feature_topk": getattr(prerouter, "feature_topk", None), + "owners": getattr(prerouter, "owners", None), + "top_k": getattr(cfg, "prerouter_top_k", None)}), + "tier config has no prerouter"), + } + small = manifest["small_file_sha256"] + tokenizer_files = [] + for entry in manifest["manifest"].get("files") or []: + name = entry["path"] + if not name.startswith(_TOKENIZER_FILE_PREFIXES): + continue + item = {"name": name, "size_bytes": entry["size_bytes"], + "sha256": small.get(name), "unavailable_reasons": {}} + if item["sha256"] is None: + item["unavailable_reasons"]["sha256"] = ( + "not hashed: larger than the small-file cap or unreadable") + tokenizer_files.append(item) + group = { + "checkpoint_path": str(root.resolve()), + "tier": engine.name, + "model_type": config.get("model_type") or None, + "architectures": _jsonable(config.get("architectures") or []), + "config_sha256": manifest["config_sha256"], + "manifest": manifest["manifest"], + "small_file_sha256": small, + "safetensors_headers": manifest["safetensors_headers"], + "adapters": adapters, + "tokenizer_files": tokenizer_files, + "unavailable_reasons": {}, + } + if group["model_type"] is None: + group["unavailable_reasons"]["model_type"] = ( + "config.json missing or has no model_type") + if group["config_sha256"] is None: + group["unavailable_reasons"]["config_sha256"] = ( + "config.json not found in the checkpoint directory") + return group + + +def _runtime_group(bench: dict, backend) -> dict: + info = measure.runtime_info(backend.name) + try: + version = str(backend.version) + except Exception as exc: # noqa: BLE001 - never fail the report here + version = None + info["unavailable_reasons"]["backend_version"] = ( + f"backend.version failed: {type(exc).__name__}: {exc}") + return { + "backend": backend.name, + "backend_version": version, + "resolved_device": bench["measurement"]["resolved_device"], + "measurement": bench["measurement"], + "execution_evidence": bench["execution_evidence"], + "power_mode": measure.power_mode(), + **info, + } + + +def _quant_paths_group(engine) -> dict: + """Resolved state of the torch backend's opt-in quantized paths + (Task 6): the batched expert gather, the capped exact weight cache + and the int4 kernel, with the bytes each one added.""" + from edge0.backends.cuda import int4pack as i4 + from edge0.backends.cuda import nn as cuda_nn + from edge0.backends.cuda import quant as cq + repacked = 0 + extra = 0 + filled = 0 + filled_bytes = 0 + reasons: dict[str, str] = {} + model = getattr(engine, "model", None) + if model is None or not hasattr(model, "modules"): + reasons["repacked_modules"] = "engine exposes no torch model" + reasons["extra_bytes"] = "engine exposes no torch model" + repacked = extra = None + filled = filled_bytes = None + else: + for module in model.modules(): + n = getattr(module, "int4pack_bytes", None) + if callable(n): + b = n() + if b: + repacked += 1 + extra += b + if getattr(module, "_dequantized", None) is not None: + filled += 1 + filled_bytes += cuda_nn.dequantized_bytes(module) + policy = dict(cuda_nn.WEIGHT_CACHE.summary()) + # Admission is deliberately conservative: a module is admitted on the + # assumption it will be used at the priced dtype, but the cache only + # fills when the activation itemsize matches. Report what is actually + # resident next to what was reserved, so the two are never confused. + policy["filled_modules"] = filled + policy["filled_bytes"] = filled_bytes + policy = _with_reasons(policy, "weight cache policy off or not finalized") + return { + "batched_gather": {"enabled": bool(cq.BATCHED), + "max_transient_bytes": int(cq.BATCHED_MAX_BYTES)}, + "weight_cache_policy": policy, + "int4pack": {"enabled": bool(i4.INT4PACK), + "repacked_modules": repacked, + "extra_bytes": extra, + "unavailable_reasons": reasons}, + } + + +def _caches_group(engine, backend) -> dict: + opts = getattr(engine.cfg, "options", None) + reasons: dict[str, str] = {} + if dataclasses.is_dataclass(opts) and not isinstance(opts, type): + layer_options = _with_reasons( + _jsonable(dataclasses.asdict(opts)), + "option unset (None): the model config value applies") + else: + layer_options = {"unavailable_reasons": {}} + reasons["layer_options"] = "engine config has no LayerOptions" + layers = getattr(engine, "_all_stream_layers", {}) or {} + first = next(iter(layers.values()), None) + shared_slots = prefetch_cap = None + if first is None: + reasons["shared_cache_slots_resolved"] = ( + "engine has no streaming layers") + reasons["prefetch_cap_resolved"] = "engine has no streaming layers" + else: + shared_slots = getattr(getattr(first, "shared_cache", None), + "slots", None) + prefetch_cap = getattr(getattr(first, "_prefetch_buf", None), + "cap", None) + if shared_slots is None: + reasons["shared_cache_slots_resolved"] = ( + "streaming layer exposes no shared_cache.slots") + if prefetch_cap is None: + reasons["prefetch_cap_resolved"] = ( + "streaming layer exposes no _prefetch_buf.cap") + threads: dict[str, Any] = { + "load_threads": getattr(opts, "load_threads", None), + "prefetch_threads": getattr(opts, "prefetch_threads", None), + "torch_num_threads": None} + thread_reasons = {} + weight_cache = None + mlx_cache_limit = None + quant_paths = None + if backend.name == "cuda": + from edge0.backends.cuda import nn as cuda_nn + weight_cache = bool(getattr(cuda_nn, "CACHE_DEQUANTIZED", False)) + quant_paths = _quant_paths_group(engine) + try: + import torch + threads["torch_num_threads"] = int(torch.get_num_threads()) + except Exception as exc: # noqa: BLE001 - reported, not raised + thread_reasons["torch_num_threads"] = f"{type(exc).__name__}: {exc}" + reasons["mlx_cache_limit_bytes"] = "backend is cuda" + else: + reasons["weight_cache"] = ( + "EDGE0_TORCH_WEIGHT_CACHE applies to the torch backend only") + reasons["quant_paths"] = ( + "the opt-in quantized paths exist in the torch backend only") + thread_reasons["torch_num_threads"] = "backend is not torch" + try: + mlx_cache_limit = int( + os.environ.get("MLX_CACHE_LIMIT_MB", "256")) * 1024 * 1024 + except ValueError: + reasons["mlx_cache_limit_bytes"] = "MLX_CACHE_LIMIT_MB not an int" + for key in ("load_threads", "prefetch_threads"): + if threads[key] is None: + thread_reasons[key] = "engine config has no LayerOptions" + threads["unavailable_reasons"] = thread_reasons + prewarm = os.environ.get( + "EDGE0_PREWARM", os.environ.get("LING_PREWARM", "0")) == "1" + try: + stats = _jsonable(engine.stats()) + except Exception as exc: # noqa: BLE001 - diagnostics only + stats = {"stats": None, "unavailable_reasons": { + "stats": f"{type(exc).__name__}: {exc}"}} + resolved_budget = getattr(engine, "memory_budget", None) + if resolved_budget is None: + reasons["memory_budget"] = ( + "no memory budget resolved (EDGE0_MEMORY_BUDGET off or " + "engine predates Task 4)") + return { + "layer_options": layer_options, + "shared_cache_slots_resolved": shared_slots, + "prefetch_cap_resolved": prefetch_cap, + "memory_budget": (resolved_budget.as_dict() + if resolved_budget is not None else None), + "weight_cache": weight_cache, + "quant_paths": quant_paths, + "prewarm": prewarm, + "threads": threads, + "mlx_cache_limit_bytes": mlx_cache_limit, + "streaming_stats_after_runs": stats, + "unavailable_reasons": reasons, + } + + +def _build_report(engine, bench: dict, args, argv: list[str], + started_utc: str, model_dir: str, probe, sampler, + env_knobs: dict, model_env: dict, backend) -> dict: + """Assemble the JSON report. Called after the timed runs (and after + the sampler stopped): hashing, git and power queries happen here.""" + requested = { + "model": args.model, + "ntok": args.ntok, + "warmup": args.warmup, + "rss_sample_interval_seconds": args.rss_sample_interval, + "json_output": args.json_output, + "probe_json": args.probe_json, + } + identity = { + "timestamp_utc": started_utc, + "finished_utc": datetime.now(timezone.utc).isoformat(), + "git": measure.git_identity(REPO_ROOT), + "run_count": len(bench["report_runs"]), + "command": {"argv": list(argv), "env_knobs": dict(env_knobs)}, + "requested": _with_reasons(requested, "not requested"), + "unavailable_reasons": {}, + } + workload = dict(bench["workload"]) + workload["prefill_chunk"] = getattr(engine, "prefill_chunk", None) + workload["think"] = getattr(engine, "think", None) + workload["model_env"] = model_env + workload = _with_reasons(workload, "not set") + if workload["seed"] is None: + workload["unavailable_reasons"]["seed"] = ( + "BENCH_SEED not set: sampling is unseeded") + memory = { + "process_peak_rss": measure.process_peak_rss(), + **sampler.metrics(), + "notes": [ + ("process RSS and device allocator peaks are separate views of " + "the same shared DRAM on Jetson; never add them"), + ("RSS includes resident file-backed pages of the mmap'd " + "checkpoint; the sampled *_rss_anon peak is the anonymous " + "footprint (Linux only)"), + ("allocator peaks cover the framework allocator only, not the " + "CUDA context, driver or non-torch memory"), + ("cuda_peak_reserved's floor is every segment cached at the " + "reset, so run 0 and run 1 are expected to be close; use " + "cuda_peak_allocated for per-run activity"), + ("process_peak_rss is a lifetime high-water mark (model load " + "included); per-run peaks are the backend metrics in runs[]"), + ], + "unavailable_reasons": {}, + } + return report.build_report( + identity=identity, + model=_model_group(engine, model_dir), + runtime=_runtime_group(bench, backend), + workload=workload, + caches=_caches_group(engine, backend), + runs=bench["report_runs"], + memory=memory, + probe=probe, + ) + + +# ---- CLI --------------------------------------------------------------------- + + +def _build_engine(model_dir: str, name): + from edge0 import AutoEngine + return AutoEngine.from_pretrained(model_dir, name=name) + + +#: Copy of ``edge0.cli.TIER_ENV`` so tier resolution needs no backend +#: import (``edge0.cli`` loads the model registry, which loads the +#: backend). ``tests/test_bench_backend.py::test_tier_env_matches_the_cli`` +#: keeps the two in sync. +_TIER_ENV = { + "edge0-35b": "EDGE0_35B_MODEL", + "edge0-8b": "EDGE0_8B_MODEL", +} + + +def _resolve_model(model: str) -> tuple[str, str | None]: + """tier name -> checkpoint dir via $EDGE0__MODEL (cli parity).""" if model.startswith("edge0-"): - from edge0.cli import TIER_ENV - env = os.environ.get(TIER_ENV.get(model, "")) + env_name = _TIER_ENV.get(model) + if env_name is None: + raise SystemExit( + f"unknown tier {model!r}; known tiers: " + f"{', '.join(sorted(_TIER_ENV))} (or pass the checkpoint " + f"directory itself)") + env = os.environ.get(env_name) if not env: raise SystemExit( - f"set {TIER_ENV.get(model, 'EDGE0__MODEL')} to the " - f"checkpoint directory for {model}, or pass the dir itself") - model = env - engine = AutoEngine.from_pretrained(model, name=name) - print(f"[bench] tier={engine.name} ntok={args.ntok} warmup={args.warmup}", - flush=True) - run_bench(engine, args.ntok, args.warmup) - engine.close() - return 0 + f"set {env_name} to the checkpoint directory for {model}, " + f"or pass the dir itself") + return env, None + return model, None + + +def _prepare_output(path: str) -> str: + """Refuse an existing destination and make its directory exist, both + before the model loads.""" + dest = Path(path).expanduser() + if dest.exists() or dest.is_symlink(): + raise FileExistsError( + f"--json-output {dest} already exists; choose a new path so a " + f"stale report is never mistaken for this run") + dest.parent.mkdir(parents=True, exist_ok=True) + return str(dest.resolve()) + + +def _load_probe(path: str) -> dict: + """Read and validate an explicitly requested probe before inference.""" + return report.probe_link(path=path, text=Path(path).read_bytes()) + + +def _parse_args(ap: argparse.ArgumentParser, argv: list[str]): + args = ap.parse_args(argv) + if args.ntok is None: + raw = os.environ.get("BENCH_NTOK", "200") + try: + args.ntok = int(raw) + except ValueError: + ap.error(f"BENCH_NTOK must be an integer, got {raw!r}") + if args.ntok <= 0: + ap.error("--ntok (or BENCH_NTOK) must be > 0: an empty timed window " + "is not a benchmark") + if args.warmup < 0: + ap.error("--warmup must be >= 0") + if not (0 <= args.rss_sample_interval < math.inf): + ap.error("--rss-sample-interval must be a finite number >= 0") + if args.json_output is not None and not args.json_output.strip(): + ap.error("--json-output needs a path") + if args.probe_json is not None and not args.probe_json.strip(): + ap.error("--probe-json needs a path") + seed = os.environ.get("BENCH_SEED") + if seed is not None: + try: + int(seed) + except ValueError: + ap.error(f"BENCH_SEED must be an integer, got {seed!r}") + temp = os.environ.get("BENCH_TEMP") + if temp is not None: + try: + value = float(temp) + except ValueError: + ap.error(f"BENCH_TEMP must be a number, got {temp!r}") + else: + if not math.isfinite(value): + ap.error(f"BENCH_TEMP must be finite, got {temp!r}") + return args + + +def main(argv: list[str] | None = None) -> int: + ap = argparse.ArgumentParser( + description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("model", help="checkpoint dir or tier name") + ap.add_argument("--ntok", type=int, default=None, + help="timed decode tokens per run (> 0; default " + "$BENCH_NTOK or 200)") + ap.add_argument("--warmup", type=int, default=10, + help="untimed greedy warmup steps per run (>= 0)") + ap.add_argument("--json-output", metavar="PATH", + help="write the machine-readable report to this NEW file") + ap.add_argument("--probe-json", metavar="PATH", + help="link an existing scripts/jetson_probe.py report") + ap.add_argument("--rss-sample-interval", type=float, default=0.25, + metavar="SECONDS", + help="process memory sampling interval (0 disables)") + argv = list(sys.argv[1:] if argv is None else argv) + args = _parse_args(ap, argv) + + started_utc = datetime.now(timezone.utc).isoformat() + model_dir, name = _resolve_model(args.model) + # Everything that can be rejected is rejected before the model loads. + try: + json_output = (_prepare_output(args.json_output) + if args.json_output is not None else None) + probe = (_load_probe(args.probe_json) + if args.probe_json is not None else None) + except (OSError, ValueError) as exc: + ap.exit(2, f"[bench] error: {exc}\n") + env_knobs = _env_snapshot() # before the engine can touch the environment + + sampler = measure.RssSampler(args.rss_sample_interval) + engine = None + stage = "engine build" + try: + sampler.start() # before the model loads: load transients are sampled + engine = _build_engine(model_dir, name) + model_env = _model_env() # resolved after the build, before the runs + print(f"[bench] tier={engine.name} ntok={args.ntok} " + f"warmup={args.warmup}", flush=True) + stage = "benchmark" + bench = run_bench(engine, args.ntok, args.warmup) + if json_output is not None: + stage = "report" + sampler.stop() + backend = _framework()[0] + doc = _build_report(engine, bench, args, argv, started_utc, + model_dir, probe, sampler, env_knobs, + model_env, backend) + stage = "write" + report.write_report(doc, json_output) + stage = "output after a successful write" + print(f"[bench] wrote {json_output}", flush=True) + return 0 + except Exception as exc: # noqa: BLE001 - reported with its stage + traceback.print_exc() + print(f"[bench] FAILED during {stage}: {type(exc).__name__}: {exc}", + file=sys.stderr, flush=True) + return 1 + finally: + try: + sampler.stop() + finally: + if engine is not None: + engine.close() if __name__ == "__main__": diff --git a/examples/benchmark_measure.py b/examples/benchmark_measure.py new file mode 100644 index 0000000..463cb64 --- /dev/null +++ b/examples/benchmark_measure.py @@ -0,0 +1,842 @@ +"""Benchmark-owned measurement: backend peaks, process memory, run identity. + +``examples/bench.py`` is the only consumer. Everything that touches a +backend (``torch.cuda``, ``mlx.core``), the process (``psutil``, +``resource``, ``/proc``) or the machine (``git``, ``nvpmodel``) lives +here, behind lazy imports and injectable hooks so the dispatch logic is +unit-tested on a host with neither CUDA nor MLX. Nothing here runs +inside a timed interval except the explicit ``synchronize`` / +``reset_peak`` / ``peak_metrics`` calls the benchmark makes at phase +boundaries. + +The pure report builder (``benchmark_report``) only sees the dicts these +functions return. +""" + +from __future__ import annotations + +import hashlib +import importlib +import json +import os +import platform +import re +import struct +import subprocess +import threading +from pathlib import Path +from typing import Any + +try: # imported as ``examples.benchmark_measure`` (tests) + from examples import benchmark_report as _report +except ImportError: # run next to bench.py as a script + import benchmark_report as _report # type: ignore[no-redef] + +memory_metric = _report.memory_metric + +_AUTO = object() + +#: What a per-run backend peak covers. The reset happens immediately +#: before prefill, the read after the synchronized final decode step. +RUN_PEAK_SCOPE = ( + "one benchmark run: statistics reset immediately before prefill and " + "read after the synchronized final decode step; covers prefill, " + "warmup and decode and includes allocations already resident at the " + "reset; an earlier initialization transient is not recovered") + +_CUDA_ALLOCATED_SCOPE = ( + RUN_PEAK_SCOPE + "; torch caching allocator only, not driver, context " + "or non-torch memory") + +#: ``max_memory_reserved`` is not reset to zero by +#: ``reset_peak_memory_stats``: its floor is every segment the allocator +#: already holds, so it is close to a lifetime figure. +CUDA_RESERVED_SCOPE = ( + "one benchmark run: statistics reset immediately before prefill and " + "read after the synchronized final decode step; peak of the segments " + "held by the torch caching allocator, whose floor is every segment " + "already cached at the reset (model load, earlier runs, any earlier " + "initialization transient) because the allocator does not release " + "segments between runs; not driver, context or non-torch memory") + + +def _bytes(value): + """Integer bytes from a backend counter; floats that are whole numbers + are accepted, anything else is left for ``memory_metric`` to reject.""" + if isinstance(value, float) and value.is_integer(): + return int(value) + return value + + +# ---- backend measurement adapters -------------------------------------------- + + +class NullMeasurement: + """No backend peak available (torch on cpu/mps, unknown backend).""" + + kind = "null" + headline = "backend_peak" + + def __init__(self, backend: str, device: str, reason: str): + self.backend = backend + self.device = device + self.reason = reason + + def synchronize(self) -> None: + return None + + def reset_peak(self) -> None: + return None + + def peak_metrics(self) -> dict: + return {self.headline: memory_metric( + None, method="none", scope=RUN_PEAK_SCOPE, + unavailable_reason=self.reason)} + + def describe(self) -> dict: + sync = "none: the benchmark synchronizes only an actual CUDA device" + if str(self.device).startswith("mps"): + sync += ("; mps execution is asynchronous, so prefill/decode wall " + "times may exclude the final step's in-flight device work") + return {"kind": self.kind, "backend": self.backend, + "resolved_device": self.device, + "synchronization": sync, + "peak_source": f"none: {self.reason}", + "headline_peak": self.headline} + + +class MlxMeasurement: + """The pre-existing MLX behavior: ``reset_peak_memory`` before + prefill, ``get_peak_memory`` after decode, no added synchronization + (the engine's ``core.eval(logits)`` already materializes every step, + so the wall time is unchanged from the original benchmark).""" + + kind = "mlx" + headline = "mlx_peak_active" + + def __init__(self, core): + self._core = core + + def synchronize(self) -> None: + return None + + def reset_peak(self) -> None: + self._core.reset_peak_memory() + + def peak_metrics(self) -> dict: + return {self.headline: memory_metric( + _bytes(self._core.get_peak_memory()), + method="mlx.core.get_peak_memory", scope=RUN_PEAK_SCOPE)} + + def describe(self) -> dict: + default = getattr(self._core, "default_device", None) + try: + device = str(default()) if callable(default) else "mlx default" + except Exception as exc: # noqa: BLE001 - never fail a report here + device = f"mlx default (query failed: {type(exc).__name__})" + return {"kind": self.kind, "backend": "mlx", + "resolved_device": device, + "synchronization": ( + "none added: MLX evaluation is forced by the engine " + "(core.eval) at every step; timing behavior unchanged"), + "peak_source": "mlx.core.get_peak_memory (reset before prefill)", + "headline_peak": self.headline} + + +class CudaMeasurement: + """Torch on an actual CUDA device: synchronize the resolved device at + phase boundaries and read the caching-allocator peaks. The headline + (human output, legacy ``peak_gib``) is ``max_memory_allocated``, the + analogue of MLX's peak active memory; ``max_memory_reserved`` is + JSON-only.""" + + kind = "cuda" + headline = "cuda_peak_allocated" + + def __init__(self, device, cuda_api): + self.device = device + self._cuda = cuda_api + self._at_reset: dict[str, Any] = {} + + def synchronize(self) -> None: + self._cuda.synchronize(self.device) + + def reset_peak(self) -> None: + self._cuda.reset_peak_memory_stats(self.device) + # What is resident right now is the floor of both peaks; record it + # (attribute reads, before the prefill clock starts). + self._at_reset = {} + for key, name in (("allocated_at_reset", "memory_allocated"), + ("reserved_at_reset", "memory_reserved")): + reader = getattr(self._cuda, name, None) + if callable(reader): + self._at_reset[key] = _bytes(reader(self.device)) + + def peak_metrics(self) -> dict: + allocated = _bytes(self._cuda.max_memory_allocated(self.device)) + reserved = _bytes(self._cuda.max_memory_reserved(self.device)) + alloc_details = {k: v for k, v in self._at_reset.items() + if k == "allocated_at_reset"} + res_details = {k: v for k, v in self._at_reset.items() + if k == "reserved_at_reset"} + return { + "cuda_peak_allocated": memory_metric( + allocated, method="torch.cuda.max_memory_allocated", + scope=_CUDA_ALLOCATED_SCOPE, **alloc_details), + "cuda_peak_reserved": memory_metric( + reserved, method="torch.cuda.max_memory_reserved", + scope=CUDA_RESERVED_SCOPE, **res_details), + } + + def describe(self) -> dict: + return {"kind": self.kind, "backend": "cuda", + "resolved_device": str(self.device), + "synchronization": ( + "torch.cuda.synchronize(device) after engine reset, " + "at the end of prefill, at the end of warmup and after " + "the final timed decode step; no per-token syncs"), + "peak_source": ( + "torch.cuda caching allocator statistics, " + "reset_peak_memory_stats before prefill"), + "headline_peak": self.headline} + + +def _torch_cuda(): + import torch + return torch.cuda + + +def make_measurement(backend_name: str, core_module, *, cuda_api=None): + """Pick the adapter for the ACTIVE backend and its RESOLVED device. + + ``EDGE0_BACKEND=cuda`` alone does not prove the device is CUDA: the + torch backend runs on cpu/mps when ``EDGE0_TORCH_DEVICE`` says so or + when no CUDA device exists, and then allocator peaks are simply + unavailable. ``resolved_device`` and the unavailable reason are both + derived from the same ``core_module.DEVICE`` object. + """ + if backend_name == "mlx": + if (callable(getattr(core_module, "reset_peak_memory", None)) + and callable(getattr(core_module, "get_peak_memory", None))): + return MlxMeasurement(core_module) + return NullMeasurement( + "mlx", "mlx default", + "mlx core namespace has no reset_peak_memory/get_peak_memory; " + "peak memory unavailable") + if backend_name == "cuda": + device = getattr(core_module, "DEVICE", None) + if device is None: + return NullMeasurement( + "cuda", "unknown", + "torch core namespace has no DEVICE; resolved device unknown") + dev_type = getattr(device, "type", str(device)) + if dev_type == "cuda": + return CudaMeasurement(device, cuda_api or _torch_cuda()) + return NullMeasurement( + "cuda", str(device), + f"torch device is {dev_type}, not cuda: torch allocator peaks " + f"are unavailable") + return NullMeasurement( + backend_name, "unknown", + f"backend {backend_name!r} exposes no peak-memory interface") + + +def execution_evidence(logits) -> dict: + """What the engine actually produced: the array class and, when the + class carries one, its device (an attribute read, no sync).""" + cls = type(logits) + array_type = f"{cls.__module__}.{cls.__name__}" + evidence: dict[str, Any] = {"array_type": array_type, "device": None, + "unavailable_reasons": {}} + device = getattr(logits, "device", None) + if device is None: + evidence["unavailable_reasons"]["device"] = ( + f"{array_type} carries no device attribute (MLX arrays are " + f"device-agnostic); see runtime.measurement.resolved_device") + else: + evidence["device"] = str(device) + return evidence + + +# ---- process memory -------------------------------------------------------------- + + +_RSS_COMPOSITION = ( + "resident set size = anonymous memory plus resident file-backed pages " + "(the mmap'd checkpoint: prod_k8 whole-layer prefill touches every " + "expert of model.safetensors each run, and those pages are reclaimable " + "page cache), so it is not the process's anonymous footprint and not " + "comparable to the 'peak anonymous memory' figure in docs/nvidia.md; " + "must not be added to device allocator peaks") + +_RSS_LIFETIME_SCOPE = ( + "process lifetime high-water mark of " + _RSS_COMPOSITION + + "; includes model load and every run; cannot be reset per run") + + +def _rusage_self(): + try: + import resource + return getattr(resource, "RUSAGE_SELF", 0) + except ImportError: + return 0 + + +def process_peak_rss(*, system=None, getrusage: Any = _AUTO, + psutil_module: Any = _AUTO) -> dict: + """Process-lifetime peak RSS with a platform-specific method and unit + normalization to bytes.""" + system = system or platform.system() + if getrusage is _AUTO: + try: + import resource + getrusage = getattr(resource, "getrusage", None) + except ImportError: + getrusage = None + if psutil_module is _AUTO: + try: + import psutil as psutil_module + except ImportError: + psutil_module = None + + if system == "Linux": + if getrusage is None: + return memory_metric(None, method="none", + scope=_RSS_LIFETIME_SCOPE, + unavailable_reason="resource module unavailable") + return memory_metric( + int(getrusage(_rusage_self()).ru_maxrss) * 1024, + method="resource.getrusage(RUSAGE_SELF).ru_maxrss (KiB x 1024)", + scope=_RSS_LIFETIME_SCOPE) + if system == "Darwin": + if getrusage is None: + return memory_metric(None, method="none", + scope=_RSS_LIFETIME_SCOPE, + unavailable_reason="resource module unavailable") + return memory_metric( + int(getrusage(_rusage_self()).ru_maxrss), + method="resource.getrusage(RUSAGE_SELF).ru_maxrss (bytes)", + scope=_RSS_LIFETIME_SCOPE) + if system == "Windows": + if psutil_module is None: + return memory_metric(None, method="none", + scope=_RSS_LIFETIME_SCOPE, + unavailable_reason="psutil unavailable") + return memory_metric( + int(psutil_module.Process().memory_info().peak_wset), + method="psutil.Process().memory_info().peak_wset", + scope=_RSS_LIFETIME_SCOPE) + return memory_metric( + None, method="none", scope=_RSS_LIFETIME_SCOPE, + unavailable_reason=f"no process peak RSS method for platform {system!r}") + + +_PROC_STATUS_FIELDS = (("VmRSS", "rss"), ("RssAnon", "rss_anon"), + ("RssFile", "rss_file"), ("VmSwap", "swap")) + + +def parse_proc_status(text: str) -> dict: + """``/proc/self/status`` -> bytes for rss / rss_anon / rss_file / swap + (``None`` where the kernel does not report a field).""" + values: dict[str, int | None] = {key: None for _, key in _PROC_STATUS_FIELDS} + wanted = dict(_PROC_STATUS_FIELDS) + for line in text.splitlines(): + name, sep, rest = line.partition(":") + if not sep or name not in wanted: + continue + parts = rest.split() + if not parts: + continue + try: + number = int(parts[0]) + except ValueError: + continue + unit = parts[1].lower() if len(parts) > 1 else "bytes" + values[wanted[name]] = number * (1024 if unit == "kb" else 1) + return values + + +def read_process_memory() -> dict | None: + """Current process memory composition: ``/proc/self/status`` on Linux + (rss, anonymous, file-backed, swap), else psutil rss only, else + ``None``. The ``method`` key names the source used.""" + try: + with open("/proc/self/status", encoding="utf-8") as fh: + parsed = parse_proc_status(fh.read()) + if parsed["rss"] is not None: + parsed["method"] = "/proc/self/status (VmRSS, RssAnon, RssFile, VmSwap)" + return parsed + except OSError: + pass + try: + import psutil + rss = int(psutil.Process().memory_info().rss) + except Exception: # noqa: BLE001 - optional dependency + return None + return {"rss": rss, "rss_anon": None, "rss_file": None, "swap": None, + "method": "psutil.Process().memory_info().rss"} + + +_SAMPLED_METRICS = ( + ("rss", "process_sampled_peak_rss", + "sampled peak of process " + _RSS_COMPOSITION), + ("rss_anon", "process_sampled_peak_rss_anon", + ("sampled peak of the process's anonymous (non-file-backed) resident " + "pages; the closest match to the 'peak anonymous memory' figure in " + "docs/nvidia.md")), + ("rss_file", "process_sampled_peak_rss_file", + ("sampled peak of the process's resident file-backed pages (mmap'd " + "checkpoint page cache, reclaimable)")), + ("swap", "process_sampled_peak_swap", + ("sampled peak of swapped-out process memory; a non-zero value marks " + "a swapping run, which cannot establish a resident-memory target")), +) + + +class RssSampler: + """Lightweight periodic process-memory sampler (one daemon thread). + + Start it BEFORE the model is built so load transients are covered. + The result is a *sampled* peak: spikes shorter than the interval are + missed, and the report says so. On Linux the anonymous, file-backed + and swap components are tracked too. + """ + + def __init__(self, interval_seconds: float, *, read_rss=None, + read_memory=None, join_timeout: float | None = None): + self.interval_seconds = float(interval_seconds) + self.join_timeout = (float(join_timeout) if join_timeout is not None + else max(5.0, 4 * self.interval_seconds)) + if read_rss is not None: + self._read = lambda: {"rss": int(read_rss())} + self._method = "injected read_rss" + elif read_memory is not None: + self._read = read_memory + self._method = "injected read_memory" + else: + self._read = None + self._method = "none" + self._peaks: dict[str, int | None] = { + key: None for key, _, _ in _SAMPLED_METRICS} + self._count = 0 + self._error: str | None = None + self._stop = threading.Event() + self._thread: threading.Thread | None = None + self._started = False + + @property + def is_running(self) -> bool: + return self._thread is not None and self._thread.is_alive() + + @property + def peak(self) -> int | None: + return self._peaks["rss"] + + @property + def sample_count(self) -> int: + return self._count + + def _sample(self) -> None: + read = self._read + if read is None: + return + try: + mem = read() + except Exception as exc: # noqa: BLE001 - sampling must never raise + self._error = f"{type(exc).__name__}: {exc}" + return + if not mem: + self._error = "memory reader returned nothing" + return + method = mem.get("method") + if method: + self._method = str(method) + self._count += 1 + for key in self._peaks: + value = mem.get(key) + if value is None: + continue + value = int(value) + current = self._peaks[key] + self._peaks[key] = value if current is None else max(current, value) + + def _loop(self) -> None: + while not self._stop.wait(self.interval_seconds): + self._sample() + + def start(self) -> None: + if self.interval_seconds <= 0 or self._started: + return + if self._read is None: + probe = read_process_memory() + if probe is None: + self._error = "no process memory reader (no /proc, no psutil)" + return + self._read = read_process_memory + self._started = True + self._sample() + self._thread = threading.Thread( + target=self._loop, name="bench-rss-sampler", daemon=True) + self._thread.start() + + def stop(self) -> None: + """Stop sampling (idempotent). The thread reference is only + dropped once the thread has actually exited, so ``is_running`` + stays true (and the metric carries a reason) if a reader hangs + past ``join_timeout``.""" + self._stop.set() + thread = self._thread + if thread is None: + return + thread.join(timeout=self.join_timeout) + if thread.is_alive(): + self._error = (f"sampler thread did not stop within " + f"{self.join_timeout:g}s (reader blocked?)") + return + self._thread = None + self._sample() + + def metrics(self) -> dict: + """All sampled peaks as memory metric objects (rss always; the + composition only where the reader supplied it).""" + out = {} + for key, name, scope in _SAMPLED_METRICS: + details = {"interval_seconds": self.interval_seconds, + "sample_count": self._count} + if self.interval_seconds <= 0: + out[name] = memory_metric( + None, method="none", scope=scope, + unavailable_reason="sampler disabled (interval 0)", + interval_seconds=self.interval_seconds) + continue + value = self._peaks[key] + if value is not None: + if self._error and self.is_running: + details["warning"] = self._error + out[name] = memory_metric(value, method=self._method, + scope=scope, **details) + elif key == "rss": + out[name] = memory_metric( + None, method=self._method, scope=scope, + unavailable_reason=self._error or "sampler never ran", + **details) + else: + out[name] = memory_metric( + None, method=self._method, scope=scope, + unavailable_reason=( + f"{key} needs /proc/self/status (Linux); the " + f"memory reader did not supply it"), + **details) + return out + + def metric(self) -> dict: + return self.metrics()["process_sampled_peak_rss"] + + +# ---- run identity ------------------------------------------------------------------ + + +def git_identity(repo_dir, *, run=subprocess.run) -> dict: + """HEAD commit and tracked-file dirty state of ``repo_dir``.""" + out: dict[str, Any] = {"commit": None, "dirty": None, + "unavailable_reasons": {}} + + def _git(*args): + return run(["git", "-C", str(repo_dir), *args], capture_output=True, + text=True, timeout=15) + + try: + head = _git("rev-parse", "HEAD") + except (OSError, subprocess.SubprocessError) as exc: + reason = f"git unavailable: {type(exc).__name__}: {exc}" + out["unavailable_reasons"] = {"commit": reason, "dirty": reason} + return out + if head.returncode != 0 or not head.stdout.strip(): + reason = (f"git rev-parse failed (exit {head.returncode}): " + f"{(head.stderr or '').strip() or 'no output'}") + out["unavailable_reasons"] = {"commit": reason, "dirty": reason} + return out + out["commit"] = head.stdout.strip() + try: + status = _git("status", "--porcelain", "--untracked-files=no") + except (OSError, subprocess.SubprocessError) as exc: + out["unavailable_reasons"]["dirty"] = ( + f"git status failed: {type(exc).__name__}: {exc}") + return out + if status.returncode != 0: + out["unavailable_reasons"]["dirty"] = ( + f"git status failed (exit {status.returncode})") + return out + out["dirty"] = bool(status.stdout.strip()) + return out + + +_NVPMODEL_STATUS = Path("/var/lib/nvpmodel/status") + + +def power_mode(*, run=subprocess.run, status_path=_NVPMODEL_STATUS) -> dict: + """Jetson power mode via ``nvpmodel -q``, else the nvpmodel status + file, else an explicit unavailable reason.""" + out: dict[str, Any] = {"value": None, "mode_id": None, "method": None, + "unavailable_reasons": {}} + reasons = [] + try: + proc = run(["nvpmodel", "-q"], capture_output=True, text=True, + timeout=10) + if proc.returncode == 0: + name = None + mode_id = None + for line in (proc.stdout or "").splitlines(): + text = line.strip() + if text.lower().startswith("nv power mode:"): + name = text.split(":", 1)[1].strip() + elif name is not None and mode_id is None and text.isdigit(): + mode_id = int(text) + if name: + out.update(value=name, mode_id=mode_id, method="nvpmodel -q") + if mode_id is None: + out["unavailable_reasons"]["mode_id"] = ( + "nvpmodel -q printed no numeric mode id") + return out + reasons.append("nvpmodel -q printed no 'NV Power Mode' line") + else: + reasons.append( + f"nvpmodel -q exited {proc.returncode}: " + f"{(proc.stderr or '').strip() or 'no output'}") + except (OSError, subprocess.SubprocessError) as exc: + reasons.append(f"nvpmodel not runnable: {type(exc).__name__}: {exc}") + try: + text = Path(status_path).read_text(errors="replace") + match = re.search(r"pmode:(\d+)", text) + if match: + mode_id = int(match.group(1)) + out.update(value=f"pmode {mode_id}", mode_id=mode_id, + method=f"{status_path} pmode field") + return out + reasons.append(f"{status_path} has no pmode field") + except OSError as exc: + reasons.append(f"{status_path}: {type(exc).__name__}") + reason = ("; ".join(reasons) + + " (not a Jetson, or nvpmodel is not on PATH)") + out["unavailable_reasons"] = {"value": reason, "mode_id": reason, + "method": reason} + return out + + +# ---- checkpoint identity ----------------------------------------------------------- + + +def _sha256_file(path: Path) -> str: + digest = hashlib.sha256() + with open(path, "rb") as fh: + for chunk in iter(lambda: fh.read(1 << 20), b""): + digest.update(chunk) + return digest.hexdigest() + + +def safetensors_header(path) -> dict: + """``__metadata__``, tensor count and a hash of the JSON header of a + safetensors file: an 8-byte length prefix plus the header, no weight + bytes touched.""" + out: dict[str, Any] = {"metadata": {}, "tensor_count": None, + "header_sha256": None, "unavailable_reasons": {}} + try: + with open(path, "rb") as fh: + (length,) = struct.unpack(" (256 << 20): + raise ValueError(f"header length {length} is implausible") + raw = fh.read(length) + header = json.loads(raw.decode("utf-8")) + if not isinstance(header, dict): + raise ValueError("header is not a JSON object") + metadata = header.get("__metadata__") + out["metadata"] = dict(metadata) if isinstance(metadata, dict) else {} + out["tensor_count"] = len([k for k in header if k != "__metadata__"]) + out["header_sha256"] = hashlib.sha256(raw).hexdigest() + except (OSError, ValueError, struct.error) as exc: + reason = f"safetensors header unreadable: {type(exc).__name__}: {exc}" + out["unavailable_reasons"] = {"tensor_count": reason, + "header_sha256": reason} + return out + + +MANIFEST_ALGORITHM = ("sha256 over 'path\\tsize\\n' lines, one per regular " + "file, sorted by relative POSIX path") + + +def checkpoint_manifest(model_dir, *, hash_limit_bytes: int = 32 << 20) -> dict: + """Identify a checkpoint directory without hashing weights. + + Every regular file contributes ``(relative path, size)`` to the + manifest digest and the retained ``files`` list. Files up to + ``hash_limit_bytes`` that are not safetensors are content-hashed + (config, chat template, tokenizer); safetensors files contribute + their header metadata, tensor count and header hash. + """ + out: dict[str, Any] = { + "manifest": {"algorithm": MANIFEST_ALGORITHM, "file_count": None, + "total_bytes": None, "sha256": None, "files": [], + "unavailable_reasons": {}}, + "small_file_sha256": {}, + "safetensors_headers": {}, + "config": {}, + "config_sha256": None, + } + root = Path(model_dir) + if not root.is_dir(): + reason = f"checkpoint directory {str(root)!r} does not exist" + out["manifest"]["unavailable_reasons"] = { + "file_count": reason, "total_bytes": reason, "sha256": reason} + return out + entries = [] + total = 0 + for path in sorted(p for p in root.rglob("*") if p.is_file()): + rel = path.relative_to(root).as_posix() + size = path.stat().st_size + entries.append((rel, size)) + total += size + if path.suffix == ".safetensors": + out["safetensors_headers"][rel] = safetensors_header(path) + elif size <= hash_limit_bytes: + try: + out["small_file_sha256"][rel] = _sha256_file(path) + except OSError: + pass + entries.sort() + digest = hashlib.sha256() + for rel, size in entries: + digest.update(f"{rel}\t{size}\n".encode()) + out["manifest"].update( + file_count=len(entries), total_bytes=total, + sha256=digest.hexdigest(), + files=[{"path": rel, "size_bytes": size} for rel, size in entries]) + config_path = root / "config.json" + if config_path.is_file(): + try: + with open(config_path, encoding="utf-8") as fh: + config = json.load(fh) + if isinstance(config, dict): + out["config"] = config + except (OSError, ValueError): + pass + out["config_sha256"] = out["small_file_sha256"].get("config.json") + return out + + +def adapter_identity(path, checkpoint_dir, *, + sha256_max_bytes: int = 256 << 20) -> dict: + """Identity of the adapter file the engine actually applied + (``engine.cfg.lora`` / ``engine.cfg.prerouter.weights_file``), which + may live outside the checkpoint directory (``artifacts/`` fallback). + Adapter files are small enough to hash; the cap keeps a misconfigured + path from hashing a weight shard.""" + out: dict[str, Any] = {"file": None, "path": None, + "inside_checkpoint_dir": None, "size_bytes": None, + "sha256": None, "safetensors_header": None, + "unavailable_reasons": {}} + reasons = out["unavailable_reasons"] + if not path: + reason = "disabled: no adapter file configured for this tier" + for key in ("file", "path", "inside_checkpoint_dir", "size_bytes", + "sha256", "safetensors_header"): + reasons[key] = reason + return out + p = Path(os.fspath(path)) + resolved = p.resolve() + out["path"] = str(resolved) + root = Path(checkpoint_dir).resolve() + try: + rel = resolved.relative_to(root).as_posix() + inside = True + except ValueError: + rel = None + inside = False + out["inside_checkpoint_dir"] = inside + out["file"] = rel if inside else str(resolved) + if not p.is_file(): + reason = f"adapter file not found: {resolved}" + for key in ("size_bytes", "sha256", "safetensors_header"): + reasons[key] = reason + return out + size = p.stat().st_size + out["size_bytes"] = size + if size <= sha256_max_bytes: + out["sha256"] = _sha256_file(p) + else: + reasons["sha256"] = (f"not hashed: {size} bytes is larger than the " + f"{sha256_max_bytes}-byte cap") + if p.suffix == ".safetensors": + out["safetensors_header"] = safetensors_header(p) + else: + reasons["safetensors_header"] = "not a safetensors file" + return out + + +# ---- runtime ----------------------------------------------------------------------- + + +def runtime_info(backend_name: str, *, importer=importlib.import_module) -> dict: + """Interpreter, platform and framework versions for the ACTIVE backend + only (the other framework is not imported).""" + info: dict[str, Any] = { + "python_version": platform.python_version(), + "platform": {"system": platform.system(), + "release": platform.release(), + "machine": platform.machine()}, + "torch_version": None, + "torch_cuda_version": None, + "mlx_version": None, + "unavailable_reasons": {}, + } + reasons = info["unavailable_reasons"] + if backend_name == "cuda": + reasons["mlx_version"] = "backend is cuda; mlx not imported" + try: + torch = importer("torch") + except Exception as exc: # noqa: BLE001 - reported, not raised + reason = f"torch import failed: {exc}" + reasons["torch_version"] = reason + reasons["torch_cuda_version"] = reason + return info + info["torch_version"] = str(getattr(torch, "__version__", "unknown")) + cuda = getattr(getattr(torch, "version", None), "cuda", None) + if cuda: + info["torch_cuda_version"] = str(cuda) + else: + reasons["torch_cuda_version"] = ( + "torch build has no CUDA runtime (CPU-only build)") + elif backend_name == "mlx": + reasons["torch_version"] = "backend is mlx; torch not imported" + reasons["torch_cuda_version"] = "backend is mlx; torch not imported" + try: + mlx = importer("mlx") + info["mlx_version"] = str(getattr(mlx, "__version__", "unknown")) + except Exception as exc: # noqa: BLE001 - reported, not raised + reasons["mlx_version"] = f"mlx import failed: {exc}" + else: + reason = f"backend {backend_name!r} has no known runtime" + for key in ("torch_version", "torch_cuda_version", "mlx_version"): + reasons[key] = reason + return info + + +__all__ = [ + "CUDA_RESERVED_SCOPE", + "MANIFEST_ALGORITHM", + "RUN_PEAK_SCOPE", + "CudaMeasurement", + "MlxMeasurement", + "NullMeasurement", + "RssSampler", + "adapter_identity", + "checkpoint_manifest", + "execution_evidence", + "git_identity", + "make_measurement", + "parse_proc_status", + "power_mode", + "process_peak_rss", + "read_process_memory", + "runtime_info", + "safetensors_header", +] diff --git a/examples/benchmark_report.py b/examples/benchmark_report.py new file mode 100644 index 0000000..1bf58bd --- /dev/null +++ b/examples/benchmark_report.py @@ -0,0 +1,546 @@ +"""Pure benchmark report builder: measured numbers -> JSON contract. + +This module is deliberately standard-library only. It never imports +``edge0``, ``torch`` or ``mlx`` and performs no measurement of its own: +``examples/bench.py`` measures, ``examples/benchmark_measure.py`` collects +backend/process metrics, and the functions here validate and assemble the +result into the versioned schema documented in ``docs/nvidia.md``. + +Conventions of the schema: + +* ``schema_version`` is the benchmark schema, independent of the probe + schema in ``scripts/jetson_probe.py``. +* A required measurement (a token count, a wall-clock duration, a byte + count) that is invalid -- negative, non-finite, wrong type -- raises + ``ValueError``; it is never written as a plausible-looking number. +* Unavailable *metadata* is ``null`` plus a reason: every group object + carries an ``unavailable_reasons`` map (``{field: reason}``) and every + ``null`` field in the group has an entry there (and vice versa). + Single-valued memory metric objects carry their own + ``unavailable_reason`` instead. +* Serialization is strict JSON (``allow_nan=False``). +""" + +from __future__ import annotations + +import hashlib +import json +import math +import os +import statistics +import tempfile +from typing import Any + +#: Benchmark report schema. Independent of ``scripts/jetson_probe.py``'s +#: ``schema_version``; bump when a field's meaning or presence changes. +SCHEMA_VERSION = 1 + +#: What ``examples/bench.py`` measures. Stated in the report so a number +#: is never mistaken for time-to-first-token or an EOS-aware completion. +PROTOCOL = { + "name": "fixed-length-sample-and-step", + "version": 1, + "description": ( + "Two sequential runs in one process. Each run: engine reset, " + "timed chat-templated prefill, untimed warmup steps, then a " + "timed fixed-length decode window. Wall time is host monotonic " + "time and includes storage, transfer and CPU sampling costs; on " + "CUDA the device is synchronized at phase boundaries only."), + "warmup": ( + "greedy argmax steps, untimed; excluded from every throughput " + "figure (expert cold-start and first-step overhead land here)"), + "timed_phase": ( + "each timed iteration samples one token from the current logits " + "and runs one engine step (a forward pass); the final step's " + "logits are computed but never sampled; the loop is fixed-length " + "and does not stop at EOS"), + "not_measured": ["TTFT", "EOS-aware completion latency"], +} + +#: Required direct fields per group (``unavailable_reasons`` is added or +#: checked separately). +REQUIRED_FIELDS = { + "identity": ("timestamp_utc", "git", "run_count", "command", + "requested"), + "model": ("checkpoint_path", "tier", "model_type", "architectures", + "config_sha256", "manifest", "small_file_sha256", "adapters", + "tokenizer_files"), + "runtime": ("backend", "backend_version", "resolved_device", + "measurement", "execution_evidence", "python_version", + "torch_version", "torch_cuda_version", "mlx_version", + "platform", "power_mode"), + "workload": ("prompt_source", "prompt_sha256", "prompt_token_ids_sha256", + "prompt_chars", "prompt_tokens", "requested_warmup_tokens", + "requested_timed_tokens", "seed", "temperature", "top_k", + "top_p", "repetition_penalty", "prefill_chunk", "think"), + "caches": ("layer_options", "shared_cache_slots_resolved", + "prefetch_cap_resolved", "weight_cache", "prewarm", + "threads", "mlx_cache_limit_bytes"), + "memory": ("process_peak_rss", "process_sampled_peak_rss", "notes"), +} + +_RUN_FIELDS = ("index", "prompt_tokens", "requested_warmup_tokens", + "warmup_tokens", "requested_timed_tokens", + "timed_decode_tokens", "total_generated_tokens", + "decode_start_context_tokens", "prefill_seconds", + "prefill_tokens_per_second", "decode_seconds", + "decode_tokens_per_second", "memory", "execution_evidence", + "unavailable_reasons") + +_SUMMARY_METRICS = ("decode_tokens_per_second", "prefill_tokens_per_second", + "prefill_seconds", "decode_seconds") + + +def _is_int(value) -> bool: + return isinstance(value, int) and not isinstance(value, bool) + + +def _require_count(name: str, value) -> int: + if not _is_int(value): + raise ValueError(f"{name} must be an int, got {value!r}") + if value < 0: + raise ValueError(f"{name} must be >= 0, got {value}") + return value + + +def _require_seconds(name: str, value) -> float: + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise ValueError(f"{name} must be a number, got {value!r}") + value = float(value) + if not math.isfinite(value): + raise ValueError(f"{name} must be finite, got {value!r}") + if value < 0: + raise ValueError(f"{name} must be >= 0, got {value!r}") + return value + + +def memory_metric(bytes_value, *, method: str, scope: str, + unavailable_reason: str | None = None, **details) -> dict: + """One memory measurement: integer bytes, how it was measured, and + what interval/process it covers. + + ``bytes_value`` is ``None`` only together with ``unavailable_reason`` + (and a measured value never carries a reason). Extra keyword details + (``interval_seconds``, ``sample_count``, ...) are kept verbatim. + """ + if not method: + raise ValueError("memory metric needs a non-empty method") + if not scope: + raise ValueError("memory metric needs a non-empty scope") + if bytes_value is None: + if not unavailable_reason: + raise ValueError( + "memory metric without bytes needs an unavailable_reason") + else: + if unavailable_reason: + raise ValueError( + "memory metric with bytes must not carry an unavailable " + "reason") + if not _is_int(bytes_value): + raise ValueError( + f"memory metric bytes must be an int, got {bytes_value!r}") + if bytes_value < 0: + raise ValueError( + f"memory metric bytes must be >= 0, got {bytes_value}") + metric = { + "bytes": bytes_value, + "unit": "bytes", + "method": method, + "scope": scope, + "unavailable_reason": unavailable_reason or None, + } + metric.update(details) + return metric + + +def _rate(tokens: int, seconds: float): + """tokens/seconds, or ``None`` when the measured duration is zero (an + unavailable rate, never a fabricated zero).""" + if seconds == 0.0: + return None + return tokens / seconds + + +def run_result(*, index: int, prompt_tokens: int, + requested_warmup_tokens: int, warmup_tokens: int, + requested_timed_tokens: int, timed_decode_tokens: int, + prefill_seconds: float, decode_seconds: float, + memory: dict, execution_evidence: dict | None = None) -> dict: + """Validate one benchmark run and derive its counts and rates. + + For the fixed-length sample-and-step loop in ``examples/bench.py``: + ``total_generated_tokens = warmup + timed_decode`` and + ``decode_start_context_tokens = prompt + warmup``. Throughput uses + only the timed decode tokens. Requested counts are kept separately + from the counts actually executed. + """ + index = _require_count("index", index) + prompt_tokens = _require_count("prompt_tokens", prompt_tokens) + requested_warmup_tokens = _require_count( + "requested_warmup_tokens", requested_warmup_tokens) + warmup_tokens = _require_count("warmup_tokens", warmup_tokens) + requested_timed_tokens = _require_count( + "requested_timed_tokens", requested_timed_tokens) + timed_decode_tokens = _require_count( + "timed_decode_tokens", timed_decode_tokens) + if timed_decode_tokens <= 0: + raise ValueError( + "timed_decode_tokens must be > 0: an empty timed window is " + "not a successful benchmark") + prefill_seconds = _require_seconds("prefill_seconds", prefill_seconds) + decode_seconds = _require_seconds("decode_seconds", decode_seconds) + + unavailable: dict[str, str] = {} + prefill_tps = _rate(prompt_tokens, prefill_seconds) + if prefill_tps is None: + unavailable["prefill_tokens_per_second"] = ( + "measured prefill duration is zero") + decode_tps = _rate(timed_decode_tokens, decode_seconds) + if decode_tps is None: + unavailable["decode_tokens_per_second"] = ( + "measured decode duration is zero") + + return { + "index": index, + "prompt_tokens": prompt_tokens, + "requested_warmup_tokens": requested_warmup_tokens, + "warmup_tokens": warmup_tokens, + "requested_timed_tokens": requested_timed_tokens, + "timed_decode_tokens": timed_decode_tokens, + "total_generated_tokens": warmup_tokens + timed_decode_tokens, + "decode_start_context_tokens": prompt_tokens + warmup_tokens, + "prefill_seconds": prefill_seconds, + "prefill_tokens_per_second": prefill_tps, + "decode_seconds": decode_seconds, + "decode_tokens_per_second": decode_tps, + "memory": dict(memory), + "execution_evidence": dict(execution_evidence or {}), + "unavailable_reasons": unavailable, + } + + +# ---- validation -------------------------------------------------------------- + + +def _walk(value, path: str, errors: list[str]) -> None: + """Strict-JSON and null-with-reason walk over a report fragment.""" + if isinstance(value, dict): + for key in value: + if not isinstance(key, str): + errors.append(f"{path}: non-string key {key!r}") + reasons = value.get("unavailable_reasons") + if "unavailable_reasons" in value: + if not isinstance(reasons, dict): + errors.append(f"{path}.unavailable_reasons must be a map") + reasons = {} + for key, reason in reasons.items(): + if key not in value: + errors.append( + f"{path}.unavailable_reasons names unknown field " + f"{key!r}") + elif value[key] is not None: + errors.append( + f"{path}.{key} has an unavailable reason but also a " + f"value {value[key]!r}") + if not isinstance(reason, str) or not reason: + errors.append( + f"{path}.unavailable_reasons[{key!r}] must be a " + f"non-empty string") + elif "bytes" in value and "unavailable_reason" in value: + # memory metric object: single value, single reason + if value["bytes"] is None and not value["unavailable_reason"]: + errors.append(f"{path}.bytes is null without a reason") + if value["bytes"] is not None and value["unavailable_reason"]: + errors.append( + f"{path}.bytes has a value and an unavailable reason") + for key, child in value.items(): + if key == "unavailable_reasons": + continue + child_path = f"{path}.{key}" + if child is None: + is_metric_null = key == "unavailable_reason" or ( + key == "bytes" and "unavailable_reason" in value) + if "unavailable_reasons" in value: + if key not in (reasons or {}): + errors.append( + f"{child_path} is null without an unavailable " + f"reason") + elif not is_metric_null: + errors.append( + f"{child_path} is null in an object without an " + f"unavailable_reasons map") + continue + _walk(child, child_path, errors) + elif isinstance(value, list): + for i, child in enumerate(value): + if child is None: + # list entries mirror per-run fields; the reason lives on + # the run object (see summarize_runs) + continue + _walk(child, f"{path}[{i}]", errors) + elif value is None or isinstance(value, (bool, str, int)): + return + elif isinstance(value, float): + if not math.isfinite(value): + errors.append(f"{path}: non-finite number {value!r}") + else: + errors.append( + f"{path}: unsupported type {type(value).__name__} " + f"(not JSON-serializable)") + + +def validate_report(report: dict) -> None: + """Raise ``ValueError`` listing every contract violation.""" + errors: list[str] = [] + if not isinstance(report, dict): + raise ValueError("report must be a dict") + for key in ("schema_version", "protocol", "identity", "model", + "runtime", "workload", "caches", "runs", "summary", + "memory", "probe"): + if key not in report: + errors.append(f"missing top-level field {key!r}") + if errors: + raise ValueError("invalid benchmark report:\n " + "\n ".join(errors)) + if report["schema_version"] != SCHEMA_VERSION: + errors.append( + f"schema_version {report['schema_version']!r} != " + f"{SCHEMA_VERSION}") + for group, fields in REQUIRED_FIELDS.items(): + obj = report[group] + if not isinstance(obj, dict): + errors.append(f"{group} must be an object") + continue + for field in fields: + if field not in obj: + errors.append(f"{group}.{field} is required") + if "unavailable_reasons" not in obj: + errors.append(f"{group}.unavailable_reasons is required") + runs = report["runs"] + if not isinstance(runs, list) or not runs: + errors.append("runs must be a non-empty list") + else: + for i, run in enumerate(runs): + if not isinstance(run, dict): + errors.append(f"runs[{i}] must be an object") + continue + for field in _RUN_FIELDS: + if field not in run: + errors.append(f"runs[{i}].{field} is required") + if run.get("index") != i: + errors.append( + f"runs[{i}].index is {run.get('index')!r}, expected {i}") + run_count = report["identity"].get("run_count") \ + if isinstance(report["identity"], dict) else None + if run_count != len(runs): + errors.append( + f"identity.run_count {run_count!r} != len(runs) " + f"{len(runs)}") + _walk(report, "report", errors) + if errors: + raise ValueError("invalid benchmark report:\n " + "\n ".join(errors)) + + +# ---- summary ----------------------------------------------------------------- + + +def _stat_block(values: list) -> dict: + """Descriptive statistics over per-run values (``None`` allowed).""" + values = list(values) + present = [v for v in values if v is not None] + block: dict[str, Any] = { + "n": len(values), "n_valid": len(present), "values": values, + "mean": None, "min": None, "max": None, "sample_stdev": None, + "unavailable_reasons": {}} + if not present: + reason = "no run has a value for this metric (n_valid = 0)" + for key in ("mean", "min", "max", "sample_stdev"): + block["unavailable_reasons"][key] = reason + return block + block["mean"] = statistics.fmean(present) + block["min"] = min(present) + block["max"] = max(present) + if len(present) >= 2: + block["sample_stdev"] = statistics.stdev(present) + else: + block["unavailable_reasons"]["sample_stdev"] = ( + f"sample standard deviation needs at least 2 runs with a value " + f"(n_valid = {len(present)})") + return block + + +def summarize_runs(runs: list[dict]) -> dict: + """Summary over every retained run. Runs are sequential within one + process, so this is descriptive, not an estimate from independent + launches (those are Task 3/7's job).""" + summary: dict[str, Any] = {"run_count": len(runs)} + for metric in _SUMMARY_METRICS: + summary[metric] = _stat_block([run.get(metric) for run in runs]) + summary["method"] = ( + "arithmetic mean, min, max and sample standard deviation (n-1) " + "over every retained run with a value (n_valid; values keeps the " + "per-run entries, null included); runs are sequential in one " + "process (run 0 is the first request after model load) and are " + "not independent launches") + return summary + + +# ---- probe link -------------------------------------------------------------- + + +_PROBE_FIELDS = ("path", "schema_version", "sha256", "ready", + "collected_at_utc") + + +def probe_link(*, path, text) -> dict: + """Link the report to an actual probe file: its schema version, the + SHA-256 of the file content as read, and the local reference. + + ``text`` is the raw file content (``bytes`` are hashed as-is; ``str`` + is hashed as UTF-8). ``None`` means no probe was supplied. + """ + link: dict[str, Any] + if text is None: + reason = "no probe file supplied (--probe-json not given)" + link = {"supplied": False} + link.update({field: None for field in _PROBE_FIELDS}) + link["unavailable_reasons"] = {f: reason for f in _PROBE_FIELDS} + return link + raw = text if isinstance(text, bytes) else text.encode("utf-8") + digest = hashlib.sha256(raw).hexdigest() + try: + data = json.loads(raw.decode("utf-8")) + except (UnicodeDecodeError, ValueError) as exc: + raise ValueError(f"probe file {path!r} is not valid JSON: {exc}") + if not isinstance(data, dict): + raise ValueError(f"probe file {path!r} is not a JSON object") + schema = data.get("schema_version") + if not _is_int(schema): + raise ValueError( + f"probe file {path!r} has no integer schema_version") + link = {"supplied": True, "path": os.fspath(path), + "schema_version": schema, "sha256": digest, + "ready": None, "collected_at_utc": None, + "unavailable_reasons": {}} + ready = data.get("ready") + if isinstance(ready, bool): + link["ready"] = ready + else: + link["unavailable_reasons"]["ready"] = ( + "probe has no boolean 'ready' field") + collected = data.get("collected_at_utc") + if isinstance(collected, str) and collected: + link["collected_at_utc"] = collected + else: + link["unavailable_reasons"]["collected_at_utc"] = ( + "probe has no 'collected_at_utc' field") + return link + + +# ---- assembly and serialization ------------------------------------------------ + + +def build_report(*, identity: dict, model: dict, runtime: dict, + workload: dict, caches: dict, runs: list[dict], + memory: dict, probe: dict | None = None) -> dict: + """Assemble and validate the full report. ``runs`` are + ``run_result`` dicts; ``probe`` is a ``probe_link`` dict (``None`` = + not supplied).""" + def _group(obj, name): + if not isinstance(obj, dict): + raise ValueError(f"{name} must be an object") + out = dict(obj) + out.setdefault("unavailable_reasons", {}) + return out + + report = { + "schema_version": SCHEMA_VERSION, + "protocol": dict(PROTOCOL), + "identity": _group(identity, "identity"), + "model": _group(model, "model"), + "runtime": _group(runtime, "runtime"), + "workload": _group(workload, "workload"), + "caches": _group(caches, "caches"), + "runs": [dict(run) for run in runs], + "summary": summarize_runs(runs) if runs else {}, + "memory": _group(memory, "memory"), + "probe": (probe if probe is not None + else probe_link(path=None, text=None)), + } + validate_report(report) + return report + + +def dumps_report(report: dict) -> str: + """Strict JSON text (``allow_nan=False``, sorted keys, trailing + newline). Validates first so a NaN never reaches the encoder.""" + validate_report(report) + return json.dumps(report, indent=2, sort_keys=True, allow_nan=False) + "\n" + + +def _place_new_file(tmp: str, dest: str) -> None: + """Move ``tmp`` to ``dest`` without ever overwriting ``dest``. + + A hard link is atomic and fails with EEXIST when the destination + appeared meanwhile; the temporary name is then unlinked. Filesystems + without hard links fall back to a rename (``os.rename`` refuses an + existing destination on Windows; POSIX ``os.replace`` is the last + resort after a fresh existence check). + """ + try: + os.link(tmp, dest) + except FileExistsError: + raise FileExistsError( + f"benchmark output {dest!r} appeared while the report was being " + f"written; refusing to overwrite it") + except (OSError, AttributeError): + if os.path.lexists(dest): + raise FileExistsError( + f"benchmark output {dest!r} appeared while the report was " + f"being written; refusing to overwrite it") + if os.name == "nt": + os.rename(tmp, dest) + else: + os.replace(tmp, dest) + return + os.unlink(tmp) + + +def write_report(report: dict, path) -> str: + """Validate, then write the report atomically to a NEW file. + + The destination must not exist (a stale success report can otherwise + be mistaken for this run's result). The text is written to a + temporary file in the destination directory, fsynced, and linked or + renamed into place without overwriting; on any failure the temporary + file is removed and the destination is left as it was. + """ + text = dumps_report(report) + path = os.fspath(path) + if os.path.lexists(path): + raise FileExistsError( + f"benchmark output {path!r} already exists; refusing to " + f"overwrite (choose a new path)") + dest_dir = os.path.dirname(os.path.abspath(path)) or os.curdir + if not os.path.isdir(dest_dir): + raise FileNotFoundError( + f"benchmark output directory {dest_dir!r} does not exist") + fd, tmp = tempfile.mkstemp(prefix=".bench-", suffix=".json.tmp", + dir=dest_dir) + try: + if os.name != "nt": + # mkstemp creates 0600; give the report the usual umask mode + mask = os.umask(0) + os.umask(mask) + os.chmod(tmp, 0o666 & ~mask) + with os.fdopen(fd, "w", encoding="utf-8", newline="\n") as fh: + fh.write(text) + fh.flush() + os.fsync(fh.fileno()) + _place_new_file(tmp, path) + except BaseException: + try: + os.unlink(tmp) + except OSError: + pass + raise + return text diff --git a/pyproject.toml b/pyproject.toml index 5186b0b..98ce6a1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -50,6 +50,10 @@ where = ["src"] [tool.pytest.ini_options] testpaths = ["tests"] +# tests import repo-root packages (``from scripts import jetson_probe``, +# ``from examples import bench``); put the root on sys.path for the plain +# ``pytest`` invocation too, not only for ``python -m pytest``. +pythonpath = ["."] markers = [ "slow: real-weight e2e tests (skip with -m 'not slow')", ] diff --git a/scripts/jetson_probe.py b/scripts/jetson_probe.py new file mode 100755 index 0000000..1e01f10 --- /dev/null +++ b/scripts/jetson_probe.py @@ -0,0 +1,332 @@ +#!/usr/bin/env python3 +"""Collect reproducible Jetson/CUDA facts before running Edge0 benchmarks. + +The probe is read-only and uses the Python standard library unless PyTorch is +installed. It prints JSON so benchmark reports can retain the exact platform +and storage context that produced them. +""" + +from __future__ import annotations + +import argparse +import json +import os +import platform +import shutil +import sys +from datetime import datetime, timezone +from pathlib import Path +from typing import Any + +_AUTO_TORCH = object() + + +def _read_text(path: Path) -> str | None: + try: + return path.read_text(errors="replace").rstrip("\x00\n") + except (OSError, UnicodeError): + return None + + +def parse_meminfo(text: str) -> dict[str, int | None]: + """Return total and available memory in bytes from Linux meminfo text.""" + values: dict[str, int] = {} + for line in text.splitlines(): + if ":" not in line: + continue + name, raw = line.split(":", 1) + fields = raw.split() + if not fields: + continue + try: + value = int(fields[0]) + except ValueError: + continue + unit = fields[1].lower() if len(fields) > 1 else "bytes" + multiplier = 1024 if unit == "kb" else 1 + values[name] = value * multiplier + return { + "total_bytes": values.get("MemTotal"), + "available_bytes": values.get("MemAvailable"), + } + + +def _unescape_mount_field(value: str) -> str: + for escaped, plain in ( + ("\\040", " "), + ("\\011", "\t"), + ("\\012", "\n"), + ("\\134", "\\"), + ): + value = value.replace(escaped, plain) + return value + + +def _mounts(mountinfo_text: str) -> list[dict[str, str]]: + mounts = [] + for line in mountinfo_text.splitlines(): + fields = line.split() + try: + separator = fields.index("-") + mounts.append( + { + "device_id": fields[2], + "mount_point": _unescape_mount_field(fields[4]), + "filesystem": fields[separator + 1], + "device": _unescape_mount_field(fields[separator + 2]), + } + ) + except (ValueError, IndexError): + continue + return mounts + + +def _sysfs_tree_is_nvme(path: Path, seen: set[Path] | None = None) -> bool: + try: + resolved = path.resolve(strict=True) + except OSError: + return False + seen = seen or set() + if resolved in seen: + return False + seen.add(resolved) + try: + slaves = list((resolved / "slaves").iterdir()) + except OSError: + slaves = [] + if slaves: + return all(_sysfs_tree_is_nvme(slave, seen) for slave in slaves) + return any(part.startswith("nvme") for part in resolved.parts) + + +def _device_is_nvme( + device: str, + device_id: str | None, + sys_dev_block: Path, +) -> bool | None: + if device_id: + sysfs_device = sys_dev_block / device_id + if sysfs_device.exists(): + return _sysfs_tree_is_nvme(sysfs_device) + if device.startswith(("/dev/mmc", "/dev/sd", "/dev/vd")): + return False + return None + + +def storage_device_for_path( + path: Path, + mountinfo_text: str, + *, + sys_dev_block: Path = Path("/sys/dev/block"), + disk_usage_fn: Any = shutil.disk_usage, + path_device_id: str | None = None, +) -> dict[str, Any]: + """Describe the longest-prefix Linux mount containing *path*.""" + target = path.resolve() + if path_device_id is None: + try: + stat = target.stat() + path_device_id = f"{os.major(stat.st_dev)}:{os.minor(stat.st_dev)}" + except OSError: + path_device_id = None + candidates = [] + for position, mount in enumerate(_mounts(mountinfo_text)): + mount_path = Path(mount["mount_point"]) + try: + target.relative_to(mount_path) + except ValueError: + continue + if path_device_id is not None and mount["device_id"] != path_device_id: + continue + candidates.append((len(mount_path.parts), position, mount)) + + selected = ( + max(candidates, key=lambda item: (item[0], item[1]))[2] + if candidates + else { + "device_id": None, + "mount_point": None, + "filesystem": None, + "device": None, + } + ) + device = selected.get("device") or "" + is_nvme = _device_is_nvme( + device, + selected.get("device_id"), + sys_dev_block, + ) + + try: + usage = disk_usage_fn(target) + total_bytes, free_bytes = usage.total, usage.free + disk_usage_available = True + except OSError: + total_bytes = free_bytes = None + disk_usage_available = False + + return { + **selected, + "is_nvme": is_nvme, + "path": str(target), + "path_exists": target.exists(), + "path_is_dir": target.is_dir(), + "disk_usage_available": disk_usage_available, + "total_bytes": total_bytes, + "free_bytes": free_bytes, + } + + +def collect_torch_info(torch_module: Any = _AUTO_TORCH) -> dict[str, Any]: + """Collect CUDA facts without making torch a probe dependency.""" + info: dict[str, Any] = { + "installed": False, + "version": None, + "cuda_available": False, + "cuda_version": None, + "device_name": None, + "compute_capability": None, + "device_memory_bytes": None, + "errors": [], + } + if torch_module is _AUTO_TORCH: + try: + import torch as torch_module # type: ignore[no-redef] + except Exception as exc: # noqa: BLE001 - vendor imports can fail broadly + info["errors"].append(f"PyTorch import failed: {type(exc).__name__}: {exc}") + return info + if torch_module is None: + info["errors"].append("PyTorch import was skipped or unavailable") + return info + + info["installed"] = True + info["version"] = str(getattr(torch_module, "__version__", "unknown")) + cuda = getattr(torch_module, "cuda", None) + info["cuda_version"] = getattr(getattr(torch_module, "version", None), "cuda", None) + try: + cuda_available = bool(cuda is not None and cuda.is_available()) + except Exception as exc: # noqa: BLE001 - CUDA driver failures vary + info["errors"].append( + f"CUDA availability check failed: {type(exc).__name__}: {exc}" + ) + return info + info["cuda_available"] = cuda_available + if cuda_available: + assert cuda is not None + try: + props = cuda.get_device_properties(0) + info.update( + { + "device_name": props.name, + "compute_capability": [props.major, props.minor], + "device_memory_bytes": props.total_memory, + } + ) + except Exception as exc: # noqa: BLE001 - CUDA driver failures vary + info["errors"].append( + f"CUDA device query failed: {type(exc).__name__}: {exc}" + ) + return info + + +def collect_probe( + model_dir: Path, + *, + fs_root: Path = Path("/"), + mountinfo_text: str | None = None, + torch_module: Any = _AUTO_TORCH, + disk_usage_fn: Any = shutil.disk_usage, + path_device_id: str | None = None, +) -> dict[str, Any]: + """Collect a JSON-serializable Edge0/Jetson readiness report.""" + model = _read_text(fs_root / "proc/device-tree/model") + l4t_release = _read_text(fs_root / "etc/nv_tegra_release") + is_jetson_orin = bool( + model and "jetson" in model.lower() and "orin" in model.lower() + ) + + meminfo_text = _read_text(fs_root / "proc/meminfo") or "" + memory = parse_meminfo(meminfo_text) + if mountinfo_text is None: + mountinfo_text = _read_text(fs_root / "proc/self/mountinfo") or "" + storage = storage_device_for_path( + model_dir, + mountinfo_text, + sys_dev_block=fs_root / "sys/dev/block", + disk_usage_fn=disk_usage_fn, + path_device_id=path_device_id, + ) + torch_info = collect_torch_info(torch_module) + + blockers = [] + if not storage["path_exists"]: + blockers.append("model path does not exist") + elif not storage["path_is_dir"]: + blockers.append("model path is not a directory") + if not storage["disk_usage_available"]: + blockers.append("model path disk usage could not be determined") + if not is_jetson_orin: + blockers.append("hardware is not a Jetson Orin") + if not torch_info["installed"]: + blockers.append("PyTorch is not installed") + if not torch_info["cuda_available"]: + blockers.append("CUDA is not available through PyTorch") + blockers.extend( + f"PyTorch/CUDA probe error: {error}" for error in torch_info["errors"] + ) + if storage["is_nvme"] is None: + blockers.append("model path NVMe backing could not be confirmed") + elif not storage["is_nvme"]: + blockers.append("model path is not on NVMe storage") + + return { + "schema_version": 1, + "collected_at_utc": datetime.now(timezone.utc).isoformat(), + "platform": { + "system": platform.system(), + "release": platform.release(), + "machine": platform.machine(), + "python": platform.python_version(), + }, + "jetson": { + "is_jetson_orin": is_jetson_orin, + "model": model, + "l4t_release": l4t_release, + }, + "memory": memory, + "storage": storage, + "torch": torch_info, + "edge0_backend": os.environ.get("EDGE0_BACKEND", "mlx"), + "ready": not blockers, + "blockers": blockers, + } + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--model-dir", + type=Path, + default=Path.cwd(), + help="checkpoint directory whose backing storage should be inspected", + ) + parser.add_argument("--output", type=Path, help="also write JSON to this path") + parser.add_argument( + "--skip-torch", + action="store_true", + help="do not import PyTorch (reported as an explicit blocker)", + ) + args = parser.parse_args(argv) + + torch_module = None if args.skip_torch else _AUTO_TORCH + report = collect_probe(args.model_dir, torch_module=torch_module) + rendered = json.dumps(report, indent=2, sort_keys=True) + "\n" + sys.stdout.write(rendered) + if args.output: + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(rendered) + return 0 if report["ready"] else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/orin_reference_check.py b/scripts/orin_reference_check.py new file mode 100644 index 0000000..59465fb --- /dev/null +++ b/scripts/orin_reference_check.py @@ -0,0 +1,426 @@ +#!/usr/bin/env python3 +"""On-device Edge0-8B correctness check (Task 3, Gate B). + +Generates >= 32 greedy tokens with the torch backend on the TEST device +(default ``cuda``), then teacher-forces the SAME token sequence through +the torch backend on the REFERENCE device (default ``cpu``) and compares +the full per-step logits and token choices. + +Why torch-CPU is the reference here: the plan's Gate B fixtures were to +be generated against MLX, which does not exist on a Jetson. The torch +implementation itself was checked against real MLX layer by layer on a +GB10 (edge0-8b backbone: 5.3e-7 max per layer, 4.3e-7 on the logits, +same argmax — docs/nvidia.md), so torch-on-CPU is the closest identified +reference data available on this device, and the comparison isolates +exactly the variable Task 3 introduces: the CUDA device. When MLX +fixtures land, point ``--reference-npz`` at them instead. + +Tolerances (DEFINED IN ADVANCE, before any Orin result was examined — +do not loosen them to make a run pass): + +* ``--logit-tol`` (default 1.0): max absolute per-step logit difference. + The dense model math runs in bfloat16 (~8 mantissa bits); reduction + order differs across devices, so logits at a scale of tens can move + by a few tenths through 24 layers. 1.0 is the a-priori bf16 bound, + NOT an observed number. +* ``--near-tie-margin`` (default 1.0): a token divergence at step i is + acceptable only if the reference's own logit gap between its top + token and the test's chosen token is within this margin (a genuine + near-tie), AND the logit tolerance holds at that step. Any other + divergence fails. + +Each phase runs in its own subprocess because the torch backend binds +``EDGE0_TORCH_DEVICE`` at import time. Exit status: 0 PASS, 1 FAIL, +2 usage/environment error. The JSON report records checkpoint / +adapter / git identity, prompt digest, per-step statistics and the +verdict; keep it in the untracked evidence set ($EDGE0_RUN_DIR). +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import subprocess +import sys +import time +from datetime import datetime, timezone +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(REPO_ROOT)) + +DEFAULT_PROMPT = "9.11 和 9.8 哪个大?请仔细比较。" # examples/bench.py edge0-8b + +#: Knobs recorded per phase as RESOLVED (inherited from the shell or set +#: by --test-env/--reference-env): the comparison is only meaningful when +#: the two phases differ in exactly what the caller intended. +RECORDED_ENV = ( + "EDGE0_BACKEND", "EDGE0_TORCH_DEVICE", "EDGE0_QMM_BATCHED", + "EDGE0_QMM_BATCHED_MAX_BYTES", "EDGE0_INT4PACK", + "EDGE0_TORCH_WEIGHT_CACHE", "EDGE0_TORCH_WEIGHT_CACHE_BYTES", + "EDGE0_MEMORY_BUDGET", "EDGE0_BUDGET_CONTEXT", "EDGE0_CACHE_SLOTS", + "EDGE0_PREDICT_PREFETCH", "EDGE0_PREWARM", "LING_PREWARM", + "LING_HIDDEN_CLIP", "PREROUTER_FEATURE_TOPK", "PREROUTER_INTRA", +) + + +def host_identity() -> dict: + """Which machine produced this phase. A report without it reads the + same whether it came from the Orin or a workstation GPU, and this + script's results gate an Orin claim.""" + import platform + out = { + "hostname_hash": hashlib.sha256( + platform.node().encode("utf-8")).hexdigest()[:16], + "machine": platform.machine(), + "system": platform.system(), + "python": platform.python_version(), + "gpu_name": None, + "compute_capability": None, + "torch_version": None, + "torch_cuda_version": None, + "jetson_model": None, + "unavailable_reasons": {}, + } + try: + import torch + out["torch_version"] = str(torch.__version__) + out["torch_cuda_version"] = getattr(torch.version, "cuda", None) + if torch.cuda.is_available(): + props = torch.cuda.get_device_properties(0) + out["gpu_name"] = props.name + out["compute_capability"] = f"{props.major}.{props.minor}" + else: + out["unavailable_reasons"]["gpu_name"] = "no CUDA device" + out["unavailable_reasons"]["compute_capability"] = "no CUDA device" + except Exception as exc: # noqa: BLE001 - identity is best-effort + reason = f"{type(exc).__name__}: {exc}" + for key in ("torch_version", "torch_cuda_version", "gpu_name", + "compute_capability"): + out["unavailable_reasons"].setdefault(key, reason) + try: + out["jetson_model"] = Path( + "/proc/device-tree/model").read_text(errors="replace").rstrip("\x00\n") + except OSError: + out["unavailable_reasons"]["jetson_model"] = ( + "/proc/device-tree/model unreadable: not a Jetson") + if out["torch_cuda_version"] is None: + out["unavailable_reasons"].setdefault( + "torch_cuda_version", "torch build has no CUDA runtime") + return out + + +# --------------------------------------------------------------- phases + +def _encode_prompt(engine, prompt: str) -> list[int]: + # Same chat templating as examples/bench.py (think=False). + tok = engine._tok + if hasattr(engine, "encode_chat"): + return [int(t) for t in engine.encode_chat( + [{"role": "user", "content": prompt}], think=False)] + return [int(t) for t in tok(tok.apply_chat_template( + [{"role": "user", "content": prompt}], + tokenize=False, add_generation_prompt=True, + enable_thinking=False))["input_ids"]] + + +def _logits_to_numpy(logits): + import numpy as np + import torch + if isinstance(logits, torch.Tensor): + return logits.detach().reshape(-1).to("cpu", torch.float32).numpy() + return np.asarray(logits, dtype="float32").reshape(-1) + + +def run_phase(model_dir: str, prompt: str, ntok: int, out_npz: str, + forced: list[int] | None) -> None: + """One engine pass on the device the environment selected. + + ``forced is None``: greedy generation, logits captured at every step. + ``forced``: teacher-forcing — capture logits, then feed forced[i] + regardless of what this device would have chosen. + """ + import numpy as np + + from edge0 import AutoEngine + + t0 = time.perf_counter() + engine = AutoEngine.from_pretrained(model_dir) + try: + ids = _encode_prompt(engine, prompt) + engine.reset() + engine.prefill(ids) + logits = engine.next_logits() + first = _logits_to_numpy(logits) + steps = [first] + tokens: list[int] = [] + for i in range(ntok): + if forced is None: + tid = int(steps[-1].argmax()) + else: + tid = int(forced[i]) + tokens.append(tid) + if i < ntok - 1: + logits = engine.step(tid) + steps.append(_logits_to_numpy(logits)) + evidence = { + "logits_class": type(logits).__name__, + "device": str(getattr(logits, "device", "unknown")), + "elapsed_s": time.perf_counter() - t0, + **host_identity(), + # every knob as this phase RESOLVED it, not just the CLI + # overrides: a value inherited from the shell would otherwise + # turn a same-device comparison into a self-comparison + "resolved_env": {k: os.environ[k] for k in RECORDED_ENV + if k in os.environ}, + } + np.savez_compressed( + out_npz, logits=np.stack(steps), tokens=np.asarray(tokens), + prompt_ids=np.asarray(ids)) + Path(out_npz + ".meta.json").write_text(json.dumps(evidence)) + finally: + engine.close() + + +# --------------------------------------------------------------- compare + +def compare(test_npz: str, ref_npz: str, logit_tol: float, + near_tie_margin: float) -> dict: + import numpy as np + + test = np.load(test_npz) + ref = np.load(ref_npz) + if not np.array_equal(test["prompt_ids"], ref["prompt_ids"]): + return {"pass": False, "error": "prompt ids differ between phases"} + if not np.array_equal(test["tokens"], ref["tokens"]): + return {"pass": False, + "error": "teacher phase did not replay the test tokens"} + a, b = test["logits"], ref["logits"] + if a.shape != b.shape: + return {"pass": False, + "error": f"logit shapes differ: {a.shape} vs {b.shape}"} + if not (np.isfinite(a).all() and np.isfinite(b).all()): + return {"pass": False, "error": "non-finite logits"} + + steps = [] + divergences = [] + ok = True + for i in range(a.shape[0]): + diff = float(np.abs(a[i] - b[i]).max()) + t_test = int(test["tokens"][i]) # what the test device chose + r_top = int(b[i].argmax()) # reference greedy choice + entry = {"step": i, "max_abs_logit_diff": diff, + "test_token": t_test, "reference_token": r_top} + if diff > logit_tol: + entry["fail"] = f"logit diff {diff:.4g} > tol {logit_tol}" + ok = False + if t_test != r_top: + gap = float(b[i][r_top] - b[i][t_test]) + entry["reference_gap"] = gap + near_tie = 0.0 <= gap <= near_tie_margin and diff <= logit_tol + entry["explained_near_tie"] = near_tie + divergences.append(entry) + if not near_tie: + entry.setdefault( + "fail", f"token divergence with reference gap " + f"{gap:.4g} > margin {near_tie_margin}") + ok = False + steps.append(entry) + diffs = [s["max_abs_logit_diff"] for s in steps] + return { + "pass": ok, + "steps_compared": len(steps), + "tokens_matched": len(steps) - len(divergences), + "divergences": divergences, + "max_abs_logit_diff": max(diffs), + "mean_abs_logit_diff_of_max": sum(diffs) / len(diffs), + "per_step": steps, + } + + +# ----------------------------------------------------------- orchestrate + +def parse_env_overrides(items: list[str] | None) -> dict[str, str]: + """``KEY=VALUE`` pairs for one phase's environment (Task 6: run the + reference path and an opt-in path on the SAME device, e.g. + ``--test-env EDGE0_QMM_BATCHED=1``). An empty value unsets the key.""" + out: dict[str, str] = {} + for item in items or (): + key, sep, value = item.partition("=") + if not sep or not key.strip(): + raise ValueError(f"expected KEY=VALUE, got {item!r}") + out[key.strip()] = value + return out + + +def _spawn(device: str, argv: list[str], + overrides: dict[str, str] | None = None) -> None: + env = dict(os.environ) + env["EDGE0_BACKEND"] = "cuda" + env["EDGE0_TORCH_DEVICE"] = device + for key, value in (overrides or {}).items(): + if value == "": + env.pop(key, None) + else: + env[key] = value + proc = subprocess.run([sys.executable, os.fspath(Path(__file__).resolve()), + "--phase-internal", *argv], env=env) + if proc.returncode != 0: + raise RuntimeError(f"phase on device {device!r} failed " + f"(exit {proc.returncode})") + + +def main(argv: list[str] | None = None) -> int: + ap = argparse.ArgumentParser(description=__doc__.split("\n", 1)[0]) + ap.add_argument("model", help="model directory or tier name") + ap.add_argument("--tokens", type=int, default=32) + ap.add_argument("--prompt", default=DEFAULT_PROMPT) + ap.add_argument("--test-device", default="cuda") + ap.add_argument("--reference-device", default="cpu") + ap.add_argument("--reference-npz", default=None, + help="use pre-generated reference logits instead of " + "running the reference device phase") + ap.add_argument("--logit-tol", type=float, default=1.0) + ap.add_argument("--near-tie-margin", type=float, default=1.0) + ap.add_argument("--test-env", action="append", metavar="KEY=VALUE", + help="environment override for the TEST phase only " + "(repeatable; empty value unsets), e.g. " + "EDGE0_QMM_BATCHED=1 to check an opt-in path " + "against the reference path on the same device") + ap.add_argument("--reference-env", action="append", metavar="KEY=VALUE", + help="environment override for the REFERENCE phase only") + ap.add_argument("--output", required=True, + help="JSON report path (must not exist)") + ap.add_argument("--phase-internal", action="store_true", + help=argparse.SUPPRESS) + ap.add_argument("--phase-out", help=argparse.SUPPRESS) + ap.add_argument("--phase-forced", help=argparse.SUPPRESS) + args = ap.parse_args(argv) + + if args.phase_internal: + forced = (json.loads(Path(args.phase_forced).read_text()) + if args.phase_forced else None) + run_phase(args.model, args.prompt, args.tokens, args.phase_out, forced) + return 0 + + if args.tokens < 32: + print("[reference-check] Gate B requires at least 32 tokens", + file=sys.stderr) + return 2 + try: + test_env = parse_env_overrides(args.test_env) + reference_env = parse_env_overrides(args.reference_env) + except ValueError as exc: + print(f"[reference-check] {exc}", file=sys.stderr) + return 2 + out = Path(args.output) + if out.exists(): + print(f"[reference-check] refusing to overwrite {out}", + file=sys.stderr) + return 2 + out.parent.mkdir(parents=True, exist_ok=True) + + from examples import benchmark_measure as measure + + model_dir = args.model + if not Path(model_dir).exists(): + env_name = f"EDGE0_{args.model.split('-')[-1].upper()}_MODEL" + model_dir = os.environ.get(env_name, model_dir) + if not Path(model_dir).exists(): + print(f"[reference-check] model not found: {args.model}", + file=sys.stderr) + return 2 + + work = out.parent / (out.stem + ".phases") + work.mkdir(exist_ok=True) + test_npz = str(work / f"test-{args.test_device}.npz") + ref_npz = args.reference_npz or str( + work / f"reference-{args.reference_device}.npz") + + common = [model_dir, "--prompt", args.prompt, + "--tokens", str(args.tokens), "--output", "unused"] + print(f"[reference-check] phase 1/2: greedy {args.tokens} tokens on " + f"{args.test_device}" + + (f" with {test_env}" if test_env else "")) + _spawn(args.test_device, common + ["--phase-out", test_npz], test_env) + + if not args.reference_npz: + import numpy as np + tokens = [int(t) for t in np.load(test_npz)["tokens"]] + forced_file = work / "forced-tokens.json" + forced_file.write_text(json.dumps(tokens)) + print(f"[reference-check] phase 2/2: teacher-forcing on " + f"{args.reference_device}" + + (f" with {reference_env}" if reference_env else "")) + _spawn(args.reference_device, + common + ["--phase-out", ref_npz, + "--phase-forced", str(forced_file)], + reference_env) + + result = compare(test_npz, ref_npz, args.logit_tol, args.near_tie_margin) + + same_device = (not args.reference_npz + and args.test_device == args.reference_device) + phase_env = { + name: json.loads(Path(p + ".meta.json").read_text()).get( + "resolved_env", {}) + for name, p in (("test", test_npz), ("reference", ref_npz)) + if Path(p + ".meta.json").exists()} + env_differences = sorted( + set(phase_env.get("test", {}).items()) + ^ set(phase_env.get("reference", {}).items())) + if same_device and not env_differences: + print("[reference-check] WARNING: same device and identical " + "resolved knobs in both phases: this compares a " + "configuration with itself and proves nothing", + file=sys.stderr) + + report = { + "schema": "edge0-reference-check/2", + "timestamp_utc": datetime.now(timezone.utc).isoformat(), + "host": host_identity(), + "same_device_comparison": same_device, + "resolved_env_differences": [list(kv) for kv in env_differences], + "git": measure.git_identity(REPO_ROOT), + "model_dir": os.path.abspath(model_dir), + "checkpoint": measure.checkpoint_manifest(model_dir), + "prompt_sha256": hashlib.sha256( + args.prompt.encode("utf-8")).hexdigest(), + "prompt_text": args.prompt, + "tokens_requested": args.tokens, + "test_device": args.test_device, + "test_env": test_env, + "reference_device": (args.reference_device if not args.reference_npz + else f"npz:{args.reference_npz}"), + "reference_env": reference_env, + "tolerances": {"logit_tol": args.logit_tol, + "near_tie_margin": args.near_tie_margin}, + "phase_evidence": { + name: json.loads(Path(p + ".meta.json").read_text()) + for name, p in (("test", test_npz), ("reference", ref_npz)) + if Path(p + ".meta.json").exists()}, + "result": {k: v for k, v in result.items() if k != "per_step"}, + "per_step": result.get("per_step"), + } + tmp = out.with_suffix(out.suffix + ".tmp") + tmp.write_text(json.dumps(report, indent=1, allow_nan=False)) + tmp.rename(out) + + verdict = "PASS" if result.get("pass") else "FAIL" + # ASCII only: this line is printed on whatever console the target has, + # and a Windows cp1252 stdout raises UnicodeEncodeError on "Δ" AFTER + # the report is already written -- turning a PASS into a crash. + print(f"[reference-check] {verdict}: " + f"{result.get('tokens_matched')}/{result.get('steps_compared')} " + f"tokens matched, max abs logit diff = " + f"{result.get('max_abs_logit_diff', float('nan')):.4g} " + f"(tol {args.logit_tol}); report: {out}") + if result.get("error"): + print(f"[reference-check] error: {result['error']}", file=sys.stderr) + return 0 if result.get("pass") else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/edge0/backends/cuda/int4pack.py b/src/edge0/backends/cuda/int4pack.py new file mode 100644 index 0000000..4713b15 --- /dev/null +++ b/src/edge0/backends/cuda/int4pack.py @@ -0,0 +1,185 @@ +"""torch's built-in int4 kernel for the dense quantized linears (Task 6, +design A, opt-in ``EDGE0_INT4PACK=1``). + +``torch.ops.aten._weight_int4pack_mm`` (the gpt-fast / torchao path) +runs a group-wise int4 x bfloat16 matmul on CUDA. It wants the weight +in its own tiled layout (``_convert_weight_to_int4pack`` of a +``[n, k/2]`` uint8 tensor, high nibble first) and computes +``w = (q - 8) * scale + zero`` from a ``[k/group, n, 2]`` bfloat16 +``scales_and_zeros`` tensor. The MLX affine layout this repo stores +(8 codes per uint32, least significant bits first, ``w = code * scale ++ bias``) maps onto it with ``zero = bias + 8 * scale``. + +That mapping is NOT exact: ``zero`` is rounded to bfloat16 and the +kernel dequantizes in bfloat16, so its weights differ from the +reference's own bfloat16-rounded weights by about one bfloat16 ulp for +roughly a third of the entries (measured on torch 2.14.0+cu130, see +``tests/test_int4pack.py`` for the pre-registered bounds). The exact +paths (the on-the-fly dequantization and both dequantized-weight +caches) therefore take precedence in ``nn.QuantizedLinear``; this path +is used only when the gate is on, nothing exact is cached for the +module, and every guard in ``is_supported`` holds. Whether the +approximation is acceptable is decided by the 32-token reference check +on the target device, not here. + +Costs: the repacked payload duplicates the int4 words it covers +(``int4pack_extra_bytes``) -- the original buffers stay for the +fall-through paths -- and the kernel accepts only 2-D bfloat16 +activations, so inputs are reshaped around the call. +""" + +from __future__ import annotations + +import os + +import torch + +INT4PACK_ENV = "EDGE0_INT4PACK" +#: ``_convert_weight_to_int4pack`` tiling; the kernel then needs +#: ``in_features % (INNER_K_TILES * 16) == 0``. +INNER_K_TILES = 8 +SUPPORTED_GROUP_SIZES = (32, 64, 128, 256) + + +def _read_int4pack_gate() -> bool: + return os.environ.get(INT4PACK_ENV, "") == "1" + + +#: Read once at import, like ``quant.BATCHED``; tests patch the attribute. +INT4PACK = _read_int4pack_gate() + + +# ---- pure repack (CPU or CUDA, exact) ------------------------------------------- + + +def codes_to_uint8(weight_u32: torch.Tensor) -> torch.Tensor: + """MLX-packed ``[..., k/8]`` uint32 (8 codes per word, least + significant bits first) -> ``[..., k/2]`` uint8 with the EVEN column's + code in the high nibble and the ODD column's in the low nibble (the + order ``_convert_weight_to_int4pack`` expects).""" + words = weight_u32.view(torch.int32) + shifts = torch.arange(8, device=words.device, dtype=torch.int32) * 4 + codes = ((words.unsqueeze(-1) >> shifts) & 0xF) + codes = codes.reshape(*words.shape[:-1], words.shape[-1] * 8) + return ((codes[..., 0::2] << 4) | codes[..., 1::2]).to(torch.uint8) + + +def uint8_to_codes(packed_u8: torch.Tensor) -> torch.Tensor: + """Inverse of ``codes_to_uint8``: ``[..., k/2]`` uint8 -> ``[..., k]`` + int64 codes in [0, 15].""" + words = packed_u8.to(torch.int64) + hi = (words >> 4) & 0xF + lo = words & 0xF + return torch.stack([hi, lo], dim=-1).reshape(*words.shape[:-1], + words.shape[-1] * 2) + + +def codes_to_uint32(codes: torch.Tensor) -> torch.Tensor: + """``[..., k]`` codes -> the MLX packing (8 per uint32, LSB first).""" + grouped = codes.reshape(*codes.shape[:-1], -1, 8).to(torch.int64) + shifts = torch.arange(8, device=codes.device, dtype=torch.int64) * 4 + words = (grouped << shifts).sum(-1) + return words.to(torch.int32).view(torch.uint32) + + +def scales_and_zeros(scales: torch.Tensor, biases: torch.Tensor) -> torch.Tensor: + """``[n, k/group]`` scales/biases -> ``[k/group, n, 2]`` bfloat16 + ``(scale, zero)`` with ``zero = bias + 8 * scale`` so that the kernel's + ``(q - 8) * scale + zero`` is the affine ``q * scale + bias``.""" + s = scales.to(torch.float32) + z = biases.to(torch.float32) + 8.0 * s + return torch.stack([s.t(), z.t()], dim=-1).to(torch.bfloat16).contiguous() + + +def zero_rounding_error(scales: torch.Tensor, biases: torch.Tensor) -> float: + """Largest |bf16(zero) - zero| the mapping introduces (reported, never + hidden: it is one of the two reasons the path is approximate).""" + z = biases.double() + 8.0 * scales.double() + return float((z.to(torch.bfloat16).double() - z).abs().max().item()) + + +def is_supported(*, in_features: int, out_features: int, group_size: int, + bits: int, device_type: str, x_dtype) -> tuple[bool, str]: + """Guarded dispatch: every combination the kernel does not handle + falls through to the reference implementation.""" + if bits != 4: + return False, f"{bits}-bit layout (kernel is int4 only)" + if group_size not in SUPPORTED_GROUP_SIZES: + return False, f"group_size {group_size} not in {SUPPORTED_GROUP_SIZES}" + if in_features % (INNER_K_TILES * 16) != 0: + return False, (f"in_features {in_features} is not a multiple of " + f"{INNER_K_TILES * 16}") + if out_features % 8 != 0: + return False, f"out_features {out_features} is not a multiple of 8" + if device_type != "cuda": + return False, f"device {device_type!r} (CUDA kernel only)" + if x_dtype is not torch.bfloat16: + return False, f"activation dtype {x_dtype} (kernel takes bfloat16)" + return True, "" + + +def int4pack_extra_bytes(out_features: int, in_features: int, + group_size: int) -> int: + """Bytes the repacked payload adds next to the original buffers.""" + return (out_features * in_features // 2 + + (in_features // group_size) * out_features * 2 * 2) + + +# ---- kernel (CUDA) -------------------------------------------------------------------- + + +class KernelUnavailable(RuntimeError): + """The device has no usable image for the kernel (e.g. a wheel whose + SASS list excludes this architecture). Every guard passed, so the + caller falls through to the reference path and stops retrying.""" + + +def repack(weight_u32: torch.Tensor, scales: torch.Tensor, + biases: torch.Tensor, group_size: int): + """Build the kernel's operands from the MLX layout.""" + packed_u8 = codes_to_uint8(weight_u32) + tiled = torch.ops.aten._convert_weight_to_int4pack(packed_u8, INNER_K_TILES) + return tiled, scales_and_zeros(scales, biases) + + +def probe(device=None) -> tuple[bool, str]: + """Run the kernel once on a tiny tensor to prove the device actually + has an image for it. + + ``is_supported`` checks shapes and dtypes; only execution proves the + build covers this architecture. The Orin's torch wheel warns that + ``sm_87`` is outside its SASS list (CC 8.x binary compatibility is + expected to carry it, but the plan's rule is that execution evidence + gates, not the warning), so the dispatch probes before it commits. + """ + try: + codes = torch.randint(0, 16, (8, 256), device=device) + packed = codes_to_uint32(codes) + scales = torch.ones(8, 4, dtype=torch.bfloat16, device=device) + biases = torch.zeros(8, 4, dtype=torch.bfloat16, device=device) + tiled, saz = repack(packed, scales, biases, 64) + x = torch.ones(1, 256, dtype=torch.bfloat16, device=device) + out = _int4pack_mm(x, tiled, 64, saz) + if device is not None and getattr(device, "type", device) == "cuda": + torch.cuda.synchronize() + if not torch.isfinite(out).all(): + return False, "kernel probe produced non-finite output" + except Exception as exc: # noqa: BLE001 - any failure disables the path + return False, f"kernel probe failed: {type(exc).__name__}: {exc}" + return True, "" + + +def _int4pack_mm(x2d: torch.Tensor, tiled: torch.Tensor, group_size: int, + saz: torch.Tensor) -> torch.Tensor: + return torch.ops.aten._weight_int4pack_mm(x2d, tiled, group_size, saz) + + +def int4pack_linear(x: torch.Tensor, tiled: torch.Tensor, group_size: int, + saz: torch.Tensor, bias: torch.Tensor | None = None): + """``F.linear`` on the repacked weight: reshape to the 2-D bfloat16 + input the kernel takes and back.""" + x2d = x.reshape(-1, x.shape[-1]).contiguous() + y = _int4pack_mm(x2d, tiled, group_size, saz) + if bias is not None: + y = y + bias.to(y.dtype) + return y.reshape(*x.shape[:-1], y.shape[-1]) diff --git a/src/edge0/backends/cuda/io.py b/src/edge0/backends/cuda/io.py index e5dd816..f1332e6 100644 --- a/src/edge0/backends/cuda/io.py +++ b/src/edge0/backends/cuda/io.py @@ -19,7 +19,7 @@ import torch from edge0.backends.base import TensorStore -from edge0.backends.cuda.core import DEVICE, _as_tensor +from edge0.backends.cuda.core import DEVICE _DTYPES = { "F64": torch.float64, "F32": torch.float32, "F16": torch.float16, @@ -338,6 +338,7 @@ def _install_quantized(model, state: dict, dtype) -> set: if isinstance(mod, torch.nn.Linear): bias = state.pop(f"{path}.bias", None) new = cnn.QuantizedLinear(w, s, b, mod.in_features, bias=bias) + cnn.WEIGHT_CACHE.register(path, new) # capped exact cache candidate elif isinstance(mod, torch.nn.Embedding): new = cnn.QuantizedEmbedding(w, s, b, mod.embedding_dim, dtype=dtype) @@ -379,8 +380,6 @@ def load_model(model_path, lazy=True, strict=False, model_config=None, import json import os - import torch - with open(os.path.join(os.fspath(model_path), "config.json")) as f: raw_config = json.load(f) @@ -417,7 +416,12 @@ def load_model(model_path, lazy=True, strict=False, model_config=None, _undo_mlx_qwen35_sanitize(state) if engine_path and hasattr(model, "sanitize"): state = model.sanitize(state) + from edge0.backends.cuda import nn as _cnn + _cnn.WEIGHT_CACHE.begin() # one model per process quantized = _install_quantized(model, state, dtype) + # Admit against the requested cap now; the 8B engine re-finalizes + # against the memory budget's headroom once that is resolved. + _cnn.WEIGHT_CACHE.finalize() for k, t in state.items(): if dtype is not None and t.is_floating_point(): state[k] = t.to(dtype) diff --git a/src/edge0/backends/cuda/nn.py b/src/edge0/backends/cuda/nn.py index 1134ee9..424150b 100644 --- a/src/edge0/backends/cuda/nn.py +++ b/src/edge0/backends/cuda/nn.py @@ -7,16 +7,33 @@ here under the contract's names rather than assuming call sites import ``torch.nn`` directly (see ``backends/__init__.py``'s enforcement that framework code only ever imports ``edge0.backends.{core,nn,io,quant}``). + +The quantized modules keep the 4-bit payload resident and dequantize on +the fly. Three opt-in alternatives exist for the dense linears, tried +in this order in ``QuantizedLinear.forward`` (exact ones first): + +* ``EDGE0_TORCH_WEIGHT_CACHE=1`` -- keep EVERY weight dequantized + (bit-identical to on-the-fly; ~1.9 GB in bf16 for edge0-8b's dense + modules, 4.1 GB in float32 as measured on the GB10); +* ``EDGE0_TORCH_WEIGHT_CACHE_BYTES=`` -- keep only the modules that + fit a byte cap (``WeightCachePolicy``: first-fit in registration + order, re-finalized by the engine against the Task 4 budget's + headroom, recorded in the bench report); +* ``EDGE0_INT4PACK=1`` -- torch's built-in int4 kernel (see + ``int4pack.py``; approximate at bfloat16 rounding, guarded). """ from __future__ import annotations import os +from typing import Any import torch import torch.nn as _tnn import torch.nn.functional as F +from edge0.backends.cuda import int4pack as _i4 + Module = _tnn.Module @@ -66,14 +83,183 @@ def _quant_params(weight, scales, in_features): #: Off by default: the whole point of the quantized layers is that the #: 4-bit payload is what stays resident. Worth turning on where memory is #: plentiful: on a GB10 it takes the edge0-8b decode step from 206 ms to -#: 124 ms a token, for 4.1 GB more resident (EDGE0_TORCH_WEIGHT_CACHE=1). +#: 124 ms a token, for a measured 4.1 GB more resident +#: (EDGE0_TORCH_WEIGHT_CACHE=1). For this tier the dequantized weights +#: themselves are 1.40 GB in bf16 (``WeightCachePolicy`` prices them from +#: the loaded modules); the difference is allocator segments the fill +#: leaves behind, which is why the fill is chunked. CACHE_DEQUANTIZED = os.environ.get("EDGE0_TORCH_WEIGHT_CACHE", "") == "1" +WEIGHT_CACHE_BYTES_ENV = "EDGE0_TORCH_WEIGHT_CACHE_BYTES" + + +def _read_weight_cache_cap() -> int | None: + """``None`` when unset (policy off); ``0`` also means off.""" + raw = os.environ.get(WEIGHT_CACHE_BYTES_ENV, "").strip() + if not raw: + return None + try: + value = int(raw) + except ValueError: + raise ValueError( + f"{WEIGHT_CACHE_BYTES_ENV} must be an integer byte count, got " + f"{raw!r}") from None + if value < 0: + raise ValueError(f"{WEIGHT_CACHE_BYTES_ENV} must be >= 0, got {value}") + return value + + +def dequantized_bytes(module, itemsize: int | None = None) -> int: + """Bytes a module's dequantized weight occupies once cached (activation + itemsize defaults to the checkpoint's scales dtype, bf16 for the + shipped tiers).""" + if itemsize is None: + itemsize = module.scales.element_size() + return module.out_features * module.in_features * itemsize + + +def fill_transient_bytes(module) -> int: + """Peak EXTRA bytes of building one cached weight, on top of the + resident result. + + The fill runs ``ROWS_PER_CHUNK`` rows at a time, so the transient is + bounded by the chunk, not by the module: ``_dequantize`` keeps four + float32-sized tensors of one chunk live at its peak, plus the + float32 casts of that chunk's scales and biases, plus the chunk's + converted result before it is written into the buffer. Measured at + edge0-8b's widest module (lm_head, 157184 x 1536) on torch + 2.14.0+cu130: 556.9 MiB of total peak against a 460.5 MiB resident + result. Whole-weight dequantization instead peaks at 3702 MiB for + the same module, which is why the fill is chunked. + """ + rows = min(QuantizedLinear.ROWS_PER_CHUNK, module.out_features) + in_features = module.in_features + itemsize = module.scales.element_size() + groups = -(-in_features // module.group_size) + chain = 4 * rows * in_features * 4 + scale_terms = 2 * rows * groups * (itemsize + 4) + chunk_result = rows * in_features * itemsize + return chain + scale_terms + chunk_result + + +_UNSET: Any = object() + + +class WeightCachePolicy: + """Capped, exact dequantized-weight cache: which ``QuantizedLinear`` + modules keep their dequantized weight. + + Candidates are registered by the loader in model order; ``finalize`` + admits them first-fit until ``cap_bytes`` is used, sets + ``module.cache_weight`` and records the decision. Re-finalizing with a + smaller cap (the engine does this once the Task 4 budget resolves) + releases weights the new selection no longer covers, so nothing stale + survives. ``None`` / ``0`` means off: no module is admitted. + """ + + def __init__(self, cap_bytes: int | None = None): + self.cap_bytes = cap_bytes + self._candidates: list[tuple[str, Any, int]] = [] + self._summary: dict | None = None + + @classmethod + def from_env(cls) -> "WeightCachePolicy": + return cls(_read_weight_cache_cap()) + + @property + def enabled(self) -> bool: + return bool(self.cap_bytes) + + def begin(self) -> None: + """A new model load starts: forget the previous model's modules.""" + self._candidates = [] + self._summary = None + + def register(self, path: str, module) -> None: + self._candidates.append((path, module, dequantized_bytes(module))) + + def candidate_bytes_total(self) -> int: + return sum(nbytes for _, _, nbytes in self._candidates) + + def max_fill_transient_bytes(self, admitted_only: bool = True) -> int: + """Largest build transient among the modules that will be filled + (zero when none is admitted). The fills happen lazily at the first + forward, AFTER the budget's observation point, so this is what the + budget must still have room for.""" + chosen = [m for _, m, _ in self._candidates + if not admitted_only or getattr(m, "cache_weight", False)] + return max((fill_transient_bytes(m) for m in chosen), default=0) + + def finalize(self, cap_bytes=_UNSET) -> dict: + cap = self.cap_bytes if cap_bytes is _UNSET else cap_bytes + enabled = bool(cap) + admitted, skipped, used = [], [], 0 + for path, module, nbytes in self._candidates: + # Each module must fit with its OWN build transient: the fills + # are sequential, so only one is ever live, but a module whose + # fill does not fit would OOM at the first forward even though + # its resident bytes fit the cap. + fill = fill_transient_bytes(module) + if enabled and used + nbytes + fill <= cap: + module.cache_weight = True + used += nbytes + admitted.append({"path": path, "bytes": nbytes, + "fill_transient_bytes": fill}) + continue + if getattr(module, "cache_weight", False): + module._dequantized = None # dropped: release, not stale + module.cache_weight = False + skipped.append({ + "path": path, "bytes": nbytes, + "reason": ("policy off" if not enabled else + f"does not fit: {used + nbytes + fill:,} " + f"(incl. {fill:,} build transient) > cap {cap:,}"), + }) + self._summary = { + "enabled": enabled, + "requested_cap_bytes": (self.cap_bytes if self.cap_bytes is not None + else cap), + "effective_cap_bytes": cap, + "admitted_bytes": used, + "max_fill_transient_bytes": self.max_fill_transient_bytes(), + "admitted": admitted, + "skipped": skipped, + "candidate_bytes_total": self.candidate_bytes_total(), + "candidate_count": len(self._candidates), + "finalized": True, + } + return self._summary + + def summary(self) -> dict: + if self._summary is not None: + return self._summary + return { + "enabled": self.enabled, + "requested_cap_bytes": self.cap_bytes, + "effective_cap_bytes": None, + "admitted_bytes": 0, + "max_fill_transient_bytes": 0, + "admitted": [], + "skipped": [], + "candidate_bytes_total": self.candidate_bytes_total(), + "candidate_count": len(self._candidates), + "finalized": False, + } + + +#: Process-wide policy, filled by ``io.load_model`` (one model per process). +WEIGHT_CACHE = WeightCachePolicy.from_env() + class QuantizedLinear(_tnn.Module): """``mlx.nn.QuantizedLinear`` layout (``weight`` packed uint32, ``scales``, ``biases``) dequantized on the fly -- the 4-bit payload - stays resident, not a bf16 copy (unless ``CACHE_DEQUANTIZED``).""" + stays resident, not a bf16 copy (unless ``CACHE_DEQUANTIZED`` or the + ``WeightCachePolicy`` admitted this module).""" + + weight: torch.Tensor # buffers (declared for type checkers) + scales: torch.Tensor + biases: torch.Tensor def __init__(self, weight, scales, biases, in_features, bias=None): super().__init__() @@ -85,24 +271,83 @@ def __init__(self, weight, scales, biases, in_features, bias=None): self.register_buffer("biases", biases) self.bias = None if bias is None else _tnn.Parameter(bias, False) self._dequantized = None # (dtype, weight), see forward + #: set by WeightCachePolicy.finalize; priced at the scales' itemsize + self.cache_weight = False + self._priced_itemsize = scales.element_size() + self._int4pack = None # (tiled weight, scales_and_zeros) + #: falsy while usable; the reason string once the kernel proved + #: unusable on this device (never retried afterwards) + self._int4pack_disabled = "" ROWS_PER_CHUNK = 4096 + def int4pack_bytes(self) -> int: + """Extra resident bytes of the repacked int4 payload (0 until the + kernel path has been used).""" + if self._int4pack is None: + return 0 + return _i4.int4pack_extra_bytes(self.out_features, self.in_features, + self.group_size) + + def _int4pack_eligible(self, x) -> bool: + if self._int4pack_disabled or x.dtype is not torch.bfloat16 \ + or x.device.type != "cuda": + return False + if self._int4pack is not None: + return True + if self.weight.device.type != "cuda" or x.shape[-1] != self.in_features: + return False + ok, _reason = _i4.is_supported( + in_features=self.in_features, out_features=self.out_features, + group_size=self.group_size, bits=self.bits, + device_type=x.device.type, x_dtype=x.dtype) + if not ok: + return False + # Shapes and dtypes are not enough: prove the device has an image + # for the kernel before committing to it (once per module). + usable, reason = _i4.probe(x.device) + if not usable: + self._int4pack_disabled = reason + return False + return True + def forward(self, x): """Dequantize ``ROWS_PER_CHUNK`` output rows at a time: the full float weight of a large projection is never materialized (edge0-8b's lm_head alone would be ~1 GB per call). Same arithmetic per element as dequantizing everything first.""" from edge0.backends.cuda.quant import _dequantize - if CACHE_DEQUANTIZED: - # Same weight every call: dequantize once. The chunking below - # exists to bound the transient, which a kept weight makes moot. + if CACHE_DEQUANTIZED or (self.cache_weight + and x.element_size() == self._priced_itemsize): + # Same weight every call: dequantize once. Fill it CHUNKED + # into a preallocated buffer -- element for element the same + # arithmetic as the on-the-fly path below, but bounding the + # build transient to one chunk: whole-weight dequantization of + # edge0-8b's lm_head peaks at 3.7 GB for a 0.46 GB result, + # which does not fit an 8 GB Orin's headroom + # (``fill_transient_bytes`` prices what this costs instead). if self._dequantized is None or self._dequantized[0] != x.dtype: - w = _dequantize(self.weight, self.scales, self.biases, - self.group_size, self.bits).to(x.dtype) + w = torch.empty(self.out_features, self.in_features, + dtype=x.dtype, device=self.weight.device) + for r in range(0, self.out_features, self.ROWS_PER_CHUNK): + sl = slice(r, r + self.ROWS_PER_CHUNK) + w[sl] = _dequantize(self.weight[sl], self.scales[sl], + self.biases[sl], self.group_size, + self.bits).to(x.dtype) self._dequantized = (x.dtype, w) b = None if self.bias is None else self.bias.to(x.dtype) return F.linear(x, self._dequantized[1], b) + if _i4.INT4PACK and self._int4pack_eligible(x): + try: + if self._int4pack is None: + self._int4pack = _i4.repack(self.weight, self.scales, + self.biases, self.group_size) + tiled, saz = self._int4pack + return _i4.int4pack_linear(x, tiled, self.group_size, saz, + bias=self.bias) + except RuntimeError as exc: # no image, or a shape the kernel + self._int4pack = None # rejects at runtime: fall through + self._int4pack_disabled = f"{type(exc).__name__}: {exc}" outs = [] for r in range(0, self.out_features, self.ROWS_PER_CHUNK): sl = slice(r, r + self.ROWS_PER_CHUNK) @@ -118,6 +363,10 @@ class QuantizedEmbedding(_tnn.Module): dequantized. Output dtype: ``dtype``, else the scales' (the checkpoint's).""" + weight: torch.Tensor # buffers (declared for type checkers) + scales: torch.Tensor + biases: torch.Tensor + def __init__(self, weight, scales, biases, embedding_dim, dtype=None): super().__init__() self.num_embeddings = weight.shape[0] diff --git a/src/edge0/backends/cuda/quant.py b/src/edge0/backends/cuda/quant.py index 3201c4e..5f65bf9 100644 --- a/src/edge0/backends/cuda/quant.py +++ b/src/edge0/backends/cuda/quant.py @@ -16,14 +16,70 @@ * ``sorted_indices`` is a kernel hint in MLX; it never changes the values. Ignored here. -Not fast: every call dequantizes the gathered experts in full. The point -is a correct, testable reference for the streaming path. +Two implementations share those semantics: + +* the reference loop: one distinct expert at a time, each dequantized in + full to float32 and multiplied -- correct, testable, and the one every + other path is compared against. ``EDGE0_QMM_BATCHED=0`` selects it; +* the batched path (Task 6, **default on this backend**): the same + float32 arithmetic with ONE dequantization over the distinct experts + and ONE batched matmul over a padded per-expert slab, instead of a + python loop of ``len(unique)`` dequantize+matmul pairs. Only affine + 4-bit takes it; 2/8-bit layouts and any call whose priced transient + exceeds ``EDGE0_QMM_BATCHED_MAX_BYTES`` (default 256 MiB) fall through + to the reference loop. The transient is bounded and priced + (``batched_transient_bytes``) so the Task 4 budget deducts it. + +The default was flipped on the Orin's Task 6 acceptance evidence, which +is what Gate D asks for before a path becomes default: measured benefit +on the target (decode +21%, three independent launches) plus the +pre-registered 32-token reference check PASSING there (32/32 token +choices identical, max |delta logit| 0.777 against the registered bound +of 1.0) -- and a second GPU generation agreeing in direction (+59% on an +RTX PRO 6000). The arithmetic is the same expression per element; it +lands within one bfloat16 ulp of the reference loop at the real decode +shapes, which is why the check passes rather than merely not crashing. """ from __future__ import annotations +import os + import torch +BATCHED_ENV = "EDGE0_QMM_BATCHED" +BATCHED_MAX_BYTES_ENV = "EDGE0_QMM_BATCHED_MAX_BYTES" +DEFAULT_BATCHED_MAX_BYTES = 256 << 20 + + +def _read_batched_gate() -> bool: + """On by default (see the module docstring for the evidence that + flipped it); ``EDGE0_QMM_BATCHED=0`` selects the reference loop. + Any other value is treated as unset, so a typo cannot silently + change the path.""" + return os.environ.get(BATCHED_ENV, "").strip() != "0" + + +def _read_batched_max_bytes() -> int: + raw = os.environ.get(BATCHED_MAX_BYTES_ENV, "").strip() + if not raw: + return DEFAULT_BATCHED_MAX_BYTES + try: + value = int(raw) + except ValueError: + raise ValueError( + f"{BATCHED_MAX_BYTES_ENV} must be an integer byte count, got " + f"{raw!r}") from None + if value < 0: + raise ValueError(f"{BATCHED_MAX_BYTES_ENV} must be >= 0, got {value}") + return value + + +#: Read once at import, like ``nn.CACHE_DEQUANTIZED``; tests patch the +#: module attributes. A malformed cap fails the import visibly. +BATCHED = _read_batched_gate() +BATCHED_MAX_BYTES = _read_batched_max_bytes() + def _dequantize(w: torch.Tensor, scales: torch.Tensor, biases: torch.Tensor, group_size: int, bits: int) -> torch.Tensor: @@ -42,14 +98,38 @@ def _dequantize(w: torch.Tensor, scales: torch.Tensor, biases: torch.Tensor, return deq.reshape(codes.shape) -def gather_qmm(x, w, scales, biases, rhs_indices, transpose=True, - group_size=64, bits=4, mode="affine", - sorted_indices=False): - """Quantized matmul over a gathered subset of experts (``mx.gather_qmm``).""" - if mode != "affine" or bits not in (2, 4, 8): - raise NotImplementedError( - f"reference gather_qmm covers affine 2/4/8-bit only " - f"(got mode={mode!r}, bits={bits!r})") +def batched_transient_bytes(n_unique: int, rows_max: int, m: int, d_in: int, + n_out: int, bits: int = 4, group_size: int = 64, + scale_itemsize: int = 2) -> int: + """Upper bound on the extra bytes ``_gather_qmm_batched`` allocates for + one call (freed on return): the gathered packed words, the int32 code + tensors and float32 dequantized copies the ``_dequantize`` chain keeps + live (four float32-sized tensors of ``n_unique * n_out * d_in``), the + gathered scales/biases and their float32 casts, and the padded + input/output slabs. Pure arithmetic, checked against allocator peaks + in ``tests/test_quant_paths.py``.""" + for name, value in (("n_unique", n_unique), ("rows_max", rows_max), + ("m", m), ("d_in", d_in), ("n_out", n_out), + ("group_size", group_size), + ("scale_itemsize", scale_itemsize)): + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise ValueError(f"{name} must be a positive int, got {value!r}") + per_expert = n_out * d_in + packed = n_unique * per_expert * bits // 8 + chain = 4 * n_unique * per_expert * 4 + groups = -(-d_in // group_size) + scale_terms = 2 * n_unique * n_out * groups * (scale_itemsize + 4) + slabs = 2 * n_unique * rows_max * m * (d_in + n_out) * 4 + return packed + chain + scale_terms + slabs + + +def _gather_qmm_reference(x, w, scales, biases, rhs_indices, transpose, + group_size, bits): + """One distinct expert at a time (the reference every other path is + compared against). Dequantizing a copy per (token, expert) pair + peaked at several GB per MoE layer during prefill, hence per expert. + An index past the last expert raises here; in MLX it silently reads + out of bounds.""" idx = rhs_indices.to(torch.long) bshape = torch.broadcast_shapes(x.shape[:-2], idx.shape) M = x.shape[-2] @@ -57,9 +137,6 @@ def gather_qmm(x, w, scales, biases, rhs_indices, transpose=True, flat = idx.expand(bshape).reshape(-1) n_out = w.shape[-2] if transpose else w.shape[-1] * (32 // bits) out = torch.empty(flat.numel(), M, n_out, dtype=torch.float32, device=x.device) - # One distinct expert at a time: dequantizing a copy per (token, expert) - # pair peaked at several GB per MoE layer during prefill. An index past - # the last expert raises here; in MLX it silently reads out of bounds. for e in torch.unique(flat).tolist(): rows = (flat == e).nonzero().squeeze(-1) deq = _dequantize(w[e], scales[e], biases[e], group_size, bits) @@ -67,6 +144,77 @@ def gather_qmm(x, w, scales, biases, rhs_indices, transpose=True, return out.reshape(*bshape, M, n_out).to(x.dtype) +def _gather_qmm_batched(x, w, scales, biases, rhs_indices, transpose, + group_size, bits): + """Same result as the reference loop with one dequantization and one + batched matmul; ``None`` when the priced transient exceeds the cap + (the caller then runs the reference loop).""" + idx = rhs_indices.to(torch.long) + bshape = torch.broadcast_shapes(x.shape[:-2], idx.shape) + M = x.shape[-2] + D = x.shape[-1] + n_out = w.shape[-2] if transpose else w.shape[-1] * (32 // bits) + flat = idx.expand(bshape).reshape(-1) + R = flat.numel() + if R == 0: + return torch.empty(*bshape, M, n_out, dtype=x.dtype, device=x.device) + uniq, inverse = torch.unique(flat, return_inverse=True) + n_experts = w.shape[0] + if int(uniq.max()) >= n_experts: + raise IndexError( + f"expert index {int(uniq.max())} out of range for {n_experts} " + f"experts") + U = uniq.numel() + counts = torch.bincount(inverse, minlength=U) + rows_max = int(counts.max()) + priced = batched_transient_bytes(U, rows_max, M, D, n_out, bits, + group_size, scales.element_size()) + if priced > BATCHED_MAX_BYTES: + return None + + xf = x.to(torch.float32).expand(*bshape, M, D).reshape(-1, M, D) + # One dequantization over the distinct experts. Index the int32 view: + # CUDA has no index kernel for uint32 (see nn.QuantizedEmbedding), and + # _dequantize reads the words as int32 anyway. + deq = _dequantize(w.view(torch.int32)[uniq], scales[uniq], biases[uniq], + group_size, bits) # [U, out, in] + wmat = deq.transpose(1, 2) if transpose else deq # [U, in, out] + order = torch.argsort(inverse, stable=True) # rows by expert + grouped = inverse[order] + if rows_max == 1: + # decode: every distinct expert has exactly one row, no padding + y = torch.matmul(xf[order], wmat) # [U, M, out] + out = torch.empty(R, M, n_out, dtype=torch.float32, device=x.device) + out[order] = y + return out.reshape(*bshape, M, n_out).to(x.dtype) + starts = torch.cumsum(counts, 0) - counts # [U] + slot = torch.arange(R, device=flat.device) - starts[grouped] + slab = xf.new_zeros(U, rows_max, M, D) + slab[grouped, slot] = xf[order] + y = torch.matmul(slab.reshape(U, rows_max * M, D), wmat) + y = y.reshape(U, rows_max, M, n_out)[grouped, slot] # [R, M, out] + out = torch.empty(R, M, n_out, dtype=torch.float32, device=x.device) + out[order] = y + return out.reshape(*bshape, M, n_out).to(x.dtype) + + +def gather_qmm(x, w, scales, biases, rhs_indices, transpose=True, + group_size=64, bits=4, mode="affine", + sorted_indices=False): + """Quantized matmul over a gathered subset of experts (``mx.gather_qmm``).""" + if mode != "affine" or bits not in (2, 4, 8): + raise NotImplementedError( + f"reference gather_qmm covers affine 2/4/8-bit only " + f"(got mode={mode!r}, bits={bits!r})") + if BATCHED and bits == 4: + out = _gather_qmm_batched(x, w, scales, biases, rhs_indices, + transpose, group_size, bits) + if out is not None: + return out + return _gather_qmm_reference(x, w, scales, biases, rhs_indices, + transpose, group_size, bits) + + def gather_sort(x, indices): """Sort token rows by expert id for ``gather_qmm(sorted_indices=True)``: returns ``(x_sorted, indices_sorted, inv_order)``, as mlx-lm's diff --git a/src/edge0/engine/ling.py b/src/edge0/engine/ling.py index fb24049..5a65cde 100644 --- a/src/edge0/engine/ling.py +++ b/src/edge0/engine/ling.py @@ -29,10 +29,179 @@ ) from edge0.prerouter.install import install_prerouter from edge0.prerouter.stager import LingPrerouterStager +from edge0.streaming import budget as _budget +from edge0.streaming.cache import PrefetchBuffer, SharedExpertCache from edge0.streaming.install import install_streaming_experts from edge0.streaming.mmap import SafetensorsMmap +def _apply_env_profile(cfg, opts, spec): + """Task 5 opt-in profile knobs, applied to the REQUESTED options + before the budget prices them (the budget may still lower them). + + ``EDGE0_CACHE_SLOTS=``: request a different shared-LRU size. + The tested default (64) is smaller than the per-step working set on + this tier (23 MoE layers x K=8 = 184 bundles), so it never hits; + 512 measured a 65% hit rate on the Orin baseline (evidence in the + plan's Task 5 status). + + ``EDGE0_PREDICT_PREFETCH=1``: enable + ``LayerOptions.predict_prefetch``; the prefetch buffer is raised to + hold one full predicted step (owners x K) and the per-layer + in-flight bound defaults to K, both priced by the budget. + """ + from dataclasses import replace as _dc_replace + env_slots = os.environ.get("EDGE0_CACHE_SLOTS", "").strip() + if env_slots: + try: + slots = int(env_slots) + except ValueError: + raise ValueError( + f"EDGE0_CACHE_SLOTS must be an integer, got " + f"{env_slots!r}") from None + if slots <= 0: + raise ValueError( + f"EDGE0_CACHE_SLOTS must be positive, got {slots} " + "(0 would disable eviction, not caching)") + opts = _dc_replace(opts, cache_slots=slots) + if os.environ.get("EDGE0_PREDICT_PREFETCH", "") == "1": + opts = _dc_replace(opts, predict_prefetch=True) + if opts.predict_prefetch: + owners_n = len(getattr(cfg.prerouter, "owners", ()) or ()) + per_layer = opts.top_k or spec.top_k + opts = _dc_replace( + opts, + prefetch_cap=max(opts.prefetch_cap, owners_n * per_layer), + max_inflight=opts.max_inflight or per_layer) + return opts + + +def _torch_quant_modules(): + """``(nn, quant)`` of the torch backend, ``(None, None)`` on MLX.""" + from edge0.backends import backend + if backend.name != "cuda": + return None, None + from edge0.backends.cuda import nn as cnn + from edge0.backends.cuda import quant as cq + return cnn, cq + + +def _resolve_memory_budget(cfg, opts, shards): + """Resolve the Task 4 memory budget when ``EDGE0_MEMORY_BUDGET`` asks + for one; return ``None`` when it is off (the default — behavior is + then exactly the pre-budget code path). + + ``EDGE0_MEMORY_BUDGET`` accepts ``auto`` (observe MemAvailable and + any tighter address-space rlimit at this point, right before cache + construction) or an explicit byte count. The declared context comes + from ``EDGE0_BUDGET_CONTEXT`` (tokens, default 1024): contexts + beyond the declaration are a different workload and must be + re-declared, not discovered as an OOM. + """ + req = os.environ.get("EDGE0_MEMORY_BUDGET", "").strip() + if not req: + return None + if req == "auto": + import psutil + available = int(psutil.virtual_memory().available) + limit = None + try: + import resource + soft, _hard = resource.getrlimit(resource.RLIMIT_AS) + if soft not in (-1, resource.RLIM_INFINITY): + limit = int(soft) + except (ImportError, OSError, ValueError): + limit = None + obs = _budget.Observation(available_bytes=available, + process_limit_bytes=limit) + else: + try: + obs = _budget.Observation(available_bytes=int(req)) + except ValueError: + raise _budget.BudgetError( + f"EDGE0_MEMORY_BUDGET must be 'auto' or a byte count, " + f"got {req!r}") from None + spec = cfg.moe_spec + # Layer 0 is dense on this tier; layer 1 carries the expert tensors. + entries = {} + for shard in shards: + entries.update(shard.entries) + bundle = _budget.bundle_bytes_from_entries( + entries, spec.key_template.format(layer=1), spec.num_experts) + if not cfg.kv_bytes_per_token: + raise _budget.BudgetError( + f"{cfg.name}: kv_bytes_per_token is not measured for this " + "tier; a memory budget cannot price the declared context") + context = int(os.environ.get("EDGE0_BUDGET_CONTEXT", "1024")) + top_k = opts.top_k or spec.top_k + # In-flight transient: per-layer bound x layers that can be + # producing at once. With predict_prefetch every owner layer may + # hold a step's predictions in flight; otherwise only the thread + # pool's builds are outstanding. + per_layer_inflight = opts.max_inflight or max(1, opts.prefetch_threads) + if opts.predict_prefetch: + producing_layers = len(getattr(cfg.prerouter, "owners", ()) or ()) + inflight_total = max(1, per_layer_inflight * producing_layers) + else: + inflight_total = per_layer_inflight + # The dense dequantized-weight cache: the full cache (measured 4.1 GB + # in float32 on the GB10; the loader now prices the actual bf16 total + # of this tier's QuantizedLinear modules) is only ever vetoed by the + # budget, never enabled; the capped cache (Task 6) is shrunk to the + # headroom. The batched expert gather's per-call transient cap is + # deducted once when that path is on. + weight_cache_requested = ( + os.environ.get("EDGE0_TORCH_WEIGHT_CACHE", "") == "1") + cnn, cq = _torch_quant_modules() + candidate_total = cnn.WEIGHT_CACHE.candidate_bytes_total() if cnn else 0 + request = _budget.dense_cache_request( + full_requested=weight_cache_requested, + cap_bytes=cnn.WEIGHT_CACHE.cap_bytes if cnn else None, + candidate_total=candidate_total, + max_fill_transient_bytes=( + cnn.WEIGHT_CACHE.max_fill_transient_bytes(admitted_only=False) + if cnn else 0)) + kernel_transient = int(cq.BATCHED_MAX_BYTES) if (cq and cq.BATCHED) else 0 + resolved = _budget.resolve_budget( + obs, + _budget.ExpertFootprint( + bundle_bytes=bundle, num_experts=spec.num_experts, + num_moe_layers=23), + _budget.WorkloadDecl(max_context_tokens=context, + kv_bytes_per_token=cfg.kv_bytes_per_token), + requested_cache_slots=opts.cache_slots, + requested_prefetch_cap=opts.prefetch_cap, + full_layer_prefill=opts.full_layer_prefill, + inflight_builds=inflight_total, + min_cache_slots=2 * top_k, + weight_cache_bytes=(request["priced_bytes"] + if request["mode"] != "off" else None), + reserves=_budget.Reserves(kernel_transient_bytes=kernel_transient), + ) + if request["mode"] == "full" and not resolved.weight_cache_permitted: + raise _budget.BudgetError( + "EDGE0_TORCH_WEIGHT_CACHE=1 rejected: the dequantized dense " + f"weights ({request['resident_bytes']:,} bytes resident plus a " + f"{request['fill_transient_bytes']:,} byte build transient) do " + f"not fit the resolved headroom of {resolved.headroom_bytes:,} " + f"bytes " + f"(budget: {resolved.as_dict()})") + if request["mode"] == "capped" and cnn is not None: + from dataclasses import replace as _dc_replace + # The policy admits each module with its own build transient, so + # the effective cap is the headroom itself. + effective = _budget.effective_capped_bytes( + cap_bytes=request["cap_bytes"], + headroom_bytes=resolved.headroom_bytes) + summary = cnn.WEIGHT_CACHE.finalize(effective) + resolved = _dc_replace(resolved, notes=resolved.notes + ( + f"capped weight cache: requested {request['cap_bytes']:,} B, " + f"effective {effective:,} B, admitted " + f"{summary['admitted_bytes']:,} B in " + f"{len(summary['admitted'])} modules",)) + return resolved + + def _get_model_classes(config): """load_model class hook: serve the active backend's port of the bailing backbone (imported here so the module itself imports without MLX).""" @@ -65,8 +234,33 @@ def load_installed(model_dir: str, cfg): spec = cfg.moe_spec opts = cfg.options n_layers = model_config["num_hidden_layers"] + opts = _apply_env_profile(cfg, opts, spec) + resolved_budget = _resolve_memory_budget(cfg, opts, shards) + shared_cache = prefetch_buffer = None + if resolved_budget is not None: + from dataclasses import replace as _dc_replace + # cache/prefetch capacities come from the budget; the in-flight + # bound stays PER-LAYER (the budget priced the total across + # producing layers, see _resolve_memory_budget). + opts = _dc_replace( + opts, cache_slots=resolved_budget.cache_slots, + prefetch_cap=resolved_budget.prefetch_cap, + max_inflight=(opts.max_inflight + or max(1, opts.prefetch_threads))) + shared_cache = SharedExpertCache(resolved_budget.cache_slots) + prefetch_buffer = PrefetchBuffer( + resolved_budget.prefetch_cap, + max_cap=resolved_budget.max_prefetch_cap) + print(f"[edge0-8b] memory budget: usable=" + f"{resolved_budget.usable_bytes:,}B cache_slots=" + f"{resolved_budget.cache_slots} prefetch_cap=" + f"{resolved_budget.prefetch_cap} headroom=" + f"{resolved_budget.headroom_bytes:,}B" + + (f" notes={list(resolved_budget.notes)}" + if resolved_budget.notes else ""), flush=True) installed = install_streaming_experts( - model, shards, spec, options=opts, num_layers=n_layers) + model, shards, spec, options=opts, num_layers=n_layers, + shared_cache=shared_cache, prefetch_buffer=prefetch_buffer) all_stream = {li: t for li, t in enumerate(installed) if t is not None} stream_layers = {li: t for li, t in all_stream.items() if t._staged_mode} @@ -78,16 +272,22 @@ def load_installed(model_dir: str, cfg): if cfg.prerouter and cfg.prerouter.weights_file: pg_state, heads = install_prerouter( model=model, spec=spec, pspec=cfg.prerouter, n_layers=n_layers) + prefetch_layers = None + if opts.predict_prefetch: + prefetch_layers = {li: t for li, t in all_stream.items() + if not t._staged_mode} pg_stager = LingPrerouterStager( model=model, spec=spec, pspec=cfg.prerouter, state=pg_state, stream_layers=stream_layers, - top_k=cfg.prerouter_top_k) + top_k=cfg.prerouter_top_k, + prefetch_layers=prefetch_layers) print(f"[edge0-8b] prerouter installed: {len(heads)} heads, " f"start={cfg.prerouter.start_layer}, " f"K={cfg.prerouter_top_k}", flush=True) installs = dict(all_stream_layers=all_stream, stream_layers=stream_layers, - pg_state=pg_state, pg_stager=pg_stager) + pg_state=pg_state, pg_stager=pg_stager, + memory_budget=resolved_budget) return model, model_config, shards, installs @@ -118,6 +318,9 @@ def _build(self): self._stream_layers = inst["stream_layers"] self._pg_state = inst["pg_state"] self._pg_stager = inst["pg_stager"] + #: Resolved Task 4 memory budget (None when EDGE0_MEMORY_BUDGET + #: is off); the bench report's caches group records it. + self.memory_budget = inst["memory_budget"] opts = cfg.options if self._tok is None: try: diff --git a/src/edge0/models/base.py b/src/edge0/models/base.py index 09a14b9..f0c4fb7 100644 --- a/src/edge0/models/base.py +++ b/src/edge0/models/base.py @@ -62,6 +62,11 @@ class ModelConfig: # acceptance profile (measured on the production Mac) target_tok_s: float = 0.0 peak_active_mem_mb: float = 0.0 + #: Measured KV-cache growth per context token (bytes), used by the + #: memory-budgeted profile (EDGE0_MEMORY_BUDGET) to price the + #: declared context. 0 = not measured for this tier; a requested + #: budget then fails loudly instead of guessing. + kv_bytes_per_token: int = 0 @classmethod def from_pretrained(cls, model_dir: str | None = None, **overrides): diff --git a/src/edge0/models/edge0_8b/__init__.py b/src/edge0/models/edge0_8b/__init__.py index 7c8f4db..fdd3b4b 100644 --- a/src/edge0/models/edge0_8b/__init__.py +++ b/src/edge0/models/edge0_8b/__init__.py @@ -69,6 +69,19 @@ def _defaults(cls, model_dir: str) -> "Ling8BConfig": port=8083, target_tok_s=33.0, peak_active_mem_mb=1400.0, + # MEASURED on the torch backend (Orin, Task 7): the marginal + # cost between 1536 and 2048 tokens is 561,320 B/token + # (548 KiB), recorded in the evidence set; 600,000 keeps ~7% + # margin. Growth is sublinear below that — most layers of + # this hybrid are KDA, whose state is fixed-size, so only the + # MLA layers' KV grows with context — which is why the + # marginal slope at the top of the range is the right + # constant to extrapolate a declared context from. + # The previous 1,100,000 came from the README's Apple/MLX + # figure and over-reserved by ~2x, rejecting contexts that + # fit. Over-reserving is safe but shrinks the envelope, so + # it is corrected from the measurement, not from the doc. + kv_bytes_per_token=600_000, ) diff --git a/src/edge0/prerouter/stager.py b/src/edge0/prerouter/stager.py index e47b65f..56d3955 100644 --- a/src/edge0/prerouter/stager.py +++ b/src/edge0/prerouter/stager.py @@ -65,12 +65,17 @@ def __init__( state: PrerouterState, stream_layers: dict[int, StreamingSwitchGLU], top_k: int, + prefetch_layers: dict[int, StreamingSwitchGLU] | None = None, ): self.model = model self.spec = spec self.pspec = pspec self.state = state self.stream_layers = stream_layers + #: Non-staged layers whose predicted expert sets are routed to + #: ``prefetch()`` (LayerOptions.predict_prefetch). Scheduling + #: only: the exact consumption path is untouched. + self.prefetch_layers = prefetch_layers or {} self.top_k = top_k self.num_experts = spec.num_experts self.records: list | None = None @@ -190,10 +195,12 @@ class LingPrerouterStager(PrerouterStager): def __init__(self, *, model, spec: MoESpec, pspec: PrerouterSpec, state: PrerouterState, - stream_layers: dict[int, StreamingSwitchGLU], top_k: int): + stream_layers: dict[int, StreamingSwitchGLU], top_k: int, + prefetch_layers: dict[int, StreamingSwitchGLU] + | None = None): super().__init__(model=model, spec=spec, pspec=pspec, state=state, stream_layers=stream_layers, - top_k=top_k) + top_k=top_k, prefetch_layers=prefetch_layers) self.pg_cache: dict[int, core.array] = {} def _head_of(self, owner: int): @@ -232,6 +239,15 @@ def _store(self, consumer: int, logits_i, inds_i, scores_i, exp = self.stream_layers.get(consumer) if exp is not None: exp.stage_experts(experts) + return + # Non-staged consumer (prod profiles): overlap the predicted + # next-step builds with the current forward. prefetch() is + # bounded (max_inflight / prefetch buffer caps) and consumption + # stays on the exact _get_bundles path — a wrong or late + # prediction is just a demand load, never a zero row. + pre = self.prefetch_layers.get(consumer) + if pre is not None: + pre.prefetch(experts) def _note(self, owner: int, cur_oh) -> None: block = self.spec.block_of(self.model, owner) diff --git a/src/edge0/streaming/__init__.py b/src/edge0/streaming/__init__.py index 8e412b0..5ea6592 100644 --- a/src/edge0/streaming/__init__.py +++ b/src/edge0/streaming/__init__.py @@ -1,9 +1,14 @@ """Streaming expert offload: mmap stores, shared LRU caches, staged -decode slots, whole-layer prefill and hot-expert pins.""" +decode slots, whole-layer prefill and hot-expert pins. + +``StreamingSwitchGLU`` and ``install_streaming_experts`` bind to the +active backend at import; they are exposed lazily (PEP 562) so the +backend-free modules (``cache``, ``options``, ``budget``, ``mmap``) +stay importable on hosts with no backend installed — the Task 4 budget +tests run everywhere. +""" from edge0.streaming.cache import PrefetchBuffer, SharedExpertCache -from edge0.streaming.install import install_streaming_experts -from edge0.streaming.layer import StreamingSwitchGLU from edge0.streaming.mmap import SafetensorsMmap from edge0.streaming.options import LayerOptions @@ -15,3 +20,19 @@ "LayerOptions", "install_streaming_experts", ] + +_LAZY = { + "StreamingSwitchGLU": ("edge0.streaming.layer", "StreamingSwitchGLU"), + "install_streaming_experts": ("edge0.streaming.install", + "install_streaming_experts"), +} + + +def __getattr__(name): + try: + module_name, attr = _LAZY[name] + except KeyError: + raise AttributeError( + f"module {__name__!r} has no attribute {name!r}") from None + import importlib + return getattr(importlib.import_module(module_name), attr) diff --git a/src/edge0/streaming/budget.py b/src/edge0/streaming/budget.py new file mode 100644 index 0000000..61f3737 --- /dev/null +++ b/src/edge0/streaming/budget.py @@ -0,0 +1,339 @@ +"""Memory-budgeted streaming profile (pure calculation). + +Task 4 of ``docs/plans/jetson-orin-nano.md``: turn an OBSERVED amount of +memory into explicit byte allowances for everything the streaming stack +retains or transiently allocates, and refuse impossible profiles before +inference instead of OOMing during it. + +Everything here is pure: the caller supplies observations (available +RAM, process limits), checkpoint-derived payload sizes and the declared +workload; this module only does arithmetic and policy. No psutil, no +torch, no engine imports — see ``engine/ling.py`` for the integration +that gathers the inputs. + +Accounting rules (Gate C): + +* Start from observed available RAM, tightened by any process limit. +* Deduct, each exactly once: an OS/application growth reserve; KV cache + at the DECLARED maximum context; the prefill transient (whole-layer + load-drop materializes every expert of one layer at once when + ``full_layer_prefill`` is on); in-flight expert builds (bounded by + the resolved ``max_inflight``, not the thread count alone); pinned + staging (Task 5 — zero today); the batched expert gather's per-call + transient cap (Task 6 — zero unless ``EDGE0_QMM_BATCHED=1``); and an + allocator/driver allowance (fraction of the observed total, applied + once). +* What remains is the retained-cache allowance, split between the + shared LRU and the prefetch buffer in payload bytes computed from + checkpoint metadata — slot counts alone are meaningless. +* ``SharedExpertCache(0)`` and ``PrefetchBuffer(0)`` DISABLE eviction + (unbounded). A calculated zero is therefore never passed through: + any allowance that cannot fit the minimum working set raises + ``BudgetError`` with the itemized deductions. +* The budget only ever LOWERS the profile's requested capacities. A + budget roomier than the tested profile keeps the tested values. +* Full dequantized-weight caching (``EDGE0_TORCH_WEIGHT_CACHE``) is + reported as permitted only when its measured payload fits in the + allowance that remains AFTER the caches — it is never turned on by + this module, only vetoed. +""" + +from __future__ import annotations + +from dataclasses import dataclass + + +class BudgetError(ValueError): + """The declared workload cannot fit the observed memory.""" + + +def _positive_int(name: str, value) -> int: + if isinstance(value, bool) or not isinstance(value, int): + raise BudgetError(f"{name} must be an int, got {value!r}") + if value <= 0: + raise BudgetError(f"{name} must be positive, got {value}") + return value + + +def _non_negative_int(name: str, value) -> int: + if isinstance(value, bool) or not isinstance(value, int): + raise BudgetError(f"{name} must be an int, got {value!r}") + if value < 0: + raise BudgetError(f"{name} must be >= 0, got {value}") + return value + + +@dataclass(frozen=True) +class Observation: + """What was actually observed on the target, at the observation + point (immediately before engine construction — allocations already + resident at that point are outside ``available_bytes`` by + definition and must not be deducted again).""" + + available_bytes: int + #: A tighter limit than system RAM when one exists (cgroup, + #: rlimit); None means no separate limit was observed. + process_limit_bytes: int | None = None + + def usable(self) -> int: + avail = _positive_int("available_bytes", self.available_bytes) + if self.process_limit_bytes is None: + return avail + return min(avail, + _positive_int("process_limit_bytes", + self.process_limit_bytes)) + + +@dataclass(frozen=True) +class ExpertFootprint: + """Per-expert payload from checkpoint metadata (see + ``bundle_bytes_from_entries``).""" + + bundle_bytes: int + num_experts: int + num_moe_layers: int + + +@dataclass(frozen=True) +class WorkloadDecl: + """The declared operating envelope. KV for ``max_context_tokens`` + is deducted up front: a context the budget did not declare is a + different workload and must be re-resolved, not discovered OOM.""" + + max_context_tokens: int + kv_bytes_per_token: int + + +@dataclass(frozen=True) +class Reserves: + os_growth_bytes: int = 512 * 1024 * 1024 + #: Allocator + driver allowance as a fraction of the usable + #: observation, deducted once. + allocator_fraction: float = 0.10 + #: Task 5's pinned staging pool; zero until it exists. + pinned_staging_bytes: int = 0 + #: Task 6's batched expert gather: the per-call transient cap + #: (``EDGE0_QMM_BATCHED_MAX_BYTES``) when that path is on, else zero. + kernel_transient_bytes: int = 0 + + +@dataclass(frozen=True) +class ResolvedBudget: + """The resolved policy, with its arithmetic shown.""" + + usable_bytes: int + deductions: tuple # ((name, bytes), ...) in application order + cache_allowance_bytes: int + cache_slots: int + cache_bytes: int + prefetch_cap: int + prefetch_bytes: int + #: Hard ceiling for ``PrefetchBuffer.set_cap`` growth + #: (``prefetch_all`` raises the cap to num_experts + 32 otherwise). + max_prefetch_cap: int + max_inflight: int + weight_cache_permitted: bool + weight_cache_bytes: int | None + headroom_bytes: int + notes: tuple = () + + def as_dict(self) -> dict: + reasons = {} + if self.weight_cache_bytes is None: + reasons["weight_cache_bytes"] = ( + "no dequantized-weight cache was requested " + "(EDGE0_TORCH_WEIGHT_CACHE unset)") + return { + "unavailable_reasons": reasons, + "usable_bytes": self.usable_bytes, + "deductions": [list(d) for d in self.deductions], + "cache_allowance_bytes": self.cache_allowance_bytes, + "cache_slots": self.cache_slots, + "cache_bytes": self.cache_bytes, + "prefetch_cap": self.prefetch_cap, + "prefetch_bytes": self.prefetch_bytes, + "max_prefetch_cap": self.max_prefetch_cap, + "max_inflight": self.max_inflight, + "weight_cache_permitted": self.weight_cache_permitted, + "weight_cache_bytes": self.weight_cache_bytes, + "headroom_bytes": self.headroom_bytes, + "notes": list(self.notes), + } + + +def dense_cache_request(*, full_requested: bool, cap_bytes, + candidate_total: int, + max_fill_transient_bytes: int = 0) -> dict: + """What the budget prices for the dense dequantized-weight cache + (Task 6). + + The price is the RESIDENT bytes the cache keeps plus the largest + single build transient: the weights are filled lazily at the first + forward, one chunked module at a time, so exactly one transient is + live but it must still fit alongside everything resident. + + ``full`` (``EDGE0_TORCH_WEIGHT_CACHE=1``) prices the loader's measured + total of every ``QuantizedLinear`` (the GB10's 4.1 GB figure only when + no model is loaded yet) and keeps its veto; ``capped`` + (``EDGE0_TORCH_WEIGHT_CACHE_BYTES``) prices ``min(cap, total)`` and is + later shrunk to the headroom instead of vetoed; ``off`` prices nothing. + """ + fill = _non_negative_int("max_fill_transient_bytes", + max_fill_transient_bytes) + if full_requested: + resident = candidate_total or 4_100_000_000 + return {"mode": "full", "priced_bytes": resident + fill, + "resident_bytes": resident, "fill_transient_bytes": fill, + "cap_bytes": None} + if cap_bytes: + resident = min(int(cap_bytes), candidate_total) + return {"mode": "capped", "priced_bytes": resident + fill, + "resident_bytes": resident, "fill_transient_bytes": fill, + "cap_bytes": int(cap_bytes)} + return {"mode": "off", "priced_bytes": 0, "resident_bytes": 0, + "fill_transient_bytes": 0, "cap_bytes": None} + + +def effective_capped_bytes(*, cap_bytes: int, headroom_bytes: int) -> int: + """The capped cache never exceeds the budget's post-cache headroom.""" + return max(0, min(int(cap_bytes), int(headroom_bytes))) + + +def bundle_bytes_from_entries(entries: dict, key_prefix: str, + num_experts: int, + projections=("gate_proj", "up_proj", + "down_proj"), + parts=("weight", "scales", "biases")) -> int: + """Per-expert payload bytes for one layer, from a safetensors header. + + ``entries`` maps tensor name to ``{"size": bytes, ...}`` (the + ``SafetensorsMmap.entries`` shape) or directly to an int size. + Expert tensors are stacked ``[num_experts, ...]``; the per-expert + payload is the summed tensor size divided by the expert count. + """ + _positive_int("num_experts", num_experts) + total = 0 + for proj in projections: + for part in parts: + name = f"{key_prefix}.{proj}.{part}" + if name not in entries: + raise BudgetError( + f"checkpoint metadata is missing {name!r}; cannot " + "compute the expert payload from slot counts alone") + meta = entries[name] + size = meta["size"] if isinstance(meta, dict) else meta + total += _positive_int(f"size of {name}", size) + if total % num_experts: + # Not fatal — round up so the budget never undercounts. + return total // num_experts + 1 + return total // num_experts + + +def resolve_budget( + observation: Observation, + experts: ExpertFootprint, + workload: WorkloadDecl, + *, + requested_cache_slots: int, + requested_prefetch_cap: int, + full_layer_prefill: bool, + inflight_builds: int, + min_cache_slots: int | None = None, + weight_cache_bytes: int | None = None, + reserves: Reserves = Reserves(), +) -> ResolvedBudget: + """Resolve the streaming memory policy for one engine build. + + Raises ``BudgetError`` when the declared workload cannot fit — with + the itemized deductions, so the failure is a review artifact rather + than a mystery OOM. ``min_cache_slots`` defaults to the larger of + 2 and the per-step working set the caller declares; the shared LRU + is never resolved below it and never to zero. + """ + bundle = _positive_int("bundle_bytes", experts.bundle_bytes) + n_experts = _positive_int("num_experts", experts.num_experts) + _positive_int("num_moe_layers", experts.num_moe_layers) + slots_req = _positive_int("requested_cache_slots", requested_cache_slots) + cap_req = _positive_int("requested_prefetch_cap", requested_prefetch_cap) + inflight = _positive_int("inflight_builds", inflight_builds) + min_slots = (2 if min_cache_slots is None + else _positive_int("min_cache_slots", min_cache_slots)) + if weight_cache_bytes is not None: + _non_negative_int("weight_cache_bytes", weight_cache_bytes) + if not 0.0 <= reserves.allocator_fraction < 1.0: + raise BudgetError( + f"allocator_fraction must be in [0, 1), got " + f"{reserves.allocator_fraction!r}") + + usable = observation.usable() + kv_bytes = (_non_negative_int("max_context_tokens", + workload.max_context_tokens) + * _positive_int("kv_bytes_per_token", + workload.kv_bytes_per_token)) + prefill_transient = bundle * n_experts if full_layer_prefill else 0 + inflight_bytes = bundle * inflight + deductions = ( + ("os_growth_reserve", _non_negative_int( + "os_growth_bytes", reserves.os_growth_bytes)), + ("allocator_allowance", int(usable * reserves.allocator_fraction)), + ("kv_cache_at_declared_context", kv_bytes), + ("full_layer_prefill_transient", prefill_transient), + ("inflight_expert_builds", inflight_bytes), + ("pinned_staging", _non_negative_int( + "pinned_staging_bytes", reserves.pinned_staging_bytes)), + ("kernel_transient", _non_negative_int( + "kernel_transient_bytes", reserves.kernel_transient_bytes)), + ) + allowance = usable - sum(d[1] for d in deductions) + if allowance <= 0: + detail = ", ".join(f"{n}={b:,}" for n, b in deductions) + raise BudgetError( + f"declared workload does not fit: usable={usable:,} bytes, " + f"deductions leave {allowance:,} ({detail})") + + # Shared LRU first (correctness path), prefetch from the remainder. + cache_slots = min(slots_req, allowance // bundle) + if cache_slots < min_slots: + raise BudgetError( + f"cache allowance {allowance:,} bytes fits only " + f"{cache_slots} expert bundles of {bundle:,} bytes — below " + f"the minimum working set of {min_slots}; a zero/starved " + f"cache is rejected, not constructed (SharedExpertCache(0) " + f"would disable eviction entirely)") + cache_bytes = cache_slots * bundle + prefetch_cap = min(cap_req, (allowance - cache_bytes) // bundle) + notes = [] + if prefetch_cap < 1: + # An explicit bounded floor, never PrefetchBuffer(0): one slot + # keeps the buffer's eviction semantics intact and costs one + # bundle, which min_cache_slots' margin covers. + prefetch_cap = 1 + notes.append("prefetch_cap floored to 1 (0 disables eviction)") + prefetch_bytes = prefetch_cap * bundle + if cache_slots < slots_req: + notes.append( + f"cache_slots lowered {slots_req} -> {cache_slots} by budget") + if prefetch_cap < cap_req: + notes.append( + f"prefetch_cap lowered {cap_req} -> {prefetch_cap} by budget") + + headroom = allowance - cache_bytes - prefetch_bytes + permitted = (weight_cache_bytes is not None + and weight_cache_bytes <= headroom) + return ResolvedBudget( + usable_bytes=usable, + deductions=deductions, + cache_allowance_bytes=allowance, + cache_slots=cache_slots, + cache_bytes=cache_bytes, + prefetch_cap=prefetch_cap, + prefetch_bytes=prefetch_bytes, + # prefetch_all may not grow the buffer past what the budget + # priced in; without a budget it grows to num_experts + 32. + max_prefetch_cap=prefetch_cap, + max_inflight=inflight, + weight_cache_permitted=permitted, + weight_cache_bytes=weight_cache_bytes, + headroom_bytes=headroom, + notes=tuple(notes), + ) diff --git a/src/edge0/streaming/cache.py b/src/edge0/streaming/cache.py index d426768..b3eb518 100644 --- a/src/edge0/streaming/cache.py +++ b/src/edge0/streaming/cache.py @@ -56,11 +56,22 @@ class PrefetchBuffer: it (it is then promoted into the shared LRU by the caller). """ - def __init__(self, cap: int = 48): - self.cap = max(0, cap) + def __init__(self, cap: int = 48, max_cap: int | None = None): + # ``max_cap`` is a hard ceiling for later ``set_cap`` growth + # (``prefetch_all`` raises the cap to num_experts + 32 on its + # own); a memory budget passes the ceiling it priced in. None + # keeps the historical unbounded-growth semantics. + self.max_cap = max_cap + self.cap = self._clamp(cap) self._buf = OrderedDict() self._lock = threading.Lock() + def _clamp(self, cap: int) -> int: + cap = max(0, cap) + if self.max_cap is not None: + cap = min(cap, self.max_cap) + return cap + def put(self, key, bundle): with self._lock: self._buf[key] = bundle @@ -78,4 +89,4 @@ def contains(self, key) -> bool: return key in self._buf def set_cap(self, cap: int): - self.cap = max(0, cap) + self.cap = self._clamp(cap) diff --git a/src/edge0/streaming/layer.py b/src/edge0/streaming/layer.py index 78a330f..2c0767a 100644 --- a/src/edge0/streaming/layer.py +++ b/src/edge0/streaming/layer.py @@ -94,6 +94,7 @@ def __init__( self.hot_decay = o.hot_decay self._pin_bonus = o.pin_bonus self._use_compile = o.use_compile + self._max_inflight = o.max_inflight self._pool = ThreadPoolExecutor(max_workers=o.load_threads) self._prefetch_pool = ThreadPoolExecutor( @@ -469,6 +470,13 @@ def prefetch(self, experts): ): continue missing.append(e) + limit = self._max_inflight + if limit is not None: + # Speculative builds only: the queue may not outrun the + # budget's priced-in transient. Demand loads + # (_get_bundles) are never dropped by this bound. + room = max(0, limit - len(self._inflight)) + missing = missing[:room] for e in missing: fut = self._prefetch_pool.submit(self._build, e) self._inflight[e] = fut diff --git a/src/edge0/streaming/options.py b/src/edge0/streaming/options.py index 34a27a3..397727c 100644 --- a/src/edge0/streaming/options.py +++ b/src/edge0/streaming/options.py @@ -42,6 +42,20 @@ class LayerOptions: selecting hot pins. cache_slots: shared LRU capacity across all layers. prefetch_cap: prefetch buffer capacity. + max_inflight: per-layer bound on QUEUED speculative prefetch + builds (None = unbounded, the historical behavior). Demand + loads are never dropped; a memory budget sets this so the + producer queue cannot outrun the priced-in transient. + predict_prefetch: route each step's prerouter predictions into + ``prefetch()`` for NON-staged layers (the prod profiles), + overlapping next-step expert builds with the current + forward. Scheduling only: consumption stays on the exact + ``_get_bundles`` path (late/wrong predictions fall back to + demand loads; no staged zero rows). Default off — on the + Orin the unoptimized decode is compute-bound and the + measured end-to-end benefit was below noise (see + docs/plans/jetson-orin-nano.md Task 5); revisit once the + Task 6 kernel work shrinks compute. load_threads: threads for on-demand expert builds. prefetch_threads: threads for eager prefetch builds. use_compile: wrap the MoE math in ``mx.compile``. @@ -69,6 +83,8 @@ class LayerOptions: pin_bonus: float = 2.0 cache_slots: int = 64 prefetch_cap: int = 48 + max_inflight: int | None = None + predict_prefetch: bool = False load_threads: int = 8 prefetch_threads: int = 4 use_compile: bool = True diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..9a398b8 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,18 @@ +"""Repo-wide pytest options. + +``--require-cuda`` turns the CUDA smoke suite's environment skips into +hard failures (Task 3, Gate B: the acceptance invocation must fail if +CUDA is required but absent, rather than quietly skipping). Without the +flag every suite keeps its current skip behavior, so plain ``pytest`` +on a host without torch or a GPU is unchanged. +""" + +from __future__ import annotations + + +def pytest_addoption(parser): + parser.addoption( + "--require-cuda", action="store_true", default=False, + help="fail (instead of skip) tests that need a working torch CUDA " + "device; the Orin acceptance contract from " + "docs/plans/jetson-orin-nano.md Task 3") diff --git a/tests/test_bench_backend.py b/tests/test_bench_backend.py new file mode 100644 index 0000000..ea4fb70 --- /dev/null +++ b/tests/test_bench_backend.py @@ -0,0 +1,187 @@ +"""``examples/bench.py`` on a real backend namespace (still no weights). + +``tests/test_bench_cli.py`` covers the benchmark with fakes only; this +module repeats the core cases through the actual ``core`` ops and the +actual sampler of whichever backend is importable, so a real-array +regression (dtype, device, argmax semantics) is caught. On a host +without MLX and with torch installed the torch backend is selected on the +CPU unless the environment already chose otherwise; with neither, the +module is skipped. A CPU/MPS result here is an instrumentation check, +not CUDA or Orin evidence. +""" + +from __future__ import annotations + +import json +import os +from types import SimpleNamespace + +import pytest + +if "EDGE0_BACKEND" not in os.environ: + try: + import mlx.core # type: ignore[import-not-found] # noqa: F401 + except ImportError: + os.environ["EDGE0_BACKEND"] = "cuda" + os.environ.setdefault("EDGE0_TORCH_DEVICE", "cpu") + +pytest.importorskip( + "edge0.backends", + reason="needs an importable edge0 backend: MLX, or torch (EDGE0_BACKEND=cuda)") + +from edge0.backends import backend, core # noqa: E402 +from edge0.config import GenerationConfig # noqa: E402 +from edge0.streaming.options import LayerOptions # noqa: E402 +from examples import bench # noqa: E402 +from examples import benchmark_report as br # noqa: E402 +from tests.test_bench_cli import ( # noqa: E402 + HUMAN_LINE, + FakeTok, + RecordingMeasurement, + _kinds, + _match, +) + +VOCAB = 8 + + +class ArrayFakeEngine: + """Fake engine whose logits are real backend arrays.""" + + name = "fake-tier" + + def __init__(self, prompt_ids=(1, 2, 3)): + self._tok = FakeTok(prompt_ids) + self.cfg = SimpleNamespace( + gen=GenerationConfig(temperature=0.7, top_p=0.95, top_k=4, + repetition_penalty=1.0), + options=LayerOptions.prod_k8(), model_dir="fake-dir", + lora="", lora_r=16, lora_alpha=32.0, prerouter=None, + prerouter_top_k=8) + self.dir = "fake-dir" + self.prefill_chunk = 2048 + self.think = False + self.events: list = [] + self.closed = 0 + self.pos = 0 + self._logits = None + self._all_stream_layers = {} + + def _make_logits(self): + row = [0.0] * VOCAB + row[self.pos % VOCAB] = 10.0 + return core.array(row, dtype=core.float32) + + def reset(self): + self.events.append("reset") + self.pos = 0 + + def prefill(self, ids): + self.events.append(("prefill", len(ids))) + self.pos += len(ids) + self._logits = self._make_logits() + + def next_logits(self): + return self._logits + + def step(self, tid): + self.events.append(("step", int(tid))) + self.pos += 1 + return self._make_logits() + + def stats(self): + return {} + + def close(self): + self.closed += 1 + + +def test_tier_env_matches_the_cli(): + from edge0.cli import TIER_ENV + assert bench._TIER_ENV == TIER_ENV + + +def test_real_core_warmup_is_greedy_and_timed_steps_sample(monkeypatch): + # At T=0 the real sampler is greedy too, so the token sequence alone + # cannot tell a sampled warmup from a greedy one: count the calls to + # the REAL ops (still executed) through the framework seam. + from edge0.sampling import sample as real_sample + monkeypatch.setenv("BENCH_TEMP", "0") + monkeypatch.delenv("BENCH_SEED", raising=False) + calls = {"argmax": 0, "sample": 0} + real_argmax = core.argmax + + def counting_argmax(*args, **kwargs): + calls["argmax"] += 1 + return real_argmax(*args, **kwargs) + + def counting_sample(*args, **kwargs): + calls["sample"] += 1 + return real_sample(*args, **kwargs) + monkeypatch.setattr(core, "argmax", counting_argmax) + monkeypatch.setattr(bench, "_framework", + lambda: (backend, core, counting_sample)) + eng = ArrayFakeEngine(prompt_ids=(1, 2, 3)) + out = bench.run_bench(eng, ntok=5, warmup=2, + measurement=RecordingMeasurement(eng)) + steps = [e[1] for e in eng.events if isinstance(e, tuple) and e[0] == "step"] + assert steps[:7] == [3, 4, 5, 6, 7, 0, 1] + assert calls == {"argmax": 2 * 2, "sample": 5 * 2} # per run: warmup, ntok + kinds = _kinds(eng.events) + assert kinds.count("sync") == 8 and kinds.count("reset_peak") == 2 + run = out["report_runs"][0] + assert run["total_generated_tokens"] == 7 + assert run["execution_evidence"]["array_type"] + + +def test_seed_is_applied_through_the_real_core(monkeypatch): + seen: list = [] + monkeypatch.setattr(core.random, "seed", seen.append) + monkeypatch.setenv("BENCH_SEED", "0") + monkeypatch.setenv("BENCH_TEMP", "0") + eng = ArrayFakeEngine() + bench.run_bench(eng, ntok=1, warmup=0, + measurement=RecordingMeasurement(eng)) + assert seen == [0, 0] + + +def test_default_measurement_dispatches_on_the_active_backend(): + eng = ArrayFakeEngine() + out = bench.run_bench(eng, ntok=1, warmup=0) + desc = out["measurement"] + assert desc["backend"] == backend.name + if backend.name == "cuda": + expected = "cuda" if core.DEVICE.type == "cuda" else "null" + assert desc["kind"] == expected + assert desc["resolved_device"] == str(core.DEVICE) + else: + assert desc["kind"] == "mlx" + + +def test_cli_report_states_the_actually_resolved_device(tmp_path, monkeypatch, + capsys): + eng = ArrayFakeEngine() + monkeypatch.setattr(bench, "_build_engine", lambda model, name=None: eng) + monkeypatch.setenv("BENCH_TEMP", "0") + model_dir = tmp_path / "ckpt" + model_dir.mkdir() + out = tmp_path / "bench.json" + assert bench.main([str(model_dir), "--ntok", "2", "--warmup", "1", + "--rss-sample-interval", "0", + "--json-output", str(out)]) == 0 + assert eng.closed == 1 + doc = json.loads(out.read_text(encoding="utf-8")) + br.validate_report(doc) + rt = doc["runtime"] + assert rt["backend"] == backend.name + if backend.name == "cuda": + assert rt["resolved_device"] == str(core.DEVICE) + assert rt["measurement"]["kind"] == ( + "cuda" if core.DEVICE.type == "cuda" else "null") + assert rt["torch_version"] + else: + assert rt["measurement"]["kind"] == "mlx" + assert rt["mlx_version"] + assert doc["caches"]["layer_options"]["cache_slots"] == 64 + lines = capsys.readouterr().out.splitlines() + assert _match(HUMAN_LINE, lines[1]) diff --git a/tests/test_bench_cli.py b/tests/test_bench_cli.py new file mode 100644 index 0000000..fd6c0f6 --- /dev/null +++ b/tests/test_bench_cli.py @@ -0,0 +1,794 @@ +"""``examples/bench.py`` end to end with fakes only: no backend, no weights. + +``bench.py`` imports the framework lazily (``bench._framework``), so this +module runs on a host with neither torch nor MLX: the engine, the array +namespace, the sampler and the measurement adapter are all fakes that +record the order of every call. That is what lets the phase +synchronization order, the count arithmetic, the memory units, the +pre-load validation, the failure paths and the cleanup be asserted +anywhere. ``tests/test_bench_backend.py`` repeats the core cases on a +real backend when one is importable. Host results are instrumentation +checks, not Orin evidence. +""" + +from __future__ import annotations + +import hashlib +import json +import os +import re +import struct +import subprocess +import sys +import time +from dataclasses import dataclass +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +import pytest + +from examples import bench +from examples import benchmark_measure as bm +from examples import benchmark_report as br + +ROOT = Path(__file__).resolve().parents[1] +VOCAB = 8 +GIB = 1024 ** 3 + +HUMAN_LINE = re.compile( + r"^prompt=(\d+) tok prefill=([\d.]+)s \((\d+) tok/s\) " + r"decode=(\d+)/([\d.]+)s tok/s=([\d.]+) " + r"peak_active=((?:[\d.]+ GiB)|n/a)$") +SUMMARY_LINE = re.compile( + r"^\[bench\] mean tok/s=([\d.]+) peak_active<=((?:[\d.]+ GiB)|n/a)$") + + +# ---- import boundary -------------------------------------------------------------- + + +def test_bench_module_imports_without_torch_mlx_or_edge0(): + code = ( + f"import sys; sys.path.insert(0, {str(ROOT)!r}); " + "from examples import bench; " + "bad = sorted(m for m in sys.modules " + "if m.split('.')[0] in ('torch', 'mlx', 'mlx_lm', 'edge0')); " + "print(bad)" + ) + proc = subprocess.run([sys.executable, "-c", code], + capture_output=True, text=True, check=True) + assert proc.stdout.strip() == "[]", proc.stdout + + +def test_help_and_argument_errors_need_no_backend(): + env = dict(os.environ, EDGE0_BACKEND="no-such-backend") + proc = subprocess.run( + [sys.executable, str(ROOT / "examples" / "bench.py"), "--help"], + capture_output=True, text=True, env=env, check=False) + assert proc.returncode == 0, proc.stderr + assert "--json-output" in proc.stdout + proc = subprocess.run( + [sys.executable, str(ROOT / "examples" / "bench.py"), "x", "--ntok", "0"], + capture_output=True, text=True, env=env, check=False) + assert proc.returncode == 2 + assert "--ntok" in proc.stderr + + +# ---- fakes ------------------------------------------------------------------------------ + + +def _argmax(logits): + return max(range(len(logits)), key=logits.__getitem__) + + +class FakeCore: + """The two ``core`` calls the benchmark loop makes.""" + + def __init__(self): + self.seeds: list[int] = [] + self.random = SimpleNamespace(seed=self.seeds.append) + + @staticmethod + def argmax(logits, axis=-1): + return SimpleNamespace(item=lambda: _argmax(logits)) + + +def fake_sample(logits, temperature=0.7, top_k=None, top_p=None, + repetition_penalty=1.0, history=(), seed=None): + """Deliberately NOT the argmax, so a warmup step that sampled (or a + timed step that took the argmax) is visible in the recorded tokens.""" + return (_argmax(logits) + 1) % VOCAB + + +class FakeTok: + def __init__(self, ids): + self.ids = list(ids) + + def apply_chat_template(self, messages, tokenize=False, + add_generation_prompt=True, + enable_thinking=False): + return "templated:" + messages[-1]["content"] + + def __call__(self, text): + return {"input_ids": list(self.ids)} + + +@dataclass +class FakeOptions: + cache_slots: int = 64 + prefetch_cap: int = 48 + load_threads: int = 8 + prefetch_threads: int = 4 + top_k: int | None = None + + +def _safetensors(path: Path, metadata: dict, n_tensors: int = 1): + header: dict = {"__metadata__": metadata} + for i in range(n_tensors): + header[f"t{i}"] = {"dtype": "F16", "shape": [1], + "data_offsets": [2 * i, 2 * i + 2]} + blob = json.dumps(header).encode() + path.write_bytes(struct.pack("