From dc1d9b3b78f51eb1a5f3705d63ed16c5b2cc9dfe Mon Sep 17 00:00:00 2001 From: KarlTaylorKnight Date: Sun, 13 Sep 2026 17:58:25 +0400 Subject: [PATCH 1/9] [verified] feat(jetson): add Orin bring-up foundation Add a fail-closed capability probe for Jetson, CUDA, memory, and NVMe-backed model storage. Document the staged Orin Nano implementation plan and add focused tests for platform, storage, error, schema, and CLI behavior. --- docs/nvidia.md | 21 ++ docs/plans/jetson-orin-nano.md | 177 ++++++++++++ scripts/jetson_probe.py | 332 ++++++++++++++++++++++ tests/test_jetson_probe.py | 494 +++++++++++++++++++++++++++++++++ 4 files changed, 1024 insertions(+) create mode 100644 docs/plans/jetson-orin-nano.md create mode 100755 scripts/jetson_probe.py create mode 100644 tests/test_jetson_probe.py diff --git a/docs/nvidia.md b/docs/nvidia.md index b86994e..6546204 100644 --- a/docs/nvidia.md +++ b/docs/nvidia.md @@ -1,5 +1,26 @@ # NVIDIA / CUDA support — status +## Jetson Orin Nano bring-up + +Jetson support is not yet claimed or benchmarked. 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). Start with the +8B tier and retain the existing correctness path while optimizing memory and +I/O. + `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.md b/docs/plans/jetson-orin-nano.md new file mode 100644 index 0000000..92fc383 --- /dev/null +++ b/docs/plans/jetson-orin-nano.md @@ -0,0 +1,177 @@ +# Jetson Orin Nano CUDA Backend Implementation Plan + +> **For Claude Code:** Execute this plan vertically with strict RED → GREEN → REFACTOR cycles. Do not optimize kernels until an Orin Nano baseline has been captured and committed. + +**Goal:** Make the Edge0 Torch/CUDA backend run reproducibly on an 8 GB Jetson Orin Nano with NVMe-backed expert streaming, then optimize it without violating MLX parity. + +**Architecture:** Build on the correctness reference in upstream PR #19 rather than porting the MLX backend again. Treat Jetson as a distinct memory/I/O profile: dense quantized weights stay resident where possible, sparse experts move through a bounded slot cache, and prerouter predictions schedule asynchronous NVMe → pinned host buffer → CUDA slot transfers. Preserve the existing backend facade and model engines; introduce Jetson behavior as measured configuration and CUDA implementation details, not model-specific forks. + +**Tech stack:** Python 3.10+, PyTorch for Jetson, CUDA, safetensors, NVMe, pytest; later CUTLASS/CUDA extensions for fused INT4 gather-GEMM. + +--- + +## Constraints and evidence + +- This branch is stacked on upstream PR #19 (`agourakis82:cuda-backend`), which already ports both shipped models and establishes CPU/MPS/CUDA parity. +- Upstream explicitly warns that the Apple path depends on unified-memory mmap semantics; an efficient CUDA implementation needs a DirectStorage-style data path, not a mechanical MLX port. +- The current Torch backend is a correctness reference. On GB10, Edge0-8B was measured at 3.2 tok/s by default and 8.1 tok/s with large caches, but the latter used about 6 GB resident memory. Those results cannot be projected onto Orin Nano. +- Orin Nano has 8 GB shared memory. The OS, CUDA context, tokenizer, KV cache, and application processes must all fit alongside Edge0. +- Use NVMe. microSD is not an acceptable performance target for expert streaming. + +## Non-goals for the first PR + +- No unverified performance claims. +- No Edge0-35B optimization before Edge0-8B is stable. +- No new model family. +- No speculative TensorRT rewrite. +- No changes to MLX behavior. + +## Acceptance gates + +### Gate A — reproducible platform evidence + +On the target Orin Nano, save a JSON probe containing: + +- Jetson model and L4T/JetPack release data; +- architecture, total/available shared memory; +- PyTorch, CUDA, CUDA availability, GPU name, compute capability, and reported device memory; +- model/checkpoint filesystem, free space, and whether it resolves to NVMe; +- explicit readiness blockers instead of guessed defaults. + +### Gate B — correctness baseline + +With `Edge0/Edge0-8B-A1B-preview` on NVMe: + +- run the CUDA backend smoke tests; +- generate at least 32 greedy tokens; +- record peak process RSS, peak device allocation, prefill time, and decode tok/s; +- compare greedy tokens or logits against the existing reference fixture; +- retain the full command and JSON result. + +### Gate C — bounded memory + +- No OOM on an 8 GB Orin Nano in the agreed power mode. +- Target peak process footprint: below 6.5 GB for a short-context Edge0-8B run, leaving explicit OS headroom. +- Cache sizes are selected from available memory at startup and can be overridden. + +### Gate D — optimized data path + +- Expert-cache misses copy into preallocated slots; no per-token CUDA allocation. +- Transfers run on a dedicated CUDA stream and are synchronized with events. +- Prerouter predictions issue prefetch before the expert is consumed. +- Dense quantized layers do not repeatedly unpack/dequantize unchanged weights. +- Every optimized path has a parity test against the correctness path. + +## Task 1: Add a read-only Jetson capability probe + +**Objective:** Capture the facts needed to reproduce and interpret every Jetson benchmark. + +**Files:** +- Create: `scripts/jetson_probe.py` +- Create: `tests/test_jetson_probe.py` +- Modify: `docs/nvidia.md` + +**TDD steps:** + +1. Test `/proc/meminfo` parsing with kB values and missing `MemAvailable`. +2. Run the focused test and verify it fails because the module is absent. +3. Implement only the memory parser; rerun until green. +4. Test mount selection and NVMe classification using synthetic `/proc/self/mountinfo`. +5. Implement longest-prefix mount matching without shelling out. +6. Test Jetson model/L4T detection from an injected filesystem root. +7. Implement platform and device-tree readers. +8. Test Torch reporting with fake unavailable and fake CUDA modules. +9. Implement Torch collection with import failure represented as data. +10. Test readiness blockers: not Jetson, no CUDA, non-NVMe storage, and a missing model path. +11. Implement `collect_probe()` and JSON CLI output. +12. Run `pytest tests/test_jetson_probe.py -q`, then the full non-slow suite. +13. Commit: `feat(jetson): add Orin capability probe`. + +## Task 2: Add a machine-readable inference benchmark result + +**Objective:** Extend the existing benchmark path to emit reproducible JSON without changing generation behavior. + +**Files:** +- Modify: `examples/bench.py` +- Create or modify: `tests/test_benchmark_report.py` +- Modify: `docs/nvidia.md` + +**TDD steps:** + +1. Extract a pure report-building function and test its schema first. +2. Include exact checkpoint path, backend, device, context/prompt tokens, generated tokens, prefill seconds/tok-s, decode seconds/tok-s, RSS peak, CUDA allocation peak, cache settings, power mode, and git SHA. +3. Add `--json-output PATH`; retain current human-readable output. +4. Verify output on CPU without real weights using a fixture/fake measurement object. +5. Run the existing benchmark manually on Orin and commit only the command/result documentation—not model files. +6. Commit: `feat(bench): emit reproducible inference reports`. + +## Task 3: Establish the unoptimized Orin baseline + +**Objective:** Learn which constraint kills the current implementation before changing it. + +**Target commands:** + +```bash +python scripts/jetson_probe.py --model-dir "$EDGE0_8B_MODEL" --output artifacts/orin-probe.json +EDGE0_BACKEND=cuda EDGE0_TORCH_DEVICE=cuda \ + pytest tests/test_nvidia_backend_smoke.py -q +EDGE0_BACKEND=cuda EDGE0_TORCH_DEVICE=cuda \ + python examples/bench.py edge0-8b --json-output artifacts/orin-baseline.json +``` + +Capture `tegrastats` beside the benchmark and note JetPack, power mode, clocks, NVMe model/filesystem, cooling, and ambient throttling. If the run OOMs, reduce context/generation—not model correctness—and record the first failing allocation. + +Commit: `docs(jetson): record Orin Nano baseline`. + +## Task 4: Add a memory-budgeted runtime profile + +**Objective:** Select safe cache settings from measured available memory rather than copying GB10 defaults. + +**Files:** +- Modify: `src/edge0/backends/cuda/backend.py` +- Modify: `src/edge0/backends/cuda/nn.py` +- Modify: model/runtime configuration at the narrowest existing override point +- Test: `tests/test_cuda_backend.py` + +Start with policy only—no new kernels. The policy must expose its budget calculation, reserve OS/KV headroom, reject impossible overrides, and print the resolved cache budget. Test the calculation as a pure function before wiring it into the engine. + +Commit: `feat(cuda): add memory-budgeted cache profile`. + +## Task 5: Replace per-token expert bundle copies + +**Objective:** Reuse bounded CUDA expert slots and overlap misses with compute. + +**Approach:** + +- Allocate fixed slots once. +- Stage safetensors ranges into reusable pinned host buffers. +- Copy on a dedicated CUDA stream. +- Use CUDA events to publish slot readiness. +- Drive prefetch from the prerouter; preserve exact fallback on a miss. +- Add counters for cache hit, predicted hit, late hit, bytes read, bytes copied, and stall time. + +First test state transitions and eviction using fake streams/events. Then add CUDA parity tests. Do not add GPUDirect Storage until the staged path is measured; Jetson support and benefit must be demonstrated first. + +Commit: `perf(cuda): add asynchronous expert slot cache`. + +## Task 6: Fuse quantized gather/dequantize/matmul + +**Objective:** Remove repeated Python/Torch unpack operations from decode. + +Prototype one kernel for the exact Edge0 INT4 affine layout: packed u32, group size 64, bf16 scales/biases, gathered expert rows. Compare against `quant.gather_qmm` over every broadcast shape already covered by `tests/test_cuda_backend.py`; require the existing numerical tolerance before benchmarking. Prefer a small PyTorch CUDA extension using CUTLASS primitives over a broad TensorRT rewrite. + +Commit: `perf(cuda): fuse INT4 expert gather matmul`. + +## Task 7: Final Orin validation + +Repeat the exact baseline commands with three or more benchmark runs. Report individual runs plus mean/variation; keep prefill and decode separate. Verify CUDA execution from runtime logs, not merely `torch.cuda.is_available()`. Document remaining thermal, storage, context-length, and quality constraints. + +## Claude handoff rules + +1. Begin by reading upstream PR #19, issue #20, `docs/nvidia.md`, `docs/streaming.md`, and the backend parity tests. +2. Do not rewrite the CUDA backend or replace the existing engines. +3. Keep the first implementation PR limited to Task 1 plus documentation. +4. Never claim Orin support until Gate B runs on the physical device. +5. Every optimization must retain a selectable correctness/reference path. +6. Do not commit checkpoints, generated artifacts, credentials, or machine identifiers. +7. Stop and report after each acceptance gate with measured evidence and the next bottleneck. 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/tests/test_jetson_probe.py b/tests/test_jetson_probe.py new file mode 100644 index 0000000..fd7d199 --- /dev/null +++ b/tests/test_jetson_probe.py @@ -0,0 +1,494 @@ +from __future__ import annotations + +import json +from pathlib import Path + +from scripts import jetson_probe +from scripts.jetson_probe import ( + collect_probe, + collect_torch_info, + parse_meminfo, + storage_device_for_path, +) + + +def _make_nvme_sysfs(root: Path, device_id: str = "259:1") -> Path: + sys_dev_block = root / "sys/dev/block" + backing = root / "sys/devices/platform/nvme/nvme0/nvme0n1" + backing.mkdir(parents=True, exist_ok=True) + sys_dev_block.mkdir(parents=True, exist_ok=True) + (sys_dev_block / device_id).symlink_to(backing) + return sys_dev_block + + +def test_parse_meminfo_converts_kib_to_bytes(): + parsed = parse_meminfo("MemTotal: 8192000 kB\nMemAvailable: 4096000 kB\n") + + assert parsed == { + "total_bytes": 8_388_608_000, + "available_bytes": 4_194_304_000, + } + + +def test_parse_meminfo_allows_missing_available_memory(): + parsed = parse_meminfo("MemTotal: 1024 kB\n") + + assert parsed == {"total_bytes": 1_048_576, "available_bytes": None} + + +def test_storage_device_uses_longest_mount_prefix(tmp_path: Path): + model_dir = tmp_path / "models" / "edge0-8b" + model_dir.mkdir(parents=True) + mountinfo = "\n".join( + [ + "30 20 8:1 / / rw,relatime - ext4 /dev/mmcblk0p1 rw", + f"31 30 259:1 / {tmp_path / 'models'} rw,relatime - ext4 /dev/nvme0n1p1 rw", + ] + ) + + storage = storage_device_for_path( + model_dir, + mountinfo, + path_device_id="259:1", + sys_dev_block=_make_nvme_sysfs(tmp_path), + ) + + assert storage["device"] == "/dev/nvme0n1p1" + assert storage["mount_point"] == str(tmp_path / "models") + assert storage["filesystem"] == "ext4" + assert storage["is_nvme"] is True + + +def test_storage_device_marks_non_nvme_media(tmp_path: Path): + mountinfo = "30 20 179:1 / / rw,relatime - ext4 /dev/mmcblk0p1 rw" + + storage = storage_device_for_path( + tmp_path, + mountinfo, + path_device_id="179:1", + ) + + assert storage["is_nvme"] is False + + +def test_storage_device_resolves_device_mapper_backing_nvme(tmp_path: Path): + sys_dev_block = tmp_path / "sys/dev/block" + mapped = tmp_path / "sys/devices/virtual/block/dm-0" + physical = tmp_path / "sys/devices/platform/nvme/nvme0/nvme0n1" + (mapped / "slaves").mkdir(parents=True) + physical.mkdir(parents=True) + sys_dev_block.mkdir(parents=True) + (sys_dev_block / "253:0").symlink_to(mapped) + (mapped / "slaves/nvme0n1").symlink_to(physical) + mountinfo = "30 20 253:0 / / rw,relatime - ext4 /dev/mapper/cryptroot rw" + + storage = storage_device_for_path( + tmp_path, + mountinfo, + path_device_id="253:0", + sys_dev_block=sys_dev_block, + ) + + assert storage["device_id"] == "253:0" + assert storage["is_nvme"] is True + + +def test_storage_device_rejects_mixed_device_mapper_backing(tmp_path: Path): + sys_dev_block = tmp_path / "sys/dev/block" + mapped = tmp_path / "sys/devices/virtual/block/dm-0" + nvme = tmp_path / "sys/devices/platform/nvme/nvme0/nvme0n1" + mmc = tmp_path / "sys/devices/platform/mmc/mmcblk0" + (mapped / "slaves").mkdir(parents=True) + nvme.mkdir(parents=True) + mmc.mkdir(parents=True) + sys_dev_block.mkdir(parents=True) + (sys_dev_block / "253:0").symlink_to(mapped) + (mapped / "slaves/nvme0n1").symlink_to(nvme) + (mapped / "slaves/mmcblk0").symlink_to(mmc) + + storage = storage_device_for_path( + tmp_path, + "30 20 253:0 / / rw - ext4 /dev/mapper/mixed rw", + path_device_id="253:0", + sys_dev_block=sys_dev_block, + ) + + assert storage["is_nvme"] is False + + +def test_storage_device_handles_escaped_mounts_and_malformed_rows(tmp_path: Path): + model_dir = tmp_path / "model files" + model_dir.mkdir() + escaped = str(model_dir).replace(" ", "\\040") + mountinfo = "\n".join( + [ + "malformed", + f"31 30 259:1 / {escaped} rw - ext4 /dev/nvme0n1p1 rw", + ] + ) + + storage = storage_device_for_path( + model_dir, + mountinfo, + path_device_id="259:1", + sys_dev_block=_make_nvme_sysfs(tmp_path), + ) + + assert storage["mount_point"] == str(model_dir) + assert storage["is_nvme"] is True + + +def test_storage_device_preserves_unknown_nvme_state(tmp_path: Path): + mountinfo = "30 20 0:42 / / rw - overlay overlay rw" + + storage = storage_device_for_path( + tmp_path, + mountinfo, + path_device_id="0:42", + sys_dev_block=tmp_path / "missing-sysfs", + ) + + assert storage["is_nvme"] is None + + +def test_storage_device_does_not_trust_network_or_pseudo_source_names(tmp_path: Path): + for mountinfo, device_id in ( + ("30 20 0:42 / / rw - nfs server:/nvme0n1 rw", "0:42"), + ("31 20 0:43 / / rw - tmpfs /dev/nvme0n1 rw", "0:43"), + ): + storage = storage_device_for_path( + tmp_path, + mountinfo, + path_device_id=device_id, + sys_dev_block=tmp_path / "missing-sysfs", + ) + + assert storage["is_nvme"] is not True + + +def test_storage_device_matches_active_device_id_for_stacked_mounts(tmp_path: Path): + mountinfo = ( + "30 20 259:1 / / rw - ext4 /dev/nvme0n1p1 rw\n" + "31 20 179:1 / / rw - ext4 /dev/mmcblk0p1 rw" + ) + + storage = storage_device_for_path( + tmp_path, + mountinfo, + path_device_id="179:1", + ) + + assert storage["device_id"] == "179:1" + assert storage["is_nvme"] is False + + +def test_storage_device_reports_disk_usage_failure(tmp_path: Path): + def fail_disk_usage(_path): + raise OSError("unavailable") + + storage = storage_device_for_path( + tmp_path, + "30 20 259:1 / / rw - ext4 /dev/nvme0n1p1 rw", + disk_usage_fn=fail_disk_usage, + ) + + assert storage["disk_usage_available"] is False + assert storage["total_bytes"] is None + assert storage["free_bytes"] is None + + +class _FakeProps: + name = "Orin" + total_memory = 8_000_000_000 + major = 8 + minor = 7 + + +class _FakeCuda: + @staticmethod + def is_available(): + return True + + @staticmethod + def get_device_properties(_index): + return _FakeProps() + + +class _FakeVersion: + cuda = "12.6" + + +class _FakeTorch: + __version__ = "2.8.0" + cuda = _FakeCuda() + version = _FakeVersion() + + +def test_collect_torch_info_reports_cuda_device(): + info = collect_torch_info(_FakeTorch()) + + assert info == { + "installed": True, + "version": "2.8.0", + "cuda_available": True, + "cuda_version": "12.6", + "device_name": "Orin", + "compute_capability": [8, 7], + "device_memory_bytes": 8_000_000_000, + "errors": [], + } + + +def test_collect_torch_info_reports_missing_torch(): + info = collect_torch_info(None) + + assert info == { + "installed": False, + "version": None, + "cuda_available": False, + "cuda_version": None, + "device_name": None, + "compute_capability": None, + "device_memory_bytes": None, + "errors": ["PyTorch import was skipped or unavailable"], + } + + +class _BrokenCuda: + @staticmethod + def is_available(): + raise RuntimeError("CUDA driver initialization failed") + + +class _BrokenTorch: + __version__ = "2.8.0" + cuda = _BrokenCuda() + version = _FakeVersion() + + +def test_collect_torch_info_turns_cuda_driver_failure_into_probe_data(): + info = collect_torch_info(_BrokenTorch()) + + assert info["installed"] is True + assert info["cuda_available"] is False + assert info["errors"] == [ + ( + "CUDA availability check failed: " + "RuntimeError: CUDA driver initialization failed" + ) + ] + + +class _BrokenPropertiesCuda(_FakeCuda): + @staticmethod + def get_device_properties(_index): + raise RuntimeError("device query failed") + + +class _BrokenPropertiesTorch: + __version__ = "2.8.0" + cuda = _BrokenPropertiesCuda() + version = _FakeVersion() + + +def test_collect_torch_info_turns_device_query_failure_into_probe_data(): + info = collect_torch_info(_BrokenPropertiesTorch()) + + assert info["cuda_available"] is True + assert info["errors"] == [ + "CUDA device query failed: RuntimeError: device query failed" + ] + + +def test_collect_probe_reports_actionable_blockers(tmp_path: Path): + root = tmp_path / "root" + (root / "proc/device-tree").mkdir(parents=True) + (root / "etc").mkdir() + (root / "proc/device-tree/model").write_bytes( + b"NVIDIA Jetson Orin Nano Developer Kit\x00" + ) + (root / "etc/nv_tegra_release").write_text("# R36 (release), REVISION: 4.3\n") + (root / "proc/meminfo").write_text( + "MemTotal: 8192000 kB\nMemAvailable: 4096000 kB\n" + ) + model_dir = tmp_path / "models" + model_dir.mkdir() + mountinfo = "30 20 179:1 / / rw,relatime - ext4 /dev/mmcblk0p1 rw" + + report = collect_probe( + model_dir=model_dir, + fs_root=root, + mountinfo_text=mountinfo, + torch_module=None, + path_device_id="179:1", + ) + + assert report["schema_version"] == 1 + assert report["jetson"]["is_jetson_orin"] is True + assert report["jetson"]["l4t_release"].startswith("# R36") + assert report["storage"]["is_nvme"] is False + assert report["ready"] is False + assert "PyTorch is not installed" in report["blockers"] + assert "CUDA is not available through PyTorch" in report["blockers"] + assert "model path is not on NVMe storage" in report["blockers"] + + +def test_collect_probe_rejects_a_missing_model_path(tmp_path: Path): + root = tmp_path / "root" + (root / "proc/device-tree").mkdir(parents=True) + (root / "etc").mkdir() + (root / "proc/device-tree/model").write_text("NVIDIA Jetson Orin Nano") + (root / "etc/nv_tegra_release").write_text("# R36") + (root / "proc/meminfo").write_text("MemTotal: 8192000 kB\n") + _make_nvme_sysfs(root) + + report = collect_probe( + model_dir=tmp_path / "missing-model", + fs_root=root, + mountinfo_text="30 20 259:1 / / rw - ext4 /dev/nvme0n1p1 rw", + torch_module=_FakeTorch(), + path_device_id="259:1", + ) + + assert report["storage"]["is_nvme"] is True + assert report["ready"] is False + assert "model path does not exist" in report["blockers"] + + +def test_collect_probe_rejects_a_model_path_that_is_not_a_directory(tmp_path: Path): + model_file = tmp_path / "model.safetensors" + model_file.write_bytes(b"") + root = tmp_path / "root" + (root / "proc/device-tree").mkdir(parents=True) + (root / "etc").mkdir() + (root / "proc/device-tree/model").write_text("NVIDIA Jetson Orin Nano") + (root / "proc/meminfo").write_text("MemTotal: 8192000 kB\n") + _make_nvme_sysfs(root) + + report = collect_probe( + model_dir=model_file, + fs_root=root, + mountinfo_text="30 20 259:1 / / rw - ext4 /dev/nvme0n1p1 rw", + torch_module=_FakeTorch(), + path_device_id="259:1", + ) + + assert report["ready"] is False + assert "model path is not a directory" in report["blockers"] + + +def test_collect_probe_blocks_incomplete_cuda_device_evidence(tmp_path: Path): + root = tmp_path / "root" + (root / "proc/device-tree").mkdir(parents=True) + (root / "etc").mkdir() + (root / "proc/device-tree/model").write_text("NVIDIA Jetson Orin Nano") + (root / "proc/meminfo").write_text("MemTotal: 8192000 kB\n") + _make_nvme_sysfs(root) + + report = collect_probe( + model_dir=tmp_path, + fs_root=root, + mountinfo_text="30 20 259:1 / / rw - ext4 /dev/nvme0n1p1 rw", + torch_module=_BrokenPropertiesTorch(), + path_device_id="259:1", + ) + + assert report["ready"] is False + assert any("CUDA device query failed" in item for item in report["blockers"]) + + +def test_collect_probe_can_report_fully_ready(tmp_path: Path): + root = tmp_path / "root" + (root / "proc/device-tree").mkdir(parents=True) + (root / "etc").mkdir() + (root / "proc/device-tree/model").write_text("NVIDIA Jetson Orin Nano") + (root / "etc/nv_tegra_release").write_text("# R36") + (root / "proc/meminfo").write_text("MemTotal: 8192000 kB\n") + _make_nvme_sysfs(root) + + report = collect_probe( + model_dir=tmp_path, + fs_root=root, + mountinfo_text="30 20 259:1 / / rw - ext4 /dev/nvme0n1p1 rw", + torch_module=_FakeTorch(), + path_device_id="259:1", + ) + + assert report["ready"] is True + assert report["blockers"] == [] + + +def test_collect_probe_blocks_when_disk_usage_is_unavailable(tmp_path: Path): + root = tmp_path / "root" + (root / "proc/device-tree").mkdir(parents=True) + (root / "etc").mkdir() + (root / "proc/device-tree/model").write_text("NVIDIA Jetson Orin Nano") + (root / "proc/meminfo").write_text("MemTotal: 8192000 kB\n") + _make_nvme_sysfs(root) + + def fail_disk_usage(_path): + raise OSError("unavailable") + + report = collect_probe( + model_dir=tmp_path, + fs_root=root, + mountinfo_text="30 20 259:1 / / rw - ext4 /dev/nvme0n1p1 rw", + torch_module=_FakeTorch(), + disk_usage_fn=fail_disk_usage, + path_device_id="259:1", + ) + + assert report["ready"] is False + assert "model path disk usage could not be determined" in report["blockers"] + + +def test_probe_report_is_json_serializable(tmp_path: Path): + root = tmp_path / "root" + (root / "proc/device-tree").mkdir(parents=True) + (root / "etc").mkdir() + (root / "proc/device-tree/model").write_text("generic arm board") + (root / "proc/meminfo").write_text("MemTotal: 1024 kB\n") + + report = collect_probe( + model_dir=tmp_path, + fs_root=root, + mountinfo_text="30 20 8:1 / / rw - ext4 /dev/sda1 rw", + torch_module=None, + ) + + json.dumps(report) + assert "hardware is not a Jetson Orin" in report["blockers"] + + +def test_cli_writes_json_and_returns_nonzero_for_blockers( + tmp_path: Path, + monkeypatch, + capsys, +): + expected = {"schema_version": 1, "ready": False, "blockers": ["test"]} + monkeypatch.setattr(jetson_probe, "collect_probe", lambda *_args, **_kw: expected) + output = tmp_path / "probe.json" + + result = jetson_probe.main( + [ + "--model-dir", + str(tmp_path), + "--output", + str(output), + "--skip-torch", + ] + ) + + assert result == 1 + assert json.loads(output.read_text()) == expected + assert json.loads(capsys.readouterr().out) == expected + + +def test_cli_returns_zero_for_a_ready_report(tmp_path: Path, monkeypatch, capsys): + expected = {"schema_version": 1, "ready": True, "blockers": []} + monkeypatch.setattr(jetson_probe, "collect_probe", lambda *_args, **_kw: expected) + + result = jetson_probe.main(["--model-dir", str(tmp_path)]) + + assert result == 0 + assert json.loads(capsys.readouterr().out) == expected From 20d62aad35d914bc5d560d8ff38170f224b4b1a4 Mon Sep 17 00:00:00 2001 From: KarlTaylorKnight <102148054+KarlTaylorKnight@users.noreply.github.com> Date: Mon, 14 Sep 2026 15:44:55 +0400 Subject: [PATCH 2/9] feat(bench): machine-readable benchmark reports, synchronized CUDA timing (Task 2) examples/bench.py could not run on the torch backend: it called core.reset_peak_memory()/get_peak_memory(), which only MLX has, and its host timers never waited for queued CUDA work. This is the portable reporting increment of docs/plans/jetson-orin-nano.md (Task 2); no physical Orin run is claimed. * examples/benchmark_report.py: pure, stdlib-only report builder (schema_version 1): validated per-run counts and timings, integer-byte memory metrics with method/scope, null-with-reason convention, strict JSON (allow_nan=False), probe linkage by content SHA-256, atomic no-clobber write. * examples/benchmark_measure.py: benchmark-owned measurement adapter (MLX peak memory kept as before; torch CUDA allocator peaks with torch.cuda.synchronize at phase boundaries on an actual cuda device; explicit unavailable values on cpu/mps), process RSS peak per OS, a sampler thread for rss/anon/file/swap, git identity, nvpmodel power mode, checkpoint manifest and adapter identity without hashing weights. * examples/bench.py: --json-output PATH (new file only), --probe-json PATH, --rss-sample-interval; synchronize after reset, after prefill, after warmup and after the final timed step; same two-run greedy-warmup / sampled fixed-length protocol and human output (peak_active=n/a when unavailable); pre-load validation; framework imported lazily; engine and sampler closed in finally; failures exit non-zero and leave no report. * tests: test_benchmark_report.py, test_benchmark_measure.py, test_bench_cli.py (fakes only) and test_bench_backend.py (real backend namespace, no weights). pyproject: pytest pythonpath = ["."] so plain `pytest` collects the repo-root scripts/examples imports. * docs: benchmark protocol and schema in docs/nvidia.md; the revised plan replaces docs/plans/jetson-orin-nano.md with its C01-C16 justifications alongside; README warmup wording corrected (it was always greedy). Host checks only (Windows torch-CPU venv, WSL Linux, Python 3.9/3.11): no CUDA execution, no MLX, no Orin evidence. Co-Authored-By: Claude Fable 5.1 --- README.md | 2 +- README_zh.md | 2 +- docs/nvidia.md | 144 ++++- docs/plans/jetson-orin-nano-review.md | 145 +++++ docs/plans/jetson-orin-nano.md | 241 ++++---- examples/bench.py | 684 +++++++++++++++++++-- examples/benchmark_measure.py | 842 ++++++++++++++++++++++++++ examples/benchmark_report.py | 546 +++++++++++++++++ pyproject.toml | 4 + tests/test_bench_backend.py | 187 ++++++ tests/test_bench_cli.py | 794 ++++++++++++++++++++++++ tests/test_benchmark_measure.py | 686 +++++++++++++++++++++ tests/test_benchmark_report.py | 527 ++++++++++++++++ 13 files changed, 4623 insertions(+), 181 deletions(-) create mode 100644 docs/plans/jetson-orin-nano-review.md create mode 100644 examples/benchmark_measure.py create mode 100644 examples/benchmark_report.py create mode 100644 tests/test_bench_backend.py create mode 100644 tests/test_bench_cli.py create mode 100644 tests/test_benchmark_measure.py create mode 100644 tests/test_benchmark_report.py 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 6546204..5e66a0d 100644 --- a/docs/nvidia.md +++ b/docs/nvidia.md @@ -17,9 +17,147 @@ 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). Start with the -8B tier and retain the existing correctness path while optimizing memory and -I/O. +[`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 will use, is described in [Benchmark reporting](#benchmark-reporting) +below. + +## 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 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 index 92fc383..24c3708 100644 --- a/docs/plans/jetson-orin-nano.md +++ b/docs/plans/jetson-orin-nano.md @@ -1,177 +1,178 @@ -# Jetson Orin Nano CUDA Backend Implementation Plan +# Jetson Orin Nano CUDA implementation plan — revised -> **For Claude Code:** Execute this plan vertically with strict RED → GREEN → REFACTOR cycles. Do not optimize kernels until an Orin Nano baseline has been captured and committed. +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:** Make the Edge0 Torch/CUDA backend run reproducibly on an 8 GB Jetson Orin Nano with NVMe-backed expert streaming, then optimize it without violating MLX parity. +## Goal and scope [C01, C02] -**Architecture:** Build on the correctness reference in upstream PR #19 rather than porting the MLX backend again. Treat Jetson as a distinct memory/I/O profile: dense quantized weights stay resident where possible, sparse experts move through a bounded slot cache, and prerouter predictions schedule asynchronous NVMe → pinned host buffer → CUDA slot transfers. Preserve the existing backend facade and model engines; introduce Jetson behavior as measured configuration and CUDA implementation details, not model-specific forks. +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. -**Tech stack:** Python 3.10+, PyTorch for Jetson, CUDA, safetensors, NVMe, pytest; later CUTLASS/CUDA extensions for fused INT4 gather-GEMM. +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. -## Constraints and evidence +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. -- This branch is stacked on upstream PR #19 (`agourakis82:cuda-backend`), which already ports both shipped models and establishes CPU/MPS/CUDA parity. -- Upstream explicitly warns that the Apple path depends on unified-memory mmap semantics; an efficient CUDA implementation needs a DirectStorage-style data path, not a mechanical MLX port. -- The current Torch backend is a correctness reference. On GB10, Edge0-8B was measured at 3.2 tok/s by default and 8.1 tok/s with large caches, but the latter used about 6 GB resident memory. Those results cannot be projected onto Orin Nano. -- Orin Nano has 8 GB shared memory. The OS, CUDA context, tokenizer, KV cache, and application processes must all fit alongside Edge0. -- Use NVMe. microSD is not an acceptable performance target for expert streaming. +## Acceptance gates -## Non-goals for the first PR +### Gate A — target and installation evidence [C08, C09] -- No unverified performance claims. -- No Edge0-35B optimization before Edge0-8B is stable. -- No new model family. -- No speculative TensorRT rewrite. -- No changes to MLX behavior. +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. -## Acceptance gates +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. -### Gate A — reproducible platform evidence +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. -On the target Orin Nano, save a JSON probe containing: +## Task 1 — retain the probe; record its limits [C09] -- Jetson model and L4T/JetPack release data; -- architecture, total/available shared memory; -- PyTorch, CUDA, CUDA availability, GPU name, compute capability, and reported device memory; -- model/checkpoint filesystem, free space, and whether it resolves to NVMe; -- explicit readiness blockers instead of guessed defaults. +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`. -### Gate B — correctness baseline +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. -With `Edge0/Edge0-8B-A1B-preview` on NVMe: +## Task 2 — functioning, machine-readable benchmark reporting [C03–C07] -- run the CUDA backend smoke tests; -- generate at least 32 greedy tokens; -- record peak process RSS, peak device allocation, prefill time, and decode tok/s; -- compare greedy tokens or logits against the existing reference fixture; -- retain the full command and JSON result. +**Scope:** reporting, backend-aware measurement, focused tests and documentation. No cache/kernel optimization and no mandatory physical-device run in this increment. -### Gate C — bounded memory +**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. -- No OOM on an 8 GB Orin Nano in the agreed power mode. -- Target peak process footprint: below 6.5 GB for a short-context Edge0-8B run, leaving explicit OS headroom. -- Cache sizes are selected from available memory at startup and can be overridden. +### Measurement behavior [C03, C04, C05] -### Gate D — optimized data path +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. -- Expert-cache misses copy into preallocated slots; no per-token CUDA allocation. -- Transfers run on a dedicated CUDA stream and are synchronized with events. -- Prerouter predictions issue prefetch before the expert is consumed. -- Dense quantized layers do not repeatedly unpack/dequantize unchanged weights. -- Every optimized path has a parity test against the correctness path. +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). -## Task 1: Add a read-only Jetson capability probe +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. -**Objective:** Capture the facts needed to reproduce and interpret every Jetson benchmark. +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. -**Files:** -- Create: `scripts/jetson_probe.py` -- Create: `tests/test_jetson_probe.py` -- Modify: `docs/nvidia.md` +### Report contract [C05, C06, C16] -**TDD steps:** +Use a benchmark schema version independent of the probe schema. Include: -1. Test `/proc/meminfo` parsing with kB values and missing `MemAvailable`. -2. Run the focused test and verify it fails because the module is absent. -3. Implement only the memory parser; rerun until green. -4. Test mount selection and NVMe classification using synthetic `/proc/self/mountinfo`. -5. Implement longest-prefix mount matching without shelling out. -6. Test Jetson model/L4T detection from an injected filesystem root. -7. Implement platform and device-tree readers. -8. Test Torch reporting with fake unavailable and fake CUDA modules. -9. Implement Torch collection with import failure represented as data. -10. Test readiness blockers: not Jetson, no CUDA, non-NVMe storage, and a missing model path. -11. Implement `collect_probe()` and JSON CLI output. -12. Run `pytest tests/test_jetson_probe.py -q`, then the full non-slow suite. -13. Commit: `feat(jetson): add Orin capability probe`. +| 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 | -## Task 2: Add a machine-readable inference benchmark result +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. -**Objective:** Extend the existing benchmark path to emit reproducible JSON without changing generation behavior. +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). -**Files:** -- Modify: `examples/bench.py` -- Create or modify: `tests/test_benchmark_report.py` -- Modify: `docs/nvidia.md` +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. -**TDD steps:** +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). -1. Extract a pure report-building function and test its schema first. -2. Include exact checkpoint path, backend, device, context/prompt tokens, generated tokens, prefill seconds/tok-s, decode seconds/tok-s, RSS peak, CUDA allocation peak, cache settings, power mode, and git SHA. -3. Add `--json-output PATH`; retain current human-readable output. -4. Verify output on CPU without real weights using a fixture/fake measurement object. -5. Run the existing benchmark manually on Orin and commit only the command/result documentation—not model files. -6. Commit: `feat(bench): emit reproducible inference reports`. +### CLI and acceptance tests [C07] -## Task 3: Establish the unoptimized Orin baseline +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. -**Objective:** Learn which constraint kills the current implementation before changing it. +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 +``` -**Target commands:** +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 -python scripts/jetson_probe.py --model-dir "$EDGE0_8B_MODEL" --output artifacts/orin-probe.json -EDGE0_BACKEND=cuda EDGE0_TORCH_DEVICE=cuda \ - pytest tests/test_nvidia_backend_smoke.py -q -EDGE0_BACKEND=cuda EDGE0_TORCH_DEVICE=cuda \ - python examples/bench.py edge0-8b --json-output artifacts/orin-baseline.json +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" ``` -Capture `tegrastats` beside the benchmark and note JetPack, power mode, clocks, NVMe model/filesystem, cooling, and ambient throttling. If the run OOMs, reduce context/generation—not model correctness—and record the first failing allocation. +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. -Commit: `docs(jetson): record Orin Nano baseline`. +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). -## Task 4: Add a memory-budgeted runtime profile +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. -**Objective:** Select safe cache settings from measured available memory rather than copying GB10 defaults. +## Task 4 — memory-budgeted profile [C11] -**Files:** -- Modify: `src/edge0/backends/cuda/backend.py` -- Modify: `src/edge0/backends/cuda/nn.py` -- Modify: model/runtime configuration at the narrowest existing override point -- Test: `tests/test_cuda_backend.py` +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. -Start with policy only—no new kernels. The policy must expose its budget calculation, reserve OS/KV headroom, reject impossible overrides, and print the resolved cache budget. Test the calculation as a pure function before wiring it into the engine. +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. -Commit: `feat(cuda): add memory-budgeted cache profile`. +**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. -## Task 5: Replace per-token expert bundle copies +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. -**Objective:** Reuse bounded CUDA expert slots and overlap misses with compute. +## Task 5 — bounded asynchronous expert transfers [C12, C13] -**Approach:** +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. -- Allocate fixed slots once. -- Stage safetensors ranges into reusable pinned host buffers. -- Copy on a dedicated CUDA stream. -- Use CUDA events to publish slot readiness. -- Drive prefetch from the prerouter; preserve exact fallback on a miss. -- Add counters for cache hit, predicted hit, late hit, bytes read, bytes copied, and stall time. +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). -First test state transitions and eviction using fake streams/events. Then add CUDA parity tests. Do not add GPUDirect Storage until the staged path is measured; Jetson support and benefit must be demonstrated first. +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). -Commit: `perf(cuda): add asynchronous expert slot cache`. +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. -## Task 6: Fuse quantized gather/dequantize/matmul +## Task 6 — conditional quantized-kernel work [C14] -**Objective:** Remove repeated Python/Torch unpack operations from decode. +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. -Prototype one kernel for the exact Edge0 INT4 affine layout: packed u32, group size 64, bf16 scales/biases, gathered expert rows. Compare against `quant.gather_qmm` over every broadcast shape already covered by `tests/test_cuda_backend.py`; require the existing numerical tolerance before benchmarking. Prefer a small PyTorch CUDA extension using CUTLASS primitives over a broad TensorRT rewrite. +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. -Commit: `perf(cuda): fuse INT4 expert gather matmul`. +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. -## Task 7: Final Orin validation +## Task 7 — final validation and evidence [C15, C16] -Repeat the exact baseline commands with three or more benchmark runs. Report individual runs plus mean/variation; keep prefill and decode separate. Verify CUDA execution from runtime logs, not merely `torch.cuda.is_available()`. Document remaining thermal, storage, context-length, and quality constraints. +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. -## Claude handoff rules +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. -1. Begin by reading upstream PR #19, issue #20, `docs/nvidia.md`, `docs/streaming.md`, and the backend parity tests. -2. Do not rewrite the CUDA backend or replace the existing engines. -3. Keep the first implementation PR limited to Task 1 plus documentation. -4. Never claim Orin support until Gate B runs on the physical device. -5. Every optimization must retain a selectable correctness/reference path. -6. Do not commit checkpoints, generated artifacts, credentials, or machine identifiers. -7. Stop and report after each acceptance gate with measured evidence and the next bottleneck. +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. diff --git a/examples/bench.py b/examples/bench.py index 971c459..65a5c53 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,45 @@ "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_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 +138,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 +254,482 @@ 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 _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 + if backend.name == "cuda": + from edge0.backends.cuda import nn as cuda_nn + weight_cache = bool(getattr(cuda_nn, "CACHE_DEQUANTIZED", False)) + 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") + 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}"}} + return { + "layer_options": layer_options, + "shared_cache_slots_resolved": shared_slots, + "prefetch_cap_resolved": prefetch_cap, + "weight_cache": weight_cache, + "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/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("