diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 0000000..c332b06 --- /dev/null +++ b/.gitattributes @@ -0,0 +1,3 @@ +# Installed kernel profiles are pinned by SHA-256 of their bytes (profile.json); +# keep them LF on every checkout so the digests match on Windows too. +embodiinfer/backend/triton/tuned/*/** text eol=lf diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 638280c..b305032 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -160,6 +160,131 @@ exclusive, so uv rejects combinations with incompatible Torch or Transformers requirements. Declare new dependencies in the right group and regenerate the lock against the public index. +### Optional kernel tuning tools + +Kernel tuning is repository developer tooling. Its Humanize2 (`hmz`) environment +requires Python 3.12+, separately from the inference runtime and its model groups. + +Real searches run on the Linux/POSIX target machine: the pinned Humanize2 release +uses POSIX journaling primitives. Task checks and artifact inspection/export also +work on Windows. Windows CPU tests adapt only that upstream file-write primitive; +they are not a claim of native Windows coding-agent support. + +```bash +uv sync --project scripts/kernel_tuning --python 3.12 --frozen +uv run --project scripts/kernel_tuning python -m scripts.kernel_tuning --help +uv run --project scripts/kernel_tuning python -m scripts.kernel_tuning generate --catalog gated_residual --output results/kernel_tuning/tasks +uv run --project scripts/kernel_tuning python -m scripts.kernel_tuning check results/kernel_tuning/tasks/gated_residual +``` + +For the end-to-end workflow, start with the +[profiling-to-model-benchmark guide](docs/proposals/0008-kernel-tuning.md#11-automated-tuning-workflow). +Use `profile` to select registered operators, then `capture-ops`, `calibrate`, +`tune-all`, `install`, `verify`, and the existing model benchmark. Selection is +manual; there is no unified `optimize` command. Only registry captures produce +exports accepted by `install`. + +Capture selected registry operators on real inputs in the prepared model runtime, +with CUDA Graphs and compilation disabled. The driver must execute at least +warmup + measured steps: + +```bash +/absolute/runtime/python -m scripts.kernel_tuning capture-ops --model pi05 --ops linear,gelu_tanh --warmup 2 --steps 8 --output results/kernel_tuning/captures/pi05 -- /absolute/preparation.py +/absolute/runtime/python -m scripts.kernel_tuning calibrate --captured results/kernel_tuning/captures/pi05 --output results/kernel_tuning/captures/pi05-calibrated +uv run --project scripts/kernel_tuning python -m scripts.kernel_tuning generate --model pi05 --captured results/kernel_tuning/captures/pi05-calibrated --list +uv run --project scripts/kernel_tuning python -m scripts.kernel_tuning tune-all --model pi05 --captured results/kernel_tuning/captures/pi05-calibrated --agent 'HARNESS/MODEL:EFFORT' --set evaluator_python=/absolute/runtime/python +uv run --project scripts/kernel_tuning python -m scripts.kernel_tuning tune-all --resume results/kernel_tuning/batches/BATCH_ID +``` + +`/absolute/preparation.py` constructs the selected policy through `make_policy` +and runs its inference entries. The shared driver +`scripts.kernel_tuning.discovery.drive` can replay recorded benchmark inputs. +Capture freezes real inputs, output layouts, aliases, numerical settings and call +counts for the selected registry operators. The default fixture budget is 1 GiB; +increase `--fixture-bytes` if necessary. Calibration freezes per-workload/output +FP64 bounds before search; every recorded, random and boundary input must pass. + +`--ops` selects operators at capture time. All resulting tasks must pass baseline +preflight before any model search starts. Model batches cannot use +`--skip-preflight`, partial task selection, estimated shapes or standalone `run` +to bypass that barrier. `--dry-run` builds tasks without evaluation; +`--preflight-only` evaluates baselines without agents. + +The full-model ATen discovery command `capture --model` has been removed. +Use `capture-ops`; old discovery artifacts must be recaptured as registry calls. +The shared capture/task/replay modules keep their existing `discovery.*` paths. +See [capture artifacts and migration](docs/proposals/0008-kernel-tuning.md#registry-capture-artifacts-and-shared-tooling). + +The independent catalog workflow remains available explicitly with +`generate --catalog`, `tune-all --catalog`, and `capture --catalog`. +Only it accepts legacy JSONL captures, catalog model groups and `--estimated`; +catalog exports need separate runtime integration. + +Run these commands from the repository root. Humanize2 is pinned to a reviewed +source revision in the tool's `pyproject.toml`; it drives an existing, separately +installed coding-agent CLI. Select the harness/model/effort explicitly when +starting a search. `run` and `resume` execute real agent turns and kernels; +`check`, `status`, and `export` do not. `check --environment` additionally imports +the task evaluator and probes dependencies/hardware, without compiling candidates. + +On the target machine, set the task's `evaluator_python` to the absolute Python +executable of its prepared GPU runtime. Install +`benchmarks/kernel_tuning/requirements.txt` there under that runtime's constraints. +In particular, preserve Thor's compatible Torch/CUDA/Triton versions; do not copy +the tooling interpreter's packages or another GPU's Torch pins over them. Inspect +the proposed dependency changes first, for example with +`uv --no-config pip install --dry-run --python /absolute/runtime/python -r benchmarks/kernel_tuning/requirements.txt -c /absolute/runtime/constraints.txt`. +For a separately prepared GPU runtime, `--no-config` avoids applying this +repository's development-only Torch upper bound. Pin that runtime's Torch, +Triton, and CUDA packages in the constraints file and inspect the installation +plan before applying it. When inheriting system packages, verify that the installer +actually reuses them; do not reinstall the GPU stack just to add the evaluator. +For a custom evaluator, install only its declared dependencies plus PyYAML. + +See the [canonical design and implementation status](docs/proposals/0008-kernel-tuning.md). +The [workflow guide](docs/proposals/0008-kernel-tuning.md#11-automated-tuning-workflow) +explains task inputs, the full tuning loop, stopping conditions, and saved results. +The controller snapshots task inputs, saves all attempts, and exports improving +solutions with their correctness and timing evidence under `results/kernel_tuning/`. +Registry exports can be installed into an optional hardware profile, followed by +whole-model verification, benchmark and review. Other exports require separate +integration before runtime use; installation alone does not certify a profile. +For an unregistered operator, the +[registration workflow](scripts/kernel_tuning/registration/README.md) creates an +isolated review branch with `register new`, then runs a resumable CPU +plan/implement/check/repair loop with `register run`. Supply a prepared model +interpreter using `--set test_python=/absolute/runtime/python` on `new` (and, if +needed, the first `run`). `finalize` exports the checked patch; commit and push +are explicit options. This registration gate does not replace target-GPU +whole-model bit-exact verification before merge. +Humanize2 agents may execute commands without interactive approval; the scratch +directory and content hashes are not an OS sandbox. Use an appropriate isolated +account/container on the target host when needed. +By default Humanize2 fences each agent session (workdir writable, rest of the +machine read-only), which needs Linux Landlock: kernel 5.13+ built with +`CONFIG_SECURITY_LANDLOCK` and listed in `/sys/kernel/security/lsm`. Some vendor +kernels, such as Jetson L4T, omit it. On such hosts pass `--unfenced` to `run`, +`tune-all` or `register run`; the choice is frozen into the run, and resume reuses +it. Unfenced agents can write anything their account can, so run them only under +a dedicated account without `sudo` or `docker` group membership, or in a container. + +CPU tests for the tooling use official Humanize2 fake agents and fake evaluators, +without a coding-agent account or GPU stack: + +```bash +uv run --project scripts/kernel_tuning python -m pytest tests/kernel_tuning --confcutdir=tests/kernel_tuning -q +uv run --project scripts/kernel_tuning ruff check scripts/kernel_tuning benchmarks/kernel_tuning tests/kernel_tuning +uv run --project scripts/kernel_tuning ruff format --check scripts/kernel_tuning benchmarks/kernel_tuning tests/kernel_tuning +``` + +CPU registry capture/replay tests also need Torch in the selected interpreter; they do +not load checkpoints or run GPU tuning. The `--confcutdir` option avoids loading the core test suite's Torch-dependent +fixtures in the separate tooling environment. Optional numerical/GPU tests skip +when their dependencies/hardware are unavailable; skips are not GPU validation. +After preparing the target GPU runtime and CUDA toolkit, exercise both the Triton +and CUDA C++ builders, byte-level correctness, and changed-input CUDA Graph replay +with `/absolute/runtime/python -m pytest tests/kernel_tuning/test_numerics.py --confcutdir=tests/kernel_tuning -q`. + ## Checks to run ```bash diff --git a/README.md b/README.md index 82f815a..4d3ccc7 100644 --- a/README.md +++ b/README.md @@ -212,6 +212,7 @@ Existing `vvla-*` command aliases and `vvla.policy.*` wire schemas remain suppor | Integrate or extend the engine | [Architecture](docs/en/architecture.md) · [Python API](docs/en/api.md) | | Inspect model behavior | [Models](docs/en/models.md) | | Evaluate performance | [Benchmarks](docs/en/benchmark.md) | +| Profile and tune kernels with an agent | [Kernel tuning workflow](docs/proposals/0008-kernel-tuning.md) | ## Contributing diff --git a/benchmarks/kernel_tuning/calibration.py b/benchmarks/kernel_tuning/calibration.py new file mode 100644 index 0000000..576d60c --- /dev/null +++ b/benchmarks/kernel_tuning/calibration.py @@ -0,0 +1,274 @@ +"""Candidate-free calibration and replay of frozen operator numerical contracts.""" + +from __future__ import annotations + +import importlib +import os +import shutil +import tempfile +from collections.abc import Iterator +from contextlib import contextmanager +from pathlib import Path +from typing import Any + +import torch +from scripts.kernel_tuning.artifacts import atomic_json +from scripts.kernel_tuning.contracts import ContractError, digest, tree_hashes +from scripts.kernel_tuning.discovery.contracts import ModelCapture +from scripts.kernel_tuning.discovery.numerics import ( + PROFILES, + SEEDS, + VERSION, + execution_precision, + implementation_digest, + policy, + profile_key, + validate_contract, +) +from scripts.kernel_tuning.discovery.tensors import ( + clone_inputs, + describe_inputs, + extent, + materialize, + restore, + tensor_outputs, +) + + +def environment(device: str) -> dict[str, Any]: + """Identity of the actual arithmetic environment, checked again at evaluation.""" + cuda = torch.device(device).type == "cuda" + return { + "torch": str(torch.__version__), + "cuda": torch.version.cuda, + "device": device, + "hardware": torch.cuda.get_device_name(device) if cuda else "cpu", + "capability": list(torch.cuda.get_device_capability(device)) if cuda else [], + "matmul_precision": torch.get_float32_matmul_precision(), + "allow_tf32": torch.backends.cuda.matmul.allow_tf32, + "cudnn_allow_tf32": torch.backends.cudnn.allow_tf32, + "fp16_reduced_reduction": torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction, + "bf16_reduced_reduction": torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction, + "deterministic_algorithms": torch.are_deterministic_algorithms_enabled(), + "num_threads": torch.get_num_threads(), + "default_dtype": str(torch.get_default_dtype()), + "autocast_enabled": torch.is_autocast_enabled("cuda" if cuda else "cpu"), + "environment": { + key: os.environ.get(key) for key in ("NVIDIA_TF32_OVERRIDE", "CUBLAS_WORKSPACE_CONFIG") + }, + } + + +@contextmanager +def captured_settings(precision: dict[str, Any], flags: dict[str, Any]) -> Iterator[None]: + """Use captured arithmetic switches and restore every changed process-global flag.""" + old = environment("cpu") + warn = torch.is_deterministic_algorithms_warn_only_enabled() + try: + torch.set_float32_matmul_precision(precision["matmul_precision"]) + torch.backends.cudnn.allow_tf32 = flags["cudnn_allow_tf32"] + torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = flags["fp16_reduced_reduction"] + torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = flags["bf16_reduced_reduction"] + torch.use_deterministic_algorithms(flags["deterministic_algorithms"]) + yield + finally: + torch.set_float32_matmul_precision(old["matmul_precision"]) + torch.backends.cudnn.allow_tf32 = old["cudnn_allow_tf32"] + torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = old["fp16_reduced_reduction"] + torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = old["bf16_reduced_reduction"] + torch.use_deterministic_algorithms(old["deterministic_algorithms"], warn_only=warn) + + +def production(identity: dict[str, Any], inputs: list[torch.Tensor]) -> list[torch.Tensor]: + """Execute only the captured production implementation, never a candidate.""" + args, kwargs = restore(identity["arguments"], inputs) + if identity["kind"] != "registry": + raise ContractError("Only registry operators can be calibrated; recapture with capture-ops") + fn = getattr(importlib.import_module(identity["module"]), identity["entry"]) + return tensor_outputs(fn(*args, **kwargs)) + + +def _arguments(identity: dict[str, Any], inputs: list[torch.Tensor]) -> dict[str, Any]: + import inspect + + args, kwargs = restore(identity["arguments"], inputs) + function = getattr(importlib.import_module(identity["module"]), identity["entry"]) + return {**dict(zip(inspect.signature(function).parameters, args)), **kwargs} + + +def validation_inputs( + identity: dict[str, Any], + case: dict[str, Any], + root: Path, + profile: str, + seed: int, +) -> list[torch.Tensor]: + """Generate a fixed validation input while retaining masks, indices, layouts and aliases.""" + inputs = materialize(case["inputs"], root, 0) + if profile == "recorded": + return inputs + if (profile, seed) not in PROFILES: + raise ContractError("Unknown numerical validation profile") + arguments = _arguments(identity, inputs) + protected = [ + value for name, value in arguments.items() if "mask" in name and isinstance(value, torch.Tensor) + ] + protected_groups = { + spec["storage"] + for value, spec in zip(inputs, case["inputs"]) + if any(value is item for item in protected) + } + specs = describe_inputs(inputs) + seen = set() + for tensor, spec in zip(inputs, specs): + group = spec["storage"] + if group in seen or group in protected_groups or not tensor.is_floating_point(): + continue + seen.add(group) + size = max(extent(item) for item in specs if item["storage"] == group) + storage = tensor.as_strided((size,), (1,), 0) + if profile == "random": + generator = torch.Generator(device=tensor.device).manual_seed(seed * 1009 + group) + storage.copy_(torch.randn(size, dtype=tensor.dtype, device=tensor.device, generator=generator)) + elif profile == "zeros": + storage.zero_() + else: + storage.copy_((torch.arange(size, device=tensor.device) % 2) * 2 - 1) + return inputs + + +def high_precision(identity: dict[str, Any], inputs: list[torch.Tensor]) -> list[torch.Tensor]: + """Evaluate the registry's Torch reference with floating inputs promoted to FP64.""" + if policy(identity) != "registry": + raise ContractError("No trusted high-precision adapter for this operator") + promoted = [x.to(torch.float64) if x.is_floating_point() else x for x in inputs] + return production(identity, promoted) + + +def derive_bound(baseline: torch.Tensor, reference: torch.Tensor) -> dict[str, Any]: + """Derive a per-output all-element bound using only baseline/reference error.""" + if ( + baseline.shape != reference.shape + or not baseline.is_floating_point() + or reference.dtype != torch.float64 + ): + raise ContractError("Calibration requires matching floating outputs and an FP64 reference") + if not torch.isfinite(baseline).all() or not torch.isfinite(reference).all(): + raise ContractError("Nonfinite baseline/reference cannot calibrate a tolerance") + info = torch.finfo(baseline.dtype) + error = (baseline.double() - reference).abs() + rms = reference.square().mean().sqrt().item() if reference.numel() else 0.0 + residual = (error - info.eps * reference.abs()).clamp_min(0).max().item() if error.numel() else 0.0 + return { + "dtype": str(baseline.dtype).removeprefix("torch."), + "reference_dtype": "float64", + "atol": max(info.eps * rms, 2 * residual, info.tiny), + "rtol": info.eps, + "baseline_max_abs": error.max().item() if error.numel() else 0.0, + "residual_max": residual, + "rms": rms, + "epsilon": info.eps, + "tiny": info.tiny, + } + + +def compare_calibrated( + actual: list[torch.Tensor], reference: list[torch.Tensor], bounds: list[dict[str, Any]] +) -> tuple[float, float]: + """Enforce frozen per-output thresholds on every finite element, with no matched-ratio escape.""" + if len(actual) != len(reference) or len(actual) != len(bounds): + raise ContractError("Calibrated output count changed") + maximum, relative = 0.0, 0.0 + for value, ref, bound in zip(actual, reference, bounds): + if ( + str(value.dtype).removeprefix("torch.") != bound["dtype"] + or value.shape != ref.shape + or value.device != ref.device + ): + raise ContractError("Calibrated output shape/dtype/device changed") + if not torch.isfinite(value).all() or not torch.isfinite(ref).all(): + raise ContractError("Calibrated comparison requires finite outputs") + error = (value.double() - ref).abs() + if not torch.all(error <= bound["atol"] + bound["rtol"] * ref.abs()): + raise ContractError("Output exceeds frozen calibrated tolerance") + if error.numel(): + maximum = max(maximum, error.max().item()) + # Use the absolute floor as the denominator floor to keep near-zero diagnostics finite. + relative = max(relative, (error / ref.abs().clamp_min(bound["atol"])).max().item()) + return maximum, relative + + +def calibrate(capture: ModelCapture, destination: Path) -> ModelCapture: + """Create a new immutable calibrated capture before any candidate source exists.""" + from scripts.kernel_tuning.discovery.tasks import _sources + + checked = ModelCapture.load(capture.root) + if checked.identity != capture.identity: + raise ContractError("Capture changed before calibration") + if capture.manifest["torch_version"] != str(torch.__version__): + raise ContractError("Calibration requires the capture's Torch version") + destination = destination.resolve() + if destination.exists(): + raise ContractError("Calibration requires a new destination; existing captures/runs are immutable") + if destination.is_relative_to(capture.root): + raise ContractError("Calibration destination must be outside its source capture") + results = {} + precision, flags = capture.manifest["precision"], capture.manifest["flags"] + with captured_settings(precision, flags), torch.no_grad(): + for op in capture.operators: + identity = op["identity"] + reference = policy(identity) if op["status"] == "ready" else None + if reference is None: + continue + _sources(op) # Verify backend source/dependency digests before executing them. + cases = {case["id"]: case for case in op["workloads"]} + device = op["workloads"][0]["outputs"][0]["device"] + contract = { + "version": VERSION, + "implementation_digest": implementation_digest(), + "reference": reference, + "operator_identity": digest(identity), + "precision": execution_precision(precision), + "flags": flags, + "profiles": [list(pair) for pair in PROFILES], + "seeds": list(SEEDS), + "environment": environment(device), + "workloads": {k: digest(v) for k, v in cases.items()}, + "cases": {}, + } + for key, case in cases.items(): + if any("fixture" not in spec for spec in case["inputs"]): + raise ContractError( + f"{op['name']}/{key}: real inputs missing; recapture with a larger --fixture-bytes" + ) + samples = {} + for profile, seed in PROFILES: + inputs = validation_inputs(identity, case, capture.root, profile, seed) + baseline = production(identity, clone_inputs(inputs)) + ref = high_precision(identity, inputs) + if len(baseline) != len(ref) or len(ref) != len(case["outputs"]): + raise ContractError("High-precision reference output count differs from capture") + bounds = [derive_bound(b, r) for b, r in zip(baseline, ref)] + compare_calibrated(baseline, ref, bounds) + samples[profile_key(profile, seed)] = bounds + contract["cases"][key] = samples + if environment(device) != contract["environment"]: + raise ContractError("Numerical environment changed during calibration") + validate_contract(contract, identity, cases, precision, flags) + results[op["id"]] = contract + destination.parent.mkdir(parents=True, exist_ok=True) + with tempfile.TemporaryDirectory(prefix="calibration-", dir=destination.parent) as scratch: + root = Path(scratch) / "capture" + shutil.copytree(capture.root, root) + # Refuse a source changed while calibration was running or while copying. + if ModelCapture.load(root).identity != capture.identity: + raise ContractError("Capture changed during calibration") + atomic_json(root / "calibration.json", results) + manifest = dict(capture.manifest) + files = tree_hashes(root) + files.pop("manifest.json") + manifest.update(files=files, numerical_policy=VERSION, calibration_source=capture.identity) + atomic_json(root / "manifest.json", manifest) + ModelCapture.load(root) + root.rename(destination) + return ModelCapture.load(destination) diff --git a/benchmarks/kernel_tuning/captures/pi05-libero10.jsonl b/benchmarks/kernel_tuning/captures/pi05-libero10.jsonl new file mode 100644 index 0000000..5a5f855 --- /dev/null +++ b/benchmarks/kernel_tuning/captures/pi05-libero10.jsonl @@ -0,0 +1,8 @@ +{"axes": {"B": 1, "H": 1024, "H3": 3072, "T": 50}, "operator": "ada_rms_norm", "scalars": {"eps": 1e-06}, "count": 74370} +{"axes": {"B": 1, "G": 1, "H": 1024, "T": 50}, "operator": "gated_residual", "count": 72360} +{"axes": {"B": 1, "D": 256, "HK": 1, "HQ": 8, "S": 50}, "operator": "rotate_qk", "count": 36180} +{"axes": {"B": 1, "I": 4096, "T": 50}, "operator": "gated_gelu", "count": 36180} +{"axes": {"B": 1, "D": 256, "HK": 1, "HQ": 8, "P": 576, "S": 50}, "operator": "split_kv_attention", "scalars": {"scaling": 0.0625}, "count": 30780} +{"axes": {"B": 1, "D": 256, "HK": 1, "HQ": 8, "P": 560, "S": 50}, "operator": "split_kv_attention", "scalars": {"scaling": 0.0625}, "count": 5400} +{"axes": {"B": 1, "D": 256, "HK": 1, "HQ": 8, "P": 0, "S": 576}, "operator": "split_kv_attention", "scalars": {"scaling": 0.0625}, "count": 3078} +{"axes": {"B": 1, "D": 256, "HK": 1, "HQ": 8, "P": 0, "S": 560}, "operator": "split_kv_attention", "scalars": {"scaling": 0.0625}, "count": 540} diff --git a/benchmarks/kernel_tuning/captures/qwen-r2r-low.jsonl b/benchmarks/kernel_tuning/captures/qwen-r2r-low.jsonl new file mode 100644 index 0000000..93b0747 --- /dev/null +++ b/benchmarks/kernel_tuning/captures/qwen-r2r-low.jsonl @@ -0,0 +1,2 @@ +{"axes": {"D": 80, "HK": 16, "HQ": 16, "T": 1980}, "operator": "rotate_half_rope", "count": 2304} +{"axes": {"D": 80, "HK": 16, "HQ": 16, "N": 46, "T": 1980}, "fixed": {"segment_offsets": [0, 64, 128, 176, 240, 304, 352, 368, 384, 396, 460, 524, 572, 636, 700, 748, 764, 780, 792, 856, 920, 968, 1032, 1096, 1144, 1160, 1176, 1188, 1252, 1316, 1364, 1428, 1492, 1540, 1556, 1572, 1584, 1648, 1712, 1760, 1824, 1888, 1936, 1952, 1968, 1980]}, "operator": "segmented_attention", "scalars": {"max_length": 64}, "count": 2016} diff --git a/benchmarks/kernel_tuning/captures/qwen-r2r-panoramic.jsonl b/benchmarks/kernel_tuning/captures/qwen-r2r-panoramic.jsonl new file mode 100644 index 0000000..8756986 --- /dev/null +++ b/benchmarks/kernel_tuning/captures/qwen-r2r-panoramic.jsonl @@ -0,0 +1,2 @@ +{"axes": {"D": 80, "HK": 16, "HQ": 16, "T": 7704}, "operator": "rotate_half_rope", "count": 2208} +{"axes": {"D": 80, "HK": 16, "HQ": 16, "N": 172, "T": 7704}, "fixed": {"segment_offsets": [0, 64, 128, 192, 256, 320, 384, 448, 512, 544, 608, 672, 736, 800, 864, 928, 992, 1056, 1088, 1104, 1120, 1136, 1152, 1168, 1184, 1200, 1216, 1224, 1288, 1352, 1416, 1480, 1544, 1608, 1672, 1736, 1768, 1832, 1896, 1960, 2024, 2088, 2152, 2216, 2280, 2312, 2328, 2344, 2360, 2376, 2392, 2408, 2424, 2440, 2448, 2512, 2576, 2640, 2704, 2768, 2832, 2896, 2960, 2992, 3056, 3120, 3184, 3248, 3312, 3376, 3440, 3504, 3536, 3552, 3568, 3584, 3600, 3616, 3632, 3648, 3664, 3672, 3736, 3800, 3864, 3928, 3992, 4056, 4120, 4184, 4216, 4280, 4344, 4408, 4472, 4536, 4600, 4664, 4728, 4760, 4776, 4792, 4808, 4824, 4840, 4856, 4872, 4888, 4896, 4960, 5024, 5088, 5152, 5216, 5280, 5344, 5408, 5440, 5504, 5568, 5632, 5696, 5760, 5824, 5888, 5952, 5984, 6000, 6016, 6032, 6048, 6064, 6080, 6096, 6112, 6120, 6184, 6248, 6296, 6360, 6424, 6472, 6488, 6504, 6516, 6580, 6644, 6692, 6756, 6820, 6868, 6884, 6900, 6912, 6976, 7040, 7088, 7152, 7216, 7264, 7280, 7296, 7308, 7372, 7436, 7484, 7548, 7612, 7660, 7676, 7692, 7704]}, "operator": "segmented_attention", "scalars": {"max_length": 64}, "count": 1932} diff --git a/benchmarks/kernel_tuning/captures/streamvln-r2r-rxr.jsonl b/benchmarks/kernel_tuning/captures/streamvln-r2r-rxr.jsonl new file mode 100644 index 0000000..ec8318b --- /dev/null +++ b/benchmarks/kernel_tuning/captures/streamvln-r2r-rxr.jsonl @@ -0,0 +1,3 @@ +{"axes": {"H": 3584, "M": 1}, "operator": "add_rms_norm", "scalars": {"eps": 1e-06}, "count": 330400} +{"axes": {"I": 18944, "M": 1}, "operator": "swiglu", "count": 165200} +{"axes": {"H": 3584, "M": 1}, "operator": "rms_norm", "scalars": {"eps": 1e-06}, "count": 5900} diff --git a/benchmarks/kernel_tuning/discovery_inputs.py b/benchmarks/kernel_tuning/discovery_inputs.py new file mode 100644 index 0000000..f5c0b65 --- /dev/null +++ b/benchmarks/kernel_tuning/discovery_inputs.py @@ -0,0 +1,176 @@ +"""Export benchmark data as end-to-end inputs for profiling, registry capture and verification. + +The output is read by ``scripts.kernel_tuning.discovery.drive``. Data selection +reuses each benchmark's own loader, so the shared driver sees the same images, states, +instructions and episode order as the measurement. Run in the model's runtime:: + + python -m benchmarks.kernel_tuning.discovery_inputs pi05 --config benchmarks/pi05-benchmark/config.yaml \\ + --output pi05-inputs.pt --serving-config pi05-serving.json + python -m benchmarks.kernel_tuning.discovery_inputs qwen --kind low --config benchmarks/qwenvl-benchmark/config.yaml \\ + --output qwen-low-inputs.pt + python -m benchmarks.kernel_tuning.discovery_inputs streamvln --config benchmarks/streamvln-benchmark/config.yaml \\ + --output streamvln-inputs.pt + +pi05 and StreamVLN produce serving requests (their deployment entry); Qwen +navigation profiles produce engine observations. Each episode keeps its frame +order so recurrent memory grows as it does in deployment. +""" + +from __future__ import annotations + +import argparse +import importlib.util +import io +import json +import sys +from pathlib import Path +from types import ModuleType +from typing import Any + +import yaml +from scripts.kernel_tuning.discovery.drive import save_inputs + +from embodiinfer.engine.serve.contracts import RawImage, RawPolicyRequest + +_ROOT = Path(__file__).resolve().parents[2] +_PI05_IMAGES = ( + ("agentview_rgb", "observation.images.image"), + ("eye_in_hand_rgb", "observation.images.image2"), +) +_PI05_STATE = ("ee_pos", "ee_ori", "gripper_states") + + +def _benchmark(directory: str) -> ModuleType: + """Import a benchmark script whose directory name is not a Python package.""" + path = _ROOT / "benchmarks" / directory / "benchmark.py" + spec = importlib.util.spec_from_file_location(f"_discovery_{directory.replace('-', '_')}", path) + module = importlib.util.module_from_spec(spec) + # Dataclasses resolve their module through sys.modules while the script executes. + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +def _config(path: Path) -> dict[str, Any]: + config = yaml.safe_load(path.read_text(encoding="utf-8")) + if not isinstance(config, dict) or config.get("schema_version") != 1: + raise ValueError("expected a benchmark schema_version: 1 mapping") + return config + + +def _png(image: Any) -> bytes: + buffer = io.BytesIO() + image.save(buffer, format="PNG") + return buffer.getvalue() + + +def pi05(args: argparse.Namespace) -> list[list[RawPolicyRequest]]: + """LIBERO frames as pi05 serving requests, plus the matching adapter config.""" + from PIL import Image + + module = _benchmark("pi05-benchmark") + samples = module.load_libero(_config(args.config)["dataset"])[: args.samples] + requests = [] + for index, sample in enumerate(samples): + raw = module.read_libero(sample) + images = [] + for source, field in _PI05_IMAGES: + image = raw[source] + if sample.image_convention == "opengl": + image = image.transpose(Image.Transpose.FLIP_TOP_BOTTOM) + images.append(RawImage(field, "image/png", _png(image))) + requests.append( + RawPolicyRequest( + session_id=f"libero-{index}", + request_id=f"libero-{index}", + step_id=0, + instruction=sample.instruction, + state={name: raw[name].tolist() for name in _PI05_STATE}, + images=tuple(images), + metadata={}, + ) + ) + serving = { + "state_fields": list(_PI05_STATE), + "image_fields": [field for _, field in _PI05_IMAGES], + "policy_kwargs": json.loads(args.policy_kwargs), + } + args.serving_config.write_text(json.dumps(serving, indent=2) + "\n", encoding="utf-8") + # pi05 is stateless: one request per session, all in a single episode. + return [requests] + + +def qwen(args: argparse.Namespace) -> list[list[Any]]: + """Consecutive navigation frames as Qwen engine observations, one episode each.""" + module = _benchmark("qwenvl-benchmark") + spec = {**_config(args.config)["datasets"][0], "episode_limit": args.episodes} + episodes = [] + for episode in module.load_navigation(spec): + # Panoramic inputs read two neighbouring frames on either side of a step. + steps = range(2, min(len(episode.frames) - 2, 2 + args.steps)) + episodes.append( + [ + module.prepare_sample( + episode, step, 0, args.kind, module.read_images(episode, step, 0, args.kind) + ).observation + for step in steps + ] + ) + return episodes + + +def streamvln(args: argparse.Namespace) -> list[list[RawPolicyRequest]]: + """Consecutive navigation frames as StreamVLN serving requests, one session per episode.""" + from embodiinfer.policies.streamvln.serving import STREAMVLN_IMAGE_FIELD + + module = _benchmark("streamvln-benchmark") + spec = {**_config(args.config)["datasets"][0], "episode_limit": args.episodes} + episodes = [] + for episode in module.load_navigation(spec): + session = f"{episode.dataset}-{episode.episode_id}" + episodes.append( + [ + RawPolicyRequest( + session_id=session, + request_id=f"{session}-{step}", + step_id=step, + instruction=episode.instruction, + state={}, + images=(RawImage(STREAMVLN_IMAGE_FIELD, "image/jpeg", path.read_bytes()),), + metadata={}, + ) + for step, path in enumerate(episode.frames[: args.steps]) + ] + ) + return episodes + + +def main(argv: list[str] | None = None) -> None: + """Write one model's recorded inference inputs.""" + parser = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + models = parser.add_subparsers(dest="model", required=True) + for name in ("pi05", "qwen", "streamvln"): + model = models.add_parser(name) + model.add_argument("--config", type=Path, required=True, help="The model benchmark's YAML config") + model.add_argument("--output", type=Path, required=True) + if name == "pi05": + model.add_argument("--samples", type=int, default=4) + model.add_argument("--serving-config", type=Path, required=True, help="Adapter config to write") + model.add_argument( + "--policy-kwargs", default="{}", help="JSON builder options for the adapter config" + ) + else: + model.add_argument("--episodes", type=int, default=2) + model.add_argument("--steps", type=int, default=8, help="Frames per episode") + if name == "qwen": + model.add_argument("--kind", choices=("low", "panoramic"), required=True) + args = parser.parse_args(argv) + episodes = {"pi05": pi05, "qwen": qwen, "streamvln": streamvln}[args.model](args) + save_inputs(args.output, "observations" if args.model == "qwen" else "requests", episodes) + print(json.dumps({"inputs": str(args.output), "episodes": [len(episode) for episode in episodes]})) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/kernel_tuning/evaluate.py b/benchmarks/kernel_tuning/evaluate.py new file mode 100644 index 0000000..16c6e52 --- /dev/null +++ b/benchmarks/kernel_tuning/evaluate.py @@ -0,0 +1,110 @@ +"""Trusted, isolated entry point for the kernel-tuning measurement protocol.""" + +from __future__ import annotations + +import argparse +import importlib.util +import sys +import traceback +from pathlib import Path +from typing import Any, Protocol + +# This is a repository measurement program, also usable with an isolated Python. +sys.path.insert(0, str(Path(__file__).resolve().parents[2])) + +from scripts.kernel_tuning.artifacts import atomic_json # noqa: E402 +from scripts.kernel_tuning.contracts import TaskPackage, digest, read_json # noqa: E402 + + +class EvaluationAdapter(Protocol): + """Task-owned extension for hardware/input semantics outside the default adapter. + + Implement create_adapter(task) in the frozen benchmark.py. Methods return JSON + data; evaluate uses the checks/rounds contract documented in the proposal. + environment must separate stable identity from transient measurement conditions. + """ + + def environment(self) -> dict[str, Any]: + """Report stable hardware/software identity and experimental conditions.""" + ... + + def evaluate(self, solutions: dict[str, dict[str, Any]]) -> dict[str, Any]: + """Validate every solution and measure baseline/incumbent/candidate pairs.""" + ... + + +def load_adapter(task: TaskPackage) -> EvaluationAdapter: + """Select the frozen task benchmark or the pinned FlashInfer implementation.""" + custom = task.root / "benchmark.py" + if custom.exists(): + spec = importlib.util.spec_from_file_location(f"kernel_benchmark_{task.identity}", custom) + if spec is None or spec.loader is None: + raise ValueError("Cannot import task benchmark.py") + module = importlib.util.module_from_spec(spec) + # Helpers can be shipped beside benchmark.py; the entire task is hashed. + sys.path.insert(0, str(task.root)) + spec.loader.exec_module(module) + return module.create_adapter(task) + from flashinfer_adapter import FlashInferAdapter + + return FlashInferAdapter(task) + + +def execute(request: dict[str, Any], output: Path) -> dict[str, Any]: + """Bind the adapter's evidence to the exact input sources and request nonce.""" + if request.get("schema_version") != 1 or request.get("operation") not in ("probe", "evaluate", "profile"): + raise ValueError("Unsupported evaluator request") + task = TaskPackage.load(Path(request["task_dir"])) + if task.identity != request["task_digest"]: + raise ValueError("Task content changed before evaluation") + solutions = {role: read_json(Path(path)) for role, path in request["solutions"].items()} + if {role: digest(solution) for role, solution in solutions.items()} != request["solution_digests"]: + raise ValueError("Solution content changed before evaluation") + adapter = load_adapter(task) + before = adapter.environment() + response = { + "schema_version": 1, + "status": "ok", + "request_digest": digest(request), + "task_digest": task.identity, + "solution_digests": request["solution_digests"], + "environment": before, + } + if request["operation"] == "evaluate": + response["measurement"] = adapter.evaluate(solutions) + elif request["operation"] == "profile": + profiler = getattr(adapter, "profile", None) + if profiler is None: + raise ValueError("This benchmark adapter does not implement profile(solutions, output_dir)") + response["profile"] = profiler(solutions, output.parent) + after = adapter.environment() + if before["identity"] != after["identity"]: + raise ValueError("Hardware/software or precision conditions changed during evaluation") + response["environment_after"] = after + task.verify() + for role, path in request["solutions"].items(): + if digest(read_json(Path(path))) != request["solution_digests"][role]: + raise ValueError("Solution file changed during evaluation") + return response + + +def main() -> int: + """Emit structured results only after successful measurement; errors exit nonzero.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--request", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + try: + response = execute(read_json(args.request), args.output) + atomic_json(args.output, response) + return 0 + except Exception as exc: + traceback.print_exc() + atomic_json( + args.output, {"schema_version": 1, "status": "error", "error": f"{type(exc).__name__}: {exc}"} + ) + return 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/benchmarks/kernel_tuning/flashinfer_adapter.py b/benchmarks/kernel_tuning/flashinfer_adapter.py new file mode 100644 index 0000000..20bd85b --- /dev/null +++ b/benchmarks/kernel_tuning/flashinfer_adapter.py @@ -0,0 +1,542 @@ +"""FlashInfer 0.1.2 building/input generation with strict KDA measurement gates. + +The upstream default evaluator permits numerical tolerances and only eager +timing. We reuse its builders, generators, and eager timer and implement the +additional frozen precision, paired timing, and graph contracts here. +""" + +from __future__ import annotations + +import importlib.metadata +import os +import platform +import random +import shutil +import subprocess +import sys +from dataclasses import asdict +from pathlib import Path +from typing import Any + +sys.path.insert(0, str(Path(__file__).resolve().parents[2])) + +from scripts.kernel_tuning.contracts import ( # noqa: E402 + FLASHINFER_VERSION, + ROLES, + ContractError, + Precision, + TaskPackage, +) + + +def compare_outputs(actual: list[Any], expected: list[Any], precision: Precision) -> tuple[float, float]: + """Compare all output tensor bytes or a predeclared elementwise tolerance. + + Byte comparison preserves signed zero and NaN payload distinctions. Shape, + device and dtype are checked first, so conversion cannot conceal an error. + """ + import torch + + if len(actual) != len(expected): + raise ContractError("Output count differs from reference") + max_abs, max_rel = 0.0, 0.0 + for value, reference in zip(actual, expected): + if ( + value.shape != reference.shape + or value.dtype != reference.dtype + or value.device != reference.device + ): + raise ContractError("Output shape, dtype, or device differs from reference") + lhs = value.detach().contiguous().reshape(-1) + rhs = reference.detach().contiguous().reshape(-1) + if precision.mode == "bit_exact": + if not torch.equal(lhs.view(torch.uint8), rhs.view(torch.uint8)): + raise ContractError("Output bytes differ from the bit-exact reference") + continue + if not torch.is_floating_point(lhs): + if not torch.equal(lhs, rhs): + raise ContractError("Non-floating outputs must match exactly") + continue + if not torch.isfinite(lhs).all().item() or not torch.isfinite(rhs).all().item(): + raise ContractError("Tolerance comparison requires finite reference and candidate outputs") + lhs, rhs = lhs.to(torch.float64), rhs.to(torch.float64) + delta = (lhs - rhs).abs() + if not torch.all(delta <= precision.atol + precision.rtol * rhs.abs()).item(): + raise ContractError("Output exceeds the fixed atol + rtol * abs(reference) threshold") + if delta.numel(): + max_abs = max(max_abs, delta.max().item()) + max_rel = max( + max_rel, (delta / rhs.abs().clamp_min(torch.finfo(torch.float64).tiny)).max().item() + ) + return max_abs, max_rel + + +class FlashInferAdapter: + """Measure pure tensor operators on one selected NVIDIA device. + + Tasks requiring aliasing, mutable input state, non-contiguous fixture creation, + statistical correctness, or other accelerators provide benchmark.py instead. + """ + + def __init__(self, task: TaskPackage) -> None: + import torch + from flashinfer_bench.data import Trace + + if importlib.metadata.version("flashinfer-bench") != FLASHINFER_VERSION: + raise ContractError(f"This adapter requires flashinfer-bench=={FLASHINFER_VERSION}") + if not torch.cuda.is_available() or torch.version.hip or not task.settings.device.startswith("cuda"): + raise ContractError("Default FlashInfer adapter requires a prepared NVIDIA CUDA environment") + self.task = task + self.cfg = task.settings + if self.cfg.language == "triton": + import triton # noqa: F401 -- verify the selected builder before opening an agent + elif not shutil.which("nvcc"): + raise ContractError("CUDA C++ tasks require nvcc on PATH before tuning") + self.device = str(torch.device(self.cfg.device)) + torch.cuda.set_device(self.device) + self.device = f"cuda:{torch.cuda.current_device()}" + self.definition = self._load_definition(task.definition) + self.traces = [Trace.model_validate(trace) for trace in task.workloads] + for trace in self.traces: + axes = {k: v.value for k, v in self.definition.axes.items() if v.type == "const"} + axes.update(trace.workload.axes) + for constraint in self.definition.constraints: + if not eval(constraint, {"__builtins__": {}, "min": min, "max": max}, axes): + raise ContractError(f"Workload {trace.workload.uuid} violates {constraint}") + torch.set_float32_matmul_precision(self.cfg.precision.matmul_precision) + torch.backends.cudnn.allow_tf32 = self.cfg.precision.allow_tf32 + self._precision_flags = self._flags() + self._current_stream = torch.cuda.current_stream(self.device).cuda_stream + self._nvcc = self._command(["nvcc", "--version"]) + self._driver = self._command( + ["nvidia-smi", "--query-gpu=uuid,driver_version", "--format=csv,noheader"] + ) + self._flush: Any = None + + def _load_definition(self, definition: dict[str, Any]) -> Any: + """Validate the upstream schema; specialized replay adapters may extend its dtype set.""" + from flashinfer_bench.data import Definition + + return Definition.model_validate(definition) + + @staticmethod + def _command(argv: list[str]) -> str: + if not shutil.which(argv[0]): + return "unavailable" + try: + result = subprocess.run(argv, capture_output=True, text=True, timeout=10, check=False) + return result.stdout.strip() if result.returncode == 0 else "unavailable" + except (OSError, subprocess.TimeoutExpired): + return "unavailable" + + @staticmethod + def _flags() -> dict[str, Any]: + import torch + + return { + "matmul_precision": torch.get_float32_matmul_precision(), + "current_device": torch.cuda.current_device(), + "default_dtype": str(torch.get_default_dtype()), + "default_device": str(torch.get_default_device()), + "autocast_enabled": torch.is_autocast_enabled("cuda"), + "autocast_dtype": str(torch.get_autocast_dtype("cuda")), + "matmul_allow_tf32": torch.backends.cuda.matmul.allow_tf32, + "cudnn_allow_tf32": torch.backends.cudnn.allow_tf32, + "fp16_reduced_reduction": torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction, + "bf16_reduced_reduction": torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction, + "deterministic_algorithms": torch.are_deterministic_algorithms_enabled(), + } + + def _check_flags(self) -> None: + import torch + + if self._flags() != self._precision_flags: + raise ContractError("Solution changed the frozen PyTorch precision settings") + if torch.cuda.current_stream(self.device).cuda_stream != self._current_stream: + raise ContractError("Solution changed the current CUDA stream") + + def environment(self) -> dict[str, Any]: + """Fingerprint the interpreter, packages, GPU, compiler, and precision flags.""" + import torch + + self._check_flags() + properties = torch.cuda.get_device_properties(self.device) + versions = {} + for name in ( + "torch", + "triton", + "flashinfer-bench", + "flashinfer-python", + "apache-tvm-ffi", + "numpy", + "safetensors", + ): + try: + versions[name] = importlib.metadata.version(name) + except importlib.metadata.PackageNotFoundError: + versions[name] = "absent" + return { + "identity": { + "hardware": properties.name, + "device": self.device, + "capability": list(torch.cuda.get_device_capability(self.device)), + "total_memory": properties.total_memory, + "gpu_uuid": str(getattr(properties, "uuid", "unavailable")), + "hostname": platform.node(), + "platform": platform.platform(), + "machine": platform.machine(), + "python": platform.python_version(), + "interpreter": sys.executable, + "packages": versions, + "cuda": torch.version.cuda, + "driver": self._driver, + "compiler": self._nvcc, + "precision_flags": self._precision_flags, + "environment": { + key: os.environ.get(key) + for key in ( + "CUDA_VISIBLE_DEVICES", + "CUDA_LAUNCH_BLOCKING", + "NVIDIA_TF32_OVERRIDE", + "CUBLAS_WORKSPACE_CONFIG", + "TORCH_CUDA_ARCH_LIST", + ) + }, + }, + "conditions": { + "timing": asdict(self.cfg.timing), + "precision": asdict(self.cfg.precision), + "seeds": list(self.cfg.seeds), + "cache_policy": ( + "cold L2 per sample; input copy/clone excluded; graph samples are device " + "time without host launch latency" + ), + }, + } + + @staticmethod + def _seed(seed: int) -> None: + import numpy as np + import torch + + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + + @staticmethod + def _clone(inputs: list[Any]) -> list[Any]: + import torch + + return [x.clone() if isinstance(x, torch.Tensor) else x for x in inputs] + + def _inputs(self, workload: Any, seed: int) -> list[Any]: + from flashinfer_bench.bench.utils import gen_inputs, load_safetensors + + self._seed(seed) + safe = ( + load_safetensors(self.definition, workload, self.task.root) + if any(x.type == "safetensors" for x in workload.inputs.values()) + else {} + ) + return gen_inputs(self.definition, workload, device=self.device, safe_tensors=safe) + + def _outputs(self, result: Any) -> list[Any]: + import torch + + names = list(self.definition.outputs) + if isinstance(result, dict): + if set(result) != set(names): + raise ContractError("Output dictionary keys differ from definition.outputs") + result = [result[name] for name in names] + values = list(result) if isinstance(result, (tuple, list)) else [result] + if len(values) != len(names): + raise ContractError("Output count differs from definition.outputs") + out = [] + for name, value, dtype in zip(names, values, self.definition.torch_output_dtypes): + if not isinstance(value, torch.Tensor): + if self.definition.outputs[name].shape is not None: + raise ContractError("Tensor output was replaced by a Python scalar") + value = torch.tensor(value, dtype=dtype, device=self.device) + out.append(value) + return out + + def _call(self, runnable: Any, inputs: list[Any]) -> tuple[list[Any], list[Any]]: + import torch + from flashinfer_bench.bench.evaluators.utils import allocate_outputs + + args = list(inputs) + if runnable.metadata.destination_passing_style: + outputs = allocate_outputs(self.definition, inputs, self.device) + # A deterministic poison catches most omitted writes, including zero outputs. + for output in outputs: + output.view(torch.uint8).fill_(0xA5) + args += outputs + runnable(*args) + else: + outputs = self._outputs(runnable(*args)) + shapes = self.definition.get_output_shapes(self.definition.get_axes_values_from_inputs(inputs)) + for value, shape, dtype in zip(outputs, shapes, self.definition.torch_output_dtypes): + if ( + tuple(value.shape) != tuple(shape or ()) + or value.dtype != dtype + or str(value.device) != self.device + ): + raise ContractError("Output violates the definition's shape, dtype, or device") + return args, outputs + + def _verify_call( + self, runnable: Any, reference: Any, inputs: list[Any], seed: int + ) -> tuple[float, float]: + import torch + + self._seed(seed) + _, expected = self._call(reference, self._clone(inputs)) + expected = [v.clone() for v in expected] + self._seed(seed) + local = self._clone(inputs) + _, outputs = self._call(runnable, local) + torch.cuda.synchronize(self.device) + self._check_flags() + # Pure operators may alias an input but must not mutate it. + compare_outputs( + [x for x in local if isinstance(x, torch.Tensor)], + [x for x in inputs if isinstance(x, torch.Tensor)], + Precision(), + ) + return compare_outputs(outputs, expected, self.cfg.precision) + + def _capture(self, runnable: Any, inputs: list[Any]) -> tuple[Any, list[Any], list[Any]]: + import torch + from flashinfer_bench.bench.evaluators.utils import allocate_outputs + + stream = torch.cuda.Stream(device=self.device) + stream.wait_stream(torch.cuda.current_stream(self.device)) + args = list(inputs) + is_dps = runnable.metadata.destination_passing_style + outputs = allocate_outputs(self.definition, inputs, self.device) if is_dps else [] + if is_dps: + args += outputs + with torch.cuda.stream(stream): + for _ in range(self.cfg.timing.warmup): + runnable(*args) + torch.cuda.current_stream(self.device).wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + result = runnable(*args) + if not is_dps: + outputs = self._outputs(result) + torch.cuda.synchronize(self.device) + return graph, args, outputs + + def _verify_graph(self, runnable: Any, reference: Any, workload: Any, seed: int) -> tuple[float, float]: + import torch + + original = self._inputs(workload, seed) + static = self._clone(original) + graph, _, outputs = self._capture(runnable, static) + error = (0.0, 0.0) + # A -> B -> A catches cached results and stale graph buffers. Task-specific + # valid perturbations of fixed fixtures belong to the custom benchmark. + changed = self._inputs(workload, (seed + 1) % 2**32) + if not any( + isinstance(value, torch.Tensor) and not torch.equal(value, original[index]) + for index, value in enumerate(changed) + ): + raise ContractError( + "Graph validation needs changed valid tensor inputs; supply benchmark.py for fixed fixtures" + ) + for values in (original, changed, original): + for dst, src in zip(static, values): + if isinstance(dst, torch.Tensor): + dst.copy_(src) + elif dst != src: + raise ContractError( + "Graph scalar arguments must stay fixed; use benchmark.py for custom semantics" + ) + self._seed(seed) + _, expected = self._call(reference, self._clone(values)) + expected = [v.clone() for v in expected] + self._seed(seed) + graph.replay() + torch.cuda.synchronize(self.device) + self._check_flags() + compare_outputs( + [x for x in static if isinstance(x, torch.Tensor)], + [x for x in values if isinstance(x, torch.Tensor)], + Precision(), + ) + current = compare_outputs(outputs, expected, self.cfg.precision) + error = tuple(max(a, b) for a, b in zip(error, current)) + return error + + def _time(self, runnable: Any, inputs: list[Any]) -> float: + from flashinfer_bench.bench.timing import time_runnable + + args, _ = self._call(runnable, self._clone(inputs)) + if self.cfg.timing.mode == "eager": + result = time_runnable( + runnable, args, self.cfg.timing.warmup, self.cfg.timing.iterations, self.device + ) + else: + graph, args, _ = self._capture(runnable, self._clone(inputs)) + result = self._time_graph(graph, args, inputs) + self._check_flags() + return float(result) + + def _time_graph(self, graph: Any, args: list[Any], inputs: list[Any]) -> float: + """Mean device time of graph replays, each from a cold L2 and freshly reset inputs. + + FlashInfer's do_bench synchronizes before every timed call. The idle GPU then + records the start event before the host has launched the replay, so each + sample also contains host launch latency and its scheduler jitter: several + microseconds, comparable to the kernels themselves. Here the L2 flush and + the input reset are queued ahead of the start event instead. The flush keeps + the device busy while the replay is enqueued, so the events bracket only the + replay's execution on the device. + """ + import torch + + if self._flush is None: + # The same 256 MiB buffer FlashInfer zeroes to evict inputs from L2. + self._flush = torch.empty(64 * 1024 * 1024, dtype=torch.int, device=self.device) + tensors = [(dst, src) for dst, src in zip(args, inputs) if isinstance(dst, torch.Tensor)] + + def prepare() -> None: + self._flush.zero_() + for dst, src in tensors: + dst.copy_(src) + + for _ in range(self.cfg.timing.warmup): + prepare() + graph.replay() + count = self.cfg.timing.iterations + starts = [torch.cuda.Event(enable_timing=True) for _ in range(count)] + ends = [torch.cuda.Event(enable_timing=True) for _ in range(count)] + torch.cuda.synchronize(self.device) + for start, end in zip(starts, ends): + prepare() + start.record() + graph.replay() + end.record() + torch.cuda.synchronize(self.device) + return sum(start.elapsed_time(end) for start, end in zip(starts, ends)) / count + + def evaluate(self, solutions: dict[str, dict[str, Any]]) -> dict[str, Any]: + """Verify all seeds before timing and alternate role order in paired rounds.""" + import torch + from flashinfer_bench.compile import BuilderRegistry + from flashinfer_bench.data import Solution + + if set(solutions) != set(ROLES): + raise ContractError("Expected baseline, incumbent, and candidate solutions") + registry = BuilderRegistry.get_instance() + reference = self._build_reference(registry) + runnables = { + role: registry.build(self.definition, Solution.model_validate(solution)) + for role, solution in solutions.items() + } + self._check_flags() + checks = [] + with torch.no_grad(): + for trace in self.traces: + for role in ROLES: + max_abs, max_rel = 0.0, 0.0 + for seed in self.cfg.seeds: + inputs = self._inputs(trace.workload, seed) + error = self._verify_call(runnables[role], reference, inputs, seed) + max_abs, max_rel = max(max_abs, error[0]), max(max_rel, error[1]) + if self.cfg.timing.mode == "cuda_graph": + error = self._verify_graph(runnables[role], reference, trace.workload, seed) + max_abs, max_rel = max(max_abs, error[0]), max(max_rel, error[1]) + checks.append( + { + "workload": trace.workload.uuid, + "role": role, + "passed": True, + "seeds": list(self.cfg.seeds), + "precision": asdict(self.cfg.precision), + "max_abs_error": max_abs, + "max_rel_error": max_rel, + "graph_replay": self.cfg.timing.mode == "cuda_graph", + **self._check_metadata(), + } + ) + rounds = [] + for round_index in range(self.cfg.timing.paired_rounds): + measurements = {} + for trace in self.traces: + row = {role: [] for role in ROLES} + for trial in range(self.cfg.timing.trials): + seed = self.cfg.seeds[trial % len(self.cfg.seeds)] + inputs = self._inputs(trace.workload, seed) + order = ROLES if (round_index + trial) % 2 == 0 else tuple(reversed(ROLES)) + for role in order: + self._seed(seed) + row[role].append(self._time(runnables[role], inputs)) + measurements[trace.workload.uuid] = row + rounds.append(measurements) + return {"checks": checks, "rounds": rounds} + + def _check_metadata(self) -> dict[str, Any]: + """Additional evidence required by specialized numerical contracts.""" + return {} + + def _build_reference(self, registry: Any) -> Any: + return registry.build_reference(self.definition) + + def profile(self, solutions: dict[str, dict[str, Any]], output: Path) -> dict[str, Any]: + """Capture optional NCU evidence in a separate worker, never promotion timing.""" + from scripts.kernel_tuning.artifacts import atomic_json + + if not shutil.which("ncu"): + raise ContractError("profile: true requires ncu on PATH") + candidate = output / "candidate.json" + atomic_json(candidate, solutions["candidate"]) + command = [ + "ncu", + "--target-processes", + "all", + "--set", + "basic", + "--force-overwrite", + "--export", + str(output / "ncu"), + sys.executable, + str(Path(__file__).resolve()), + "--profile-once", + str(self.task.root), + str(candidate), + ] + with (output / "ncu.log").open("w", encoding="utf-8") as log: + subprocess.run( + command, + stdout=log, + stderr=subprocess.STDOUT, + check=True, + timeout=self.cfg.search.evaluation_timeout_seconds, + ) + return { + "command": command, + "report": str(output / "ncu.ncu-rep"), + "workload": self.traces[0].workload.uuid, + } + + +if __name__ == "__main__": + # NCU child: one representative workload, excluded from optimization metrics. + sys.path.insert(0, str(Path(__file__).resolve().parents[2])) + from flashinfer_bench.compile import BuilderRegistry + from flashinfer_bench.data import Solution + from scripts.kernel_tuning.contracts import read_json + + if len(sys.argv) != 4 or sys.argv[1] != "--profile-once": + raise SystemExit("Internal NCU worker; use tuning.yaml profile: true") + adapter = FlashInferAdapter(TaskPackage.load(Path(sys.argv[2]))) + runnable = BuilderRegistry.get_instance().build( + adapter.definition, Solution.model_validate(read_json(Path(sys.argv[3]))) + ) + adapter._call(runnable, adapter._inputs(adapter.traces[0].workload, adapter.cfg.seeds[0])) + import torch + + torch.cuda.synchronize(adapter.device) diff --git a/benchmarks/kernel_tuning/replay_adapter.py b/benchmarks/kernel_tuning/replay_adapter.py new file mode 100644 index 0000000..806f97a --- /dev/null +++ b/benchmarks/kernel_tuning/replay_adapter.py @@ -0,0 +1,324 @@ +"""Evaluate registry calls with captured layouts, aliases, and mutable state.""" + +from __future__ import annotations + +from typing import Any + +from scripts.kernel_tuning.contracts import ContractError, Precision, TaskPackage, read_json + +from .flashinfer_adapter import FlashInferAdapter, compare_outputs + + +class ReplayAdapter(FlashInferAdapter): + """Use the frozen production call as oracle, resetting storage before each sample. + + The normal adapter remains suitable for pure, contiguous tensor tasks. This + adapter handles input views and declared mutations, and rejects changes to + the output layout/alias contract as well as incorrect values. + """ + + def __init__(self, task: TaskPackage) -> None: + import torch + + self.replay = read_json(task.root / "replay.json") + if self.replay["torch_version"] != str(torch.__version__): + raise ContractError("Capture and evaluator Torch versions differ; recapture on the target") + super().__init__(task) + flags = self.replay["flags"] + torch.backends.cudnn.allow_tf32 = flags["cudnn_allow_tf32"] + torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = flags["fp16_reduced_reduction"] + torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = flags["bf16_reduced_reduction"] + torch.use_deterministic_algorithms(flags["deterministic_algorithms"]) + self._precision_flags = self._flags() + self.case: dict[str, Any] = {} + self.numerical = read_json(task.root / "numerics.json") if "numerics.json" in task.hashes else None + if self.numerical is not None: + from .calibration import environment + + if environment(self.device) != self.numerical["environment"]: + raise ContractError( + "Calibration and evaluator environments differ; recalibrate on the target" + ) + if self.cfg.profile: + raise ContractError("Registry tasks need a replay-aware NCU adapter before profile=true") + + def _load_definition(self, definition: dict[str, Any]) -> Any: + from .replay_schema import ReplayDefinition + + return ReplayDefinition.model_validate(definition) + + def _build_reference(self, registry: Any) -> Any: + from flashinfer_bench.data import Solution + + # Build a package, not a standalone source string: backend calls may use + # snapshotted relative imports. Both reference and baseline are immutable. + return registry.build(self.definition, Solution.model_validate(self.task.baseline())) + + def _inputs(self, workload: Any, seed: int) -> list[Any]: + from scripts.kernel_tuning.discovery.tensors import materialize + + self._seed(seed) + self.case = self.replay["cases"][workload.uuid] + if getattr(self, "numerical", None) is not None: + from scripts.kernel_tuning.discovery.numerics import profile_key + + from .calibration import validation_inputs + + profile = "recorded" if seed == 0 else "random" + self._sample = profile_key(profile, seed) + return validation_inputs( + self.replay["numerical_identity"], self.case, self.task.root, profile, seed + ) + return materialize(self.case["inputs"], self.task.root, seed) + + @staticmethod + def _clone(inputs: list[Any]) -> list[Any]: + from scripts.kernel_tuning.discovery.tensors import clone_inputs + + return clone_inputs(inputs) + + @staticmethod + def _storage_views(inputs: list[Any]) -> list[Any]: + from scripts.kernel_tuning.discovery.tensors import describe_inputs, extent + + specs = describe_inputs(inputs) + buffers = {} + for tensor, spec in zip(inputs, specs): + group = spec["storage"] + if group not in buffers: + size = max(extent(item) for item in specs if item["storage"] == group) + buffers[group] = tensor.as_strided((size,), (1,), 0) + return list(buffers.values()) + + @classmethod + def _reset(cls, destination: list[Any], source: list[Any]) -> None: + for dst, src in zip(cls._storage_views(destination), cls._storage_views(source)): + dst.copy_(src) + + def _check_outputs(self, inputs: list[Any], outputs: list[Any]) -> None: + from scripts.kernel_tuning.discovery.tensors import storage_key, tensor_spec + + if len(outputs) != len(self.case["outputs"]): + raise ContractError("Output count differs from the captured call") + for tensor, spec, aliases in zip(outputs, self.case["outputs"], self.case["output_aliases"]): + if tensor_spec(tensor) != spec: + raise ContractError("Output layout, dtype, or device differs from the captured call") + actual = [i for i, value in enumerate(inputs) if storage_key(value) == storage_key(tensor)] + if actual != aliases: + raise ContractError("Output aliases differ from the captured call") + groups = [ + [j for j, other in enumerate(outputs) if storage_key(value) == storage_key(other)] + for value in outputs + ] + if groups != self.case["output_groups"]: + raise ContractError("Output-to-output aliases differ from the captured call") + + def _call(self, runnable: Any, inputs: list[Any]) -> tuple[list[Any], list[Any]]: + if runnable.metadata.destination_passing_style: + raise ContractError("Registry tasks require return-value style") + outputs = self._outputs(runnable(*inputs)) + self._check_outputs(inputs, outputs) + return list(inputs), outputs + + def _verify_call( + self, runnable: Any, reference: Any, inputs: list[Any], seed: int + ) -> tuple[float, float]: + if getattr(self, "numerical", None) is not None: + from scripts.kernel_tuning.discovery.numerics import PROFILES, profile_key + + from .calibration import validation_inputs + + error = (0.0, 0.0) + for profile, profile_seed in PROFILES: + if profile_seed != seed: + continue + self._sample = profile_key(profile, seed) + values = validation_inputs( + self.replay["numerical_identity"], self.case, self.task.root, profile, seed + ) + current = self._verify_transition(runnable, reference, values, seed) + error = tuple(max(a, b) for a, b in zip(error, current)) + return error + return self._verify_transition(runnable, reference, inputs, seed) + + def _verify_transition( + self, runnable: Any, reference: Any, inputs: list[Any], seed: int + ) -> tuple[float, float]: + import torch + + expected_inputs = self._clone(inputs) + self._seed(seed) + _, expected = self._call(reference, expected_inputs) + expected = [value.clone() for value in expected] + self._seed(seed) + local = self._clone(inputs) + _, outputs = self._call(runnable, local) + torch.cuda.synchronize(self.device) + self._check_flags() + # Check every byte of backing storage, including aliases and untouched + # regions. Only declared mutable storage may use the task's tolerance. + error = self._check_state(local, expected_inputs, inputs) + current = self._compare(outputs, expected, inputs) + return tuple(max(a, b) for a, b in zip(error, current)) + + def _compare(self, actual: list[Any], expected: list[Any], inputs: list[Any]) -> tuple[float, float]: + if getattr(self, "numerical", None) is None: + return compare_outputs(actual, expected, self.cfg.precision) + from .calibration import compare_calibrated, high_precision + + reference = high_precision(self.replay["numerical_identity"], inputs) + bounds = self.numerical["cases"][self.case["id"]][self._sample] + return compare_calibrated(actual, reference, bounds) + + def _check_metadata(self) -> dict[str, Any]: + if self.numerical is None: + return {} + from scripts.kernel_tuning.discovery.numerics import PROFILES + + return { + "numerical_contract": self.task.hashes["numerics.json"], + "validation_profiles": [list(pair) for pair in PROFILES], + } + + def _check_state( + self, actual: list[Any], expected: list[Any], original: list[Any] + ) -> tuple[float, float]: + from scripts.kernel_tuning.discovery.tensors import describe_inputs + + groups = describe_inputs(original) + if describe_inputs(actual) != groups or describe_inputs(expected) != groups: + raise ContractError("Input layout or aliasing changed during the operator call") + mutable = {groups[i]["storage"] for i in self.replay["mutates"]} + error = (0.0, 0.0) + for index, (lhs, rhs, old) in enumerate( + zip(self._storage_views(actual), self._storage_views(expected), self._storage_views(original)) + ): + if index in mutable: + current = compare_outputs([lhs], [rhs], self.cfg.precision) + error = tuple(max(a, b) for a, b in zip(error, current)) + else: + compare_outputs([rhs, lhs], [old, old], Precision()) + return error + + def _capture(self, runnable: Any, inputs: list[Any]) -> tuple[Any, list[Any], list[Any]]: + import torch + + saved = self._clone(inputs) + stream = torch.cuda.Stream(device=self.device) + stream.wait_stream(torch.cuda.current_stream(self.device)) + with torch.cuda.stream(stream): + for _ in range(self.cfg.timing.warmup): + self._reset(inputs, saved) + self._call(runnable, inputs) + self._reset(inputs, saved) + torch.cuda.current_stream(self.device).wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + _, outputs = self._call(runnable, inputs) + torch.cuda.synchronize(self.device) + return graph, inputs, outputs + + def _verify_graph(self, runnable: Any, reference: Any, workload: Any, seed: int) -> tuple[float, float]: + import torch + + if getattr(self, "numerical", None) is not None: + return self._verify_calibrated_graph(runnable, reference, workload, seed) + original = self._inputs(workload, seed) + changed = self._inputs(workload, (seed + 1) % 2**32) + if not any(not torch.equal(a, b) for a, b in zip(original, changed)): + raise ContractError( + "Graph validation requires changed valid inputs; supply a domain-specific adapter" + ) + static = self._clone(original) + graph, _, outputs = self._capture(runnable, static) + error = (0.0, 0.0) + for values in (original, changed, original): + expected_inputs = self._clone(values) + self._seed(seed) + _, expected = self._call(reference, expected_inputs) + expected = [value.clone() for value in expected] + self._reset(static, values) + self._seed(seed) + graph.replay() + torch.cuda.synchronize(self.device) + self._check_flags() + self._check_outputs(static, outputs) + state_error = self._check_state(static, expected_inputs, values) + current = compare_outputs(outputs, expected, self.cfg.precision) + error = tuple(max(a, b, c) for a, b, c in zip(error, current, state_error)) + return error + + def _verify_calibrated_graph( + self, + runnable: Any, + reference: Any, + workload: Any, + seed: int, + ) -> tuple[float, float]: + import torch + from scripts.kernel_tuning.discovery.numerics import PROFILES, profile_key + + from .calibration import validation_inputs + + original = self._inputs(workload, 0) + static = self._clone(original) + graph, _, outputs = self._capture(runnable, static) + profiles = [("recorded", 0), *(pair for pair in PROFILES if pair[1] == seed), ("recorded", 0)] + error, changed = (0.0, 0.0), False + for profile, profile_seed in profiles: + values = validation_inputs( + self.replay["numerical_identity"], self.case, self.task.root, profile, profile_seed + ) + changed |= any(not torch.equal(a, b) for a, b in zip(original, values)) + self._sample = profile_key(profile, profile_seed) + expected_inputs = self._clone(values) + _, expected = self._call(reference, expected_inputs) + self._reset(static, values) + graph.replay() + torch.cuda.synchronize(self.device) + self._check_flags() + self._check_outputs(static, outputs) + state = self._check_state(static, expected_inputs, values) + current = self._compare(outputs, expected, values) + error = tuple(max(a, b, c) for a, b, c in zip(error, state, current)) + if not changed: + raise ContractError("Calibrated graph validation requires changed inputs") + return error + + def _time(self, runnable: Any, inputs: list[Any]) -> float: + import torch + + local = self._clone(inputs) + if self.cfg.timing.mode == "cuda_graph": + graph, _, _ = self._capture(runnable, local) + execute = graph.replay + else: + + def execute() -> Any: + return runnable(*local) + + if self._flush is None: + self._flush = torch.empty(64 * 1024 * 1024, dtype=torch.int, device=self.device) + pairs = list(zip(self._storage_views(local), self._storage_views(inputs))) + + def prepare() -> None: + for dst, src in pairs: + dst.copy_(src) + # Reset copies touch L2 too; evict them before the timed call. + self._flush.zero_() + + for _ in range(self.cfg.timing.warmup): + prepare() + execute() + count = self.cfg.timing.iterations + starts = [torch.cuda.Event(enable_timing=True) for _ in range(count)] + ends = [torch.cuda.Event(enable_timing=True) for _ in range(count)] + torch.cuda.synchronize(self.device) + for start, end in zip(starts, ends): + prepare() + start.record() + execute() + end.record() + torch.cuda.synchronize(self.device) + self._check_flags() + return sum(start.elapsed_time(end) for start, end in zip(starts, ends)) / count diff --git a/benchmarks/kernel_tuning/replay_schema.py b/benchmarks/kernel_tuning/replay_schema.py new file mode 100644 index 0000000..1898849 --- /dev/null +++ b/benchmarks/kernel_tuning/replay_schema.py @@ -0,0 +1,42 @@ +"""Opt-in replay definitions retaining real float64 tensors with FlashInfer builders.""" + +from __future__ import annotations + +from functools import cached_property +from typing import TYPE_CHECKING, Literal + +from flashinfer_bench.data.definition import Definition, DType, TensorSpec +from flashinfer_bench.data.utils import NonEmptyString +from flashinfer_bench.utils import dtype_str_to_torch_dtype + +if TYPE_CHECKING: + import torch + + +class ReplayTensorSpec(TensorSpec): + """Keep upstream tensor validation while admitting native double precision.""" + + dtype: DType | Literal["float64"] + + +def _dtype(value: DType | Literal["float64"]) -> torch.dtype: + import torch + + return torch.float64 if value == "float64" else dtype_str_to_torch_dtype(value) + + +class ReplayDefinition(Definition): + """A validated Definition subclass for the local replay adapter, without global patches.""" + + inputs: dict[NonEmptyString, ReplayTensorSpec] + outputs: dict[NonEmptyString, ReplayTensorSpec] + + @cached_property + def torch_input_dtypes(self) -> list[torch.dtype]: + """Resolve input dtypes without narrowing doubles to an upstream enum member.""" + return [_dtype(spec.dtype) for spec in self.inputs.values()] + + @cached_property + def torch_output_dtypes(self) -> list[torch.dtype]: + """Resolve output dtypes using the same declared precision as replay metadata.""" + return [_dtype(spec.dtype) for spec in self.outputs.values()] diff --git a/benchmarks/kernel_tuning/requirements.txt b/benchmarks/kernel_tuning/requirements.txt new file mode 100644 index 0000000..67e1490 --- /dev/null +++ b/benchmarks/kernel_tuning/requirements.txt @@ -0,0 +1,4 @@ +# Install into the already prepared target-machine runtime. Do not replace its +# Torch/CUDA/Triton pins with those from a different GPU's environment. +flashinfer-bench==0.1.2 +PyYAML>=6.0.3,<7 diff --git a/benchmarks/pi05-benchmark/benchmark.py b/benchmarks/pi05-benchmark/benchmark.py index 82e08f2..488bb5f 100644 --- a/benchmarks/pi05-benchmark/benchmark.py +++ b/benchmarks/pi05-benchmark/benchmark.py @@ -397,6 +397,12 @@ def engine_graph_stats(core: Any) -> dict[str, Any]: } +def graph_capture_count(stats: dict[str, Any]) -> int: + """Count engine graphs plus graphs captured by the PI0.5 native runtime.""" + native = stats.get("pi05_native") or {} + return stats["graphs"]["capture_count"] + native.get("capture_count", 0) + + def run( default: Path, build: Callable[[dict[str, Any], torch.device], tuple[LiberoInference, dict[str, Any], RuntimeStats]], @@ -428,12 +434,15 @@ def run( for index in range(config["warmup_calls"]): sample = samples[warmup_indices[index % len(warmup_indices)]] infer(sample, read_libero(sample)) + warm_layouts = getattr(infer, "warm_native_layouts", None) + if warm_layouts is not None: + details["warmed_native_layouts"] = warm_layouts(samples) torch.cuda.synchronize(device) details["warmup_seconds_including_data_io"] = time.perf_counter() - started details["warmup_sample_ids"] = [samples[index].sample_id for index in warmup_indices] before = runtime_stats() details["runtime_before_measurement"] = before - if config["cuda_graph"] and not before["graphs"]["capture_count"]: + if config["cuda_graph"] and not graph_capture_count(before): raise RuntimeError("CUDA Graph was requested but warmup captured no graph") torch.cuda.reset_peak_memory_stats(device) torch.manual_seed(config["seed"]) @@ -449,7 +458,7 @@ def run( ) after = runtime_stats() details["runtime_after_measurement"] = after - if after["graphs"]["capture_count"] != before["graphs"]["capture_count"]: + if graph_capture_count(after) != graph_capture_count(before): raise RuntimeError("CUDA Graph capture occurred during measurement; extend warmup") write_report(config, device, rows, details, Path(config["output"])) @@ -509,7 +518,13 @@ def build( ), **{ k: config[k] - for k in ("native_inference", "prefix_cuda_graph", "denoise_attention", "prefix_attention") + for k in ( + "native_inference", + "prefix_cuda_graph", + "denoise_attention", + "prefix_attention", + "kernel_profile", + ) if k in config }, load_device=str(device), @@ -527,7 +542,7 @@ def build( ), ) - def infer(sample: LiberoSample, raw: dict[str, Any]) -> dict[str, Any]: + def prepare(sample: LiberoSample, raw: dict[str, Any]) -> tuple[Any, torch.Tensor]: state = processor.prepare_state( torch.from_numpy( np.concatenate((raw["ee_pos"], raw["ee_ori"], raw["gripper_states"])).astype(np.float32) @@ -544,7 +559,10 @@ def infer(sample: LiberoSample, raw: dict[str, Any]) -> dict[str, Any]: images[feature] = torch.from_numpy(processor.resize_image(image, 224, 224)) batch = processor.prepare(state, images, sample.instruction) batch.request_ids = [sample.sample_id] - batch = policy.pad(batch, 1).to(device, core.dtype) + return policy.pad(batch, 1).to(device, core.dtype), state + + def infer(sample: LiberoSample, raw: dict[str, Any]) -> dict[str, Any]: + batch, state = prepare(sample, raw) prefix, normalized, timing = timed_model( lambda: policy.encode_prefix(batch), lambda prefix: policy.decoder.integrate( @@ -555,6 +573,26 @@ def infer(sample: LiberoSample, raw: dict[str, Any]) -> dict[str, Any]: actions = policy.finalize_actions(normalized, prefix)[0].float().cpu() return {**output_record(processor.restore_actions(actions, state)), "model_timing_ms": timing} + def warm_native_layouts(samples: list[LiberoSample]) -> list[list[Any]]: + """Capture every native language-bucket graph before timing; return warmed layouts.""" + if not policy.native_inference: + return [] + warmed: set[tuple[int, bool]] = set() + for sample in samples: + raw = read_libero(sample) + batch, _ = prepare(sample, raw) + length = int(batch.masks.sum()) + bucket = next( + (n for n in (16, 32, 48, 64, 96, 128, 160, 200) if n >= length), batch.tokens.shape[1] + ) + layout = (bucket, length == bucket) + if layout not in warmed: + infer(sample, raw) + warmed.add(layout) + return [list(layout) for layout in sorted(warmed)] + + infer.warm_native_layouts = warm_native_layouts # type: ignore[attr-defined] + def runtime_stats() -> dict[str, Any]: from embodiinfer.policies.pi05 import modeling_pi05 diff --git a/benchmarks/pi05-benchmark/validate_nonquant.py b/benchmarks/pi05-benchmark/validate_nonquant.py index c5993ac..319322a 100644 --- a/benchmarks/pi05-benchmark/validate_nonquant.py +++ b/benchmarks/pi05-benchmark/validate_nonquant.py @@ -77,10 +77,10 @@ def ephemeral_packed_mlp_weights() -> Any: original = modeling_pi05._fused_linear_pair - def bounded(inputs: Any, first: Any, second: Any) -> Any: + def bounded(inputs: Any, first: Any, second: Any, **kwargs: Any) -> Any: if torch.is_grad_enabled() or torch.cuda.is_current_stream_capturing(): raise RuntimeError("ephemeral packed weights are only for untimed eager inference") - result = original(inputs, first, second) + result = original(inputs, first, second, **kwargs) first._vvla_fused_pair_weight = None return result diff --git a/docs/proposals/0008-kernel-tuning.md b/docs/proposals/0008-kernel-tuning.md new file mode 100644 index 0000000..ba3a157 --- /dev/null +++ b/docs/proposals/0008-kernel-tuning.md @@ -0,0 +1,935 @@ +# 0008 — Agent-driven kernel optimization: profiling to whole-model benchmarks + +- Status: Accepted for the existing tuning contracts; end-to-end extensions are partially implemented. +- Author: EmbodiInfer maintainers +- Date: 2026-10-02 + +This is the canonical document for the former proposals 0008 and 0009. It combines +the Humanize2 search/evaluation contract with the hotspot-to-model workflow. +Implementation status below is checked against the current source; consolidation +does not approve pending numerical or runtime changes. +For executable examples, see [automated tuning workflow](#11-automated-tuning-workflow). + +## 1. Summary + +The workflow is **profile → select → register if needed → capture → calibrate → +tune → install → verify → benchmark**. Start with the model's production latency, +choose worthwhile registered operators, let a coding agent improve their kernels, +and measure the resulting model with a hardware profile. A faster isolated kernel +is an intermediate result; success requires whole-model correctness and latency +evidence under fixed conditions. + +Humanize2 (`hmz`) owns agent execution, budgets, traces and resume. Repository +scripts prepare tasks, validate evaluator evidence and install exported sources. +Formal measurement lives in `benchmarks/`; model-neutral dispatch lives in +`layers/`; concrete kernels and profile loading live in `backend/`. The engine +has no tuning scheduler or model-name branches. + +| Stage | Current entry point | Current scope | +|---|---|---| +| Profile | `profile --output DIR -- DRIVER` | GPU attribution by stage, registry operator, attention backend and module | +| Select | Read `hotspots.md` and choose operators | Manual; automatic selection and fusion planning are future work | +| Register | `register new/run/resume/status/finalize` | Agent-assisted registration of one caller-selected operator; pi0.5 routing convention | +| Capture | `capture-ops --model MODEL --ops OPS ...` | Real arguments and counts of selected registered operators, in eager mode | +| Calibrate | `calibrate --captured DIR --output NEW_DIR` | Freeze supported numerical contracts before search | +| Tune | `tune-all`, or `run` for a prepared task | Humanize2 candidate loop with baseline preflight, correctness and paired timing | +| Install | `install EXPORT --profile NAME` | Python/Triton exports from registry captures only | +| Verify | `verify --profile NAME --output REPORT -- DRIVER_ARGS` | Whole-model FP32-reference/production-Torch/candidate comparison | +| Benchmark | Existing model benchmark scripts | Production end-to-end latency; separate configurations for each comparison row | + +There is **no `optimize`, `select`, or `benchmark` subcommand** in +`scripts.kernel_tuning` today. The stages are invoked separately. `tune-all` +resumes search batches, and `register resume` resumes registration; neither +resumes the entire end-to-end workflow. pi0.5 is the current integrated policy. +Qwen R2R/StreamVLN registry integration and committed hardware profiles remain +outstanding; their benchmark-input/catalog support does not imply profile support. + +## 2. Motivation and current gap + +An isolated operator can improve without reducing model latency: its contribution +may be small, it may already be near launch overhead, or its tuned layout may +never be selected in production. Profiling and a final model benchmark connect +search effort to the actual inference path. + +The registry workflow uses an explicit Torch reference and stable model call +sites so verified kernels can reach inference without editing the policy on every +tuning run. The independent catalog supports investigations of hand-maintained +kernels, but its exports require separate runtime integration. Full-model ATen +discovery has been removed; it did not produce installable registry call sites. + +Current boundaries are visible in `scripts/kernel_tuning/__main__.py` (commands), +`regions.py` (registry capture), `install.py` (accepted exports), and +`hotspots.py` (attribution). The remaining gap is automated region selection, +new fusion integration, shape-range coverage and orchestration across stages. + +## 3. Goals and non-goals + +- Optimize a selected model on a prepared target machine with explicit agent, + checkpoint, inputs, dtype, precision settings and execution configuration. +- Freeze workload coverage, numerical gates, timing and acceptance rules before + search; retain all attempts and export exact verified source with evidence. +- Select installed kernels through optional hardware profiles with layout checks + and fallback; evaluate the complete profile at model level. +- Keep the existing catalog and custom-task workflows usable, with + their original contracts and artifact formats. + +Robot/simulator communication, task success metrics, deployment orchestration, +automatic merge, automatic GPU-stack installation and changes to model mathematics +are outside this workflow. Installing a kernel does not make it the default. +Multi-GPU, broader batch sizes and unobserved shape ranges need their own coverage +and measurements; a B=1 capture establishes none of these automatically. + +## 4. Design + +### 4.1 Freeze the experiment and profile the production path + +Record checkpoint/revision, evaluation input identity, dtype, attention semantics, +matmul precision and TF32 flags, seeds/noise, batch shape, decode/denoise schedule, +compile settings and CUDA Graph mode. Use the existing model benchmark for the +uninstrumented baseline latency, then run `profile` with the same production +settings to explain where GPU time goes. + +`hotspots.py` instruments model-neutral inference entries and uses +`torch.profiler` correlation IDs to attribute kernels to policy/decoder stages, +registry calls, attention backends and module ranges. Repeated layer indices fold +to `layers.*`. CUDA Graph captures record node labels; replay uses those labels +when node counts agree, otherwise attribution stays at the replay-time range. +Compiled regions are attributed to their enclosing range; Python ranges are +skipped inside tracing. The implementation does not automatically rerun eagerly. + +Outputs are `hotspots.json`, `hotspots.md`, and optional `trace.json`. They report +GPU time/share, kernel counts, small-kernel fractions and graph attribution +coverage. GPU activity sums and profiler-instrumented times are not an +uninstrumented end-to-end latency measurement. Inspect attribution coverage before +selecting regions; any separate eager diagnostic run must be labelled as such. + +### 4.2 Select and, if necessary, register an operator + +Selection is currently manual. Use time share, call frequency, kernel size and +production dispatch coverage to choose operators. A rough estimate after tuning +is `sum(call_count * (baseline_latency - candidate_latency))`; compilation, +overlap and graph execution mean this is not a measured model speedup. + +An existing registered operator proceeds to capture. For an unregistered +mathematical operator, `register` drives an agent through plan, implementation, +CPU checks and repair. The caller supplies the operator and exact source +expressions; the tool does not select hotspots or fuse adjacent operations. +Its initial integration convention is pi0.5's `_op` helper and explicit +`KernelSelection` parameters. The full specification, commands and gate details +are maintained in the [registration guide](../../scripts/kernel_tuning/registration/README.md). + +Registration preserves the original expression, ATen arithmetic order, casts and +shapes. Its CPU gates cover file scope, reference/registration/routing structure, +Ruff, existing and new tests, and new tests failing by assertion on BASE. +Evidence binds source bytes; interruptions and repairs preserve journaled state. +The output is a branch, `patch.diff` and `SUMMARY.md`, followed by source review +and the target-GPU bit-exact integration gate in Section 6. CPU checks do not +prove whole-model equivalence. Commit/push are explicit finalization options; +merging is separate. + +New fusion regions require a separate reviewed code change with a bit-exact +Torch default before tuning. Automatic fusion selection and implementation remain +planned, as described in Section 7. + +### 4.3 Capture real workloads and calibrate + +`capture-ops` records registry calls on real model inputs, with selected operators +specified by `--ops`. Run capture eagerly with CUDA Graphs and compilation +disabled: replayed/compiled regions do not invoke the registry from Python. +The driver must execute at least `warmup + steps` inference calls. Every requested +operator must be observed, and the fixture budget must hold its required inputs. + +The result uses the model-capture format (`manifest.json`, `operators.json`, +`coverage.json`, `fixtures/`) with `capture: registry`. Each task uses the Torch +reference in `embodiinfer/backend/torch/ops.py` as baseline and the implementation +seen during capture as a starting point. Workload weights are observed call +counts. Each case's output contract (shapes, strides, offsets, dtypes and +aliases) is taken from the call the model actually made, not from the Torch +reference: downstream operators' recorded inputs come from that call, so a tuned +kernel must return the same layout or the next tuned kernel's layout gate stops +matching. The reference must agree on everything except strides and offsets; +where its layout differs, the case also records `reference_outputs` and the +generated baseline copies the reference result into the production layout +(values unchanged; the copy counts toward baseline time). A production output +that aliases an input, has an offset or is not dense cannot be rebuilt this way +and is reported as a capture error. Registry tasks currently default to CUDA +Graph kernel timing, with an +uncompiled Torch reference; this is not a per-region Inductor comparison. +Production compilation is restored for model verification and benchmarking. + +Calibrate supported tolerance contracts in the model runtime before generating +tasks. `calibrate` writes a new immutable capture with `calibration.json`; +generated tasks carry `numerics.json`. See Section 6 for exact/tolerance selection. +The capture's layouts must also occur in the compiled production path. Otherwise +the installed kernel falls back even if its isolated benchmark improves. + +### 4.4 Search, evaluate and promote + +A task contains `README.md`, `definition.json`, `workloads.jsonl`, `baseline.py`, +`tuning.yaml`, and optionally trusted `benchmark.py`; replay tasks also include +layout/fixture contracts. Task and solution schemas reuse NVlabs KDA/FlashInfer +formats. Humanize2 is pinned in the isolated tooling environment. + +The agent reads the task and recent candidate source/evidence, writes a hypothesis +and executable plan, and produces a candidate. A separate evaluator compiles it, +checks every workload and times baseline, incumbent and candidate. Feedback drives +the next attempt, including rejected attempts. Agent turns do source/JSON checks; +the evaluator owns GPU compilation and validation. + +`tune-all` generates tasks and preflights all baselines before opening any model +batch agent. Missing fixtures, unsupported contracts or a failed baseline block +the batch. A later individual search failure is recorded and other tasks may +continue. Each task runs serially in its own `run`/`resume` process. + +Promotion requires at least 3% weighted latency improvement and at most 5% +regression on every workload, against both the initial baseline and incumbent, +in every paired round. These defaults are configurable before the task freezes. +All checks must pass; a process exit code or printed PASS is insufficient. +Defaults stop at 20 attempts, two active hours or five consecutive +failures/non-improvements. Interruptions consume attempts; downtime is excluded. +Exhausted runs need a new task/run. Changed task/tool/environment identity +prevents resume. + +CUDA Graph measurements flush a 256 MiB buffer and refresh inputs before the +start event; events bracket device replay, without per-sample synchronization. +Eager timing includes launch cost. Compare only under the same timing mode, +warmup, iterations and paired-round protocol. Optional NCU profiling is separate +from promotion. [Appendix A](#appendix-a-task-owned-benchmark-interface) specifies +custom evaluators; [Appendix B](#appendix-b-archive-and-reproduction) specifies +archives and reproduction. + +### 4.5 Install and resolve a hardware profile + +`install` accepts an exported registry task from `capture-ops`, with a +Python/Triton solution and an operator already registered in the current tree. +Catalog exports and CUDA C++ solutions are not installable through +this command. It writes `embodiinfer/backend/triton/tuned///` and +`profile.json`; overwriting an installed operator requires `--replace`. + +The manifest records device name, compute capability, Torch/Triton match keys, +source digest, verified layout count and kernel-search provenance. CUDA version +and full experiment conditions belong in the archived evidence as well. The +installer generates `applies` from verified argument shapes, dtypes, strides and +constants. Size-one and empty-tensor strides are normalized; other layouts fall +back to the policy's preferred implementation. + +`embodiinfer/layers/ops.py` exposes `register_op`, `register_lazy_op` and +`KernelSelection`. A policy prepares its selection before compiled execution, +then resolves a cached callable. `kernel_profile` supports: + +| Value | Behavior | +|---|---| +| `none` (default) | Preserve existing Torch/Triton preferences | +| `torch` | Force registered operators to their Torch references | +| `auto` | Select a profile matching device/capability/Torch/Triton, otherwise use `none` | +| Explicit name | Select that profile; a missing or incompatible profile raises | + +Only operators present in a profile are overridden. Profile loading verifies +source digests. Backend modules remain pure tensor code with no engine/policy +imports. Registration and profile loading are runtime contracts; search is +repository tooling. + +### 4.6 Verify the complete model, then benchmark + +After installation, `verify` runs the same recorded inputs through an FP32 Torch +reference, a production-dtype Torch baseline and the selected profile, with the +same per-input seeds. It applies the model gate in Section 6 and writes a JSON +report. Its arguments after `--` are driver arguments, not a Python module name. + +Installation precedes this check so the candidate can be exercised. **`install` +does not certify model correctness, and `verify` does not automatically remove a +failed install or update `profile.json`.** Retain the report with input/checkpoint +identity and kernel evidence; do not publish a failed profile as verified. + +Next use the model's existing benchmark under the frozen production settings. +Compare `kernel_profile: torch`, existing Triton with `kernel_profile: none`, +and the candidate profile; record attention choices separately because the +registry switch does not replace every backend in the model. Measure end-to-end +mean/P50/P95, available stage splits, and per-operator ablations using separate +profile configurations. Confirm that measured calls actually select tuned kernels. +Profiler timings and weighted kernel gains cannot substitute for this run. + +### 4.7 Alternatives and compatibility + +- A separate agent controller would duplicate Humanize2's sessions, budgets and + journals; reuse its Flow instead. +- Full ATen discovery did not establish installable registry call sites and has + been removed. Capture selected registered operators with `capture-ops`; shared + replay contracts and the inference driver remain available (Appendix C). +- Existing Triton is a useful starting point/comparison row; the registry task's + baseline remains its frozen Torch reference. Inductor production performance + is measured at model level, not assumed from kernel timing. +- Uncached runtime dispatch adds launch-path overhead; prepare selections before + tracing and resolve cached callables. +- Generated policy changes and default profile promotion require ordinary review; + neither search nor installation automatically merges them. + +## 5. Model-agnosticism verdict + +Task orchestration, evaluation, registry dispatch and hardware matching are +model-neutral. Model inputs use existing drivers; model-specific call sites and +fused-region semantics stay in the owning policy. Reusable mathematical +references/kernels live under `backend/`, which does not import `models`, +`policies` or `engine`. The inference engine does not acquire agent dependencies, +training adapters, environment control or benchmark loops. + +## 6. Losslessness and precision criterion + +Bit-exact comparison is the default. Registration/routing changes must preserve +outputs against BASE with `kernel_profile: none`, including whole-model target-GPU +outputs (`max |delta action| = 0`). Compare the same checkpoint, dtype, attention, +matmul precision/TF32, batch, seeds/noise and decode/denoise schedule. The tolerant +`verify` command below is not a substitute for this integration gate. + +Calibrated tolerances are an explicit contract for supported arithmetic whose +accumulation/rounding differs. This document preserves those existing contracts; +it does not relax assertions or authorize approximate kernels merely because a +candidate fails. + +### Operator gate + +Registry capture records execution settings separately from acceptance rules. +Floating registry calls use their declared Torch reference for FP64 calibration. +Mutation, non-floating outputs, FP64 and unsupported dtypes retain bit-exact +validation. Existing hand-written/catalog contracts retain their own rules. + +Before task generation, `calibrate --captured CAPTURE --output CALIBRATED` runs +in the matching model environment, using only the captured production operator +and a trusted FP64 mathematical reference. Each shape/layout/argument signature +gets separate output bounds for real recorded inputs, seeded random inputs, +zeros, and alternating-sign cancellation inputs. Attention masks, scale, causal +and GQA semantics are preserved. Real-input calibration requires complete input +fixtures; capture must fail explicitly if its fixture budget cannot hold them. +No candidate source is accepted by the calibration API. + +For each floating output, let B be the production baseline, R the FP64 reference, +and u the output dtype's epsilon. Fix rtol=u and +atol=max(u*RMS(R), 2*max(max(abs(B-R)-u*abs(R), 0)), dtype.tiny). +Every candidate element must satisfy abs(C-R) <= atol+rtol*abs(R). The one-epsilon +scale floor allows rounding near zero; the factor of two budgets measured +baseline error. These are explicit engineering policy constants, not a proof of +an error bound for unseen inputs. Non-floating outputs, layout, aliases, mutation +constraints remain exact. Nonfinite reference/baseline/output values +fail calibration/validation. FP64 outputs retain exact validation: there is no +higher-precision oracle in this implementation. + +The immutable calibrated capture records reference identity, policy version, +seeds, profiles, dtype, layouts, scalar arguments, numerical flags, environment, +observed baseline error and each derived threshold. Generated tasks copy this +evidence into `numerics.json`; task/run hashes and the all-task preflight bind it +before any agent starts. Evaluation uses those thresholds without recalibrating. +Missing coverage, changed settings/seeds, edited evidence, or a different target +environment require a new calibration/batch. Calibration failure never widens a +threshold in response to a candidate failure. + +Implementation belongs to `scripts/kernel_tuning/discovery/numerics.py` (static +policy/contract), `benchmarks/kernel_tuning/calibration.py` (trusted reference and +measurement), registry capture/task generation, and the replay evaluator. +Calibration itself changes no runtime API or model-level action gate. CPU tests +exercise selection, calibration, per-workload bounds, attention +semantics, precision restoration, real/random/boundary coverage, and tampering; +CUDA/real-checkpoint validation remains a separate target-machine requirement. + +A single dtype-wide 1e-2 threshold was rejected because shape and accumulation +behavior matter. Online recalibration during candidate evaluation was rejected +because acceptance must be frozen. Full fixtures and FP64 references increase +preparation storage/time, and unsupported numerical references stay exact until +an explicit adapter is added. Existing incomplete captures must be recaptured; +an old task/run is never silently migrated to a looser contract. + +Default correctness compares output shapes, dtypes, and exact tensor bytes against +`definition.json`'s reference on identical seeded inputs. A task may explicitly +declare fixed absolute/relative tolerance with a reason. The agent cannot change +that contract, workloads, dtypes, or float32 matmul precision. Graph tasks must +also replay with changed inputs and compare against fresh reference outputs. +The evaluator command and full numerical contract are saved with every run. + +For registry captures, B is the registered Torch reference in the production +dtype. Unknown numerical semantics retain the exact contract. Dtype alone is +never sufficient to choose a tolerance. + +### Model gate + +`scripts/kernel_tuning/verify.py` uses FP32 Torch output R, production-dtype +Torch output B and candidate-profile output C. For each continuous output name, +compute max-absolute error and RMSE against R for every input, then require: + +| Statistic of candidate error | Limit | +|---|---| +| Mean max-absolute error | `max(1e-4, 1.25 * mean baseline max-absolute error)` | +| Mean RMSE | `max(1e-5, 1.25 * mean baseline RMSE)` | +| Worst max-absolute error | `max(1e-4, 2 * worst baseline max-absolute error)` | + +Discrete outputs must equal B on every input. Mismatched output structures and +incomparable baseline/reference runs fail explicitly. Freeze the evaluation set, +seed base (default 1000), checkpoint and settings before candidate verification; +thresholds derive from the baseline and fixed policy constants, never candidate +error. Record the actual sample count: `--samples` can shorten the run and does +not establish full-set coverage. For pi0.5, the intended integration evaluation +includes its frozen cross-task references plus 200 benchmark frames; the CLI does +not assemble or enforce that set automatically. + +When using Inductor, preserve the same compiler configuration/cache conditions; +use `TORCHINDUCTOR_DETERMINISTIC=1` for reproducible parity as in the pi0.5 +validation workflow. Any unavailable deterministic behavior must be reported, +not hidden by widening thresholds. FP32 calibration, operator correctness and +whole-model parity each have distinct reference paths and evidence. + +## 7. Implementation status and remaining design + +Implemented components are the Humanize2 Flow/archive, registry +capture, calibration, evaluator, batch tuning, GPU attribution, registry/profile +loader, install, model verification, the CPU registration controller, and an +opt-in `--unfenced` agent mode for hosts without Landlock (Section 11.1). The +pi0.5 call sites are integrated. One hardware profile is committed: `orin_pi05` +for Jetson AGX Orin (Section 9.1); `kernel_profile` still defaults to `none`, so +it is used only when selected by name or by `auto` on matching hardware. + +Remaining work from the merged design: + +1. Rank recoverable latency using measured bandwidth/compute/launch floors, + rather than GPU time alone; detect adjacent memory-bound fusion candidates. +2. Add an agent selection flow producing a reviewable `plan.json` (proposed + default: at most five regions and a 1% end-to-end contribution threshold). + These are design targets, not currently enforced CLI options. +3. Integrate new fused Torch regions through reviewed, bit-exact code changes; + extend registry use to Qwen R2R and StreamVLN. +4. Cover declared shape ranges, batches and layouts instead of only captured + signatures; validate fallback and dispatch coverage. +5. Add a resumable cross-stage `optimize` orchestrator and bind model verification, + evaluation-set identity and benchmark evidence into profile provenance. + Today these artifacts must be retained together by the operator. + +The proposed roofline bound is `max(bytes / measured bandwidth, FLOPs / measured +peak)`; recoverable time subtracts the larger of that bound and launch floors. +Fusion estimates account for eliminated intermediate traffic. These estimates +must be checked by model benchmarks before becoming performance claims. + +## 8. Test plan + +CPU tests under `tests/kernel_tuning/` cover task validation, Humanize2 fake-agent +execution, failed/interrupted attempts, budgets/resume, evidence tampering, +preflight barriers, replay layouts/aliases/input integrity, calibration, hotspot attribution, +installation, model-gate arithmetic and registration repair. Fake-agent sessions +also check that runs keep Humanize2's default permission unless `unfenced` is +recorded, and that batches and registration freeze that choice. Registry tests +also cover preference, fallback, profile mismatch and digest checks. + +GPU and checkpoint tests must use their markers and matching isolated model +runtime. Required evidence includes production-baseline correctness, changed-input +CUDA Graph replay, registration bit-exact integration, profile hit/fallback paths, +and whole-model verification. CPU/fake-agent passes prove neither GPU kernel +correctness nor speed. Test future selection/fusion logic when implemented; +do not describe those planned tests as existing coverage. + +## 9. Benchmark plan + +Follow the owning model's benchmark protocol, for pi0.5 +[`benchmarks/pi05-benchmark/README.md`](../../benchmarks/pi05-benchmark/README.md). +A target experiment is one prepared GPU, B=1, BF16, ten denoise steps, fixed +LIBERO-10 inputs, production compile settings and full-loop CUDA Graphs. The +standard config selects 1,600 frames; shortened runs must report their actual +coverage. Record GPU/driver/CUDA/Torch/Triton, checkpoint/revision, source digests, +dtype, attention, precision, seeds, batch, warmup, repetitions, step counts and +all runtime switches for each row. + +Report Torch, existing Triton and tuned-profile rows with mean/P50/P95 and stage +splits, plus operator ablations. An additional Inductor `max-autotune` row is +useful only where it compiles and passes the same numerical gate. Keep raw +conditions and measurements with the report. A win on one GPU, shape or batch +is not evidence for another. + +### 9.1 `orin_pi05` on Jetson AGX Orin + +Measured 2026-10-15 on Jetson AGX Orin 32 GB (sm_87, `nvpmodel` MAXN), L4T +R36.4.7, CUDA 12.6, Torch 2.10.0, Triton 3.6.0, checkpoint `pi05_base` +(`config.json` SHA256 `36786971…`). Inputs are 200 recorded observations from +the first two LIBERO-10 task files (dataset revision `f13aa24a…`, input file +SHA256 `94e5cad5…`). B=1, ten denoise steps, seed 1000 with one FP32 CUDA noise +draw per input rounded to BF16. Mixed precision: the SigLIP vision tower and +multimodal projector run in FP32, everything else in BF16. `native_inference`, +`native_embeddings`, Triton prefix and denoise attention, `compile_backend: +none`, full denoise-loop CUDA Graph, `TORCHINDUCTOR_DETERMINISTIC=1`. Ten warmup +calls, then the 200 inputs three times (600 timed calls). Rows differ only in +`kernel_profile`. All rows ran the source tree that adds this profile, in which +pi0.5 also routes its projections through the `linear` registry operator and its +SigLIP tanh GELU through `gelu_tanh`. + +| `kernel_profile` | Mean (ms) | P50 (ms) | P95 (ms) | vs `torch` | +|---|---|---|---|---| +| `torch` | 825.5 | 825.3 | 827.5 | — | +| `none` (existing Triton) | 609.8 | 609.7 | 611.1 | −26.1% | +| `orin_pi05` | 537.8 | 537.5 | 539.9 | −34.9% (−11.8% vs `none`) | + +Deviations from the protocol above: an experiment-local evaluator replaced +`benchmarks/pi05-benchmark/benchmark.py`, which does not pass +`native_embeddings` and fixes camera keys for the LIBERO fine-tuned checkpoint; +200 inputs instead of 1,600; no prefill/decode split. The prefix encoder runs +eagerly instead of under Inductor because the tuned kernels' layout gates call +`Tensor.storage_offset()`, which Dynamo cannot trace under `fullgraph`; the +whole denoise loop is still captured in one CUDA Graph. + +Every captured layout dispatched to its tuned kernel: no `rotate_qk` or +`split_kv_attention` call fell back. `linear` falls back for the layouts that +were not tuned (FP32 vision layers and the fused gate/up projections). Against +`torch`, the first action of every chunk (6,400 values) differs by at most +0.0156 (RMSE 0.0016). The `none` row differs from `torch` by a similar amount +(0.0195 / 0.0016). The whole-model `verify` gate (Section 6) was not run for +this profile. + +Operators, from the promoted attempt's paired evaluator rounds against the +task's Torch baseline (calls per inference in parentheses): + +| Operator (run) | Workload | Torch (ms) | Tuned (ms) | +|---|---|---|---| +| `split_kv_attention` (`20261015T055303Z-8eeab7204c`) | prefix, 672 keys (18) | 3.62 | 0.49 | +| | denoise, 672 + 50 keys (180) | 0.908 | 0.070 | +| `linear` (`20261015T090245Z-031b8d3bd3`) | prefix `down_proj`, 672×16384→2048 (18) | 6.45 | 3.77 | +| | expert `down_proj`, 50×4096→1024 (180) | 0.076 | 0.078 | +| `rotate_qk` prefix (`20261015T134002Z-6a72e1be9e`) | 672 tokens (18) | 0.434 | 0.049 | +| `rotate_qk` decode (`20261015T150031Z-03be0c2142`) | 50 tokens (180) | 0.071 | 0.013 | + +The tuned expert `down_proj` is about 3% slower than Torch. It was promoted +because the weighted mean over both workloads improved. +`rotate_qk` combines two separately tuned exports with disjoint verified +layouts. Its package dispatches each call to the variant whose `applies` +accepts it. The decode export was tuned after capture recorded the production +output layout (a contiguous query). Replacing the existing Triton decode +rotation with it changed end-to-end latency by less than 1 ms and left actions +bit-identical. + +## 10. Risks and limitations + +- Captured layouts are finite; a production compile can change strides and cause + fallback. Check dispatch coverage and benchmark the actual production path. +- Region selection and fusion remain manual. Large unregistered GEMMs cannot be + optimized by `capture-ops` until their routing/reference contract exists. +- Profile environment matching is strict; upgrades require new validation. +- CUPTI/graph attribution may be incomplete. Diagnostic eager execution changes + the experiment and must be distinguished from production timing. +- Finite numerical tests do not prove correctness on every input or task success. + Thermal state, other GPU users and clocks can affect timing despite pairing. +- Profile installation and verification are separate; failed verification has + no automatic rollback or publication gate in the current installer. +- Humanize2/FlashInfer use pinned optional dependencies. Real Humanize2 execution + requires POSIX and Python 3.12+ in a separate tool environment; model runtimes + retain their own dependency pins and Python support floor. +- Agents and kernels execute with process permissions. Workspace bounds and + artifact digests are integrity checks, not an OS security sandbox. Humanize2's + default session fence needs Linux Landlock; `--unfenced` (frozen per run or + batch) drops it on hosts without Landlock and leaves isolation to a dedicated + account or container. + +## 11. Automated tuning workflow + +The following pi0.5 example connects the currently implemented stages. Run these +commands from the repository root, replacing paths, checkpoint, agent and profile +names with values for the target machine. Prepare real model weights and benchmark +data first. Random-weight skeletons support structural exploration but do not +establish real-model correctness or performance. + +### 11.1 Prepare environments and freeze inputs + +Use Python 3.12+ for the tooling environment and the model's separate runtime for +model execution and GPU evaluation. Keep the incompatible pi0.5 and Qwen dependency +groups in separate environments. See +[CONTRIBUTING](../../CONTRIBUTING.md#optional-kernel-tuning-tools) for installation +and evaluator dependencies. + +```bash +uv sync --project scripts/kernel_tuning --python 3.12 --frozen +uv run --project scripts/kernel_tuning python -m scripts.kernel_tuning --help +``` + +Copy `benchmarks/pi05-benchmark/config.yaml` into the ignored results directory, +set the local checkpoint, dataset and output paths, and save it as +`results/kernel_tuning/pi05-benchmark.yaml`. Export the inputs and serving +configuration. Set `compile_backend`, `kernel_profile` and attention options in +that configuration's `policy_kwargs` to match the experiment; the exporter does +not copy every model switch automatically. The input count must cover the +subsequent profiling warmup and measurement window. + +```bash +/absolute/runtime/python -m benchmarks.kernel_tuning.discovery_inputs pi05 \ + --config results/kernel_tuning/pi05-benchmark.yaml \ + --samples 200 \ + --output results/kernel_tuning/pi05-inputs.pt \ + --serving-config results/kernel_tuning/pi05-serving.json +``` + +#### Agent hosts without Landlock (`--unfenced`) + +By default Humanize2 fences every coding-agent session: the session workdir is +writable, the rest of the machine is read-only, and the network stays on. On +Linux the fence is enforced with Landlock, which requires kernel 5.13 or later +built with `CONFIG_SECURITY_LANDLOCK=y` and listed as active. Check the target +host before starting agents: + +```bash +uname -r # 5.13 or later +cat /sys/kernel/security/lsm # must list landlock +zcat /proc/config.gz 2>/dev/null | grep LANDLOCK || grep LANDLOCK /boot/config-$(uname -r) +``` + +If Landlock is compiled in but inactive, add it to the kernel's `lsm=` boot +parameter. Some vendor kernels omit it entirely; for example, Jetson L4T R36 +ships with `# CONFIG_SECURITY_LANDLOCK is not set`, and enabling it means +rebuilding the kernel and its out-of-tree modules. Without Landlock, every agent +session is refused with +`HarnessSandboxed: ... fencing it from outside needs Landlock`. + +On such hosts, pass `--unfenced` when the run is created. Its sessions are then +granted every Humanize2 scope, which needs no fence: + +| Command | Where to pass it | On resume | +| --- | --- | --- | +| `run TASK --agent ... --unfenced` | when the run is created | `resume RUN` reuses it; no flag is accepted | +| `tune-all ... --agent ... --unfenced` | when the batch is created; every operator run inherits it | `tune-all --resume BATCH` rejects `--unfenced` | +| `register run RUN --agent ... --unfenced` | when the reviewed sites are confirmed | `register resume RUN` reuses it; another `register run` is refused | + +The choice is recorded as `unfenced` in the run's `manifest.json` (and the batch's +`batch.json`, whose `summary.md` notes it), and a warning is printed to stderr +before agents start. Runs created without the flag keep Humanize2's default +fence exactly, so the flag cannot widen an existing run. + +An unfenced agent can read and write everything its account can, including the +repository checkout, the run archive and the evaluator source. Digests still stop +a tuning run whose task, controller/evaluator source or winning evidence changed, +and a registration run whose specification, BASE checkout or context changed, but +they are integrity checks, not a sandbox. Run unfenced agents only under a +dedicated account without `sudo`, `docker` group membership or access to other +users' files, or inside a container that mounts only what the run needs. Give +that account the device groups the evaluator needs (for example `video` and +`render` on Jetson), its own tooling environment and GPU runtime, and its own +coding-agent login. + +A refused session is recorded as a failed attempt, not as a run error. A tuning +run can therefore still finish with `status: completed` and +`stop_reason: max_candidates`, and registration retries it as agent feedback +until its attempts are `exhausted`. The batch `summary.md` shows only the stop +reason. After starting on a new host, check the first attempt for +`HarnessSandboxed` before leaving a long run unattended: +`attempts[].decision.reasons` in the run's `manifest.json` (for a batch, under +`runs/`), or `attempts/*/feedback.json` for registration. + +### 11.2 Profile and select operators + +Profile with the production configuration. This example uses BF16, ten denoise +steps and a full-loop CUDA Graph. The compile and attention settings in +`pi05-serving.json` must match the baseline. + +```bash +/absolute/runtime/python -m scripts.kernel_tuning profile \ + --output results/kernel_tuning/profile-pi05 --warmup 10 --steps 20 --trace \ + -- -m scripts.kernel_tuning.discovery.drive --policy pi05 \ + --checkpoint /absolute/checkpoint --device cuda:0 --dtype bfloat16 \ + --num-steps 10 --capture-full-loop \ + --adapter-config results/kernel_tuning/pi05-serving.json \ + --inputs results/kernel_tuning/pi05-inputs.pt +``` + +Read `hotspots.md` and manually select registered operators with room for +improvement. The commands below use `gated_residual` as an example; this does not +imply that it is a worthwhile tuning target. For an unregistered operator, follow +the [registration guide](../../scripts/kernel_tuning/registration/README.md), +complete code review and CPU checks, and verify whole-model bit-exact parity on +the target GPU before capturing workloads. + +### 11.3 Capture workloads and calibrate + +Copy the serving configuration to `pi05-eager.json`. Set +`policy_kwargs.compile_backend` and `policy_kwargs.kernel_profile` to `"none"`, +while preserving the checkpoint, dtype, attention settings and other mathematical +semantics. Also disable CUDA Graphs during capture so Python registry calls can +be recorded. + +```bash +/absolute/runtime/python -m scripts.kernel_tuning capture-ops \ + --model pi05 --ops gated_residual --warmup 2 --steps 8 \ + --fixture-bytes 8589934592 --output results/kernel_tuning/captures/pi05-ops \ + -- -m scripts.kernel_tuning.discovery.drive --policy pi05 \ + --checkpoint /absolute/checkpoint --device cuda:0 --dtype bfloat16 \ + --num-steps 10 --no-cuda-graph \ + --adapter-config results/kernel_tuning/pi05-eager.json \ + --inputs results/kernel_tuning/pi05-inputs.pt + +/absolute/runtime/python -m scripts.kernel_tuning calibrate \ + --captured results/kernel_tuning/captures/pi05-ops \ + --output results/kernel_tuning/captures/pi05-ops-calibrated +``` + +`--fixture-bytes` sets a storage limit, 8 GiB in this example; it does not +preallocate that space. If calibration fails, inspect input coverage and the +reference contract. Do not relax thresholds to make a candidate pass. +`capture-ops --ops` selects the scope of this capture. Subsequent batch tuning +must cover every task in that capture, which does not imply coverage of every +ATen operation in the model. + +### 11.4 Run the agent, resume and export + +```bash +uv run --project scripts/kernel_tuning python -m scripts.kernel_tuning tune-all \ + --model pi05 --captured results/kernel_tuning/captures/pi05-ops-calibrated \ + --agent 'HARNESS/MODEL:EFFORT' --set evaluator_python=/absolute/runtime/python + +# Resume with the printed Batch path; omit agent, selection options and --set. +uv run --project scripts/kernel_tuning python -m scripts.kernel_tuning tune-all \ + --resume results/kernel_tuning/batches/BATCH_ID +``` + +Use `--dry-run` to generate tasks only, or `--preflight-only` to measure baselines +without starting agents. On a host whose kernel lacks Landlock, add `--unfenced` +when creating the batch and run it under an isolated account; see +[Agent hosts without Landlock](#agent-hosts-without-landlock---unfenced). If a +preflight-only batch will later resume search, specify its agent (and +`--unfenced`, if needed) when creating it. A normal run starts agents only after +every baseline passes. Default limits apply per task: 20 attempts, two active hours, +or five consecutive failures/non-improvements. + +Inspect the batch's `summary.md` / `summary.json` and each task's run evidence. +Promoted results are automatically exported to `exports/`; use the +actual directory names in the batch. For individual runs, use `status RUN` to +inspect progress and `export RUN DESTINATION` to export the result. Proceed to +installation only when a candidate has been promoted. + +### 11.5 Install and verify the whole model + +Restore the production serving configuration. Install only a registry export +produced by the preceding stage. Profile names must be Python identifiers; +`EXPORT_DIRECTORY` below denotes the actual export directory. + +```bash +/absolute/runtime/python -m scripts.kernel_tuning install EXPORT_DIRECTORY --profile my_gpu + +TORCHINDUCTOR_DETERMINISTIC=1 /absolute/runtime/python -m scripts.kernel_tuning verify \ + --profile my_gpu --output results/kernel_tuning/verify-my_gpu.json \ + --seed 1000 -- --policy pi05 --checkpoint /absolute/checkpoint \ + --device cuda:0 --dtype bfloat16 --num-steps 10 --capture-full-loop \ + --adapter-config results/kernel_tuning/pi05-serving.json \ + --inputs results/kernel_tuning/pi05-inputs.pt +``` + +`verify` switches internally between FP32 Torch, production-dtype Torch and the +selected profile. It does not automatically roll back an installation that fails +verification or write verification results into `profile.json`. After the report +passes, confirm that production layouts select the tuned kernels, and retain the +checkpoint identity, input digest and evidence from each stage. This tolerance +gate does not replace the bit-exact gate for registration changes. + +### 11.6 Benchmark the whole model and prepare deliverables + +Create three configurations from the same pi0.5 benchmark configuration, setting +`kernel_profile` to `torch`, `none` (the existing Triton path), and `my_gpu`, +respectively. Give each a distinct `output` path. Keep all other experimental +conditions identical, including attention selection, inputs, seed, warmup, +repeats, compilation, CUDA Graphs and denoise steps. `kernel_profile` controls +registered operators only. + +```bash +TORCHINDUCTOR_DETERMINISTIC=1 /absolute/runtime/python \ + benchmarks/pi05-benchmark/benchmark.py --config results/kernel_tuning/bench-torch.yaml +TORCHINDUCTOR_DETERMINISTIC=1 /absolute/runtime/python \ + benchmarks/pi05-benchmark/benchmark.py --config results/kernel_tuning/bench-triton.yaml +TORCHINDUCTOR_DETERMINISTIC=1 /absolute/runtime/python \ + benchmarks/pi05-benchmark/benchmark.py --config results/kernel_tuning/bench-profile.yaml +``` + +Deliver reviewable kernel/profile sources, candidate and baseline operator +evidence, the whole-model correctness report, and all three benchmark +configurations with raw results. Profiles containing multiple operators should +also include per-operator ablations. Report hardware and software versions, +checkpoint, dataset coverage and timing boundaries. Only model-level measurements +establish end-to-end gains. Installation, committing and enabling a profile by +default are separate actions. Keep raw experiment data in the results directory +and formal reports under `benchmarks/`; keep run logs out of `docs/`. + +## Appendix A. Task-owned benchmark interface + +`benchmark.py` is trusted task code, hashed and snapshotted before search. It +exports `create_adapter(task)`, where `task` is the import-safe `TaskPackage` from +`scripts.kernel_tuning.contracts`. The returned object implements: + +```python +def environment(self) -> dict: + # Actual device, compiler/runtime/library versions, and numerical switches. + # Fields under identity must remain stable on resume. + return {"identity": {...}, "conditions": {...}} + +def evaluate(self, solutions: dict) -> dict: + # Maps baseline/incumbent/candidate to official Solution objects. + # Check all three against definition.reference using identical inputs. + # Alternate their timing order within each paired round. + return {"checks": [...], "rounds": [...]} + +# Optional, only called when tuning.yaml enables profile; never promotion timing. +def profile(self, solutions: dict, output_dir: Path) -> dict: + ... +``` + +The controller requires one check per `(workload UUID, role)`: + +```json +{ + "workload": "a-declared-uuid", + "role": "candidate", + "passed": true, + "seeds": [0, 1, 2], + "precision": { + "mode": "bit_exact", "atol": 0.0, "rtol": 0.0, "reason": "", + "matmul_precision": "highest", "allow_tf32": false + }, + "max_abs_error": 0.0, + "max_rel_error": 0.0, + "graph_replay": false +} +``` + +`seeds` and the complete `precision` object must equal the task settings (use +`dataclasses.asdict(task.settings.precision)`). Every check must pass; bit-exact +errors must be zero. For `cuda_graph`, `graph_replay` must be true **after checking +changed valid input tensors and fresh reference outputs**. A tolerance task uses +`abs(actual-reference) <= atol + rtol*abs(reference)` for every floating element, +exact integer/bool outputs, and finite values. The default adapter rejects input +mutation and incorrect output shape/dtype/device. Entrypoints must launch on the +caller's current stream or join any helper-stream work before returning, while +preserving its device, stream, and numerical settings. Input mutation, constrained +random values, non-contiguous inputs, statistical operators, fixed graph fixtures +requiring domain-specific perturbation, and other accelerators use the custom +adapter. It must implement the same frozen measurement contract. + +Each element of `rounds` maps **every** UUID to +`{"baseline": [ms, ...], "incumbent": [ms, ...], "candidate": [ms, ...]}`. +List lengths must equal `timing.trials`; the number of rounds must equal +`timing.paired_rounds` (at least two). All timings must be finite and positive. +Within each round the gate averages trials per workload, then computes +`sum(weight * workload_latency) / sum(weight)` for each role. Candidate weighted +latency must improve by the configured fraction against both comparison roles, +and each workload must meet both regression caps. Every round must pass. + +The worker attaches a nonce, request digest, task/source digests, environment +before/after, and exit status. Controller-side validation rejects mismatches and +incomplete evidence. A custom adapter remains responsible for truthful measurements +and numerical checks: the controller cannot infer GPU correctness from a boolean. +The default adapter supplies these checks for deterministic pure tensor operators. + +## Appendix B. Archive and reproduction + +`run` creates a unique directory under `results/kernel_tuning/`; `resume`, +`status`, and `export` take that directory. It contains the frozen task and tool +source, baseline Solution, every scratch attempt, candidate Solutions/source trees/plans, +worker requests/commands/stdout/stderr/results, decisions, and environment data. +`best.json` appears only after promotion. Winning evidence is hashed and checked +before export; source export writes the exact UTF-8 bytes embedded in the measured +Solution. Optional NCU profiling runs separately on the first workload; profile +failure never changes promotion. + +Humanize2 owns resumable FlowState and the agent/session trace. The manifest is an +atomically written mirror. The latest framework journal in the dedicated workspace +wins on resume, including after an unclean exit. Attempt reservation is journaled +before scratch preparation. Interrupted attempts count toward candidate/patience +budgets; active elapsed time is checkpointed every second (an unclean kill may +lose the last heartbeat interval, excluding a stalled host). Downtime is not +charged. Exhausted runs require a new task/run. Task, hardware/software, evaluator, +and installed orchestration runtime identities must match to resume. + +`humanize.json` locates the original framework epic and copied flow/resume journals +under `humanize/`; CLI-managed sessions remain in the original epic. The tooling +`uv.lock` records the isolated environment, and the manifest records the actual +installed hmz source digest and origin. Exports include the task, tool source/lock, +baseline, winning evidence, source tree, and checksums. Export remains available +after a repository tool update if archived source/evidence hashes still match. + +To reproduce a measurement on a prepared target, use the saved tool's +`benchmarks/kernel_tuning/evaluate.py` with the archived request. Adjust absolute +task/solution/output locations when moving the bundle; retain task/solution digests, +nonce, runtime, and numerical/timing conditions. Use the evaluator Python recorded +in the task. The new response's request digest reflects relocated paths. This +executes evaluation, not agent search. Actual Thor compilation and performance +are not established by CPU/fake-agent tests. + +## Appendix C. Other task sources + +### Registry capture artifacts and shared tooling + +`capture-ops` selects one canonical policy and observes registry calls during its +measured inference window. `--ops` chooses the operator scope; subsequent +generation, calibration and tuning cover every recorded workload of that scope. +It records no arbitrary ATen operations. Run the driver eagerly with compilation +and CUDA Graphs disabled so registry calls reach Python. Capture does not certify +unvisited paths, other shapes or checkpoint-level action parity. + +`scripts.kernel_tuning.discovery.drive` remains the shared driver for `profile`, +`capture-ops` and `verify`. It constructs the policy and engine using serving +arguments and replays recorded requests or observations, retaining recurrent +episode sessions. `benchmarks.kernel_tuning.discovery_inputs` exports these +inputs using the pi05, Qwen navigation and StreamVLN benchmark loaders. Those +historical module paths are retained; they do not enable full-model discovery. + +The capture contains: + +- `manifest.json`: `capture: registry`, policy identity/configuration, revision, + Torch version, numerical settings, capture window, completeness and file hashes. +- `operators.json`: registry/Torch-reference/starting-implementation identities, + argument trees, observed shapes/strides/dtypes, aliases and call counts. +- `coverage.json`: operator statuses and explicit capture errors. +- `fixtures/`: real input storage, bounded by `--fixture-bytes` (default 1 GiB). + +Calibration writes a new capture with `calibration.json`; generated tasks carry +the per-operator bounds in `numerics.json`. Inputs, seeds, masks and numerical +settings follow Section 6. First-observed real fixtures represent each structural +case; random and boundary profiles provide finite additional validation. + +Tasks use the frozen Torch reference as their baseline and include +`replay.json` for storage/layout/alias contracts. Production output layouts are +preserved, including the reference relayout described in Section 4.3. Replay +checks every output and rejects input mutation, resets storage outside timed +events, and validates changed-input graph replay. Registry tasks default to +CUDA Graph timing; eager timing remains available. NCU profiling requires a +replay-aware adapter. + +The shared `discovery.contracts`, `tasks`, `tensors` and `numerics` modules remain +in place for registry capture, calibration, task generation, batch preflight and +evaluation. The `embodiinfer-replay-v1` float64 definition extension also remains; +it preserves double precision without changing the ordinary FlashInfer schema. + +The full-model `capture --model` entry, ATen/backend observer, discovery-only +task/oracle branches, RNG-transition replay and greedy-workspace recipes have +been removed. New model captures and replay tasks must identify registry +operators. Old discovery artifacts are rejected with `capture-ops` migration +guidance; recapture through registered call sites. `capture --catalog` remains +an independent, explicit JSONL workflow. + +Calibration binds the numerical-policy/oracle source digest. This cleanup changes +that digest without changing registry reference arithmetic or threshold formulas: +existing raw registry captures can be recalibrated into a new directory before +generating new tasks/batches. Do not edit old hashes or widen thresholds to reuse +old evidence. Archived runs retain their snapshotted tooling for reproduction. + +### Generated catalog tasks and batch tuning + +The legacy catalog in `scripts/kernel_tuning/operators.py` supplies hand-maintained +operator shapes, entry points, Torch references and fixed numerical contracts. +`generate --catalog` snapshots production backend sources as the baseline; +`capture --catalog` records legacy JSONL workload shapes/counts, not registry +fixtures or full-model coverage. `tune-all --catalog` uses the same Humanize2 +search, promotion and export loop. These exports require separate integration; +`install` accepts registry exports only. + +The catalog covers pi0.5 norm/activation/RoPE/split-KV kernels, Qwen2.5-VL vision +RoPE/segmented attention, and StreamVLN norm/activation kernels. `--estimated` +explicitly opts into estimated shapes. Existing fixed tolerances remain defined +by the catalog and must pass baseline preflight on the target runtime; they are +not silently replaced by calibrated registry contracts. + +```bash +uv run --project scripts/kernel_tuning python -m scripts.kernel_tuning generate \ + --catalog gated_residual --output results/kernel_tuning/tasks +uv run --project scripts/kernel_tuning python -m scripts.kernel_tuning check \ + results/kernel_tuning/tasks/gated_residual +``` + +Catalog capture must load EmbodiInfer's backend source from the checkout and run +the relevant path; importing an unrelated copy or disabling the path that calls +the target kernel gives incomplete coverage. Use recorded workloads for real +measurement. Synthetic/estimated workloads and skeleton checkpoints do not prove +real-checkpoint performance or model correctness. diff --git a/embodiinfer/backend/torch/ops.py b/embodiinfer/backend/torch/ops.py new file mode 100644 index 0000000..4dc041d --- /dev/null +++ b/embodiinfer/backend/torch/ops.py @@ -0,0 +1,101 @@ +"""Torch reference implementations of registry operators. + +Each function is the default implementation of the operator with the same name +in :mod:`embodiinfer.layers.ops` and defines that operator's mathematics. The +arithmetic order is part of the contract: policies compile these functions with +Inductor, and changing an expression changes their rounding. +""" + +from __future__ import annotations + +import torch +from torch.nn import functional as F + + +def linear(x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None = None) -> torch.Tensor: + """Apply the native Torch affine projection without mutating any input. + + ``x`` is ``[..., in_features]``, ``weight`` is ``[out_features, in_features]``, + and optional ``bias`` is ``[out_features]``. The result is + ``[..., out_features]`` and retains ``x.dtype`` for same-dtype operands + outside autocast. Preserve the exact ``F.linear`` path, operand layouts, + bias handling, and autograd semantics; do not cast or reshape operands. + """ + return F.linear(x, weight, bias) + + +def gated_residual(residual: torch.Tensor, update: torch.Tensor, gate: torch.Tensor) -> torch.Tensor: + """Return ``residual + update * gate``.""" + return residual + update * gate + + +def ada_rms_norm(x: torch.Tensor, modulation: torch.Tensor, eps: float) -> tuple[torch.Tensor, torch.Tensor]: + """Adaptive RMSNorm over ``[batch, tokens, hidden]`` with ``[batch, 3 * hidden]`` modulation. + + The modulation holds scale, shift and residual gate. Variance and the affine + step are FP32; the normalized output and the gate return in ``x.dtype``. + """ + var = torch.mean(torch.square(x.float()), dim=-1, keepdim=True) + normed = x * torch.rsqrt(var + eps) + if x.ndim == 3: + modulation = modulation.unsqueeze(1) + scale, shift, gate = modulation.chunk(3, dim=-1) + normed = normed * (1 + scale.float()) + shift.float() + return normed.to(x.dtype), gate.to(x.dtype) + + +def gated_gelu(gate: torch.Tensor, up: torch.Tensor) -> torch.Tensor: + """Return tanh-approximate ``GELU(gate) * up``.""" + return F.gelu(gate, approximate="tanh") * up + + +def gelu_tanh(x: torch.Tensor) -> torch.Tensor: + """Return exactly ``F.gelu(x, approximate="tanh")`` without mutating ``x``. + + The floating tensor result retains the input's shape, dtype, and device. + Preserve native GELU rounding and autograd, including for noncontiguous + inputs; do not cast, reshape, or expand the formula into separate operations. + """ + return F.gelu(x, approximate="tanh") + + +def _rotate_half(x: torch.Tensor) -> torch.Tensor: + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def rotate_qk( + query: torch.Tensor, key: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + """Rotate-half RoPE for BHSD query/key with ``[batch, tokens, head_dim]`` tables.""" + cos = cos.unsqueeze(1) + sin = sin.unsqueeze(1) + return (query * cos) + (_rotate_half(query) * sin), (key * cos) + (_rotate_half(key) * sin) + + +def split_kv_attention( + query: torch.Tensor, + prefix_key: torch.Tensor, + prefix_value: torch.Tensor, + suffix_key: torch.Tensor, + suffix_value: torch.Tensor, + mask: torch.Tensor | None, + scaling: float, +) -> torch.Tensor: + """Non-causal BHSD attention over cached (prefix) and current (suffix) K/V. + + ``mask`` is a shared floating padding mask ``[batch, 1, 1, prefix + suffix]``: + zero keeps a key, any other value removes it. Query heads may be a multiple + of the KV heads (grouped-query attention). + """ + key = torch.cat((prefix_key, suffix_key), dim=2) + value = torch.cat((prefix_value, suffix_value), dim=2) + return F.scaled_dot_product_attention( + query, + key, + value, + attn_mask=None if mask is None else mask == 0, + scale=scaling, + enable_gqa=query.shape[1] != key.shape[1], + ) diff --git a/embodiinfer/backend/triton/tuned/__init__.py b/embodiinfer/backend/triton/tuned/__init__.py new file mode 100644 index 0000000..6fd4693 --- /dev/null +++ b/embodiinfer/backend/triton/tuned/__init__.py @@ -0,0 +1,163 @@ +"""Verified, hardware-specific operator kernels grouped into kernel profiles. + +Each subdirectory is one profile (for example ``thor`` or ``rtx4090``) holding a +``profile.json`` manifest and one self-contained module or package per tuned +operator (``scripts.kernel_tuning install`` writes packages). +Profiles are produced by the kernel-tuning tooling and selected at runtime through +:class:`embodiinfer.layers.ops.KernelSelection`; nothing here is used by default. + +``profile.json``:: + + { + "schema": 1, + "match": {"device_name": "NVIDIA Thor", "capability": [11, 0], + "torch": "2.13.0+cu132", "triton": "3.7.1"}, + "ops": { + "gated_residual": {"module": "gated_residual", "function": "run", + "applies": "applies", "sha256": ""} + } + } + +``applies`` names an optional predicate taking the operator's arguments and +returning whether the kernel supports them; unsupported calls fall back to the +policy's preferred implementation. Modules may import only Torch and Triton. +""" + +from __future__ import annotations + +import hashlib +import importlib.util +import json +import sys +from collections.abc import Callable, Mapping +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +PROFILE_ROOT = Path(__file__).resolve().parent +_SCHEMA = 1 +_MATCH_KEYS = ("device_name", "capability", "torch", "triton") + + +@dataclass(frozen=True) +class TunedKernel: + """One tuned operator implementation and its applicability predicate.""" + + fn: Callable[..., Any] + applies: Callable[..., bool] | None + + +@dataclass(frozen=True) +class KernelProfile: + """A validated profile manifest; kernels load on first use.""" + + name: str + root: Path + match: Mapping[str, Any] + ops: Mapping[str, Mapping[str, Any]] + _loaded: dict[str, TunedKernel] = field(default_factory=dict, compare=False, repr=False) + + def matches(self, environment: Mapping[str, Any]) -> bool: + """Return whether every match key equals the current environment.""" + return all(_normalize(self.match[key]) == _normalize(environment.get(key)) for key in _MATCH_KEYS) + + def kernel(self, op: str) -> TunedKernel: + """Import ``op``'s module after checking its recorded digest.""" + if op in self._loaded: + return self._loaded[op] + entry = self.ops[op] + path = self.root / entry["module"] + if source_digest(path) != entry["sha256"]: + raise ValueError(f"kernel profile {self.name!r}: {path.name} does not match its recorded sha256") + name = f"embodiinfer_tuned_{self.name}_{path.stem}" + if path.is_dir(): + spec = importlib.util.spec_from_file_location( + name, path / "__init__.py", submodule_search_locations=[str(path)] + ) + else: + spec = importlib.util.spec_from_file_location(name, path) + if spec is None or spec.loader is None: + raise ImportError(f"kernel profile {self.name!r}: cannot import {path}") + module = importlib.util.module_from_spec(spec) + # Packages import their own files relatively, which needs the parent registered. + sys.modules[name] = module + try: + spec.loader.exec_module(module) + except BaseException: + sys.modules.pop(name, None) + raise + applies = entry.get("applies") + kernel = TunedKernel( + getattr(module, entry["function"]), getattr(module, applies) if applies else None + ) + self._loaded[op] = kernel + return kernel + + +def source_digest(path: Path) -> str: + """SHA-256 of a module file, or of every file (by relative path) in a package directory.""" + if path.is_file(): + return hashlib.sha256(path.read_bytes()).hexdigest() + files = { + item.relative_to(path).as_posix(): hashlib.sha256(item.read_bytes()).hexdigest() + for item in sorted(path.rglob("*")) + if item.is_file() and "__pycache__" not in item.parts + } + return hashlib.sha256(json.dumps(files, sort_keys=True).encode()).hexdigest() + + +def _normalize(value: Any) -> Any: + return list(value) if isinstance(value, tuple) else value + + +def load_profile(root: Path) -> KernelProfile: + """Read and validate one profile directory without importing its kernels.""" + manifest = json.loads((root / "profile.json").read_text(encoding="utf-8")) + if manifest.get("schema") != _SCHEMA: + raise ValueError(f"kernel profile {root.name!r}: unsupported schema {manifest.get('schema')!r}") + match = manifest["match"] + missing = [key for key in _MATCH_KEYS if key not in match] + if missing: + raise ValueError(f"kernel profile {root.name!r}: match is missing {missing}") + for op, entry in manifest["ops"].items(): + for key in ("module", "function", "sha256"): + if key not in entry: + raise ValueError(f"kernel profile {root.name!r}: operator {op!r} is missing {key!r}") + return KernelProfile(root.name, root, match, manifest["ops"]) + + +def available_profiles(root: Path = PROFILE_ROOT) -> dict[str, KernelProfile]: + """Return every profile directory under ``root`` by name.""" + return { + path.name: load_profile(path) + for path in sorted(root.iterdir()) + if path.is_dir() and (path / "profile.json").is_file() + } + + +def current_environment() -> dict[str, Any]: + """Describe the active CUDA device and Torch/Triton versions for profile matching.""" + import torch + + environment: dict[str, Any] = {"torch": torch.__version__} + try: + import triton + + environment["triton"] = triton.__version__ + except ImportError: + environment["triton"] = None + if torch.cuda.is_available(): + environment["device_name"] = torch.cuda.get_device_name() + environment["capability"] = list(torch.cuda.get_device_capability()) + return environment + + +__all__ = [ + "KernelProfile", + "PROFILE_ROOT", + "TunedKernel", + "available_profiles", + "current_environment", + "load_profile", + "source_digest", +] diff --git a/embodiinfer/backend/triton/tuned/orin_pi05/linear/__init__.py b/embodiinfer/backend/triton/tuned/orin_pi05/linear/__init__.py new file mode 100644 index 0000000..e539cda --- /dev/null +++ b/embodiinfer/backend/triton/tuned/orin_pi05/linear/__init__.py @@ -0,0 +1,62 @@ +"""Tuned ``linear`` kernel for the ``orin_pi05`` profile. + +Generated by ``scripts.kernel_tuning install`` from run 20261015T090245Z-031b8d3bd3; do not edit, +reinstall instead. The kernel applies only to the argument layouts verified by +the evaluator. Strides of size-one dimensions and of empty tensors address no +other element, so they are compared as zero. Like the evaluator, the kernel +receives only the tensor arguments; constants are part of the verified layout. +""" + +import torch + +from .candidate import run as run + +_VERIFIED = frozenset([(('tensor', (1, 50, 4096), (0, 4096, 1), 'bfloat16'), ('tensor', (1024, 4096), (4096, 1), 'bfloat16'), ('constant', 'None')), (('tensor', (1, 672, 16384), (0, 16384, 1), 'bfloat16'), ('tensor', (2048, 16384), (16384, 1), 'bfloat16'), ('constant', 'None'))]) +_ARGUMENTS = {'tuple': [{'tuple': [{'tensor': 0}, {'tensor': 1}, {'constant': None}]}, {'dict': {}}]} +_TENSORS = 2 + + +def _strides(shape, stride): + if 0 in shape: + return (0,) * len(shape) + return tuple(0 if size == 1 else step for size, step in zip(shape, stride)) + + +def _leaves(value, out): + if isinstance(value, torch.Tensor): + dtype = str(value.dtype).removeprefix("torch.") + shape = tuple(value.shape) + out.append(("tensor", shape, _strides(shape, value.stride()), dtype)) + elif isinstance(value, (tuple, list)): + for item in value: + _leaves(item, out) + elif isinstance(value, dict): + for key in sorted(value): + _leaves(value[key], out) + else: + out.append(("constant", repr(value))) + return out + + +def applies(*args, **kwargs): + """Whether these arguments match a layout the evaluator verified.""" + return tuple(_leaves((args, kwargs), [])) in _VERIFIED + + +def _bind(recipe, value, tensors): + kind, item = next(iter(recipe.items())) + if kind == "tensor": + tensors[item] = value + elif kind in ("tuple", "list"): + for entry, element in zip(item, value): + _bind(entry, element, tensors) + elif kind == "dict": + for key, entry in item.items(): + _bind(entry, value[key], tensors) + + +def call(*args, **kwargs): + """Call the kernel as the evaluator did: with the captured tensor arguments x0..xN.""" + tensors = {} + _bind(_ARGUMENTS, (args, kwargs), tensors) + return run(*(tensors[index] for index in range(_TENSORS))) diff --git a/embodiinfer/backend/triton/tuned/orin_pi05/linear/candidate.py b/embodiinfer/backend/triton/tuned/orin_pi05/linear/candidate.py new file mode 100644 index 0000000..a5ca760 --- /dev/null +++ b/embodiinfer/backend/triton/tuned/orin_pi05/linear/candidate.py @@ -0,0 +1,144 @@ +"""Fixed split-K BF16 projections for the two captured Orin layouts.""" + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _linear_split_k( + X, + W, + Partials, + M: tl.constexpr, + N: tl.constexpr, + K: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, + SPLIT_K: tl.constexpr, +): + """Write disjoint FP32 output planes for contiguous, exhaustive K slices.""" + tl.static_assert(SPLIT_K == 2 or SPLIT_K == 8) + tl.static_assert(K % (SPLIT_K * BLOCK_K) == 0) + tl.static_assert(N % BLOCK_N == 0) + + tiles_m = tl.cdiv(M, BLOCK_M) + tiles_n = N // BLOCK_N + pid = tl.program_id(0) + split = pid // (tiles_m * tiles_n) + tile = pid % (tiles_m * tiles_n) + tile_m = tile % tiles_m + tile_n = tile // tiles_m + + rows = tile_m * BLOCK_M + tl.arange(0, BLOCK_M) + cols = tile_n * BLOCK_N + tl.arange(0, BLOCK_N) + ks = split * (K // SPLIT_K) + tl.arange(0, BLOCK_K) + x_ptrs = X + rows[:, None] * K + ks[None, :] + w_ptrs = W + cols[None, :] * K + ks[:, None] + + acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) + # The fixed stage count pipelines this regular loop without full unrolling. + for _ in range(K // (SPLIT_K * BLOCK_K)): + a = tl.load(x_ptrs, mask=rows[:, None] < M, other=0) + b = tl.load(w_ptrs) + acc = tl.dot(a, b, acc, input_precision="ieee", out_dtype=tl.float32) + x_ptrs += BLOCK_K + w_ptrs += BLOCK_K + + offsets = split * M * N + rows[:, None] * N + cols[None, :] + tl.store(Partials + offsets, acc, mask=rows[:, None] < M) + + +@triton.jit +def _reduce_partials( + Partials, + Y, + ELEMENTS: tl.constexpr, + SPLIT_K: tl.constexpr, + BLOCK: tl.constexpr, +): + """Sum two or eight FP32 planes in a fixed tree and round to BF16 once.""" + tl.static_assert(SPLIT_K == 2 or SPLIT_K == 8) + + offsets = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + valid = offsets < ELEMENTS + p0 = tl.load(Partials + offsets, mask=valid, other=0) + p1 = tl.load(Partials + ELEMENTS + offsets, mask=valid, other=0) + if SPLIT_K == 8: + p2 = tl.load(Partials + 2 * ELEMENTS + offsets, mask=valid, other=0) + p3 = tl.load(Partials + 3 * ELEMENTS + offsets, mask=valid, other=0) + p4 = tl.load(Partials + 4 * ELEMENTS + offsets, mask=valid, other=0) + p5 = tl.load(Partials + 5 * ELEMENTS + offsets, mask=valid, other=0) + p6 = tl.load(Partials + 6 * ELEMENTS + offsets, mask=valid, other=0) + p7 = tl.load(Partials + 7 * ELEMENTS + offsets, mask=valid, other=0) + total = ((p0 + p1) + (p2 + p3)) + ((p4 + p5) + (p6 + p7)) + else: + total = p0 + p1 + tl.store(Y + offsets, total.to(tl.bfloat16), mask=valid) + + +def run(x0: torch.Tensor, x1: torch.Tensor) -> torch.Tensor: + """Return a fresh BF16 ``x0 @ x1.T`` for either captured batch-one layout. + + Supports contiguous, zero-offset CUDA BF16 shapes + ``(1, 672, 16384)/(2048, 16384)`` and + ``(1, 50, 4096)/(1024, 4096)``. The result has shape ``(1, M, N)``, + strides ``(M*N, N, 1)``, and no alias with either read-only input. + Unsupported metadata raises ValueError. Both kernels use the input + device's current stream; the caller's device and stream are preserved. + """ + if not isinstance(x0, torch.Tensor) or not isinstance(x1, torch.Tensor): + raise ValueError("Both linear inputs must be tensors.") + if x0.layout != torch.strided or x1.layout != torch.strided: + raise ValueError("Only dense strided inputs are supported.") + if x0.ndim != 3 or x1.ndim != 2: + raise ValueError("Linear inputs must have ranks three and two.") + if x0.dtype != torch.bfloat16 or x1.dtype != torch.bfloat16: + raise ValueError("Both linear inputs must have dtype bfloat16.") + if x0.device.type != "cuda" or x1.device != x0.device: + raise ValueError("Both linear inputs must be on the same CUDA device.") + + if tuple(x0.shape) == (1, 672, 16384) and tuple(x1.shape) == (2048, 16384): + m, n, k = 672, 2048, 16384 + block_m, block_n, block_k, split_k, warps = 128, 128, 64, 8, 8 + elif tuple(x0.shape) == (1, 50, 4096) and tuple(x1.shape) == (1024, 4096): + m, n, k = 50, 1024, 4096 + block_m, block_n, block_k, split_k, warps = 64, 64, 128, 2, 4 + else: + raise ValueError("Unsupported linear shapes; expected a captured workload.") + + if tuple(x0.stride()) != (m * k, k, 1) or tuple(x1.stride()) != (k, 1): + raise ValueError("Linear input strides must match the captured layouts.") + if x0.storage_offset() != 0 or x1.storage_offset() != 0: + raise ValueError("Linear inputs must have zero storage offsets.") + + # Device context restores the caller's device without changing any stream. + with torch.cuda.device(x0.device): + output = torch.empty((1, m, n), dtype=torch.bfloat16, device=x0.device) + partials = torch.empty((split_k, m, n), dtype=torch.float32, device=x0.device) + grid = (split_k * triton.cdiv(m, block_m) * (n // block_n),) + _linear_split_k[grid]( + x0, + x1, + partials, + M=m, + N=n, + K=k, + BLOCK_M=block_m, + BLOCK_N=block_n, + BLOCK_K=block_k, + SPLIT_K=split_k, + num_warps=warps, + num_stages=2, + ) + _reduce_partials[(triton.cdiv(m * n, 1024),)]( + partials, + output, + ELEMENTS=m * n, + SPLIT_K=split_k, + BLOCK=1024, + num_warps=4, + num_stages=1, + ) + return output diff --git a/embodiinfer/backend/triton/tuned/orin_pi05/profile.json b/embodiinfer/backend/triton/tuned/orin_pi05/profile.json new file mode 100644 index 0000000..24055ca --- /dev/null +++ b/embodiinfer/backend/triton/tuned/orin_pi05/profile.json @@ -0,0 +1,267 @@ +{ + "schema": 1, + "match": { + "device_name": "Orin", + "capability": [ + 8, + 7 + ], + "torch": "2.10.0", + "triton": "3.6.0" + }, + "ops": { + "split_kv_attention": { + "module": "split_kv_attention", + "function": "call", + "applies": "applies", + "sha256": "6c828c87177744d1fdb0bebe603eb6294c9f9ae22675ce8208ea0119a9527020", + "verified_layouts": 2, + "evidence": { + "run_id": "20261015T055303Z-8eeab7204c", + "task_digest": "118acee038bd715427cecd7905d26a40e01f381724e91d83f2110dacc60e2aee", + "capture_identity": "fd2cf1e09eddabe9ddbe5046e8553f1c47aef6eb22a79e77a94eb77894ead4c8", + "solution_digest": "b9f6c56fd719fcfde302e88d0c58d7d906f7e1c44e43a2fe2e5fc199fe7135e4", + "decision": { + "promoted": true, + "reasons": [], + "rounds": [ + { + "mean_latency_ms": { + "baseline": 1.1550895679683397, + "incumbent": 0.12545788023643423, + "candidate": 0.10845637841387229 + }, + "improvement": { + "baseline": 0.9061056549886138, + "incumbent": 0.1355156152050505 + }, + "workloads": { + "3495d08ab28f5f64d8e7": { + "baseline": 3.621381514231364, + "incumbent": 0.4931022825241089, + "candidate": 0.4931384532352288 + }, + "2ce4af2d9552a240245c": { + "baseline": 0.9084603733420372, + "incumbent": 0.08869344000766675, + "candidate": 0.06998817093173663 + } + } + }, + { + "mean_latency_ms": { + "baseline": 1.1474958672306754, + "incumbent": 0.1251538822727673, + "candidate": 0.10773358529444897 + }, + "improvement": { + "baseline": 0.9061141844855186, + "incumbent": 0.13919102357809066 + }, + "workloads": { + "3495d08ab28f5f64d8e7": { + "baseline": 3.599048257271449, + "incumbent": 0.4887091305057208, + "candidate": 0.4896779734591643 + }, + "2ce4af2d9552a240245c": { + "baseline": 0.9023406282265981, + "incumbent": 0.08879835744947195, + "candidate": 0.06953914647797743 + } + } + } + ] + }, + "export_checksums": "d32a5b2f462f1d022832e8820751446a7721ca35af569d6502c906eee0739ced" + } + }, + "linear": { + "module": "linear", + "function": "call", + "applies": "applies", + "sha256": "5e1f1f051298eae832364b342455bef984a5d016fcce3716b0d682f94fa90d88", + "verified_layouts": 2, + "evidence": { + "run_id": "20261015T090245Z-031b8d3bd3", + "task_digest": "9bb461df61b63bd03ae23ea0c49bd671f0775e6eade79706554ec8fde235c0f5", + "capture_identity": "fd2cf1e09eddabe9ddbe5046e8553f1c47aef6eb22a79e77a94eb77894ead4c8", + "solution_digest": "cd7fc9616b97b764e704cefa5bb3b5185ae4aaae13fcdd760dc1512d85011e4f", + "decision": { + "promoted": true, + "reasons": [], + "rounds": [ + { + "mean_latency_ms": { + "baseline": 0.655083080725688, + "incumbent": 0.6551149717605476, + "candidate": 0.41346485340234007 + }, + "improvement": { + "baseline": 0.3688360063515731, + "incumbent": 0.3688667314514277 + }, + "workloads": { + "51946d28f5e3610d7f2d": { + "baseline": 6.4458313288688664, + "incumbent": 6.446487516562144, + "candidate": 3.767086400985718 + }, + "72036e17c4985117b154": { + "baseline": 0.07600825591137012, + "incumbent": 0.07597771728038788, + "candidate": 0.07810269864400228 + } + } + }, + { + "mean_latency_ms": { + "baseline": 0.6541345326828234, + "incumbent": 0.6540909990934712, + "candidate": 0.41351395188272 + }, + "improvement": { + "baseline": 0.367845708761527, + "incumbent": 0.3678036351886447 + }, + "workloads": { + "51946d28f5e3610d7f2d": { + "baseline": 6.433224392731984, + "incumbent": 6.4332088839213055, + "candidate": 3.766006750822067 + }, + "72036e17c4985117b154": { + "baseline": 0.0762255466779073, + "incumbent": 0.07617921061068773, + "candidate": 0.07826467198878527 + } + } + } + ] + }, + "export_checksums": "e4a6f4b241c77fa4e9a7ddb7ba889525e999154d8e43df02fbf130ed53828c38" + } + }, + "rotate_qk": { + "module": "rotate_qk", + "function": "call", + "applies": "applies", + "sha256": "3ea76fc8d3c60989ca5bfca2306f3cecb99d967eec5793cdc550f6b74bde568d", + "verified_layouts": 2, + "variants": { + "prefix": { + "module": "rotate_qk", + "function": "call", + "applies": "applies", + "sha256": "ed078ad58fcc0a59ae995ac7f2be97f1ac8183d810f7d110dbc340ce1520c159", + "verified_layouts": 1, + "evidence": { + "run_id": "20261015T134002Z-6a72e1be9e", + "task_digest": "48779c22e182ea79470ef35f83725479f1bb7aace8d6f13617699915d2bd4ff1", + "capture_identity": "78e584ba5990b937a77a6e987b0fad7930e43f00f1cb07920e97e298ff7ba036", + "solution_digest": "196ce9b53159469ef3582539be2a0837cb8370628066c9a23783f9155caffa31", + "decision": { + "promoted": true, + "reasons": [], + "rounds": [ + { + "mean_latency_ms": { + "baseline": 0.43428119434913, + "incumbent": 0.05630680534119407, + "candidate": 0.048808021304508044 + }, + "improvement": { + "baseline": 0.8876119391316079, + "incumbent": 0.13317722416050326 + }, + "workloads": { + "8137fa5317f3535e9950": { + "baseline": 0.43428119434913, + "incumbent": 0.05630680534119407, + "candidate": 0.048808021304508044 + } + } + }, + { + "mean_latency_ms": { + "baseline": 0.43342285868525504, + "incumbent": 0.05633865599210064, + "candidate": 0.048818613316863775 + }, + "improvement": { + "baseline": 0.8873649316398536, + "incumbent": 0.13347927001118498 + }, + "workloads": { + "8137fa5317f3535e9950": { + "baseline": 0.43342285868525504, + "incumbent": 0.05633865599210064, + "candidate": 0.048818613316863775 + } + } + } + ] + }, + "export_checksums": "423ff10730bdefece32a757988372f80fe1ae9ecdfd7a2bf02ceba8e9c3f29c7" + } + }, + "decode": { + "module": "rotate_qk", + "function": "call", + "applies": "applies", + "sha256": "d92c907f6df514554593163a544a55f30feeb7589e19a32669300a279cecbd7a", + "verified_layouts": 1, + "evidence": { + "run_id": "20261015T150031Z-03be0c2142", + "task_digest": "17c367ea111108fb08c962b6aeb659e8f77ff47220afd9bd9276a86a607ac55a", + "capture_identity": "78e584ba5990b937a77a6e987b0fad7930e43f00f1cb07920e97e298ff7ba036", + "solution_digest": "5275fe7a66824864bd0b57f84ef6c596e5f3adfbb585386bd0d04c546b8cd45f", + "decision": { + "promoted": true, + "reasons": [], + "rounds": [ + { + "mean_latency_ms": { + "baseline": 0.07087955191979806, + "incumbent": 0.07073959452162186, + "candidate": 0.012650250654667615 + }, + "improvement": { + "baseline": 0.8215246807854869, + "incumbent": 0.8211715696108349 + }, + "workloads": { + "326542b83ce4e9399c9c": { + "baseline": 0.07087955191979806, + "incumbent": 0.07073959452162186, + "candidate": 0.012650250654667615 + } + } + }, + { + "mean_latency_ms": { + "baseline": 0.07078119465708732, + "incumbent": 0.07084533325831095, + "candidate": 0.012763989326544106 + }, + "improvement": { + "baseline": 0.8196697669715575, + "incumbent": 0.8198330258394719 + }, + "workloads": { + "326542b83ce4e9399c9c": { + "baseline": 0.07078119465708732, + "incumbent": 0.07084533325831095, + "candidate": 0.012763989326544106 + } + } + } + ] + }, + "export_checksums": "0efd881c36796d4d4b2d76a6afc44b969f51d23a5e475546f2568e3402423aac" + } + } + } + } + } +} diff --git a/embodiinfer/backend/triton/tuned/orin_pi05/rotate_qk/__init__.py b/embodiinfer/backend/triton/tuned/orin_pi05/rotate_qk/__init__.py new file mode 100644 index 0000000..94ab531 --- /dev/null +++ b/embodiinfer/backend/triton/tuned/orin_pi05/rotate_qk/__init__.py @@ -0,0 +1,27 @@ +"""Composite ``rotate_qk`` for the orin_pi05 experiment profile. + +Built from two separately tuned and verified exports with disjoint layouts: +``prefix`` (run 20261015T134002Z-6a72e1be9e) and ``decode`` (run 20261015T150031Z-03be0c2142). Each call goes to +the variant whose generated ``applies`` accepts it; otherwise ``applies`` is false +and the caller falls back to its preferred implementation. +""" + +from . import decode as _decode +from . import prefix as _prefix + +_VARIANTS = (_prefix, _decode) +if _prefix._VERIFIED & _decode._VERIFIED: + raise ImportError("rotate_qk variants overlap; refuse an ambiguous composite") + + +def applies(*args, **kwargs): + """Whether any variant verified this exact argument layout.""" + return any(variant.applies(*args, **kwargs) for variant in _VARIANTS) + + +def call(*args, **kwargs): + """Run the variant that verified this layout.""" + for variant in _VARIANTS: + if variant.applies(*args, **kwargs): + return variant.call(*args, **kwargs) + raise ValueError("rotate_qk composite called with an unverified layout") diff --git a/embodiinfer/backend/triton/tuned/orin_pi05/rotate_qk/decode/__init__.py b/embodiinfer/backend/triton/tuned/orin_pi05/rotate_qk/decode/__init__.py new file mode 100644 index 0000000..15f5ff2 --- /dev/null +++ b/embodiinfer/backend/triton/tuned/orin_pi05/rotate_qk/decode/__init__.py @@ -0,0 +1,62 @@ +"""Tuned ``rotate_qk`` kernel for the ``orin_pi05_rope2dec`` profile. + +Generated by ``scripts.kernel_tuning install`` from run 20261015T150031Z-03be0c2142; do not edit, +reinstall instead. The kernel applies only to the argument layouts verified by +the evaluator. Strides of size-one dimensions and of empty tensors address no +other element, so they are compared as zero. Like the evaluator, the kernel +receives only the tensor arguments; constants are part of the verified layout. +""" + +import torch + +from .candidate import run as run + +_VERIFIED = frozenset([(('tensor', (1, 8, 50, 256), (0, 256, 2560, 1), 'bfloat16'), ('tensor', (1, 1, 50, 256), (0, 0, 2560, 1), 'bfloat16'), ('tensor', (1, 50, 256), (0, 256, 1), 'bfloat16'), ('tensor', (1, 50, 256), (0, 256, 1), 'bfloat16'))]) +_ARGUMENTS = {'tuple': [{'tuple': [{'tensor': 0}, {'tensor': 1}, {'tensor': 2}, {'tensor': 3}]}, {'dict': {}}]} +_TENSORS = 4 + + +def _strides(shape, stride): + if 0 in shape: + return (0,) * len(shape) + return tuple(0 if size == 1 else step for size, step in zip(shape, stride)) + + +def _leaves(value, out): + if isinstance(value, torch.Tensor): + dtype = str(value.dtype).removeprefix("torch.") + shape = tuple(value.shape) + out.append(("tensor", shape, _strides(shape, value.stride()), dtype)) + elif isinstance(value, (tuple, list)): + for item in value: + _leaves(item, out) + elif isinstance(value, dict): + for key in sorted(value): + _leaves(value[key], out) + else: + out.append(("constant", repr(value))) + return out + + +def applies(*args, **kwargs): + """Whether these arguments match a layout the evaluator verified.""" + return tuple(_leaves((args, kwargs), [])) in _VERIFIED + + +def _bind(recipe, value, tensors): + kind, item = next(iter(recipe.items())) + if kind == "tensor": + tensors[item] = value + elif kind in ("tuple", "list"): + for entry, element in zip(item, value): + _bind(entry, element, tensors) + elif kind == "dict": + for key, entry in item.items(): + _bind(entry, value[key], tensors) + + +def call(*args, **kwargs): + """Call the kernel as the evaluator did: with the captured tensor arguments x0..xN.""" + tensors = {} + _bind(_ARGUMENTS, (args, kwargs), tensors) + return run(*(tensors[index] for index in range(_TENSORS))) diff --git a/embodiinfer/backend/triton/tuned/orin_pi05/rotate_qk/decode/candidate.py b/embodiinfer/backend/triton/tuned/orin_pi05/rotate_qk/decode/candidate.py new file mode 100644 index 0000000..6c18366 --- /dev/null +++ b/embodiinfer/backend/triton/tuned/orin_pi05/rotate_qk/decode/candidate.py @@ -0,0 +1,117 @@ +"""Four-token paired-half RoPE for the captured pi05 BF16 layout on Orin.""" + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _rotate_qk_paired_kernel( + Q, + K, + Cos, + Sin, + QOut, + KOut, + BLOCK_T: tl.constexpr, + HALF_D: tl.constexpr, +): + head = tl.program_id(1) + tokens = tl.program_id(0) * BLOCK_T + tl.arange(0, BLOCK_T) + dims = tl.arange(0, HALF_D) + valid = tokens[:, None] < 50 + + # The tensor pointers already include the recorded storage offsets. + if head < 8: + input_base = Q + head * 256 + output_base = QOut + head * 12800 + else: + input_base = K + output_base = KOut + + input_offsets = tokens[:, None] * 2560 + dims[None, :] + a = tl.load(input_base + input_offsets, mask=valid, other=0.0).to(tl.float32) + b = tl.load(input_base + input_offsets + HALF_D, mask=valid, other=0.0).to(tl.float32) + + table_offsets = tokens[:, None] * 256 + dims[None, :] + c0 = tl.load(Cos + table_offsets, mask=valid, other=0.0).to(tl.float32) + c1 = tl.load(Cos + table_offsets + HALF_D, mask=valid, other=0.0).to(tl.float32) + s0 = tl.load(Sin + table_offsets, mask=valid, other=0.0).to(tl.float32) + s1 = tl.load(Sin + table_offsets + HALF_D, mask=valid, other=0.0).to(tl.float32) + + # Round each product to BF16 before addition, as in the Torch reference. + lo_first = (a * c0).to(tl.bfloat16, fp_downcast_rounding="rtne").to(tl.float32) + lo_second = ((-b) * s0).to(tl.bfloat16, fp_downcast_rounding="rtne").to(tl.float32) + hi_first = (b * c1).to(tl.bfloat16, fp_downcast_rounding="rtne").to(tl.float32) + hi_second = (a * s1).to(tl.bfloat16, fp_downcast_rounding="rtne").to(tl.float32) + lo = (lo_first + lo_second).to(tl.bfloat16, fp_downcast_rounding="rtne") + hi = (hi_first + hi_second).to(tl.bfloat16, fp_downcast_rounding="rtne") + + tl.store(output_base + table_offsets, lo, mask=valid) + tl.store(output_base + table_offsets + HALF_D, hi, mask=valid) + + +def run( + x0: torch.Tensor, + x1: torch.Tensor, + x2: torch.Tensor, + x3: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Return independent BF16 Q/K rotations in the frozen production layouts. + + Accept Q [1, 8, 50, 256], K [1, 1, 50, 256], and Cos/Sin [1, 50, 256] + with the captured strides on one CUDA device. Inputs are read-only, including + aliased Q/K views; storage offsets are handled by their tensor pointers. + Unsupported metadata raises ValueError. Arithmetic preserves BF16 product + rounding, and the launch uses the current stream of the input device. + """ + expected_shapes = ( + (1, 8, 50, 256), + (1, 1, 50, 256), + (1, 50, 256), + (1, 50, 256), + ) + expected_strides = ( + (128000, 256, 2560, 1), + (128000, 256, 2560, 1), + (12800, 256, 1), + (12800, 256, 1), + ) + for tensor, shape, stride in zip((x0, x1, x2, x3), expected_shapes, expected_strides): + if ( + tensor.shape != shape + or tensor.stride() != stride + or tensor.dtype != torch.bfloat16 + or tensor.device.type != "cuda" + or tensor.device != x0.device + ): + raise ValueError("rotate_qk requires the captured BF16 CUDA shapes and strides") + + # Restore the caller's device on exit and leave its stream selection intact. + with torch.cuda.device(x0.device): + qout = torch.empty_strided( + (1, 8, 50, 256), + (102400, 12800, 256, 1), + dtype=x0.dtype, + device=x0.device, + ) + kout = torch.empty_strided( + (1, 1, 50, 256), + (12800, 12800, 256, 1), + dtype=x1.dtype, + device=x1.device, + ) + _rotate_qk_paired_kernel[(13, 9)]( + x0, + x1, + x2, + x3, + qout, + kout, + BLOCK_T=4, + HALF_D=128, + num_warps=4, + num_stages=1, + enable_fp_fusion=False, + ) + return qout, kout diff --git a/embodiinfer/backend/triton/tuned/orin_pi05/rotate_qk/prefix/__init__.py b/embodiinfer/backend/triton/tuned/orin_pi05/rotate_qk/prefix/__init__.py new file mode 100644 index 0000000..b859b6c --- /dev/null +++ b/embodiinfer/backend/triton/tuned/orin_pi05/rotate_qk/prefix/__init__.py @@ -0,0 +1,62 @@ +"""Tuned ``rotate_qk`` kernel for the ``orin_pi05`` profile. + +Generated by ``scripts.kernel_tuning install`` from run 20261015T134002Z-6a72e1be9e; do not edit, +reinstall instead. The kernel applies only to the argument layouts verified by +the evaluator. Strides of size-one dimensions and of empty tensors address no +other element, so they are compared as zero. Like the evaluator, the kernel +receives only the tensor arguments; constants are part of the verified layout. +""" + +import torch + +from .candidate import run as run + +_VERIFIED = frozenset([(('tensor', (1, 8, 672, 256), (0, 256, 2560, 1), 'bfloat16'), ('tensor', (1, 1, 672, 256), (0, 0, 2560, 1), 'bfloat16'), ('tensor', (1, 672, 256), (0, 256, 1), 'bfloat16'), ('tensor', (1, 672, 256), (0, 256, 1), 'bfloat16'))]) +_ARGUMENTS = {'tuple': [{'tuple': [{'tensor': 0}, {'tensor': 1}, {'tensor': 2}, {'tensor': 3}]}, {'dict': {}}]} +_TENSORS = 4 + + +def _strides(shape, stride): + if 0 in shape: + return (0,) * len(shape) + return tuple(0 if size == 1 else step for size, step in zip(shape, stride)) + + +def _leaves(value, out): + if isinstance(value, torch.Tensor): + dtype = str(value.dtype).removeprefix("torch.") + shape = tuple(value.shape) + out.append(("tensor", shape, _strides(shape, value.stride()), dtype)) + elif isinstance(value, (tuple, list)): + for item in value: + _leaves(item, out) + elif isinstance(value, dict): + for key in sorted(value): + _leaves(value[key], out) + else: + out.append(("constant", repr(value))) + return out + + +def applies(*args, **kwargs): + """Whether these arguments match a layout the evaluator verified.""" + return tuple(_leaves((args, kwargs), [])) in _VERIFIED + + +def _bind(recipe, value, tensors): + kind, item = next(iter(recipe.items())) + if kind == "tensor": + tensors[item] = value + elif kind in ("tuple", "list"): + for entry, element in zip(item, value): + _bind(entry, element, tensors) + elif kind == "dict": + for key, entry in item.items(): + _bind(entry, value[key], tensors) + + +def call(*args, **kwargs): + """Call the kernel as the evaluator did: with the captured tensor arguments x0..xN.""" + tensors = {} + _bind(_ARGUMENTS, (args, kwargs), tensors) + return run(*(tensors[index] for index in range(_TENSORS))) diff --git a/embodiinfer/backend/triton/tuned/orin_pi05/rotate_qk/prefix/candidate.py b/embodiinfer/backend/triton/tuned/orin_pi05/rotate_qk/prefix/candidate.py new file mode 100644 index 0000000..ccc8b8c --- /dev/null +++ b/embodiinfer/backend/triton/tuned/orin_pi05/rotate_qk/prefix/candidate.py @@ -0,0 +1,149 @@ +"""Packed BF16 rotate-half RoPE for the captured Orin query/key layouts.""" + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _rotate_packed_bf16(lo, hi, cos_lo, cos_hi, sin_lo, sin_hi): + """Rotate two adjacent lanes, rounding every product and sum to BF16.""" + out_lo, out_hi = tl.inline_asm_elementwise( + asm=""" + { + .reg .b32 zero, one, neg_hi; + .reg .b32 product_a, product_b, product_c, product_d; + .reg .b32 result_lo, result_hi; + mov.b32 zero, 0x00000000; + mov.b32 one, 0x3f803f80; + xor.b32 neg_hi, $3, 0x80008000; + + fma.rn.bf16x2 product_a, $2, $4, zero; + fma.rn.bf16x2 product_b, neg_hi, $6, zero; + fma.rn.bf16x2 product_c, $3, $5, zero; + fma.rn.bf16x2 product_d, $2, $7, zero; + fma.rn.bf16x2 result_lo, product_a, one, product_b; + fma.rn.bf16x2 result_hi, product_c, one, product_d; + + mov.b32 $0, result_lo; + mov.b32 $1, result_hi; + } + """, + constraints="=r,=r,r,r,r,r,r,r", + args=[lo, hi, cos_lo, cos_hi, sin_lo, sin_hi], + dtype=(tl.uint32, tl.uint32), + is_pure=True, + pack=1, + ) + return out_lo, out_hi + + +@triton.jit +def _rotate_qk_token_packed(Q, K, Cos, Sin, Qout, Kout): + """Process one token using 32-bit words containing adjacent BF16 values.""" + token = tl.program_id(0) + heads = tl.arange(0, 8) + pairs = tl.arange(0, 64) + + # Only pointer interpretation changes; all tensor elements remain BF16. + q_words = Q.to(tl.pointer_type(tl.uint32)) + k_words = K.to(tl.pointer_type(tl.uint32)) + cos_words = Cos.to(tl.pointer_type(tl.uint32)) + sin_words = Sin.to(tl.pointer_type(tl.uint32)) + qout_words = Qout.to(tl.pointer_type(tl.uint32)) + kout_words = Kout.to(tl.pointer_type(tl.uint32)) + + table_offsets = token * 128 + pairs + cos_lo = tl.load(cos_words + table_offsets) + cos_hi = tl.load(cos_words + table_offsets + 64) + sin_lo = tl.load(sin_words + table_offsets) + sin_hi = tl.load(sin_words + table_offsets + 64) + + # Ten input head slots per token, of which eight contain Q and one K. + query_offsets = token * 1280 + heads[:, None] * 128 + pairs[None, :] + query_lo = tl.load(q_words + query_offsets) + query_hi = tl.load(q_words + query_offsets + 64) + query_out_lo, query_out_hi = _rotate_packed_bf16( + query_lo, + query_hi, + tl.broadcast_to(cos_lo[None, :], (8, 64)), + tl.broadcast_to(cos_hi[None, :], (8, 64)), + tl.broadcast_to(sin_lo[None, :], (8, 64)), + tl.broadcast_to(sin_hi[None, :], (8, 64)), + ) + query_out_offsets = token * 1024 + heads[:, None] * 128 + pairs[None, :] + tl.store(qout_words + query_out_offsets, query_out_lo) + tl.store(qout_words + query_out_offsets + 64, query_out_hi) + + # K's pointer already includes its BF16 storage offset of 2048 elements. + key_offsets = token * 1280 + pairs + key_lo = tl.load(k_words + key_offsets) + key_hi = tl.load(k_words + key_offsets + 64) + key_out_lo, key_out_hi = _rotate_packed_bf16( + key_lo, key_hi, cos_lo, cos_hi, sin_lo, sin_hi + ) + key_out_offsets = token * 128 + pairs + tl.store(kout_words + key_out_offsets, key_out_lo) + tl.store(kout_words + key_out_offsets + 64, key_out_hi) + + +def run( + x0: torch.Tensor, + x1: torch.Tensor, + x2: torch.Tensor, + x3: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Return BF16 rotated query/key with the captured output strides. + + Accept the recorded query, key, cosine, and sine layouts on one CUDA + device, with four-byte-aligned logical data pointers. Each output owns + fresh storage; inputs, including shared query/key storage, are read-only. + Packed BF16 arithmetic rounds every product and sum to nearest even, + subject to the frozen FP64 error bounds. Unsupported metadata raises + before allocation or launch. Use the input device's current stream and + restore the caller's device without changing numerical settings. + """ + expected_inputs = ( + ("x0", x0, (1, 8, 672, 256), (1720320, 256, 2560, 1)), + ("x1", x1, (1, 1, 672, 256), (1720320, 256, 2560, 1)), + ("x2", x2, (1, 672, 256), (172032, 256, 1)), + ("x3", x3, (1, 672, 256), (172032, 256, 1)), + ) + for name, tensor, shape, strides in expected_inputs: + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"{name} must be a torch.Tensor") + if tensor.layout != torch.strided or tensor.dtype != torch.bfloat16: + raise ValueError(f"{name} must be a strided BF16 tensor") + if tensor.device.type != "cuda" or tensor.device != x0.device: + raise ValueError("rotate_qk inputs must share one CUDA device") + if tuple(tensor.shape) != shape or tuple(tensor.stride()) != strides: + raise ValueError(f"{name} has an unsupported shape or stride") + if tensor.data_ptr() % 4 != 0: + raise ValueError(f"{name} must have a four-byte-aligned data pointer") + + # The device guard restores the original device and leaves streams intact. + with torch.cuda.device(x0.device): + query_out = torch.empty_strided( + (1, 8, 672, 256), + (1376256, 256, 2048, 1), + dtype=torch.bfloat16, + device=x0.device, + ) + key_out = torch.empty_strided( + (1, 1, 672, 256), + (172032, 172032, 256, 1), + dtype=torch.bfloat16, + device=x1.device, + ) + _rotate_qk_token_packed[(672,)]( + x0, + x1, + x2, + x3, + query_out, + key_out, + num_warps=4, + num_stages=1, + enable_fp_fusion=False, + ) + return query_out, key_out diff --git a/embodiinfer/backend/triton/tuned/orin_pi05/split_kv_attention/__init__.py b/embodiinfer/backend/triton/tuned/orin_pi05/split_kv_attention/__init__.py new file mode 100644 index 0000000..8b5f3c7 --- /dev/null +++ b/embodiinfer/backend/triton/tuned/orin_pi05/split_kv_attention/__init__.py @@ -0,0 +1,62 @@ +"""Tuned ``split_kv_attention`` kernel for the ``orin_pi05`` profile. + +Generated by ``scripts.kernel_tuning install`` from run 20261015T055303Z-8eeab7204c; do not edit, +reinstall instead. The kernel applies only to the argument layouts verified by +the evaluator. Strides of size-one dimensions and of empty tensors address no +other element, so they are compared as zero. Like the evaluator, the kernel +receives only the tensor arguments; constants are part of the verified layout. +""" + +import torch + +from .candidate import run as run + +_VERIFIED = frozenset([(('tensor', (1, 8, 50, 256), (0, 12800, 256, 1), 'bfloat16'), ('tensor', (1, 1, 672, 256), (0, 0, 256, 1), 'bfloat16'), ('tensor', (1, 1, 672, 256), (0, 0, 256, 1), 'bfloat16'), ('tensor', (1, 1, 50, 256), (0, 0, 256, 1), 'bfloat16'), ('tensor', (1, 1, 50, 256), (0, 0, 2560, 1), 'bfloat16'), ('tensor', (1, 1, 1, 722), (0, 0, 0, 1), 'float32'), ('constant', '0.0625')), (('tensor', (1, 8, 672, 256), (0, 256, 2048, 1), 'bfloat16'), ('tensor', (1, 1, 0, 256), (0, 0, 0, 0), 'bfloat16'), ('tensor', (1, 1, 0, 256), (0, 0, 0, 0), 'bfloat16'), ('tensor', (1, 1, 672, 256), (0, 0, 256, 1), 'bfloat16'), ('tensor', (1, 1, 672, 256), (0, 0, 2560, 1), 'bfloat16'), ('tensor', (1, 1, 1, 672), (0, 0, 0, 1), 'float32'), ('constant', '0.0625'))]) +_ARGUMENTS = {'tuple': [{'tuple': [{'tensor': 0}, {'tensor': 1}, {'tensor': 2}, {'tensor': 3}, {'tensor': 4}, {'tensor': 5}, {'constant': 0.0625}]}, {'dict': {}}]} +_TENSORS = 6 + + +def _strides(shape, stride): + if 0 in shape: + return (0,) * len(shape) + return tuple(0 if size == 1 else step for size, step in zip(shape, stride)) + + +def _leaves(value, out): + if isinstance(value, torch.Tensor): + dtype = str(value.dtype).removeprefix("torch.") + shape = tuple(value.shape) + out.append(("tensor", shape, _strides(shape, value.stride()), dtype)) + elif isinstance(value, (tuple, list)): + for item in value: + _leaves(item, out) + elif isinstance(value, dict): + for key in sorted(value): + _leaves(value[key], out) + else: + out.append(("constant", repr(value))) + return out + + +def applies(*args, **kwargs): + """Whether these arguments match a layout the evaluator verified.""" + return tuple(_leaves((args, kwargs), [])) in _VERIFIED + + +def _bind(recipe, value, tensors): + kind, item = next(iter(recipe.items())) + if kind == "tensor": + tensors[item] = value + elif kind in ("tuple", "list"): + for entry, element in zip(item, value): + _bind(entry, element, tensors) + elif kind == "dict": + for key, entry in item.items(): + _bind(entry, value[key], tensors) + + +def call(*args, **kwargs): + """Call the kernel as the evaluator did: with the captured tensor arguments x0..xN.""" + tensors = {} + _bind(_ARGUMENTS, (args, kwargs), tensors) + return run(*(tensors[index] for index in range(_TENSORS))) diff --git a/embodiinfer/backend/triton/tuned/orin_pi05/split_kv_attention/candidate.py b/embodiinfer/backend/triton/tuned/orin_pi05/split_kv_attention/candidate.py new file mode 100644 index 0000000..a04bd17 --- /dev/null +++ b/embodiinfer/backend/triton/tuned/orin_pi05/split_kv_attention/candidate.py @@ -0,0 +1,270 @@ +"""Specialized noncausal split-KV attention for the two captured Orin layouts.""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _partial_attention( + Q, + PK, + PV, + SK, + SV, + Mask, + Part, + Maxima, + Denominators, + QH: tl.constexpr, + QS: tl.constexpr, + PKS: tl.constexpr, + PVS: tl.constexpr, + SKS: tl.constexpr, + SVS: tl.constexpr, + MS: tl.constexpr, + P: tl.constexpr, + S: tl.constexpr, + D: tl.constexpr, + SCALE: tl.constexpr, + SPLITS: tl.constexpr, + PART_LENGTH: tl.constexpr, + BM: tl.constexpr, + BN: tl.constexpr, + PACK_QUERY_HEADS: tl.constexpr, +): + partition = tl.program_id(2) + dims = tl.arange(0, D) + lanes = tl.arange(0, BN) + if PACK_QUERY_HEADS: + # All eight query heads share K/V. Pack their independent query rows + # into full tiles, mapping back to each row's head for Q and scratch. + packed_rows = tl.program_id(0) * BM + tl.arange(0, BM) + head = packed_rows // S + rows = packed_rows % S + active_rows = packed_rows < 8 * S + q = tl.load( + Q + head[:, None] * QH + rows[:, None] * QS + dims[None, :], + active_rows[:, None], + other=0, + ) + else: + head = tl.program_id(0) + rows = tl.program_id(1) * BM + tl.arange(0, BM) + active_rows = rows < S + q = tl.load( + Q + head * QH + rows[:, None] * QS + dims[None, :], + active_rows[:, None], + other=0, + ) + + maximum = tl.full((BM,), -float("inf"), tl.float32) + denominator = tl.zeros((BM,), tl.float32) + numerator = tl.zeros((BM, D), tl.float32) + first = partition * PART_LENGTH + end = tl.minimum(first + PART_LENGTH, P + S) + + for tile_start in range(first, end, BN): + columns = tile_start + lanes + # P and partition boundaries are multiples of BN. A tile always reads + # just one K/V source, including the final, partially occupied tile. + if P == 0: + k_ptrs = SK + columns[None, :] * SKS + dims[:, None] + v_ptrs = SV + columns[:, None] * SVS + dims[None, :] + else: + if tile_start < P: + k_ptrs = PK + columns[None, :] * PKS + dims[:, None] + v_ptrs = PV + columns[:, None] * PVS + dims[None, :] + else: + suffix_columns = columns - P + k_ptrs = SK + suffix_columns[None, :] * SKS + dims[:, None] + v_ptrs = SV + suffix_columns[:, None] * SVS + dims[None, :] + + in_bounds = columns < P + S + k = tl.load(k_ptrs, in_bounds[None, :], other=0) + padding = tl.load(Mask + columns * MS, in_bounds, other=1.0) + keep = in_bounds & (padding == 0.0) + scores = tl.dot(q, k, out_dtype=tl.float32) + scores = scores * (SCALE * 1.4426950408889634) + scores = tl.where( + active_rows[:, None] & keep[None, :], scores, -float("inf") + ) + + next_maximum = tl.maximum(maximum, tl.max(scores, axis=1)) + safe_maximum = tl.where(next_maximum == -float("inf"), 0.0, next_maximum) + correction = tl.exp2(maximum - safe_maximum) + probability = tl.exp2(scores - safe_maximum[:, None]) + denominator = denominator * correction + tl.sum(probability, axis=1) + + v = tl.load(v_ptrs, in_bounds[:, None], other=0) + probability_hi = probability.to(tl.bfloat16) + probability_lo = (probability - probability_hi.to(tl.float32)).to(tl.bfloat16) + numerator = numerator * correction[:, None] + numerator = tl.dot(probability_hi, v, acc=numerator, out_dtype=tl.float32) + numerator = tl.dot(probability_lo, v, acc=numerator, out_dtype=tl.float32) + maximum = next_maximum + + state_indices = (head * SPLITS + partition) * S + rows + if SPLITS == 1: + safe_denominator = tl.where(denominator > 0.0, denominator, 1.0) + result = numerator / safe_denominator[:, None] + tl.store( + Part + state_indices[:, None] * D + dims[None, :], + result, + active_rows[:, None], + ) + else: + tl.store( + Part + state_indices[:, None] * D + dims[None, :], + numerator, + active_rows[:, None], + ) + tl.store(Maxima + state_indices, maximum, active_rows) + tl.store(Denominators + state_indices, denominator, active_rows) + + +@triton.jit +def _merge_attention( + Part, + Maxima, + Denominators, + Out, + S: tl.constexpr, + D: tl.constexpr, + SPLITS: tl.constexpr, + ROWS: tl.constexpr, +): + rows = tl.program_id(0) * ROWS + tl.arange(0, ROWS) + head = tl.program_id(1) + partitions = tl.arange(0, SPLITS) + dims = tl.arange(0, D) + active_rows = rows < S + state_indices = (head * SPLITS + partitions[None, :]) * S + rows[:, None] + + maxima = tl.load(Maxima + state_indices, active_rows[:, None], other=-float("inf")) + denominators = tl.load(Denominators + state_indices, active_rows[:, None], other=0.0) + row_maximum = tl.max(maxima, axis=1) + safe_maximum = tl.where(row_maximum == -float("inf"), 0.0, row_maximum) + weights = tl.exp2(maxima - safe_maximum[:, None]) + denominator = tl.sum(weights * denominators, axis=1) + + partial = tl.load( + Part + state_indices[:, :, None] * D + dims[None, None, :], + active_rows[:, None, None], + other=0.0, + ) + numerator = tl.sum(partial * weights[:, :, None], axis=1) + safe_denominator = tl.where(denominator > 0.0, denominator, 1.0) + result = numerator / safe_denominator[:, None] + tl.store( + Out + (head * S + rows[:, None]) * D + dims[None, :], + result, + active_rows[:, None], + ) + + +def run( + x0: torch.Tensor, + x1: torch.Tensor, + x2: torch.Tensor, + x3: torch.Tensor, + x4: torch.Tensor, + x5: torch.Tensor, +) -> torch.Tensor: + """Return BF16 attention for the captured single-batch, eight-query-head calls. + + Arguments are Q, prefix K/V, suffix K/V, and a shared FP32 padding mask. + A zero mask element keeps its key; any other value removes it. The scale + is 0.0625, all KV tensors have one head, and fully masked rows return zero. + Inputs may have the recorded strided layouts and storage aliases; they are + read without modification. The result is fresh, contiguous BHSD storage. + The 50-query call packs rows across query heads that share the same K/V. + """ + if tuple(x0.shape) not in ((1, 8, 672, 256), (1, 8, 50, 256)): + raise ValueError("Expected a recorded query shape (1, 8, 672 or 50, 256)") + size = x0.shape[2] + prefix = 0 if size == 672 else 672 + if x1.shape != (1, 1, prefix, 256) or x2.shape != x1.shape: + raise ValueError("Prefix K/V shapes do not match the captured attention call") + if x3.shape != (1, 1, size, 256) or x4.shape != x3.shape: + raise ValueError("Suffix K/V shapes do not match the captured attention call") + if x0.device.type != "cuda": + raise ValueError("This attention implementation requires CUDA tensors") + for tensor in (x0, x1, x2, x3, x4): + if tensor.dtype != torch.bfloat16 or tensor.device != x0.device: + raise ValueError("Q/K/V must be BF16 tensors on the same CUDA device") + if tensor.stride(-1) != 1: + raise ValueError("Q/K/V head dimensions must have unit stride") + if x5.shape != (1, 1, 1, prefix + size): + raise ValueError("The padding mask must have shape (1, 1, 1, prefix + suffix)") + if x5.dtype != torch.float32 or x5.device != x0.device: + raise ValueError("The padding mask must be FP32 on the Q/K/V device") + + # A 32-row packed tile reuses each K/V tile across more short-call queries. + # active_rows masks the final tile's 16 padded query rows. + # Two short-call partitions halve scratch traffic and repeated Q loads. + # Their 384-key boundary and the 672-key prefix both align to BN=32. + splits, block_m, part_length = (1, 32, 672) if size == 672 else (2, 32, 384) + pack_query_heads = size == 50 + output = torch.empty((1, 8, size, 256), device=x0.device, dtype=torch.bfloat16) + if splits == 1: + # The SPLITS == 1 specialization never accesses the two state pointers. + partial = output + maxima = output + denominators = output + else: + partial = torch.empty((8, splits, size, 256), device=x0.device, dtype=torch.float32) + maxima = torch.empty((8, splits, size), device=x0.device, dtype=torch.float32) + denominators = torch.empty((8, splits, size), device=x0.device, dtype=torch.float32) + + partial_grid = ( + (triton.cdiv(8 * size, block_m), 1, splits) + if pack_query_heads + else (8, triton.cdiv(size, block_m), splits) + ) + _partial_attention[partial_grid]( + x0, + x1, + x2, + x3, + x4, + x5, + partial, + maxima, + denominators, + QH=x0.stride(1), + QS=x0.stride(2), + PKS=x1.stride(2), + PVS=x2.stride(2), + SKS=x3.stride(2), + SVS=x4.stride(2), + MS=x5.stride(3), + P=prefix, + S=size, + D=256, + SCALE=0.0625, + SPLITS=splits, + PART_LENGTH=part_length, + BM=block_m, + BN=32, + PACK_QUERY_HEADS=pack_query_heads, + num_warps=4, + # Overlap K/V loads with attention compute across two stages for both calls. + num_stages=2, + ) + if splits > 1: + _merge_attention[(triton.cdiv(size, 4), 8)]( + partial, + maxima, + denominators, + output, + S=size, + D=256, + SPLITS=splits, + ROWS=4, + num_warps=4, + num_stages=1, + ) + return output diff --git a/embodiinfer/layers/__init__.py b/embodiinfer/layers/__init__.py index 0b25087..6a02631 100644 --- a/embodiinfer/layers/__init__.py +++ b/embodiinfer/layers/__init__.py @@ -22,6 +22,7 @@ resolve_linear_backend, ) from .normalization import NormalizationBackend, NormQuantBackend, norm_quant_backends, normalization_backends +from .ops import KernelSelection, available_ops, get_op, register_lazy_op, register_op from .quantization import ( ActivationQuantizer, EncodedActivation, @@ -67,4 +68,9 @@ "parse_quantization_config", "register_linear", "resolve_linear_backend", + "KernelSelection", + "available_ops", + "get_op", + "register_lazy_op", + "register_op", ] diff --git a/embodiinfer/layers/ops.py b/embodiinfer/layers/ops.py new file mode 100644 index 0000000..2b2d382 --- /dev/null +++ b/embodiinfer/layers/ops.py @@ -0,0 +1,180 @@ +"""Tensor operator registry and kernel-profile selection. + +An operator is named by its mathematics (``gated_residual``, ``ada_rms_norm``), +never by a model. Each operator has a ``torch`` reference implementation and may +have further implementations such as ``triton``. Callers ask for an operator by +name with their preferred implementation instead of importing a backend, so a +verified tuned kernel can replace it without changing model code. + +``kernel_profile`` selects tuned kernels (see :mod:`embodiinfer.backend.triton.tuned`): + +* ``none`` (default): the caller's preferred implementation, exactly as before. +* ``torch``: every registered operator uses its Torch reference. +* ``auto``: the profile matching this device and Torch/Triton versions, else ``none``. +* a profile name: that profile; a mismatched environment raises. + +A tuned kernel replaces only the operators its profile contains, and only for +calls its ``applies`` predicate accepts; other calls use the preferred +implementation. +""" + +from __future__ import annotations + +import importlib +import logging +from collections.abc import Callable +from typing import Any + +logger = logging.getLogger(__name__) + +OpFunction = Callable[..., Any] +TorchWrapper = Callable[[OpFunction], OpFunction] + +_REGISTRY: dict[str, dict[str, OpFunction]] = {} +_LAZY_REGISTRY: dict[str, dict[str, tuple[str, str]]] = {} +_PROFILE_MODES = ("none", "torch", "auto") + + +def register_op(op: str, impl: str, fn: OpFunction) -> None: + """Register ``fn`` as implementation ``impl`` of operator ``op``.""" + if impl in _REGISTRY.get(op, {}) or impl in _LAZY_REGISTRY.get(op, {}): + raise ValueError(f"operator {op!r} implementation {impl!r} is already registered") + _REGISTRY.setdefault(op, {})[impl] = fn + + +def register_lazy_op(op: str, impl: str, module: str, function: str) -> None: + """Register an implementation without importing its backend until first use.""" + if impl in _REGISTRY.get(op, {}) or impl in _LAZY_REGISTRY.get(op, {}): + raise ValueError(f"operator {op!r} implementation {impl!r} is already registered") + _LAZY_REGISTRY.setdefault(op, {})[impl] = (module, function) + + +def available_ops() -> dict[str, list[str]]: + """Return every registered operator with its implementation names.""" + names = set(_REGISTRY) | set(_LAZY_REGISTRY) + return {op: sorted(set(_REGISTRY.get(op, {})) | set(_LAZY_REGISTRY.get(op, {}))) for op in sorted(names)} + + +def get_op(op: str, impl: str = "torch") -> OpFunction: + """Return implementation ``impl`` of ``op`` without any profile override.""" + if impl in _REGISTRY.get(op, {}): + return _REGISTRY[op][impl] + if impl not in _LAZY_REGISTRY.get(op, {}): + known = available_ops() + if op not in known: + raise KeyError(f"unknown operator {op!r}; registered: {sorted(known)}") + raise KeyError(f"operator {op!r} has no {impl!r} implementation; registered: {known[op]}") + module, function = _LAZY_REGISTRY[op].pop(impl) + fn = getattr(importlib.import_module(module), function) + _REGISTRY.setdefault(op, {})[impl] = fn + return fn + + +class KernelSelection: + """Resolve operators for one policy under one ``kernel_profile`` setting. + + Resolution is cached per ``(op, prefer, torch_wrapper)``, so call sites may + resolve on every call. The profile is chosen at construction from the current + CUDA device; rebuild the selection after moving the policy to another device. + """ + + def __init__(self, kernel_profile: Any = "none") -> None: + from ..backend.triton.tuned import KernelProfile + + self.mode = kernel_profile.name if isinstance(kernel_profile, KernelProfile) else kernel_profile + self.profile: KernelProfile | None = ( + kernel_profile if isinstance(kernel_profile, KernelProfile) else _select_profile(kernel_profile) + ) + self._cache: dict[tuple[str, str, TorchWrapper | None], OpFunction] = {} + + def prepare(self, torch_wrappers: tuple[TorchWrapper | None, ...] = (None,)) -> None: + """Resolve every registered operator and implementation ahead of time. + + Call this outside ``torch.compile`` regions: afterwards ``resolve`` only + reads its cache, which graph tracing can follow without importing modules. + """ + for op, impls in available_ops().items(): + for impl in impls: + for wrapper in torch_wrappers: + self.resolve(op, impl, wrapper) + + def resolve( + self, op: str, prefer: str = "torch", torch_wrapper: TorchWrapper | None = None + ) -> OpFunction: + """Return the callable for ``op``. + + ``prefer`` is the implementation the caller would use without a profile; + ``torch_wrapper`` (for example an Inductor compile helper) is applied to + the Torch reference only. + """ + key = (op, prefer, torch_wrapper) + cached = self._cache.get(key) + if cached is not None: + return cached + impl = "torch" if self.mode == "torch" else prefer + base = get_op(op, impl) + if impl == "torch" and torch_wrapper is not None: + base = torch_wrapper(base) + resolved = base + if self.profile is not None and op in self.profile.ops: + tuned = self.profile.kernel(op) + resolved = tuned.fn if tuned.applies is None else _dispatch(tuned.fn, tuned.applies, base) + self._cache[key] = resolved + return resolved + + +def _dispatch(tuned: OpFunction, applies: Callable[..., bool], fallback: OpFunction) -> OpFunction: + def call(*args: Any, **kwargs: Any) -> Any: + return tuned(*args, **kwargs) if applies(*args, **kwargs) else fallback(*args, **kwargs) + + return call + + +def _select_profile(name: str) -> Any: + from ..backend.triton.tuned import available_profiles, current_environment + + if not isinstance(name, str): + raise TypeError(f"kernel_profile must be a string; got {type(name).__name__}") + if name in ("none", "torch"): + return None + profiles = available_profiles() + environment = current_environment() + if name == "auto": + for profile in profiles.values(): + if profile.matches(environment): + logger.info("kernel_profile=auto selected %r", profile.name) + return profile + logger.info("kernel_profile=auto found no profile for %s", environment) + return None + if name not in profiles: + choices = sorted(_PROFILE_MODES) + sorted(profiles) + raise ValueError(f"unknown kernel_profile {name!r}; expected one of {choices}") + profile = profiles[name] + if not profile.matches(environment): + raise RuntimeError( + f"kernel_profile {name!r} was verified for {dict(profile.match)}, not this environment {environment}" + ) + return profile + + +register_lazy_op("linear", "torch", "embodiinfer.backend.torch.ops", "linear") +register_lazy_op("gelu_tanh", "torch", "embodiinfer.backend.torch.ops", "gelu_tanh") + + +for _op, _module in ( + ("gated_residual", "norm"), + ("ada_rms_norm", "norm"), + ("gated_gelu", "activation"), + ("rotate_qk", "rotary"), + ("split_kv_attention", "split_kv_attention"), +): + register_lazy_op(_op, "torch", "embodiinfer.backend.torch.ops", _op) + register_lazy_op(_op, "triton", f"embodiinfer.backend.triton.{_module}", _op) + +__all__ = [ + "KernelSelection", + "available_ops", + "get_op", + "register_lazy_op", + "register_op", +] diff --git a/embodiinfer/policies/pi05/modeling_pi05.py b/embodiinfer/policies/pi05/modeling_pi05.py index 6239a39..9d52421 100644 --- a/embodiinfer/policies/pi05/modeling_pi05.py +++ b/embodiinfer/policies/pi05/modeling_pi05.py @@ -33,6 +33,8 @@ from __future__ import annotations +import functools +import sys from collections.abc import Mapping from dataclasses import dataclass from pathlib import Path @@ -41,9 +43,10 @@ import torch import torch.nn.functional as F +from torch import nn from ...engine.parallel import TensorParallelContext -from ...layers import get_attention_backend, get_split_kv_attention_backend +from ...layers import KernelSelection, get_attention_backend, get_split_kv_attention_backend from ...layers.linear import QuantizationConfig, parse_quantization_config from ...models.linear import QuantizedLinear, pack_quantized_linears, quantize_linear from ...types import Observation @@ -91,30 +94,68 @@ def _inference_helper(fn, use_inductor: bool): return compiled -# ---- attention math (reimplemented to match HF Gemma exactly) --------------- -def _rotate_half(x: torch.Tensor) -> torch.Tensor: - x1 = x[..., : x.shape[-1] // 2] - x2 = x[..., x.shape[-1] // 2 :] - return torch.cat((-x2, x1), dim=-1) +_DEFAULT_KERNELS = KernelSelection("none") -def _apply_rope_impl( - q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor -) -> tuple[torch.Tensor, torch.Tensor]: - cos = cos.unsqueeze(1) # [B, 1, seq, hd] broadcast over heads - sin = sin.unsqueeze(1) - return (q * cos) + (_rotate_half(q) * sin), (k * cos) + (_rotate_half(k) * sin) +def _op(kernels: KernelSelection | None, op: str, use_triton: bool): + """Resolve a registry operator with the policy's preferred implementation. + + Torch references run uncompiled. With ``compile_backend='inductor'`` the + former per-helper ``_inference_helper`` wrappers were first created while + Dynamo traced the prefix encoder, which stores the raw function, so these + helpers never ran compiled outside that graph. Keeping them uncompiled + preserves those numerics; the prefix encoder graph still inlines them. + """ + selection = _DEFAULT_KERNELS if kernels is None else kernels + return selection.resolve(op, "triton" if use_triton else "torch") + +def _is_gelu_tanh(activation: Any) -> bool: + """Recognize plain tanh GELU modules whose calls can use the tensor operator.""" + if type(activation) is nn.GELU: + if activation.approximate != "tanh": + return False + else: + # A loaded SigLIP activation already has this optional dependency in + # sys.modules. Compare the class itself, never its name or a subclass. + activations = sys.modules.get("transformers.activations") + if activations is None or type(activation) is not getattr(activations, "PytorchGELUTanh", None): + return False + # GELUTanh(use_gelu_tanh_python=True) ("gelu_python_tanh") evaluates a + # Python formula that is not bit-identical to F.gelu; only the native + # partial may use the operator. + act = activation.__dict__.get("act") + if not ( + isinstance(act, functools.partial) + and act.func is F.gelu + and not act.args + and act.keywords == {"approximate": "tanh"} + ): + return False + # Module hooks and instance overrides may change values or gradients. + return ( + "forward" not in activation.__dict__ + and "_call_impl" not in activation.__dict__ + and getattr(activation, "_compiled_call_impl", None) is None + and not activation._forward_hooks + and not activation._forward_pre_hooks + and not activation._backward_hooks + and not activation._backward_pre_hooks + ) + +# ---- attention math (reimplemented to match HF Gemma exactly) --------------- def _apply_rope( q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, use_inductor: bool, + use_triton: bool = False, + kernels: KernelSelection | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: - fn = _inference_helper(_apply_rope_impl, use_inductor) - return fn(q, k, cos, sin) + """Rotate-half Q/K RoPE in BHSD layout through the ``rotate_qk`` operator.""" + return _op(kernels, "rotate_qk", use_triton)(q, k, cos, sin) def _gated_residual_impl(x: torch.Tensor, y: torch.Tensor, gate: torch.Tensor | None) -> torch.Tensor: @@ -127,11 +168,10 @@ def _gated_residual( gate: torch.Tensor | None, use_inductor: bool, use_triton: bool = False, + kernels: KernelSelection | None = None, ) -> torch.Tensor: - if use_triton and gate is not None: - from ...backend.triton.norm import gated_residual - - return gated_residual(x, y, gate) + if gate is not None: + return _op(kernels, "gated_residual", use_triton)(x, y, gate) fn = _inference_helper(_gated_residual_impl, use_inductor) return fn(x, y, gate) @@ -160,18 +200,14 @@ def _rmsnorm( use_inductor: bool, modulation: torch.Tensor | None = None, use_triton: bool = False, + kernels: KernelSelection | None = None, ): """PiGemmaRMSNorm forward (fp32 variance); returns (output, adarms gate).""" - fn = _inference_helper(_rmsnorm_math, use_inductor) if modulation is None: if cond is None or norm.dense is None: - return fn(x, norm.weight, norm.eps, None) + return _inference_helper(_rmsnorm_math, use_inductor)(x, norm.weight, norm.eps, None) modulation = _adarms_projection(norm, cond) - if use_triton: - from ...backend.triton.norm import ada_rms_norm - - return ada_rms_norm(x, modulation, norm.eps) - return fn(x, None, norm.eps, modulation) + return _op(kernels, "ada_rms_norm", use_triton)(x, modulation, norm.eps) def _adarms_projection(norm: Any, cond: torch.Tensor) -> torch.Tensor: @@ -217,7 +253,9 @@ def _packed_quantized_projection( return packed -def _fused_linear_pair(x: torch.Tensor, first: Any, second: Any) -> tuple[torch.Tensor, torch.Tensor]: +def _fused_linear_pair( + x: torch.Tensor, first: Any, second: Any, *, kernels: KernelSelection | None = None +) -> tuple[torch.Tensor, torch.Tensor]: """Evaluate two bias-free projections as one GEMM and split the result.""" if torch.is_grad_enabled() or first.training or second.training: return first(x), second(x) @@ -235,7 +273,7 @@ def _fused_linear_pair(x: torch.Tensor, first: Any, second: Any) -> tuple[torch. if weight is None or weight.device != first.weight.device or weight.dtype != first.weight.dtype: weight = torch.cat((first.weight, second.weight), dim=0).contiguous() first._vvla_fused_pair_weight = weight - fused = F.linear(x, weight) + fused = _op(kernels, "linear", False)(x, weight) return fused.split((first.out_features, second.out_features), dim=-1) @@ -312,46 +350,37 @@ def _mlp( use_triton: bool = False, *, fuse_projections: bool = True, + kernels: KernelSelection | None = None, approximate: str = "tanh", ) -> torch.Tensor: if fuse_projections: - gate, up = _fused_linear_pair(x, mlp.gate_proj, mlp.up_proj) + gate, up = _fused_linear_pair(x, mlp.gate_proj, mlp.up_proj, kernels=kernels) else: # The eager reference keeps the two native GEMM shapes. Besides matching # LeRobot's eager path, this preserves its rounding on selectively cast # (e.g. bf16) checkpoints. gate, up = mlp.gate_proj(x), mlp.up_proj(x) - if use_triton and approximate == "tanh": - from ...backend.triton.activation import gated_gelu - - activated = gated_gelu(gate, up) + if approximate == "none": + activated = F.gelu(gate) * up else: - if approximate == "none": - activated = F.gelu(gate) * up - else: - fn = _inference_helper(_gated_gelu_impl, use_inductor) - activated = fn(gate, up) + activated = _op(kernels, "gated_gelu", use_triton)(gate, up) + if type(mlp.down_proj) is nn.Linear: + return _op(kernels, "linear", False)(activated, mlp.down_proj.weight, mlp.down_proj.bias) return mlp.down_proj(activated) -def _gated_gelu_impl(gate: torch.Tensor, up: torch.Tensor) -> torch.Tensor: - return F.gelu(gate, approximate="tanh") * up - - def _prepare_compiled_inference_helpers() -> None: """Stabilize helper lookups before tracing the compiled prefix tower. The prefix tower is itself wrapped in ``torch.compile``. Lazily adding helper wrappers while Dynamo traces that outer function invalidates its global-dictionary guard and needlessly consumes one cache entry per helper. - Creating the four lazy wrappers up front preserves the same compiled - operators and numerics while keeping that guard stable across layouts. + Creating the remaining lazy wrappers up front keeps that guard stable across + layouts. RoPE and tanh-GELU resolve through registry operators instead. """ for helper in ( - _apply_rope_impl, _gated_residual_impl, _rmsnorm_math, - _gated_gelu_impl, ): _inference_helper(helper, True) @@ -468,6 +497,7 @@ def __init__( request_cuda_graph: bool = False, denoise_attention: str = "sdpa", prefix_attention: str = "sdpa", + kernel_profile: str = "none", checkpoint_config: Mapping[str, Any] | None = None, optimizations: Pi05OptimizationConfig | None = None, tensor_parallel_size: int = 1, @@ -514,6 +544,10 @@ def __init__( self.request_cuda_graph = request_cuda_graph self.denoise_attention = denoise_attention self.prefix_attention = prefix_attention + if not isinstance(kernel_profile, str): + raise TypeError("PI0.5 kernel_profile must be a profile name, 'auto', 'torch' or 'none'") + self.kernel_profile = kernel_profile + self._kernel_selection: KernelSelection | None = None self.attention = attention self._attn = get_attention_backend(attention) self._native_attention = None @@ -620,14 +654,39 @@ def _embed_image(self, image: torch.Tensor) -> torch.Tensor: x = layer.layer_norm1(h) a = layer.self_attn b, s, _ = x.shape - q, k, v = [ - proj(x).view(b, s, a.num_heads, a.head_dim).transpose(1, 2) - for proj in (a.q_proj, a.k_proj, a.v_proj) - ] + qkv = [] + for proj in (a.q_proj, a.k_proj, a.v_proj): + if type(proj) is nn.Linear: + projected = _op(self._kernels, "linear", False)(x, proj.weight, proj.bias) + else: + projected = proj(x) + qkv.append(projected.view(b, s, a.num_heads, a.head_dim).transpose(1, 2)) + q, k, v = qkv y = self._vision_attn.attend(q, k, v, scaling=a.scale) - h = h + a.out_proj(y.transpose(1, 2).reshape(b, s, -1).contiguous()) + if type(a.out_proj) is nn.Linear: + h = h + _op(self._kernels, "linear", False)( + y.transpose(1, 2).reshape(b, s, -1).contiguous(), + a.out_proj.weight, + a.out_proj.bias, + ) + else: + h = h + a.out_proj(y.transpose(1, 2).reshape(b, s, -1).contiguous()) x = layer.layer_norm2(h) - h = h + layer.mlp.fc2(layer.mlp.activation_fn(layer.mlp.fc1(x))) + if type(layer.mlp.fc1) is nn.Linear: + activated = _op(self._kernels, "linear", False)(x, layer.mlp.fc1.weight, layer.mlp.fc1.bias) + else: + activated = layer.mlp.fc1(x) + activation_fn = layer.mlp.activation_fn + if _is_gelu_tanh(activation_fn): + activated = _op(self._kernels, "gelu_tanh", False)(activated) + else: + activated = activation_fn(activated) + if type(layer.mlp.fc2) is nn.Linear: + h = h + _op(self._kernels, "linear", False)( + activated, layer.mlp.fc2.weight, layer.mlp.fc2.bias + ) + else: + h = h + layer.mlp.fc2(activated) h = vision.post_layernorm(h) h = pg.multi_modal_projector.linear(h) # Preserve both operations in LeRobot's embedding path: eliminating @@ -660,6 +719,16 @@ def _native_enabled(self) -> bool: and self._m.action_in_proj.weight.device.type == "cuda" ) + @property + def _kernels(self) -> KernelSelection: + """Operator selection for ``kernel_profile``, chosen on first use on the current device.""" + selection = self.__dict__.get("_kernel_selection") + if selection is None: + selection = KernelSelection(self.__dict__.get("kernel_profile", "none")) + selection.prepare() + self._kernel_selection = selection + return selection + def _prepare_native_attention(self) -> None: """Resolve optional kernels before native execution or graph capture.""" if self._native_attention is None and "triton" in (self.denoise_attention, self.prefix_attention): @@ -689,6 +758,7 @@ def _clear_inference_caches(self) -> None: runtime = getattr(self, "_runtime", None) if runtime is not None: runtime.clear() + self._kernel_selection = None operators = getattr(self, "_fused_ops", None) if operators is not None: operators.clear() @@ -798,19 +868,25 @@ def _attn_sublayer( v = v.view(B, S, -1, hd).transpose(1, 2) if openpi_rlinf: q, k = operators.rotate(q, cos, sin), operators.rotate(k, cos, sin) - elif native_expert and q.dtype in (torch.bfloat16, torch.float16): - from ...backend.triton.rotary import rotate_qk - - q, k = rotate_qk(q, k, cos, sin) else: - q, k = _apply_rope(q, k, cos, sin, self.compile_backend == "inductor") + q, k = _apply_rope( + q, + k, + cos, + sin, + self.compile_backend == "inductor", + native_expert and q.dtype in (torch.bfloat16, torch.float16), + self._kernels, + ) if collected is not None: # prefix pass: cache this layer's K/V - collected.append((k, v)) + # Contiguous like Inductor's outputs, so the cache layout (which + # layout-gated tuned kernels match) is the same with or without compile. + collected.append((k.contiguous(), v.contiguous())) if native_expert and self.denoise_attention == "triton": # All action queries share the same prefix/suffix key validity. key_padding_mask = None if mask is None else mask[:, :, :1, :] - out = self._native_attention.attend_split( - q, *prefix_kv, k, v, key_padding_mask=key_padding_mask, scaling=attn.scaling + out = _op(self._kernels, "split_kv_attention", True)( + q, *prefix_kv, k, v, key_padding_mask, attn.scaling ) return attn.o_proj(out.transpose(1, 2).reshape(B, S, -1)) if prefix_kv is not None: # denoise pass: attend [prefix ++ suffix] @@ -837,7 +913,9 @@ def _attn_sublayer( if native_prefix: # Prefix preparation supplies a shared key mask; masked query outputs are unused. key_padding_mask = None if mask is None else mask[:, :, :1, :] - out = self._native_attention.attend(q, k, v, attn_mask=key_padding_mask, scaling=attn.scaling) + out = _op(self._kernels, "split_kv_attention", True)( + q, k[:, :, :0], v[:, :, :0], k, v, key_padding_mask, attn.scaling + ) return attn.o_proj(out.transpose(1, 2).reshape(B, S, -1)) # K/V stay at num_kv_heads; the attention backend does the GQA repeat # (materialized for 'eager', broadcast for 'eager_bc'). @@ -908,6 +986,7 @@ def _tower_forward( self.compile_backend == "inductor", None if modulations is None else modulations[2 * i], fused, + self._kernels, ) attn = self._attn_sublayer( layer.self_attn, @@ -926,6 +1005,7 @@ def _tower_forward( gate, self.compile_backend == "inductor", fused, + self._kernels, ) residual = hidden h, gate = _rmsnorm( @@ -935,6 +1015,7 @@ def _tower_forward( self.compile_backend == "inductor", None if modulations is None else modulations[2 * i + 1], fused, + self._kernels, ) hidden = _gated_residual( residual, @@ -944,10 +1025,12 @@ def _tower_forward( self.compile_backend == "inductor", fused, fuse_projections=self.attention != "eager", + kernels=self._kernels, ), gate, self.compile_backend == "inductor", fused, + self._kernels, ) hidden, _ = _rmsnorm( tower.norm, @@ -956,6 +1039,7 @@ def _tower_forward( self.compile_backend == "inductor", None if modulations is None else modulations[-1], fused, + self._kernels, ) return hidden, collected @@ -1020,6 +1104,7 @@ def _encode_prefix_impl( key_mask = prefix_pad_masks[:, None, :].expand(-1, length, -1) mask4d = self._attention_mask_4d(key_mask) if self.compile_backend == "inductor": + self._kernels # noqa: B018 - resolve operators before graph tracing reads them policy_key = (id(self), native_prefix) prefix_encoder = _COMPILED_PREFIX_ENCODERS.get(policy_key) if prefix_encoder is None: @@ -1318,6 +1403,7 @@ def _build_pi05( request_cuda_graph: bool = False, denoise_attention: str = "sdpa", prefix_attention: str = "sdpa", + kernel_profile: str = "none", low_cpu_mem_usage: bool = False, checkpoint_config: Mapping[str, Any] | None = None, device_type: str | None = None, @@ -1366,6 +1452,7 @@ def _build_pi05( request_cuda_graph=request_cuda_graph, denoise_attention=denoise_attention, prefix_attention=prefix_attention, + kernel_profile=kernel_profile, low_cpu_mem_usage=low_cpu_mem_usage, checkpoint_config=checkpoint_config, optimizations=optimizations, diff --git a/pyproject.toml b/pyproject.toml index 8ca62c6..0a28f83 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -147,6 +147,11 @@ vvla-wireless-serve = "embodiinfer.engine.serve.wireless_server:main" [tool.setuptools.packages.find] include = ["embodiinfer*", "vvla"] +# Repository-only tooling has its own Python >=3.12 environment and does not +# enter the inference package, model groups, or core dependency resolution. +[tool.embodiinfer.kernel-tuning] +tooling-project = "scripts/kernel_tuning" +evaluator-requirements = "benchmarks/kernel_tuning/requirements.txt" [tool.setuptools.package-data] "embodiinfer.backend.cuda" = ["csrc/*.cu", "csrc/*.cuh", "csrc/*.cpp"] "embodiinfer.policies.pi05.inference" = ["calibration/*.json"] @@ -169,6 +174,10 @@ markers = [ line-length = 110 target-version = "py310" src = ["embodiinfer", "tests"] +# Installed kernel-profile packages are generated by `scripts.kernel_tuning install` +# from evaluator-verified sources and pinned by digest in profile.json; editing or +# reformatting them would invalidate that evidence, so they are not linted. +extend-exclude = ["embodiinfer/backend/triton/tuned/*/*/"] [tool.ruff.lint] # E/W pycodestyle, F pyflakes, I isort, UP pyupgrade, B bugbear, C4 comprehensions, SIM simplify. diff --git a/scripts/kernel_tuning/__init__.py b/scripts/kernel_tuning/__init__.py new file mode 100644 index 0000000..e175126 --- /dev/null +++ b/scripts/kernel_tuning/__init__.py @@ -0,0 +1 @@ +"""Repository tooling for bounded Humanize2 kernel tuning; no runtime imports.""" diff --git a/scripts/kernel_tuning/__main__.py b/scripts/kernel_tuning/__main__.py new file mode 100644 index 0000000..4d91594 --- /dev/null +++ b/scripts/kernel_tuning/__main__.py @@ -0,0 +1,481 @@ +"""Developer commands for checking, running, resuming, and exporting tuning jobs.""" + +from __future__ import annotations + +import argparse +import asyncio +import json +import os +import shutil +import sys +from pathlib import Path + +from .artifacts import ( + REPOSITORY, + UNFENCED_HELP, + RunStore, + atomic_json, + run_lock, + runtime_identity, + warn_unfenced, +) +from .contracts import HMZ_REVISION, ContractError, TaskPackage, read_json +from .discovery.contracts import ModelCapture +from .runner import evaluate + + +def _execute(store: RunStore, *, resume: bool) -> dict: + if "replay.json" in store.task.hashes: + from .batch import verify_model_run + + verify_model_run(store) + if os.name != "posix": + raise ContractError( + "This pinned Humanize2 release requires POSIX; run tuning on the Linux target machine" + ) + if sys.version_info < (3, 12): + raise ContractError("Humanize2 requires a separate Python 3.12+ tooling environment") + try: + from hmz.flows import BudgetExceeded + from hmz.sdk import Hmz + except ImportError as exc: + raise ContractError("Install the optional scripts/kernel_tuning tooling environment first") from exc + with run_lock(store.root): + store.verify() + if runtime_identity() != store.manifest["tool_runtime"]: + raise ContractError("Humanize2/Python tooling environment changed; start a new run") + remaining = store.task.settings.search.max_seconds - store.manifest["elapsed_seconds"] + if remaining <= 0 or store.manifest["status"] == "completed": + raise ContractError("Run budget already exhausted; start a new run with a new task contract") + warn_unfenced(store.manifest) + epic = store.root / "humanize.json" + running = Hmz(workspace=store.root / "workspace").run( + flow=Path(__file__).with_name("flow.py"), + task=f"Optimize {store.task.definition['name']} under its frozen task contract", + agents={"coder": store.manifest["agent"]}, + params={"run_dir": str(store.root)}, + budget={"duration": remaining, "graceful": False}, + # A killed invocation may be newer than humanize.json. The dedicated + # workspace lets hmz select its newest journal instead of a stale path. + resume=resume, + ) + try: + return running.run() + except BudgetExceeded: + stopped = RunStore(store.root) + stopped.recover() + stopped.manifest.update(status="completed", stop_reason="max_seconds") + stopped.save() + return stopped.manifest + finally: + archived = None + if running.epic: + archived = store.root / "humanize" / running.epic.name + archived.mkdir(parents=True, exist_ok=True) + # Preserve flow/resume journals, without copying agent account + # directories. CLI-managed sessions stay at the recorded epic. + for path in running.epic.glob("*.jsonl"): + if path.is_file() and not path.is_symlink(): + shutil.copyfile(path, archived / path.name) + atomic_json( + epic, + { + "epic": str(running.epic) if running.epic else None, + "archived": str(archived) if archived else None, + "expected_hmz_revision": HMZ_REVISION, + }, + ) + + +def _selection(parser: argparse.ArgumentParser) -> None: + parser.add_argument("operators", nargs="*", help="Catalog operator names (default: all)") + parser.add_argument( + "--catalog", action="store_true", help="Use the legacy hand-maintained operator catalog" + ) + parser.add_argument("--model", action="append", help="One canonical target policy name, e.g. pi05") + parser.add_argument( + "--set", + action="append", + default=[], + metavar="KEY=VALUE", + help="Override tuning.yaml, e.g. evaluator_python=/abs/python or search.max_candidates=10", + ) + parser.add_argument("--hardware-notes", type=Path, help="Target-hardware notes copied into each task") + shapes = parser.add_mutually_exclusive_group() + shapes.add_argument( + "--captured", + type=Path, + help="Completed model capture directory (legacy JSONL only with --catalog)", + ) + shapes.add_argument("--estimated", action="store_true", help="Use the catalog's estimated shapes") + + +def _model_capture(args: argparse.Namespace) -> ModelCapture: + if not args.model or len(args.model) != 1: + raise ContractError( + "Model tuning requires exactly one --model; use --catalog for legacy operator selection" + ) + if ( + args.operators + or args.estimated + or getattr(args, "force", False) + or getattr(args, "skip_preflight", False) + ): + raise ContractError( + "Model tuning requires complete coverage: no operator filter, --estimated, --force, or --skip-preflight" + ) + if args.captured is None: + raise ContractError("Model tuning requires --captured MODEL_CAPTURE; run capture-ops --model first") + return ModelCapture.load(args.captured, args.model[0], require_ready=not getattr(args, "list", False)) + + +def _generate(args: argparse.Namespace) -> None: + from .capture import apply_file + from .generate import parse_overrides, render + from .operators import select + + if not args.catalog: + from .discovery.tasks import render_model + + capture = _model_capture(args) + if args.list: + for op in capture.operators: + print( + f"{op['status']:20} {op['name']} ({len(op['workloads'])} cases) {op.get('reason') or ''}" + ) + return + tasks = render_model( + capture, args.output, overrides=parse_overrides(args.set), hardware_notes=args.hardware_notes + ) + for task in tasks: + print(f"{task.definition['description']}: {task.root} ({len(task.workloads)} workloads)") + return + + operators, notes = apply_file( + select(args.operators or None, args.model), args.captured, estimated=args.estimated + ) + for note in notes: + print(note) + if args.list: + for operator in operators: + print(f"{operator.name:22} {','.join(operator.models):10} {operator.summary}") + return + overrides = parse_overrides(args.set) + for operator in operators: + task = render( + operator, + args.output / operator.name, + overrides=overrides, + hardware_notes=args.hardware_notes, + force=args.force, + ) + print( + f"{operator.name}: {task.root} ({len(task.workloads)} workloads, {task.settings.precision.mode})" + ) + + +def _tune_all(args: argparse.Namespace) -> int: + from . import batch + from .capture import apply_file + from .generate import parse_overrides + from .operators import select + + if args.resume: + if ( + args.operators + or args.model + or args.set + or args.hardware_notes + or args.agent + or args.unfenced + or args.captured + or args.estimated + or args.catalog + ): + raise ContractError( + "--resume continues the recorded batch; omit selection, --set, --agent, and --unfenced" + ) + root = args.resume.resolve(strict=True) + else: + agent = args.agent or "" + if not (args.dry_run or args.preflight_only) and "/" not in agent: + raise ContractError("Choose an explicit Humanize2 agent: --agent harness/model:effort") + options = { + "agent": agent, + "overrides": parse_overrides(args.set), + "hardware_notes": args.hardware_notes, + "unfenced": args.unfenced, + } + if args.catalog: + operators, notes = apply_file( + select(args.operators or None, args.model), args.captured, estimated=args.estimated + ) + for note in notes: + print(note, flush=True) + root = batch.create(operators, args.output, **options) + else: + root = batch.create_model(_model_capture(args), args.output, **options) + if args.skip_preflight and read_json(root / "batch.json").get("model"): + raise ContractError("Model batches cannot skip the all-baseline preflight") + print(f"Batch: {root}", flush=True) + if args.dry_run: + batch.write_summary(root) + print((root / "summary.md").read_text(encoding="utf-8")) + return 0 + if not args.skip_preflight: + results = batch.preflight(root) + items = read_json(root / "batch.json")["items"] + for name, error in results.items(): + timings = ", ".join( + f"{work} {us:.1f} us" for work, us in items[name].get("baseline_us", {}).items() + ) + print(f"preflight {name}: {f'ok ({timings})' if error is None else error}", flush=True) + if args.preflight_only: + batch.write_summary(root) + return ( + 0 + if all( + i["status"] != "preflight_failed" for i in read_json(root / "batch.json")["items"].values() + ) + else 1 + ) + record = batch.execute(root) + print((root / "summary.md").read_text(encoding="utf-8")) + return 0 if all(item["status"] == "completed" for item in record["items"].values()) else 1 + + +def main(argv: list[str] | None = None) -> int: + """Run only the requested command; checking/status/export never invoke an agent.""" + parser = argparse.ArgumentParser(description=__doc__) + commands = parser.add_subparsers(dest="command", required=True) + from .registration.cli import add_parser + + add_parser(commands) + check = commands.add_parser("check", help="Validate a task without executing its code") + check.add_argument("task", type=Path) + check.add_argument( + "--environment", action="store_true", help="Also probe the selected evaluator environment (no tuning)" + ) + run = commands.add_parser("run", help="Start real agent-driven tuning on this machine") + run.add_argument("task", type=Path) + run.add_argument("--agent", help="Explicit Humanize2 harness/model:effort spec; overrides tuning.yaml") + run.add_argument("--output", type=Path, default=REPOSITORY / "results/kernel_tuning") + run.add_argument("--unfenced", action="store_true", help=UNFENCED_HELP) + for name in ("resume", "status"): + commands.add_parser(name).add_argument("run", type=Path) + export = commands.add_parser("export", help="Export best verified source and evidence; no execution") + export.add_argument("run", type=Path) + export.add_argument("destination", type=Path) + generate = commands.add_parser( + "generate", help="Write tasks from a registry capture or catalog; no execution" + ) + _selection(generate) + generate.add_argument("--output", type=Path, default=REPOSITORY / "results/kernel_tuning/tasks") + generate.add_argument( + "--force", action="store_true", help="Replace previously generated task directories" + ) + generate.add_argument("--list", action="store_true", help="List captured registry or catalog operators") + tune_all = commands.add_parser("tune-all", help="Generate, preflight, and tune every selected operator") + _selection(tune_all) + tune_all.add_argument("--agent", help="Explicit Humanize2 harness/model:effort spec for every operator") + tune_all.add_argument("--output", type=Path, default=REPOSITORY / "results/kernel_tuning/batches") + tune_all.add_argument("--resume", type=Path, metavar="BATCH", help="Continue an existing batch directory") + tune_all.add_argument("--unfenced", action="store_true", help=UNFENCED_HELP) + mode = tune_all.add_mutually_exclusive_group() + mode.add_argument("--dry-run", action="store_true", help="Only generate and validate the tasks") + mode.add_argument("--preflight-only", action="store_true", help="Stop after measuring the baselines") + mode.add_argument("--skip-preflight", action="store_true", help="Let each run measure its baseline") + capture = commands.add_parser( + "capture", help="Record legacy catalog workload shapes; use capture-ops for registry tuning" + ) + capture.add_argument("operators", nargs="*", help="Catalog operators to record (default: all)") + capture.add_argument("--catalog", action="store_true", required=True, help="Record catalog-only JSONL") + capture.add_argument("--output", type=Path, required=True, help="Destination JSONL file") + capture.add_argument("script", nargs=argparse.REMAINDER, help="-- SCRIPT [ARGS] or -- -m MODULE [ARGS]") + capture_ops = commands.add_parser( + "capture-ops", help="Record operator-registry calls as a tunable model capture (model runtime)" + ) + capture_ops.add_argument("--model", required=True, help="Canonical policy name built by the driver") + capture_ops.add_argument("--ops", help="Comma-separated registry operators (default: every one called)") + capture_ops.add_argument("--warmup", type=int, default=2, help="Inference steps skipped before recording") + capture_ops.add_argument("--steps", type=int, default=8, help="Inference steps recorded") + capture_ops.add_argument( + "--fixture-bytes", type=int, default=1 << 30, help="Maximum stored real-input bytes" + ) + capture_ops.add_argument("--output", type=Path, required=True, help="New capture directory") + capture_ops.add_argument( + "script", nargs=argparse.REMAINDER, help="-- SCRIPT [ARGS] or -- -m MODULE [ARGS]" + ) + install = commands.add_parser( + "install", help="Install an exported registry solution into a kernel profile (model runtime)" + ) + install.add_argument("export", type=Path, help="Directory written by export") + install.add_argument("--profile", required=True, help="Kernel profile name, e.g. thor or rtx4090") + install.add_argument("--replace", action="store_true", help="Overwrite an installed operator") + verify = commands.add_parser( + "verify", help="Model-level numerical gate for a kernel profile (model runtime)" + ) + verify.add_argument("--profile", required=True, help="Installed kernel profile to check") + verify.add_argument("--output", type=Path, required=True, help="New gate report (JSON)") + verify.add_argument("--samples", type=int, help="Use only the first N recorded inputs") + verify.add_argument("--seed", type=int, default=1000, help="Base of the per-input seeds") + verify.add_argument("script", nargs=argparse.REMAINDER, help="-- discovery.drive arguments") + hotspots = commands.add_parser( + "profile", help="Rank where end-to-end inference spends GPU time (model runtime)" + ) + hotspots.add_argument("--output", type=Path, required=True, help="New profile directory") + hotspots.add_argument("--warmup", type=int, default=10, help="Inference steps skipped before profiling") + hotspots.add_argument("--steps", type=int, default=20, help="Inference steps profiled") + hotspots.add_argument("--no-module-ranges", action="store_true", help="Skip per-nn.Module attribution") + hotspots.add_argument("--trace", action="store_true", help="Also write a Chrome trace.json") + hotspots.add_argument("script", nargs=argparse.REMAINDER, help="-- SCRIPT [ARGS] or -- -m MODULE [ARGS]") + calibration = commands.add_parser( + "calibrate", help="Freeze baseline/FP64 numerical bounds in the model runtime before generation" + ) + calibration.add_argument("--captured", type=Path, required=True) + calibration.add_argument("--output", type=Path, required=True, help="New calibrated capture directory") + skeleton = commands.add_parser( + "skeleton", help="Random-weight copy of a Hugging Face checkpoint for shape capture (model runtime)" + ) + skeleton.add_argument("repo", help="Model repository, e.g. org/name; HF_ENDPOINT selects a mirror") + skeleton.add_argument("--revision", required=True, help="Immutable commit to copy") + skeleton.add_argument("--output", type=Path, required=True) + raw = sys.argv[1:] if argv is None else argv + if "--" in raw and raw[0] in ("capture", "capture-ops", "profile", "verify"): + split = raw.index("--") + raw, command = raw[:split], raw[split + 1 :] + else: + command = None + args = parser.parse_args(raw) + try: + if args.command == "register": + from .registration.cli import dispatch + + return dispatch(args) + if args.command == "capture": + from .capture import run as run_capture + from .operators import select + + rows = run_capture(command or args.script, args.output, select(args.operators or None)) + calls = sum(row["count"] for row in rows if "axes" in row) + print(f"Captured {calls} calls in {sum('axes' in row for row in rows)} cases -> {args.output}") + for row in rows: + if "skipped" in row: + print(f" skipped {row['count']} {row['operator']} calls: {row['skipped']}") + elif args.command == "capture-ops": + from .regions import run as run_op_capture + + ops = {name.strip() for name in args.ops.split(",") if name.strip()} if args.ops else None + result = run_op_capture( + command or args.script, + args.output, + args.model, + ops=ops, + warmup=args.warmup, + steps=args.steps, + fixture_bytes=args.fixture_bytes, + ) + print(json.dumps(result, indent=2, ensure_ascii=False)) + elif args.command == "install": + from .install import install as install_solution + + result = install_solution(args.export, args.profile, replace=args.replace) + print(json.dumps(result, indent=2, ensure_ascii=False)) + elif args.command == "verify": + from .verify import run as run_verify + + report = run_verify( + args.profile, args.output, command or args.script, samples=args.samples, seed=args.seed + ) + summary = {key: report[key] for key in ("passed", "inputs", "failed", "worst_ratio")} + summary["outputs"] = {name: item["ratio"] for name, item in report["outputs"].items()} + print(json.dumps(summary, indent=2)) + if not report["passed"]: + return 1 + elif args.command == "profile": + from .hotspots import render + from .hotspots import run as run_profile + + summary = run_profile( + command or args.script, + args.output, + warmup=args.warmup, + steps=args.steps, + modules=not args.no_module_ranges, + trace=args.trace, + ) + print(render(summary, limit=15)) + print(f"Report: {args.output / 'hotspots.md'}") + elif args.command == "calibrate": + from benchmarks.kernel_tuning.calibration import calibrate + + result = calibrate(ModelCapture.load(args.captured), args.output) + print(json.dumps({"capture": str(result.root), "identity": result.identity}, indent=2)) + elif args.command == "skeleton": + from .skeleton import build + + summary = build(args.repo, args.revision, args.output) + tensors = sum(summary["synthesized"].values()) + print( + f"{args.output}: {len(summary['copied'])} files copied, " + f"{tensors} random tensors in {len(summary['synthesized'])} safetensors files" + ) + elif args.command == "generate": + _generate(args) + elif args.command == "tune-all": + return _tune_all(args) + elif args.command == "check": + task = TaskPackage.load(args.task) + result = { + "task": task.definition["name"], + "digest": task.identity, + "workloads": len(task.workloads), + "validation": "structural; no task code executed", + } + if args.environment: + import tempfile + + with tempfile.TemporaryDirectory(prefix="kernel-preflight-") as directory: + store = RunStore.create(task, Path(directory), agent="") + result["environment"] = asyncio.run( + evaluate(store, store.root / "probe", operation="probe") + )["environment"] + print(json.dumps(result, indent=2, ensure_ascii=False)) + elif args.command == "run": + task = TaskPackage.load(args.task) + agent = args.agent or task.settings.agent + if not agent or "/" not in agent: + raise ContractError("Choose an explicit Humanize2 agent: --agent harness/model:effort") + model_batch = None + if "replay.json" in task.hashes: + from .batch import verify_model_task + + model_batch = verify_model_task(task) + store = RunStore.create(task, args.output, agent, unfenced=args.unfenced) + if model_batch is not None: + store.manifest["model_batch"] = str(model_batch) + item = read_json(model_batch / "batch.json")["items"][task.root.name] + store.manifest["environment"] = item["preflight_environment"] + store.save() + print(f"Run archive: {store.root}", flush=True) + print(json.dumps(_execute(store, resume=False), indent=2, ensure_ascii=False)) + else: + store = RunStore(args.run) + if args.command == "resume": + print(json.dumps(_execute(store, resume=True), indent=2, ensure_ascii=False)) + elif args.command == "status": + print(json.dumps(store.manifest, indent=2, ensure_ascii=False)) + else: + with run_lock(store.root): + print(store.export(args.destination)) + return 0 + except (ContractError, OSError, ValueError, ImportError) as exc: + print(f"kernel_tuning: {exc}", file=sys.stderr) + return 2 + except KeyboardInterrupt: + print("Interrupted; the run archive can be resumed.", file=sys.stderr) + return 130 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/kernel_tuning/artifacts.py b/scripts/kernel_tuning/artifacts.py new file mode 100644 index 0000000..e90241e --- /dev/null +++ b/scripts/kernel_tuning/artifacts.py @@ -0,0 +1,352 @@ +"""Durable run snapshots and verified, source-complete operator exports.""" + +from __future__ import annotations + +import contextlib +import hashlib +import importlib.metadata +import json +import os +import shutil +import sys +import tempfile +import uuid +from collections.abc import Callable, Iterator +from datetime import datetime, timezone +from pathlib import Path +from typing import Any + +from .contracts import ContractError, TaskPackage, digest, read_json, tree_hashes + +REPOSITORY = Path(__file__).resolve().parents[2] +UNFENCED_HELP = ( + "Grant the coding agent every Humanize2 scope so no Landlock fence is needed; " + "run only under a dedicated unprivileged account or container" +) + + +def warn_unfenced(manifest: dict[str, Any]) -> None: + """Say on stderr, before any agent starts, that the run relies on outside isolation.""" + if manifest.get("unfenced"): + print( + "kernel_tuning: WARNING: agent sessions are unfenced; they may write anywhere this " + "account can. Use a dedicated unprivileged account or container.", + file=sys.stderr, + flush=True, + ) + + +def atomic_json(path: Path, value: Any) -> None: + """Replace a complete JSON document atomically on the same filesystem.""" + path.parent.mkdir(parents=True, exist_ok=True) + fd, name = tempfile.mkstemp(prefix=f".{path.name}.", dir=path.parent) + try: + with os.fdopen(fd, "w", encoding="utf-8", newline="\n") as handle: + json.dump(value, handle, indent=2, ensure_ascii=False, allow_nan=False) + handle.write("\n") + handle.flush() + os.fsync(handle.fileno()) + os.replace(name, path) + finally: + if os.path.exists(name): + os.unlink(name) + + +def tool_files() -> dict[str, str]: + """List source files needed to reproduce this controller and evaluator version.""" + hashes = {} + for folder in (REPOSITORY / "scripts/kernel_tuning", REPOSITORY / "benchmarks/kernel_tuning"): + for directory, subdirs, names in os.walk(folder): + subdirs[:] = [ + name for name in subdirs if name not in ("tasks", "__pycache__", ".venv", ".ruff_cache") + ] + for name in names: + path = Path(directory) / name + if path.suffix not in (".py", ".toml", ".txt", ".md", ".lock"): + continue + hashes[path.relative_to(REPOSITORY).as_posix()] = hashlib.sha256( + path.read_bytes() + ).hexdigest() + hashes["scripts/__init__.py"] = hashlib.sha256( + (REPOSITORY / "scripts/__init__.py").read_bytes() + ).hexdigest() + return hashes + + +def tool_identity() -> str: + """Bind resumption to the orchestration, prompts, and evaluator source version.""" + return digest(tool_files()) + + +def runtime_identity() -> dict[str, Any]: + """Record the installed orchestration runtime, including actual hmz source bytes.""" + packages = {} + for name in ("hmz", "pydantic", "pyyaml", "psutil"): + try: + distribution = importlib.metadata.distribution(name) + except importlib.metadata.PackageNotFoundError: + packages[name] = {"version": "absent"} + continue + packages[name] = {"version": distribution.version} + if name == "hmz": + sources = {} + for item in distribution.files or []: + if str(item).startswith("hmz/") and str(item).endswith(".py"): + sources[str(item)] = hashlib.sha256( + Path(distribution.locate_file(item)).read_bytes() + ).hexdigest() + packages[name]["source_digest"] = digest(sources) + packages[name]["origin"] = distribution.read_text("direct_url.json") + return {"python": sys.version, "interpreter": sys.executable, "packages": packages} + + +@contextlib.contextmanager +def run_lock(root: Path) -> Iterator[None]: + """Hold an OS lock; process death releases it without stale-PID guessing.""" + path = root / ".run.lock" + with path.open("a+b") as handle: + handle.seek(0, os.SEEK_END) + if handle.tell() == 0: + handle.write(b"0") + handle.flush() + handle.seek(0) + try: + if os.name == "nt": + import msvcrt + + msvcrt.locking(handle.fileno(), msvcrt.LK_NBLCK, 1) + else: + import fcntl + + fcntl.flock(handle.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB) + except OSError as exc: + raise ContractError("Another process is using this tuning run") from exc + try: + yield + finally: + handle.seek(0) + if os.name == "nt": + msvcrt.locking(handle.fileno(), msvcrt.LK_UNLCK, 1) + else: + fcntl.flock(handle.fileno(), fcntl.LOCK_UN) + + +class RunStore: + """Persist a run without modifying the input task or engine source tree.""" + + def __init__(self, root: Path) -> None: + self.root = root.resolve(strict=True) + self.manifest = read_json(self.root / "manifest.json") + self.task = TaskPackage.load(self.root / "task") + if self.task.identity != self.manifest["task_digest"]: + raise ContractError("Task snapshot no longer matches its manifest") + + @classmethod + def create(cls, task: TaskPackage, output: Path, agent: str, *, unfenced: bool = False) -> RunStore: + """Snapshot a task and baseline under a fresh, collision-resistant run ID. + + ``unfenced`` is frozen into the manifest: every start and resume of the run + opens agent sessions without Humanize2's Landlock fence. + """ + task.verify() + output = output.resolve() + if output == task.root or task.root in output.parents: + raise ContractError("Results must be outside the input task directory") + run_id = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ-") + uuid.uuid4().hex[:10] + root = output / run_id + root.mkdir(parents=True, exist_ok=False) + snapshot = root / "task" + snapshot.mkdir() + for name in task.hashes: + target = snapshot / name + target.parent.mkdir(parents=True, exist_ok=True) + shutil.copyfile(task.root / name, target) + if tree_hashes(snapshot) != task.hashes: + raise ContractError("Task changed while being snapshotted") + (root / "workspace").mkdir() + baseline = task.baseline() + atomic_json(root / "baseline/solution.json", baseline) + source_hashes = tool_files() + for name in source_hashes: + destination = root / "tool" / name + destination.parent.mkdir(parents=True, exist_ok=True) + shutil.copyfile(REPOSITORY / name, destination) + if tree_hashes(root / "tool") != source_hashes: + raise ContractError("Tool source changed while being snapshotted") + manifest = { + "schema_version": 1, + "run_id": run_id, + "created_at": datetime.now(timezone.utc).isoformat(), + "task_source": str(task.root), + "task_digest": task.identity, + "task_hashes": task.hashes, + "tool_digest": digest(source_hashes), + "tool_runtime": runtime_identity(), + "agent": agent, + "unfenced": unfenced, + "baseline_digest": digest(baseline), + "environment": None, + "elapsed_seconds": 0.0, + "attempts": [], + "best": None, + "streak": 0, + "status": "created", + "stop_reason": None, + "pending": None, + } + atomic_json(root / "manifest.json", manifest) + return cls(root) + + def verify(self, *, check_tool: bool = True) -> None: + """Check frozen source, task, baseline, and any incumbent before another turn.""" + self.task.verify() + if check_tool and tool_identity() != self.manifest["tool_digest"]: + raise ContractError("Tuning/evaluation source changed; start a new run") + if digest(tree_hashes(self.root / "tool")) != self.manifest["tool_digest"]: + raise ContractError("Archived evaluator/controller source changed") + if digest(read_json(self.root / "baseline/solution.json")) != self.manifest["baseline_digest"]: + raise ContractError("Archived baseline changed") + if self.manifest["best"]: + self.solution_path(self.manifest["best"]) + best = next(a for a in self.manifest["attempts"] if a["id"] == self.manifest["best"]) + actual = { + k: v + for k, v in tree_hashes(self.root / "candidates" / best["id"]).items() + if not k.startswith("profile/") + } + if actual != best["evidence_hashes"]: + raise ContractError("Archived winning evidence changed") + + def save(self) -> None: + """Publish the state mirror and best pointer after durable candidate evidence.""" + atomic_json(self.root / "manifest.json", self.manifest) + if self.manifest["best"]: + best = self.manifest["best"] + atomic_json( + self.root / "best.json", + { + "candidate": best, + "solution": str(self.solution_path(best).relative_to(self.root)), + "task_digest": self.task.identity, + }, + ) + + def solution_path(self, candidate: str | None = None) -> Path: + """Resolve an archived solution and verify its recorded content digest.""" + if candidate is None: + return self.root / "baseline/solution.json" + if not candidate.isascii() or not candidate.isdigit(): + raise ContractError("Candidate IDs must contain only ASCII digits") + attempt = next((a for a in self.manifest["attempts"] if a["id"] == candidate), None) + if attempt is None or not attempt.get("solution_digest"): + raise ContractError(f"No archived solution for {candidate}") + path = self.root / "candidates" / candidate / "solution.json" + if digest(read_json(path)) != attempt["solution_digest"]: + raise ContractError(f"Archived solution changed: {candidate}") + return path + + def begin_attempt(self, *, persist: Callable[[], None] | None = None) -> tuple[str, Path, Path]: + """Allocate a counted attempt, with scratch separated from measurement records.""" + self.verify() + candidate = f"{len(self.manifest['attempts']) + 1:04d}" + archive = self.root / "candidates" / candidate + scratch = self.root / "workspace" / candidate + self.manifest["attempts"].append({"id": candidate, "status": "running", "solution_digest": None}) + self.manifest["pending"] = candidate + # Reserve the attempt in the framework journal before creating directories; + # a crash halfway through copying large fixtures must still be resumable. + (persist or self.save)() + archive.mkdir(parents=True, exist_ok=False) + scratch.mkdir(parents=True, exist_ok=False) + shutil.copytree(self.task.root, scratch / "task") + shutil.copyfile(self.solution_path(self.manifest["best"]), scratch / "incumbent.json") + feedback = [] + for previous in self.manifest["attempts"][-5:]: + record = dict(previous) + if previous.get("solution_digest"): + # A rejected kernel is not the incumbent. Give the next agent + # its verified source as well as its timings, so it can learn + # from failed launch configurations instead of guessing them. + record["solution"] = read_json(self.solution_path(previous["id"])) + feedback.append(record) + atomic_json(scratch / "feedback.json", feedback) + atomic_json(scratch / "environment.json", self.manifest["environment"]) + return candidate, archive, scratch + + def finish_attempt( + self, candidate: str, *, status: str, decision: dict[str, Any], solution: dict[str, Any] | None = None + ) -> None: + """Write evidence before changing which solution future attempts may inherit.""" + archive = self.root / "candidates" / candidate + atomic_json(archive / "decision.json", decision) + attempt = next(a for a in self.manifest["attempts"] if a["id"] == candidate) + attempt.update(status=status, decision=decision) + if solution is not None: + atomic_json(archive / "solution.json", solution) + from .contracts import validate_solution + + validate_solution(solution, self.task) + for source in solution["sources"]: + path = archive / "sources" / source["path"] + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(source["content"].encode("utf-8")) + attempt["solution_digest"] = digest(solution) + attempt["evidence_hashes"] = { + k: v for k, v in tree_hashes(archive).items() if not k.startswith("profile/") + } + if status == "promoted": + if solution is None: + raise ContractError("Cannot promote without an archived solution") + self.manifest["best"] = candidate + self.manifest["streak"] = 0 + else: + self.manifest["streak"] += 1 + self.manifest["pending"] = None + + def recover(self) -> None: + """Keep an interrupted attempt as failed evidence rather than replaying its writes.""" + if candidate := self.manifest["pending"]: + self.finish_attempt( + candidate, + status="interrupted", + decision={"promoted": False, "reasons": ["Interrupted before durable completion"]}, + ) + + def export(self, destination: Path) -> Path: + """Export the best source, frozen task, and evidence into a new directory.""" + self.verify(check_tool=False) + best = self.manifest["best"] + if not best: + raise ContractError("No candidate has passed promotion; all attempts remain in the run archive") + destination = destination.resolve() + if destination.exists(): + raise ContractError("Export destination already exists; choose a new directory") + if destination == self.root or self.root in destination.parents: + raise ContractError("Export outside the run directory") + solution = read_json(self.solution_path(best)) + from .contracts import validate_solution + + validate_solution(solution, self.task) + destination.parent.mkdir(parents=True, exist_ok=True) + staging = Path(tempfile.mkdtemp(prefix=".kernel-export-", dir=destination.parent)) + try: + for source in solution["sources"]: + path = staging / "sources" / source["path"] + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(source["content"].encode("utf-8")) + atomic_json(staging / "solution.json", solution) + atomic_json(staging / "manifest.json", self.manifest) + shutil.copytree(self.root / "candidates" / best, staging / "evidence") + shutil.copytree(self.root / "baseline", staging / "baseline") + shutil.copytree(self.task.root, staging / "task") + shutil.copytree(self.root / "tool", staging / "tool") + if (self.root / "humanize").exists(): + shutil.copytree(self.root / "humanize", staging / "humanize") + if (self.root / "humanize.json").exists(): + shutil.copyfile(self.root / "humanize.json", staging / "humanize.json") + atomic_json(staging / "checksums.json", tree_hashes(staging)) + staging.rename(destination) + finally: + if staging.exists(): + shutil.rmtree(staging) + return destination diff --git a/scripts/kernel_tuning/batch.py b/scripts/kernel_tuning/batch.py new file mode 100644 index 0000000..16bfb81 --- /dev/null +++ b/scripts/kernel_tuning/batch.py @@ -0,0 +1,403 @@ +"""One-command tuning of every selected catalog operator, serially on one machine. + +A batch directory holds freshly generated tasks, one ordinary run archive per +operator, logs, exports of promoted kernels, and a summary. Each operator runs +in its own ``run``/``resume`` subprocess, so one failure or crash does not stop +the others, and an interrupted batch continues where it stopped. +""" + +from __future__ import annotations + +import asyncio +import re +import shutil +import subprocess +import sys +from collections.abc import Callable +from datetime import datetime, timezone +from pathlib import Path +from typing import Any +from uuid import uuid4 + +from .artifacts import REPOSITORY, RunStore, atomic_json, run_lock, tool_identity +from .contracts import ContractError, TaskPackage, digest, read_json, relative_path, validate_measurement +from .discovery.contracts import ModelCapture +from .generate import render +from .operators import Operator +from .runner import evaluate + +Launcher = Callable[[list[str], Path], int] +ARCHIVE = re.compile(r"^Run archive: (.+)$", re.MULTILINE) +DONE = ("completed", "preflight_failed") + + +def launch(argv: list[str], log: Path) -> int: + """Run one tuning command from the repository root, appending its output to ``log``.""" + log.parent.mkdir(parents=True, exist_ok=True) + with log.open("ab") as handle: + handle.write(f"\n$ {' '.join(argv)}\n".encode()) + handle.flush() + return subprocess.run( + argv, cwd=REPOSITORY, stdout=handle, stderr=subprocess.STDOUT, check=False + ).returncode + + +def create( + operators: list[Operator], + output: Path, + *, + agent: str, + overrides: dict[str, Any], + hardware_notes: Path | None, + unfenced: bool = False, +) -> Path: + """Generate every task into a new batch directory and record the plan.""" + if not operators: + raise ContractError("No operators selected") + root = output.resolve() / (datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ") + "-batch") + root.mkdir(parents=True, exist_ok=False) + items = {} + for operator in operators: + task = render( + operator, root / "tasks" / operator.name, overrides=overrides, hardware_notes=hardware_notes + ) + items[operator.name] = { + "task": task.root.relative_to(root).as_posix(), + "task_digest": task.identity, + "status": "pending", + "run": None, + "exit_code": None, + "error": None, + } + atomic_json( + root / "batch.json", + { + "schema_version": 1, + "created_at": datetime.now(timezone.utc).isoformat(), + "agent": agent, + "unfenced": unfenced, + "overrides": overrides, + "items": items, + }, + ) + return root + + +def preflight(root: Path, names: list[str] | None = None) -> dict[str, str | None]: + """Measure each production baseline against its own contract before any agent runs. + + This is the run's first evaluator step, done for every operator up front so + a broken runtime or a contract the production kernel cannot meet is found + before hours of agent time are spent on earlier operators. + """ + batch = read_json(root / "batch.json") + _verify_model(root, batch) + if batch.get("model"): + if names is not None: + raise ContractError("Model preflight must include every captured registry task") + batch["preflight_complete"] = False + atomic_json(root / "batch.json", batch) + results: dict[str, str | None] = {} + for name, item in batch["items"].items(): + if (names is not None and name not in names) or ( + not batch.get("model") and item["status"] != "pending" + ): + continue + task = TaskPackage.load(root / item["task"]) + try: + store = RunStore.create(task, root / "preflight", agent="preflight/none") + report = asyncio.run(evaluate(store, store.root / "baseline/evaluation")) + validate_measurement(task, report["measurement"]) + if batch.get("model"): + if not report.get("environment", {}).get("identity"): + raise ContractError("Model preflight requires the evaluator environment identity") + receipt = { + "task_digest": task.identity, + "tool_digest": tool_identity(), + "environment": report["environment"], + "measurement": report["measurement"], + } + path = store.root / "baseline/preflight.json" + atomic_json(path, receipt) + item.update( + preflight_receipt=path.relative_to(root).as_posix(), + preflight_receipt_digest=digest(receipt), + preflight_environment=report["environment"], + ) + results[name] = None + item["preflight_passed"] = True + if item["status"] == "preflight_failed": + item.update(status="pending", error=None) + # Baselines already at the device's launch floor leave an agent nothing to win. + rounds = report["measurement"]["rounds"] + item["baseline_us"] = { + work: 1000 + * sum(sum(r[work]["baseline"]) / len(r[work]["baseline"]) for r in rounds) + / len(rounds) + for work in task.workload_ids + } + except ContractError as exc: + results[name] = str(exc) + item.update(status="preflight_failed", error=str(exc), preflight_passed=False) + atomic_json(root / "batch.json", batch) + if batch.get("model"): + batch["preflight_complete"] = bool(results) and all(error is None for error in results.values()) + batch["preflight_tool_digest"] = tool_identity() + atomic_json(root / "batch.json", batch) + write_summary(root) + return results + + +def create_model( + capture: ModelCapture, + output: Path, + *, + agent: str, + overrides: dict[str, Any], + hardware_notes: Path | None, + unfenced: bool = False, +) -> Path: + """Freeze one model capture and all tasks in a batch requiring a complete preflight.""" + from .discovery.tasks import render_model + + capture.require_ready() + root = output.resolve() / ( + datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ") + f"-{uuid4().hex[:8]}-model" + ) + root.mkdir(parents=True, exist_ok=False) + # Verify before copying, then verify the copy, to reject a stale capture. + checked = ModelCapture.load(capture.root, capture.manifest["target"]["model"]) + if checked.identity != capture.identity: + raise ContractError("Capture changed before batch creation") + shutil.copytree(capture.root, root / "capture") + frozen = ModelCapture.load(root / "capture", capture.manifest["target"]["model"]) + if frozen.identity != capture.identity: + raise ContractError("Capture changed during batch creation") + tasks = render_model(frozen, root / "tasks", overrides=overrides, hardware_notes=hardware_notes) + names = {op["id"]: op["name"] for op in capture.operators} + items = { + task.root.name: { + "name": names[task.root.name], + "task": task.root.relative_to(root).as_posix(), + "task_digest": task.identity, + "status": "pending", + "run": None, + "exit_code": None, + "error": None, + "preflight_passed": False, + } + for task in tasks + } + atomic_json( + root / "batch.json", + { + "schema_version": 2, + "created_at": datetime.now(timezone.utc).isoformat(), + "agent": agent, + "unfenced": unfenced, + "overrides": overrides, + "model": capture.manifest["target"], + "capture_identity": capture.identity, + "preflight_complete": False, + "items": items, + }, + ) + return root + + +def _verify_model(root: Path, batch: dict[str, Any], *, require_preflight: bool = False) -> None: + if batch.get("schema_version") != 2: + if batch.get("model") or (root / "tasks/model.json").exists(): + raise ContractError("Unsupported model batch schema") + return + capture = ModelCapture.load(root / "capture", batch["model"]["model"]) + if capture.identity != batch["capture_identity"] or capture.manifest["target"] != batch["model"]: + raise ContractError("Batch target/capture identity changed") + expected = {op["id"] for op in capture.operators if op["status"] == "ready"} + if set(batch["items"]) != expected: + raise ContractError("Model batch has task coverage gaps") + index = read_json(root / "tasks/model.json") + if index["capture_identity"] != capture.identity or set(index["tasks"]) != expected: + raise ContractError("Generated task index differs from the model capture") + for name, item in batch["items"].items(): + if item["task"] != f"tasks/{name}": + raise ContractError("Model task path changed") + task = TaskPackage.load(root / item["task"]) + if task.identity != item["task_digest"] or task.identity != index["tasks"][name]: + raise ContractError("Model task changed after generation; prepare a new batch") + replay = read_json(task.root / "replay.json") + if replay["capture_identity"] != capture.identity or replay["target"] != batch["model"]: + raise ContractError("Task belongs to a different target model") + if require_preflight and ( + not batch.get("preflight_complete") + or not all(item.get("preflight_passed") for item in batch["items"].values()) + ): + raise ContractError( + "Model batch blocked: every task must pass baseline preflight before any agent starts" + ) + if require_preflight and batch.get("preflight_tool_digest") != tool_identity(): + raise ContractError("Tuning tools changed after model preflight; rerun the batch preflight") + if require_preflight: + for item in batch["items"].values(): + receipt = read_json(root / relative_path(item["preflight_receipt"])) + if ( + digest(receipt) != item["preflight_receipt_digest"] + or receipt["task_digest"] != item["task_digest"] + or receipt["tool_digest"] != batch["preflight_tool_digest"] + or receipt["environment"] != item["preflight_environment"] + ): + raise ContractError("Model preflight evidence changed; rerun the batch preflight") + + +def verify_model_task(task: TaskPackage) -> Path: + """Require a generated model task to belong to a fully preflighted model batch.""" + root = task.root.parent.parent + if not (root / "batch.json").is_file(): + raise ContractError( + "Start captured registry tasks with tune-all; standalone run cannot bypass model preflight" + ) + record = read_json(root / "batch.json") + if record.get("schema_version") != 2: + raise ContractError("Discovered tasks require a model batch") + _verify_model(root, record, require_preflight=True) + if record["items"].get(task.root.name, {}).get("task_digest") != task.identity: + raise ContractError("Task does not belong to this model batch") + return root + + +def verify_model_run(store: RunStore) -> None: + """Revalidate the entire model before an archived operator search resumes.""" + if not store.manifest.get("model_batch"): + raise ContractError("Model run has no preflighted parent batch") + root = Path(store.manifest["model_batch"]) + record = read_json(root / "batch.json") + if record.get("schema_version") != 2: + raise ContractError("Model run requires a model batch") + _verify_model(root, record, require_preflight=True) + if store.task.identity not in {item["task_digest"] for item in record["items"].values()}: + raise ContractError("Run task no longer belongs to the model batch") + item = next(item for item in record["items"].values() if item["task_digest"] == store.task.identity) + if store.manifest["environment"]["identity"] != item["preflight_environment"]["identity"]: + raise ContractError("Model evaluator environment changed; start a new model batch") + + +def _state(root: Path, item: dict[str, Any]) -> dict[str, Any] | None: + if not item["run"]: + return None + return read_json(root / item["run"] / "manifest.json") + + +def execute(root: Path, *, launcher: Launcher = launch) -> dict[str, Any]: + """Start or resume every unfinished operator in order; return the batch record.""" + batch = read_json(root / "batch.json") + _verify_model(root, batch, require_preflight=True) + for name, item in batch["items"].items(): + if item["status"] in DONE: + continue + _verify_model(root, batch, require_preflight=True) + log = root / "logs" / f"{name}.log" + command = [sys.executable, "-m", "scripts.kernel_tuning"] + if item["run"]: + command += ["resume", str(root / item["run"])] + else: + command += [ + "run", + str(root / item["task"]), + "--agent", + batch["agent"], + "--output", + str(root / "runs"), + ] + if batch.get("unfenced"): + command.append("--unfenced") + item["status"] = "running" + atomic_json(root / "batch.json", batch) + try: + code = launcher(command, log) + except KeyboardInterrupt: + code = 130 + if not item["run"] and (found := ARCHIVE.findall(log.read_text(encoding="utf-8", errors="replace"))): + item["run"] = Path(found[-1].strip()).resolve().relative_to(root).as_posix() + manifest = _state(root, item) + item["exit_code"] = code + if manifest and manifest["status"] == "completed": + item.update(status="completed", error=None) + if manifest["best"]: + _export(root, name, item) + else: + item["status"] = "interrupted" if code == 130 else "failed" + item["error"] = (manifest or {}).get("stop_reason") or f"exit {code}; see logs/{name}.log" + atomic_json(root / "batch.json", batch) + write_summary(root) + if code == 130: + raise KeyboardInterrupt + write_summary(root) + return batch + + +def _export(root: Path, name: str, item: dict[str, Any]) -> None: + destination = root / "exports" / name + if destination.exists(): + return + store = RunStore(root / item["run"]) + with run_lock(store.root): + store.export(destination) + + +def _result(root: Path, name: str, item: dict[str, Any]) -> dict[str, Any]: + manifest = _state(root, item) or {} + attempts = manifest.get("attempts", []) + row: dict[str, Any] = { + "operator": name, + "status": item["status"], + "stop_reason": manifest.get("stop_reason") or item.get("error"), + "attempts": len(attempts), + "promoted": sum(a["status"] == "promoted" for a in attempts), + "best": manifest.get("best"), + "run": item["run"], + "export": f"exports/{name}" if (root / "exports" / name).is_dir() else None, + } + best = next((a for a in attempts if a["id"] == manifest.get("best")), None) + if best: + rounds = best["decision"]["rounds"] + # The weakest paired round is the improvement the evidence supports. + row["improvement_vs_baseline"] = min(r["improvement"]["baseline"] for r in rounds) + row["baseline_ms"] = rounds[-1]["mean_latency_ms"]["baseline"] + row["best_ms"] = rounds[-1]["mean_latency_ms"]["candidate"] + return row + + +def write_summary(root: Path) -> list[dict[str, Any]]: + """Write summary.json and a human-readable summary.md for the whole batch.""" + batch = read_json(root / "batch.json") + rows = [_result(root, name, item) for name, item in batch["items"].items()] + atomic_json(root / "summary.json", rows) + lines = [ + "# Kernel tuning batch", + "", + f"Agent: `{batch['agent']}`. Improvements are weighted mean latency against the production", + "baseline, from the weakest paired round. Exported kernels still need review and", + "integration before runtime use.", + *(["", "Agents ran **unfenced** (no Humanize2 Landlock fence)."] if batch.get("unfenced") else []), + "", + "| operator | status | attempts | promoted | improvement | baseline ms | best ms | detail |", + "|---|---|---|---|---|---|---|---|", + ] + if batch.get("model"): + coverage = read_json(root / "capture/coverage.json")["operators"] + lines[2:2] = [ + f"Model: `{batch['model']['model']}`. Capture: `{batch['capture_identity']}`.", + f"Coverage: `{coverage}`. All-baseline preflight passed: `{batch['preflight_complete']}`.", + "", + ] + for row in rows: + improvement = f"{row['improvement_vs_baseline']:.1%}" if "improvement_vs_baseline" in row else "-" + baseline = f"{row['baseline_ms']:.4f}" if "baseline_ms" in row else "-" + best = f"{row['best_ms']:.4f}" if "best_ms" in row else "-" + detail = row["export"] or row["stop_reason"] or "" + lines.append( + f"| {row['operator']} | {row['status']} | {row['attempts']} | {row['promoted']} | " + f"{improvement} | {baseline} | {best} | {str(detail).replace('|', '/')} |" + ) + (root / "summary.md").write_text("\n".join(lines) + "\n", encoding="utf-8", newline="\n") + return rows diff --git a/scripts/kernel_tuning/capture.py b/scripts/kernel_tuning/capture.py new file mode 100644 index 0000000..3ef2f49 --- /dev/null +++ b/scripts/kernel_tuning/capture.py @@ -0,0 +1,394 @@ +"""Record the shapes of production kernel calls while a model actually runs. + +``python -m scripts.kernel_tuning capture --catalog --output calls.jsonl -- SCRIPT [ARGS]`` +executes an unmodified inference or benchmark script in this interpreter with +each catalog operator's production function wrapped. Every call is converted to +the operator's task axes, scalars, and fixed integer inputs, so captured traffic +can replace the catalog's estimated workloads (``generate --captured``). + +Run it in the model's own runtime with its production configuration. Calls +captured into a CUDA Graph are counted on every replay. Disable compilation: +compiled code may call kernels through traced graphs this cannot observe. +Shapes depend on the configuration and inputs, never on the weight values. +""" + +from __future__ import annotations + +import dataclasses +import functools +import importlib +import inspect +import json +import re +import runpy +import sys +from collections import Counter +from collections.abc import Callable +from pathlib import Path +from typing import Any + +from .artifacts import REPOSITORY +from .contracts import ContractError +from .operators import Operator, Workload, catalog + +Bound = dict[str, Any] + + +class Skip(Exception): + """A call outside the task's semantics; it is counted but yields no workload.""" + + +def _bf16(*tensors: Any) -> None: + for tensor in tensors: + if str(tensor.dtype) != "torch.bfloat16": + raise Skip(f"dtype {tensor.dtype}") + + +def _ada_rms_norm(a: Bound) -> dict[str, Any]: + x, modulation = a["x"], a["modulation"] + _bf16(x, modulation) + batch, tokens, hidden = x.shape + return { + "axes": {"B": batch, "T": tokens, "H": hidden, "H3": modulation.shape[1]}, + "scalars": {"eps": a["eps"]}, + } + + +def _gated_residual(a: Bound) -> dict[str, Any]: + _bf16(a["residual"], a["update"], a["gate"]) + batch, tokens, hidden = a["residual"].shape + return {"axes": {"B": batch, "T": tokens, "H": hidden, "G": a["gate"].shape[1]}} + + +def _gated_gelu(a: Bound) -> dict[str, Any]: + _bf16(a["gate"], a["up"]) + batch, tokens, width = a["gate"].shape + return {"axes": {"B": batch, "T": tokens, "I": width}} + + +def _rotate_qk(a: Bound) -> dict[str, Any]: + _bf16(a["query"], a["key"], a["cos"], a["sin"]) + batch, heads, size, width = a["query"].shape + return {"axes": {"B": batch, "HQ": heads, "HK": a["key"].shape[1], "S": size, "D": width}} + + +def _split_kv_attention(a: Bound) -> dict[str, Any]: + # A padding mask keeps the task's shapes and work; it only changes which keys count. + _bf16(a["query"], a["prefix_key"]) + batch, heads, size, width = a["query"].shape + axes = {"B": batch, "HQ": heads, "HK": a["prefix_key"].shape[1], "P": a["prefix_key"].shape[2]} + return {"axes": {**axes, "S": size, "D": width}, "scalars": {"scaling": float(a["scaling"])}} + + +def _rotate_half_rope(a: Bound) -> dict[str, Any]: + q, k, cos = a["q"], a["k"], a["cos"] + _bf16(q, k, cos, a["sin"]) + if cos.shape[0] != q.shape[0]: + raise Skip("broadcast cos/sin rows") + tokens, heads, width = q.shape + return {"axes": {"T": tokens, "HQ": heads, "HK": k.shape[1], "D": width}} + + +def _segmented_attention(a: Bound) -> dict[str, Any]: + q, k, offsets = a["q"], a["k"], a["q_segment_offsets"] + _bf16(q, k, a["v"]) + if a["kv_segment_offsets"] is not None and a["kv_segment_offsets"] is not offsets: + raise Skip("separate key/value segments") + if k.shape[0] != q.shape[0]: + raise Skip("cross-length attention") + if a["scaling"] is not None and abs(a["scaling"] - q.shape[-1] ** -0.5) > 1e-12: + raise Skip("non-default scaling") + if str(offsets.dtype) != "torch.int32": + raise Skip(f"offset dtype {offsets.dtype}") + key = (offsets.data_ptr(), offsets.numel()) if hasattr(offsets, "data_ptr") else None + if _capturing(q): + # Production warms each layout eagerly before capture; reuse the values read then. + if key not in _OFFSETS: + raise Skip("offsets unreadable during CUDA Graph capture") + bounds = _OFFSETS[key] + else: + bounds = [int(x) for x in offsets.tolist()] + if key is not None: + _OFFSETS[key] = bounds + longest = max(end - start for start, end in zip(bounds, bounds[1:])) + bound = a["max_query_length"] or longest + if (a["max_key_length"] or longest) != bound: + raise Skip("different query and key bounds") + tokens, heads, width = q.shape + return { + "axes": {"T": tokens, "HQ": heads, "HK": k.shape[1], "D": width, "N": len(bounds)}, + "scalars": {"max_length": int(bound)}, + "fixed": {"segment_offsets": bounds}, + } + + +def _rms_norm(a: Bound) -> dict[str, Any]: + _bf16(a["inputs"]) + width = a["inputs"].shape[-1] + return {"axes": {"M": a["inputs"].numel() // width, "H": width}, "scalars": {"eps": float(a["epsilon"])}} + + +def _add_rms_norm(a: Bound) -> dict[str, Any]: + _bf16(a["residual"], a["update"]) + width = a["residual"].shape[-1] + return { + "axes": {"M": a["residual"].numel() // width, "H": width}, + "scalars": {"eps": float(a["epsilon"])}, + } + + +def _swiglu(a: Bound) -> dict[str, Any]: + _bf16(a["packed"]) + width = a["packed"].shape[-1] + return {"axes": {"M": a["packed"].numel() // width, "I": width // 2}} + + +#: Production arguments -> task axes/scalars/fixed inputs, keyed by catalog name. +MAPPINGS: dict[str, Callable[[Bound], dict[str, Any]]] = { + "ada_rms_norm": _ada_rms_norm, + "gated_residual": _gated_residual, + "gated_gelu": _gated_gelu, + "rotate_qk": _rotate_qk, + "split_kv_attention": _split_kv_attention, + "rotate_half_rope": _rotate_half_rope, + "segmented_attention": _segmented_attention, + "rms_norm": _rms_norm, + "add_rms_norm": _add_rms_norm, + "swiglu": _swiglu, +} + + +#: Segment offsets seen in eager calls, keyed by device address, for calls made during capture. +_OFFSETS: dict[tuple[int, int], list[int]] = {} + + +def _capturing(tensor: Any) -> bool: + if not getattr(tensor, "is_cuda", False): + return False + import torch + + return torch.cuda.is_current_stream_capturing() + + +class Recorder: + """Convert production calls to task cases; counts identical cases instead of storing calls. + + Kernels launched while a CUDA Graph is captured run on every replay without + reaching Python again. Their cases are attached to the graph during capture + and counted once per replay, so counts follow execution with or without graphs. + """ + + def __init__(self) -> None: + self.cases: Counter[str] = Counter() + self.skipped: Counter[tuple[str, str]] = Counter() + self.graphs: dict[int, Counter[str]] = {} + self._capturing: list[Counter[str]] = [] + + def record(self, name: str, bound: Bound) -> None: + """Count one call. A mapping failure must never break the model being measured.""" + try: + case = MAPPINGS[name](bound) + except Skip as exc: + self.skipped[name, str(exc)] += 1 + return + except Exception as exc: # noqa: BLE001 -- recording is best effort by design + self.skipped[name, f"unmapped call: {type(exc).__name__}: {exc}"] += 1 + return + key = json.dumps({"operator": name, **case}, sort_keys=True) + (self._capturing[-1] if self._capturing else self.cases)[key] += 1 + + def begin_graph(self) -> None: + """Start attaching recorded calls to the graph being captured.""" + self._capturing.append(Counter()) + + def end_graph(self, graph: int) -> None: + """Bind the calls recorded since :meth:`begin_graph` to ``graph``.""" + self.graphs[graph] = self._capturing.pop() + + def replay(self, graph: int) -> None: + """Count one execution of every call captured in ``graph``.""" + self.cases.update(self.graphs.get(graph, Counter())) + + def rows(self) -> list[dict[str, Any]]: + """One row per distinct case with its call count, plus counted skips.""" + rows = [{**json.loads(key), "count": count} for key, count in self.cases.most_common()] + rows += [ + {"operator": name, "skipped": reason, "count": count} + for (name, reason), count in self.skipped.most_common() + ] + return rows + + +def wrap(name: str, function: Callable[..., Any], recorder: Recorder) -> Callable[..., Any]: + """Return a drop-in replacement that records each call before running the original.""" + signature = inspect.signature(function) + + @functools.wraps(function) + def recorded(*args: Any, **kwargs: Any) -> Any: + bound = signature.bind(*args, **kwargs) + bound.apply_defaults() + recorder.record(name, dict(bound.arguments)) + return function(*args, **kwargs) + + recorded.__wrapped_kernel__ = function # type: ignore[attr-defined] + return recorded + + +def install(operators: list[Operator], recorder: Recorder) -> list[tuple[Any, str, Any]]: + """Wrap each production function wherever an imported module already holds it. + + Lazy ``from ... import`` inside model code reads the defining module at call + time; modules importing the backend package later receive the wrapper too. + Returns what was replaced so the caller can restore it. + """ + replaced = [] + for operator in operators: + module = importlib.import_module(operator.source.removesuffix(".py").replace("/", ".")) + original = getattr(module, operator.entry) + wrapper = wrap(operator.name, original, recorder) + for holder in list(sys.modules.values()): + for attribute, value in list(getattr(holder, "__dict__", {}).items()): + if value is original: + replaced.append((holder, attribute, original)) + setattr(holder, attribute, wrapper) + return replaced + + +def install_graph_hooks(recorder: Recorder) -> list[tuple[Any, str, Any]]: + """Attribute captured kernel calls to CUDA Graphs and count them on each replay.""" + try: + import torch + except ImportError: + return [] + graph_class = torch.cuda.CUDAGraph + begin, end, replay = graph_class.capture_begin, graph_class.capture_end, graph_class.replay + + def capture_begin(self: Any, *args: Any, **kwargs: Any) -> Any: + recorder.begin_graph() + return begin(self, *args, **kwargs) + + def capture_end(self: Any, *args: Any, **kwargs: Any) -> Any: + try: + return end(self, *args, **kwargs) + finally: + recorder.end_graph(id(self)) + + def replayed(self: Any, *args: Any, **kwargs: Any) -> Any: + recorder.replay(id(self)) + return replay(self, *args, **kwargs) + + graph_class.capture_begin, graph_class.capture_end, graph_class.replay = ( + capture_begin, + capture_end, + replayed, + ) + return [ + (graph_class, "capture_begin", begin), + (graph_class, "capture_end", end), + (graph_class, "replay", replay), + ] + + +def run(command: list[str], output: Path, operators: list[Operator]) -> list[dict[str, Any]]: + """Run ``SCRIPT [ARGS]`` or ``-m MODULE [ARGS]`` as ``__main__`` and write the captured cases.""" + if not command: + raise ContractError("Give the script to run after --, e.g. -- benchmark.py --config config.yaml") + recorder = Recorder() + _OFFSETS.clear() + replaced = install(operators, recorder) + install_graph_hooks(recorder) + argv = sys.argv + try: + if command[0] == "-m": + sys.argv = command[1:] + runpy.run_module(command[1], run_name="__main__", alter_sys=True) + else: + script = Path(command[0]).resolve() + sys.argv = [str(script), *command[1:]] + sys.path.insert(0, str(script.parent)) + runpy.run_path(str(script), run_name="__main__") + except SystemExit as exc: + if exc.code not in (None, 0): + raise + finally: + sys.argv = argv + for holder, attribute, original in replaced: + setattr(holder, attribute, original) + rows = recorder.rows() + output.parent.mkdir(parents=True, exist_ok=True) + output.write_text("".join(json.dumps(row) + "\n" for row in rows), encoding="utf-8") + return rows + + +#: Captures committed with the repository; they replace catalog estimates by default. +COMMITTED = REPOSITORY / "benchmarks/kernel_tuning/captures" + + +def load(path: Path) -> list[dict[str, Any]]: + """Read a capture file written by :func:`run`, or every ``*.jsonl`` in a directory.""" + files = sorted(path.glob("*.jsonl")) if path.is_dir() else [path] + rows = [ + json.loads(line) + for file in files + for line in file.read_text(encoding="utf-8").splitlines() + if line.strip() + ] + if not rows: + raise ContractError(f"{path} contains no captured calls") + return rows + + +def apply(operator: Operator, rows: list[dict[str, Any]], *, top: int = 8) -> tuple[Operator, str]: + """Replace an operator's estimated workloads with its most frequent captured cases. + + Weights are call counts, so the promotion criterion follows real traffic. + Returns the operator unchanged when the capture saw no usable call of it. + """ + cases = [row for row in rows if row["operator"] == operator.name and "axes" in row] + skipped = sum(row["count"] for row in rows if row["operator"] == operator.name and "skipped" in row) + if not cases: + return operator, f"{operator.name}: no captured calls ({skipped} skipped); keeping estimated shapes" + cases.sort(key=lambda row: -row["count"]) + kept, total = cases[:top], sum(row["count"] for row in cases) + variable = [name for name, value in operator.axes.items() if value is None] + workloads = [] + for index, row in enumerate(kept): + if set(row["axes"]) != set(variable): + raise ContractError( + f"Capture for {operator.name} binds {sorted(row['axes'])}, expected {variable}" + ) + label = "-".join(f"{name}{row['axes'][name]}" for name in variable) + workloads.append( + Workload( + re.sub(r"[^A-Za-z0-9-]", "", f"captured{index + 1}-{label}")[:64], + row["axes"], + scalars=row.get("scalars", {}), + fixed={name: tuple(values) for name, values in row.get("fixed", {}).items()}, + weight=float(row["count"]), + note=f"{row['count']} of {total} captured calls", + ) + ) + share = sum(row["count"] for row in kept) / total + summary = f"{operator.name}: {len(kept)} captured shapes covering {share:.0%} of {total} calls" + if skipped: + summary += f"; {skipped} calls outside the task skipped" + return dataclasses.replace(operator, workloads=tuple(workloads)), summary + + +def apply_file( + operators: list[Operator], path: Path | None = None, *, estimated: bool = False +) -> tuple[list[Operator], list[str]]: + """Apply captured shapes: ``path``, else the committed captures, unless ``estimated``.""" + if estimated: + return operators, ["using the catalog's estimated shapes"] + if path is None: + if not COMMITTED.is_dir() or not any(COMMITTED.glob("*.jsonl")): + return operators, [] + path = COMMITTED + rows = load(path) + if unknown := {row["operator"] for row in rows} - catalog().keys(): + raise ContractError(f"Capture names operators outside the catalog: {sorted(unknown)}") + applied = [apply(operator, rows) for operator in operators] + return [operator for operator, _ in applied], [summary for _, summary in applied] + + +__all__ = ["MAPPINGS", "Recorder", "apply", "apply_file", "install", "load", "run", "wrap"] diff --git a/scripts/kernel_tuning/contracts.py b/scripts/kernel_tuning/contracts.py new file mode 100644 index 0000000..4d41bc9 --- /dev/null +++ b/scripts/kernel_tuning/contracts.py @@ -0,0 +1,548 @@ +"""Frozen task packages, portable solutions, and promotion from measured evidence. + +This module intentionally imports neither Humanize2 nor GPU libraries. Official +FlashInfer schemas are validated again by the default evaluator before building. +""" + +from __future__ import annotations + +import ast +import hashlib +import json +import math +from dataclasses import asdict, dataclass, field, fields +from pathlib import Path, PurePosixPath, PureWindowsPath +from typing import Any, TypeVar + +HMZ_REVISION = "b70d4427f11308442aef20f9b2233dfbbc34a39e" +FLASHINFER_VERSION = "0.1.2" +ROLES = ("baseline", "incumbent", "candidate") + + +class ContractError(ValueError): + """A task, solution, or measurement does not satisfy the frozen contract.""" + + +def read_json(path: Path) -> Any: + """Read JSON while rejecting non-finite numbers and duplicate keys.""" + return parse_json(path.read_text(encoding="utf-8")) + + +def parse_json(text: str) -> Any: + """Parse strict JSON, also used for individual official Trace JSONL records.""" + + def pairs(items: list[tuple[str, Any]]) -> dict[str, Any]: + result: dict[str, Any] = {} + for key, value in items: + if key in result: + raise ContractError(f"Duplicate JSON key: {key}") + result[key] = value + return result + + def invalid(value: str) -> None: + raise ContractError(f"Non-finite JSON number: {value}") + + return json.loads(text, object_pairs_hook=pairs, parse_constant=invalid) + + +def digest(value: Any) -> str: + """Hash JSON data independently of whitespace and dictionary order.""" + return hashlib.sha256(json.dumps(value, sort_keys=True, allow_nan=False).encode()).hexdigest() + + +def relative_path(value: str) -> str: + """Validate a portable relative file path before joining it to any workspace.""" + if not isinstance(value, str) or not value or "\\" in value or ":" in value: + raise ContractError(f"Expected a portable relative path: {value!r}") + path = PurePosixPath(value) + if ( + path.is_absolute() + or PureWindowsPath(value).is_absolute() + or any(part in ("", ".", "..") for part in value.split("/")) + ): + raise ContractError(f"Unsafe relative path: {value!r}") + reserved = {"CON", "PRN", "AUX", "NUL", "CONIN$", "CONOUT$"} + reserved.update(f"{prefix}{index}" for prefix in ("COM", "LPT") for index in "123456789¹²³") + if any(part.split(".")[0].upper() in reserved or part.endswith((".", " ")) for part in path.parts): + raise ContractError(f"Non-portable path: {value!r}") + return value + + +def tree_hashes(root: Path) -> dict[str, str]: + """Hash every task/artifact file, refusing symlinks and generated Python caches.""" + result = {} + for path in sorted(root.rglob("*")): + if path.is_symlink(): + raise ContractError(f"Symlinks are not portable task inputs: {path}") + if "__pycache__" in path.parts or path.suffix == ".pyc": + continue + if path.is_file(): + name = relative_path(path.relative_to(root).as_posix()) + result[name] = hashlib.sha256(path.read_bytes()).hexdigest() + return result + + +def _positive(value: Any, name: str, *, integer: bool = False, zero: bool = False) -> None: + if isinstance(value, bool) or not isinstance(value, int if integer else (int, float)): + raise ContractError(f"{name} must be {'an integer' if integer else 'a number'}") + if not math.isfinite(value) or value < 0 or (value == 0 and not zero): + raise ContractError(f"{name} must be finite and {'nonnegative' if zero else 'positive'}") + + +@dataclass(frozen=True) +class Precision: + """Numerical conditions fixed before the first agent turn.""" + + mode: str = "bit_exact" + atol: float = 0.0 + rtol: float = 0.0 + reason: str = "" + matmul_precision: str = "highest" + allow_tf32: bool = False + + def __post_init__(self) -> None: + if self.mode not in ("bit_exact", "tolerance"): + raise ContractError("precision.mode must be bit_exact or tolerance") + for key in ("atol", "rtol"): + _positive(getattr(self, key), f"precision.{key}", zero=True) + if self.mode == "bit_exact" and (self.atol or self.rtol): + raise ContractError("bit_exact cannot declare nonzero tolerances") + if self.mode == "tolerance" and (not self.reason.strip() or not (self.atol or self.rtol)): + raise ContractError("tolerance requires a reason and a nonzero atol or rtol") + if self.matmul_precision not in ("highest", "high", "medium") or type(self.allow_tf32) is not bool: + raise ContractError("Invalid matmul precision or allow_tf32") + if self.allow_tf32 != (self.matmul_precision != "highest"): + raise ContractError("Use highest with allow_tf32: false, or high/medium with allow_tf32: true") + + +@dataclass(frozen=True) +class Timing: + """Timing protocol; each paired round must independently pass promotion.""" + + mode: str = "eager" + warmup: int = 10 + iterations: int = 100 + trials: int = 3 + paired_rounds: int = 2 + + def __post_init__(self) -> None: + if self.mode not in ("eager", "cuda_graph"): + raise ContractError("timing.mode must be eager or cuda_graph") + for key in ("warmup", "iterations", "trials", "paired_rounds"): + _positive(getattr(self, key), f"timing.{key}", integer=True) + if self.paired_rounds < 2: + raise ContractError("Promotion requires at least two paired measurement rounds") + + +@dataclass(frozen=True) +class Search: + """Whole-run and per-operation limits; failures consume attempts and patience.""" + + max_candidates: int = 20 + max_seconds: float = 7200 + patience: int = 5 + agent_timeout_seconds: float = 600 + evaluation_timeout_seconds: float = 300 + + def __post_init__(self) -> None: + for key in ("max_candidates", "patience"): + _positive(getattr(self, key), f"search.{key}", integer=True) + for key in ("max_seconds", "agent_timeout_seconds", "evaluation_timeout_seconds"): + _positive(getattr(self, key), f"search.{key}") + + +@dataclass(frozen=True) +class Acceptance: + """Latency reduction and workload regression limits, expressed as fractions.""" + + min_improvement: float = 0.03 + max_regression: float = 0.05 + + def __post_init__(self) -> None: + _positive(self.min_improvement, "acceptance.min_improvement") + _positive(self.max_regression, "acceptance.max_regression", zero=True) + if self.min_improvement >= 1 or self.max_regression >= 1: + raise ContractError("Acceptance fractions must be less than 1") + + +T = TypeVar("T") + + +def _settings(cls: type[T], data: Any) -> T: + if not isinstance(data, dict) or set(data) - {f.name for f in fields(cls)}: + raise ContractError(f"Unknown or malformed {cls.__name__} settings: {data!r}") + try: + return cls(**data) + except (TypeError, AttributeError) as exc: + raise ContractError(f"Invalid {cls.__name__}: {exc}") from exc + + +@dataclass(frozen=True) +class Settings: + """Project settings beside, rather than inside, the official task schemas.""" + + language: str + schema_version: int = 1 + agent: str = "" + evaluator_python: str = "python" + device: str = "cuda:0" + target_hardware: tuple[str, ...] = ("cuda",) + seeds: tuple[int, ...] = (0, 1, 2) + weights: dict[str, float] = field(default_factory=dict) + precision: Precision = field(default_factory=Precision) + timing: Timing = field(default_factory=Timing) + search: Search = field(default_factory=Search) + acceptance: Acceptance = field(default_factory=Acceptance) + profile: bool = False + + @classmethod + def load(cls, path: Path) -> Settings: + """Parse YAML with unknown settings rejected instead of silently ignored.""" + import yaml + + class UniqueLoader(yaml.SafeLoader): + pass + + def mapping(loader: Any, node: Any) -> dict[str, Any]: + pairs = loader.construct_pairs(node, deep=True) + if len({key for key, _ in pairs}) != len(pairs): + raise ContractError("Duplicate YAML setting") + return dict(pairs) + + UniqueLoader.add_constructor(yaml.resolver.BaseResolver.DEFAULT_MAPPING_TAG, mapping) + try: + data = yaml.load(path.read_text(encoding="utf-8"), Loader=UniqueLoader) + except yaml.YAMLError as exc: + raise ContractError(f"Invalid tuning.yaml: {exc}") from exc + return cls.load_mapping(data) + + @classmethod + def load_mapping(cls, data: Any) -> Settings: + """Validate parsed tuning.yaml data, e.g. settings rendered by the task generator.""" + if not isinstance(data, dict): + raise ContractError("tuning.yaml must contain a mapping") + data = dict(data) + for key, kind in ( + ("precision", Precision), + ("timing", Timing), + ("search", Search), + ("acceptance", Acceptance), + ): + data[key] = _settings(kind, data.get(key, {})) + for key in ("seeds", "target_hardware"): + if key in data: + if not isinstance(data[key], list): + raise ContractError(f"{key} must be a list") + data[key] = tuple(data[key]) + return _settings(cls, data) + + def __post_init__(self) -> None: + if self.language not in ("triton", "cuda") or self.schema_version != 1: + raise ContractError("Use schema_version: 1 and language: triton or cuda") + for key in ("evaluator_python", "device"): + if not isinstance(getattr(self, key), str) or not getattr(self, key).strip(): + raise ContractError(f"{key} must be a nonempty string") + if not isinstance(self.agent, str) or type(self.profile) is not bool: + raise ContractError("agent must be a string; profile must be a boolean") + if not self.target_hardware or any(not isinstance(x, str) or not x for x in self.target_hardware): + raise ContractError("target_hardware must list at least one hardware label") + if not self.seeds or len(set(self.seeds)) != len(self.seeds): + raise ContractError("seeds must be nonempty and unique") + for seed in self.seeds: + _positive(seed, "seed", integer=True, zero=True) + if seed >= 2**32: + raise ContractError("Seeds must fit uint32") + if not isinstance(self.weights, dict): + raise ContractError("weights must map workload UUIDs to positive weights") + for value in self.weights.values(): + _positive(value, "workload weight") + + +@dataclass(frozen=True) +class TaskPackage: + """A self-contained task snapshot with a content identity and declared workloads.""" + + root: Path + definition: dict[str, Any] + workloads: tuple[dict[str, Any], ...] + settings: Settings + hashes: dict[str, str] + + @property + def identity(self) -> str: + """Identity of the complete task, including data files and custom evaluator.""" + return digest(self.hashes) + + @property + def workload_ids(self) -> tuple[str, ...]: + """Stable measurement keys in the task's declared order.""" + return tuple(t["workload"]["uuid"] for t in self.workloads) + + @classmethod + def load(cls, root: Path) -> TaskPackage: + """Validate the task structurally without importing or executing its code.""" + try: + return cls._load(root) + except (KeyError, TypeError, AttributeError, SyntaxError) as exc: + raise ContractError(f"Malformed task package: {exc}") from exc + + @classmethod + def _load(cls, root: Path) -> TaskPackage: + root = root.resolve(strict=True) + hashes = tree_hashes(root) + required = {"README.md", "definition.json", "workloads.jsonl", "baseline.py", "tuning.yaml"} + if missing := required - hashes.keys(): + raise ContractError(f"Missing task files: {sorted(missing)}") + definition = read_json(root / "definition.json") + for key in ("name", "op_type", "reference"): + if not isinstance(definition.get(key), str) or not definition[key]: + raise ContractError(f"Definition requires {key}") + axes = definition.get("axes") + if not isinstance(axes, dict): + raise ContractError("Definition axes must be a mapping") + for axis in axes.values(): + if axis.get("type") not in ("var", "const"): + raise ContractError("Axis must be var or const") + if axis["type"] == "const": + _positive(axis.get("value"), "constant axis", integer=True, zero=True) + for kind in ("inputs", "outputs"): + if not isinstance(definition.get(kind), dict) or ( + not definition[kind] and not (kind == "inputs" and "replay.json" in hashes) + ): + raise ContractError(f"Definition requires {kind}") + for spec in definition[kind].values(): + shape = spec.get("shape") + if ( + shape is not None and (not isinstance(shape, list) or set(shape) - axes.keys()) + ) or not spec.get("dtype"): + raise ContractError(f"Invalid {kind} tensor spec") + bound = {axis for spec in definition["inputs"].values() for axis in spec.get("shape") or ()} + if ( + unbound := {name for name, axis in axes.items() if axis["type"] == "var"} - bound + ) and "replay.json" not in hashes: + # FlashInfer infers variable axes from input shapes when sizing outputs. + raise ContractError(f"Variable axes must appear in an input shape: {sorted(unbound)}") + for source in (definition["reference"], (root / "baseline.py").read_text(encoding="utf-8")): + functions = [ + n for n in ast.parse(source).body if isinstance(n, ast.FunctionDef) and n.name == "run" + ] + if len(functions) != 1: + raise ContractError("Reference and baseline must define run") + args = functions[0].args + if [a.arg for a in args.posonlyargs + args.args] != list(definition["inputs"]): + raise ContractError("run argument order must match definition.inputs") + traces = [] + for line in (root / "workloads.jsonl").read_text(encoding="utf-8").splitlines(): + if not line.strip(): + continue + trace = parse_json(line) + work = trace.get("workload", {}) + if ( + trace.get("definition") != definition["name"] + or not isinstance(work.get("uuid"), str) + or not work["uuid"] + ): + raise ContractError("Workload must identify this definition and have a UUID") + if trace.get("solution") is not None or trace.get("evaluation") is not None: + raise ContractError("Input workloads must be unevaluated traces") + variables = {name for name, axis in axes.items() if axis["type"] == "var"} + if not isinstance(work.get("axes"), dict) or set(work["axes"]) != variables: + raise ContractError("Workload axes must bind exactly the variable axes") + for value in work["axes"].values(): + _positive(value, "workload axis", integer=True, zero=True) + if not isinstance(work.get("inputs"), dict) or set(work["inputs"]) - definition["inputs"].keys(): + raise ContractError("Invalid workload inputs") + for spec in work["inputs"].values(): + if spec.get("type") not in ("random", "scalar", "safetensors"): + raise ContractError("Unsupported workload input source") + if spec["type"] == "scalar": + value = spec.get("value") + if not isinstance(value, (bool, int, float)) or not math.isfinite(value): + raise ContractError("Scalar input must contain a finite numeric value") + if spec["type"] == "safetensors" and relative_path(spec["path"]) not in hashes: + raise ContractError("Safetensors data must be included inside the task package") + traces.append(trace) + settings = Settings.load(root / "tuning.yaml") + ids = [t["workload"]["uuid"] for t in traces] + if not ids or len(ids) != len(set(ids)): + raise ContractError("Workloads must be nonempty and have distinct UUIDs") + if settings.weights and set(settings.weights) != set(ids): + raise ContractError("Explicit weights must cover exactly every workload UUID") + if "replay.json" in hashes: + from .discovery.contracts import validate_replay + + validate_replay(read_json(root / "replay.json"), definition, traces, hashes) + if "benchmark.py" not in hashes: + raise ContractError("Captured layouts require the replay benchmark adapter") + replay = read_json(root / "replay.json") + if "numerical_identity" in replay or "numerics.json" in hashes: + from .discovery.numerics import SEEDS, task_precision, validate_contract + + if "numerics.json" not in hashes or "numerical_identity" not in replay: + raise ContractError("Missing calibrated numerical contract or operator identity") + numerical = read_json(root / "numerics.json") + validate_contract( + numerical, + replay["numerical_identity"], + replay["cases"], + asdict(settings.precision), + replay["flags"], + ) + if ( + asdict(settings.precision) != task_precision(numerical, asdict(settings.precision)) + or settings.seeds != SEEDS + ): + raise ContractError("Frozen calibration precision/seeds changed") + elif "numerics.json" in hashes: + raise ContractError("Numerical calibration requires a model replay task") + return cls(root, definition, tuple(traces), settings, hashes) + + def verify(self) -> None: + """Refuse a changed frozen task before trusting another measurement.""" + if tree_hashes(self.root) != self.hashes: + raise ContractError("Frozen task changed during tuning") + + def baseline(self) -> dict[str, Any]: + """Represent baseline.py and its helper sources as an official Solution.""" + return { + "name": f"{self.definition['name']}_baseline", + "definition": self.definition["name"], + "author": "task-owner", + "spec": { + "language": "python", + "target_hardware": list(self.settings.target_hardware), + "entry_point": "baseline.py::run", + "destination_passing_style": False, + }, + "sources": [ + {"path": path, "content": (self.root / path).read_text(encoding="utf-8")} + for path in self.hashes + if Path(path).suffix in (".py", ".cu", ".cuh", ".h", ".hpp", ".cpp") + and path != "benchmark.py" + ], + } + + +def validate_solution(solution: Any, task: TaskPackage) -> dict[str, Any]: + """Validate candidate identity and source paths before any filesystem write/build.""" + if not isinstance(solution, dict) or set(solution) - { + "name", + "definition", + "author", + "description", + "spec", + "sources", + }: + raise ContractError("Expected an official FlashInfer Solution object") + for key in ("name", "author"): + if not isinstance(solution.get(key), str) or not solution[key]: + raise ContractError(f"Solution requires {key}") + if solution.get("definition") != task.definition["name"]: + raise ContractError("Candidate solves a different definition") + spec = solution.get("spec", {}) + if spec.get("language") != task.settings.language: + raise ContractError("Candidate language differs from the frozen task") + if spec.get("target_hardware") != list(task.settings.target_hardware): + raise ContractError("Candidate target_hardware differs from the frozen task") + if type(spec.get("destination_passing_style")) is not bool: + raise ContractError("Declare destination_passing_style explicitly") + if "replay.json" in task.hashes and spec["destination_passing_style"]: + raise ContractError("Discovered tasks require return-value style") + if spec.get("binding") not in (None, "torch", "tvm-ffi"): + raise ContractError("CUDA binding must be torch or tvm-ffi") + entry = spec.get("entry_point", "").split("::") + if len(entry) != 2 or not entry[1].isidentifier(): + raise ContractError("entry_point must be file::symbol") + relative_path(entry[0]) + sources = solution.get("sources") + if not isinstance(sources, list) or not sources: + raise ContractError("Solution requires sources") + paths = [] + for source in sources: + path = relative_path(source["path"]) + if not isinstance(source.get("content"), str) or not source["content"]: + raise ContractError("Source contents must be nonempty UTF-8 text") + paths.append(path) + if len({p.casefold() for p in paths}) != len(paths) or entry[0] not in paths: + raise ContractError("Source paths must be distinct and include the entry point") + for path in paths: + if any(parent.as_posix() in paths for parent in PurePosixPath(path).parents if str(parent) != "."): + raise ContractError("Source file/directory collision") + return solution + + +def validate_measurement(task: TaskPackage, report: dict[str, Any]) -> None: + """Require complete correctness evidence and finite timings for every role/case.""" + try: + _validate_measurement(task, report) + except (KeyError, TypeError, AttributeError) as exc: + raise ContractError(f"Malformed measurement report: {exc}") from exc + + +def _validate_measurement(task: TaskPackage, report: dict[str, Any]) -> None: + cfg = task.settings + expected = {(work, role) for work in task.workload_ids for role in ROLES} + seen = set() + for check in report.get("checks", []): + key = (check.get("workload"), check.get("role")) + if key not in expected or key in seen or check.get("passed") is not True: + raise ContractError(f"Failed, unexpected, or duplicate correctness check: {key}") + seen.add(key) + if check.get("seeds") != list(cfg.seeds) or check.get("precision") != asdict(cfg.precision): + raise ContractError("Correctness evidence used a different precision/seed contract") + if "numerics.json" in task.hashes: + from .discovery.numerics import PROFILES + + if check.get("numerical_contract") != task.hashes["numerics.json"] or check.get( + "validation_profiles" + ) != [list(pair) for pair in PROFILES]: + raise ContractError("Missing frozen calibration/profile correctness evidence") + if cfg.timing.mode == "cuda_graph" and check.get("graph_replay") is not True: + raise ContractError("Missing changed-input graph replay check") + for field_name in ("max_abs_error", "max_rel_error"): + _positive(check.get(field_name), field_name, zero=True) + if cfg.precision.mode == "bit_exact" and check[field_name] != 0: + raise ContractError("bit_exact report contains nonzero error") + if seen != expected: + raise ContractError("Incomplete correctness evidence") + rounds = report.get("rounds", []) + if len(rounds) != cfg.timing.paired_rounds: + raise ContractError("Missing paired remeasurement rounds") + for measurements in rounds: + if set(measurements) != set(task.workload_ids): + raise ContractError("Missing or unexpected timed workloads") + for row in measurements.values(): + if set(row) != set(ROLES): + raise ContractError("Missing baseline/incumbent/candidate measurements") + for samples in row.values(): + if not isinstance(samples, list) or len(samples) != cfg.timing.trials: + raise ContractError("Missing timing trials") + for value in samples: + _positive(value, "latency_ms") + + +def promotion(task: TaskPackage, report: dict[str, Any]) -> dict[str, Any]: + """Decide solely from fixed thresholds, requiring every paired round to pass.""" + validate_measurement(task, report) + cfg = task.settings + weights = cfg.weights or dict.fromkeys(task.workload_ids, 1.0) + decisions = [] + reasons = [] + for index, measurements in enumerate(report["rounds"]): + means = { + work: {role: math.fsum(values) / len(values) for role, values in row.items()} + for work, row in measurements.items() + } + aggregate = { + role: math.fsum(weights[w] * means[w][role] for w in means) / math.fsum(weights.values()) + for role in ROLES + } + improvements = {role: 1 - aggregate["candidate"] / aggregate[role] for role in ROLES[:2]} + decisions.append({"mean_latency_ms": aggregate, "improvement": improvements, "workloads": means}) + for role in ROLES[:2]: + if improvements[role] + 1e-12 < cfg.acceptance.min_improvement: + reasons.append(f"round {index + 1}: insufficient improvement over {role}") + for work in means: + if means[work]["candidate"] > means[work][role] * (1 + cfg.acceptance.max_regression + 1e-12): + reasons.append(f"round {index + 1}: {work} regresses against {role}") + return {"promoted": not reasons, "reasons": reasons, "rounds": decisions} diff --git a/scripts/kernel_tuning/discovery/__init__.py b/scripts/kernel_tuning/discovery/__init__.py new file mode 100644 index 0000000..a5238ea --- /dev/null +++ b/scripts/kernel_tuning/discovery/__init__.py @@ -0,0 +1,5 @@ +"""Shared registry capture, replay and driver utilities; the historical module paths are retained.""" + +from .contracts import ModelCapture + +__all__ = ["ModelCapture"] diff --git a/scripts/kernel_tuning/discovery/contracts.py b/scripts/kernel_tuning/discovery/contracts.py new file mode 100644 index 0000000..ce0b5e3 --- /dev/null +++ b/scripts/kernel_tuning/discovery/contracts.py @@ -0,0 +1,158 @@ +"""Portable, immutable model captures, independent of Torch and coding agents.""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from ..contracts import ContractError, digest, read_json, relative_path, tree_hashes + +SCHEMA = "embodiinfer-model-capture-v1" +DISPOSITIONS = {"ready", "blocked"} +REPLAY_SCHEMA = "embodiinfer-replay-v1" + + +def validate_replay( + replay: dict[str, Any], definition: dict[str, Any], traces: list[dict[str, Any]], hashes: dict[str, str] +) -> None: + """Validate captured tensor layouts and bind every replay case to its task axes.""" + if ( + replay.get("schema_version") != 1 + or not replay.get("capture_identity") + or not replay.get("target", {}).get("model") + ): + raise ContractError("Invalid model replay identity") + schema = replay.get("definition_schema") + if schema not in (None, REPLAY_SCHEMA) or ( + any( + spec["dtype"] == "float64" for kind in ("inputs", "outputs") for spec in definition[kind].values() + ) + and schema != REPLAY_SCHEMA + ): + raise ContractError("float64 requires the explicit replay definition schema") + if replay.get("kind") != "registry" or replay.get("rng") or replay.get("undefined_outputs"): + raise ContractError("Only registry replay tasks are supported; recapture with capture-ops") + if set(replay.get("cases", {})) != {trace["workload"]["uuid"] for trace in traces}: + raise ContractError("Replay cases must cover every workload exactly") + if not isinstance(replay.get("mutates"), list) or any( + type(index) is not int or not 0 <= index < len(definition["inputs"]) for index in replay["mutates"] + ): + raise ContractError("Invalid declared input mutations") + for trace in traces: + work = trace["workload"] + case = replay["cases"][work["uuid"]] + if "rng" in case: + raise ContractError("Registry workloads cannot contain RNG replay state") + for kind in ("inputs", "outputs"): + specs = case[kind] + if len(specs) != len(definition[kind]): + raise ContractError("Replay tensor count differs from the definition") + for spec, tensor in zip(specs, definition[kind].values()): + shape = [ + work["axes"].get(axis, definition["axes"][axis].get("value")) for axis in tensor["shape"] + ] + if spec["shape"] != shape or spec["dtype"] != tensor["dtype"]: + raise ContractError("Replay tensor shape/dtype differs from workload axes") + if len(spec["stride"]) != len(shape) or any( + type(n) is not int or n < 0 for n in [spec["offset"], *spec["stride"]] + ): + raise ContractError("Invalid replay tensor layout") + if not isinstance(spec.get("device"), str): + raise ContractError("Replay tensor must specify a device") + if kind == "inputs": + if type(spec.get("storage")) is not int or spec["storage"] < 0: + raise ContractError("Invalid replay storage group") + if "fixture" in spec and relative_path(spec["fixture"]) not in hashes: + raise ContractError("Replay fixture missing from the task") + for key, limit in (("output_aliases", len(case["inputs"])), ("output_groups", len(case["outputs"]))): + if len(case[key]) != len(case["outputs"]) or any( + not isinstance(group, list) or any(type(i) is not int or not 0 <= i < limit for i in group) + for group in case[key] + ): + raise ContractError("Invalid output alias contract") + + +@dataclass(frozen=True) +class ModelCapture: + """A registry capture bound to one policy and every case of its selected operators.""" + + root: Path + manifest: dict[str, Any] + operators: tuple[dict[str, Any], ...] + + @property + def identity(self) -> str: + """Content identity used by generated tasks and resumable model batches.""" + return digest(self.manifest) + + @classmethod + def load(cls, root: Path, model: str | None = None, *, require_ready: bool = True) -> ModelCapture: + """Verify capture contents and optionally enforce complete device-task coverage.""" + root = root.resolve() + if not (root / "manifest.json").is_file(): + raise ContractError("Model tuning requires a model capture directory; recapture legacy JSONL") + manifest = read_json(root / "manifest.json") + if manifest.get("schema") != SCHEMA: + raise ContractError("Unsupported model capture schema") + if manifest.get("capture") != "registry": + raise ContractError( + "Full-model discovery captures are no longer supported; recapture with capture-ops" + ) + target = manifest.get("target", {}) + if not target.get("model") or not target.get("type"): + raise ContractError("Capture has no verified target model") + if model is not None and target["model"] != model: + raise ContractError(f"Capture model {target['model']!r} differs from requested model {model!r}") + hashes = tree_hashes(root) + hashes.pop("manifest.json", None) + if hashes != manifest.get("files"): + raise ContractError("Model capture content changed") + operators = read_json(root / "operators.json") + if not isinstance(operators, list) or not operators: + raise ContractError("Capture contains no operator observations") + identities = set() + for op in operators: + if op.get("status") not in DISPOSITIONS or not op.get("id") or op["id"] in identities: + raise ContractError("Malformed or duplicate captured operator") + if op["id"] != digest(op["identity"])[:20] or op["name"] != op["identity"]["name"]: + raise ContractError("Captured operator identity differs from its content") + if ( + op["identity"].get("kind") != "registry" + or op["identity"].get("rng") + or op["identity"].get("undefined_outputs") + ): + raise ContractError("Only registry operators are supported; recapture with capture-ops") + identities.add(op["id"]) + cases = op.get("workloads", []) + if not cases or len({case["id"] for case in cases}) != len(cases): + raise ContractError(f"Missing or duplicate workloads: {op['id']}") + if any(type(case.get("count")) is not int or case["count"] < 1 for case in cases): + raise ContractError("Captured workload counts must be positive") + for case in cases: + if case["id"] != digest({k: v for k, v in case.items() if k not in {"id", "count"}})[:20]: + raise ContractError("Captured workload identity differs from its content") + for tensor in case.get("inputs", []): + if "fixture" in tensor and relative_path(tensor["fixture"]) not in hashes: + raise ContractError("Captured fixture must be included in the capture directory") + if "rng" in case: + raise ContractError("Registry workloads cannot contain RNG replay state") + capture = cls(root, manifest, tuple(operators)) + if require_ready: + capture.require_ready() + return capture + + def require_ready(self) -> None: + """Reject incomplete runs and any selected-operator capture gaps.""" + gaps = list(self.manifest.get("errors", [])) + gaps += [ + f"{op['name']}: {op.get('reason', 'unsupported')}" + for op in self.operators + if op["status"] == "blocked" + ] + if not self.manifest.get("complete"): + gaps.insert(0, "preparation did not complete") + if gaps: + raise ContractError("Model capture is incomplete; no tuning may start:\n" + "\n".join(gaps)) + if not any(op["status"] == "ready" for op in self.operators): + raise ContractError("Model capture has no tunable device computation") diff --git a/scripts/kernel_tuning/discovery/drive.py b/scripts/kernel_tuning/discovery/drive.py new file mode 100644 index 0000000..d9cc9a6 --- /dev/null +++ b/scripts/kernel_tuning/discovery/drive.py @@ -0,0 +1,117 @@ +"""Run recorded inputs through one model's production inference path. + +This shared inference driver serves profiling, registry capture and verification:: + + python -m scripts.kernel_tuning capture-ops --model POLICY --output CAPTURE -- \\ + -m scripts.kernel_tuning.discovery.drive --policy POLICY --inputs INPUTS.pt --no-cuda-graph [serving options] + +The policy and engine are constructed with the same arguments as the serving +transports. ``requests`` inputs go through the policy's serving adapter, exactly +as an HTTP request would; ``observations`` inputs go through ``policy.collate`` +and ``EngineCore.execute``. Recurrent policies keep one session per episode, so +memory grows the way it does in deployment. The calling tool controls profiling +or registry observation; the driver itself chooses no operators or code paths. +""" + +from __future__ import annotations + +import argparse +from collections.abc import Sequence +from pathlib import Path +from typing import Any + +import torch + +from embodiinfer.engine.config import EngineConfig +from embodiinfer.engine.core import EngineCore +from embodiinfer.engine.serve.batching import BatchedServingAdapter +from embodiinfer.engine.serve.contracts import RawPolicyRequest +from embodiinfer.engine.serve.factory import ( + _parse_adapter_config, + add_policy_arguments, + build_serving_adapter, +) +from embodiinfer.policies.factory import make_policy +from embodiinfer.types import Observation, SessionKey + +_KINDS = {"requests": RawPolicyRequest, "observations": Observation} + + +def save_inputs(path: Path, kind: str, episodes: Sequence[Sequence[Any]]) -> None: + """Write episodes of serving requests or engine observations for :func:`main`.""" + if kind not in _KINDS: + raise ValueError(f"input kind must be one of {sorted(_KINDS)}") + torch.save({"kind": kind, "episodes": [list(episode) for episode in episodes]}, path) + + +def load_inputs(path: Path) -> tuple[str, list[list[Any]]]: + """Read and validate inputs written by :func:`save_inputs`.""" + value = torch.load(path, map_location="cpu", weights_only=False) + kind = value.get("kind") if isinstance(value, dict) else None + episodes = value.get("episodes") if kind in _KINDS else None + if ( + not isinstance(episodes, list) + or not episodes + or not all(isinstance(episode, list) and episode for episode in episodes) + or not all(isinstance(item, _KINDS[kind]) for episode in episodes for item in episode) + ): + raise ValueError("inputs must contain non-empty episodes of one kind: requests or observations") + return kind, episodes + + +def serve(args: argparse.Namespace, episodes: list[list[RawPolicyRequest]]) -> None: + """Answer every request through the serving adapter, then end its sessions.""" + adapter = build_serving_adapter(args) + try: + for episode in episodes: + for request in episode: + adapter.infer(request) + for session_id in dict.fromkeys(request.session_id for request in episode): + adapter.reset(session_id) + finally: + if isinstance(adapter, BatchedServingAdapter): + adapter.shutdown() + + +def execute(args: argparse.Namespace, episodes: list[list[Observation]]) -> None: + """Execute every observation through the engine, one session per recurrent episode.""" + config = _parse_adapter_config(args.adapter_config) + policy_kwargs = dict(config.get("policy_kwargs", {})) + if args.checkpoint is not None: + policy_kwargs["checkpoint"] = args.checkpoint + policy = make_policy(args.policy, **policy_kwargs) + core = EngineCore( + policy, + EngineConfig( + device=args.device, + dtype=args.dtype, + max_batch_size=args.max_batch, + max_wait_ms=args.max_wait_ms, + num_steps=args.num_steps, + use_cuda_graph=not args.no_cuda_graph, + capture_full_loop=args.capture_full_loop, + ), + ) + for index, episode in enumerate(episodes): + key = SessionKey(env_id="discovery", episode_id=index) if policy.is_recurrent else None + for step, observation in enumerate(episode): + batch = policy.collate([observation], [f"{index}-{step}"]) + core.execute(batch, session_ids=None if key is None else [key]) + if key is not None: + core.reset_sessions([key]) + + +def main(argv: list[str] | None = None) -> None: + """Parse serving options and drive all recorded inputs.""" + parser = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + add_policy_arguments(parser) + parser.add_argument("--inputs", type=Path, required=True, help="Episodes written by save_inputs") + args = parser.parse_args(argv) + kind, episodes = load_inputs(args.inputs) + (serve if kind == "requests" else execute)(args, episodes) + + +if __name__ == "__main__": + main() diff --git a/scripts/kernel_tuning/discovery/numerics.py b/scripts/kernel_tuning/discovery/numerics.py new file mode 100644 index 0000000..a9d3dba --- /dev/null +++ b/scripts/kernel_tuning/discovery/numerics.py @@ -0,0 +1,114 @@ +"""Static numerical policies and validation of frozen, candidate-free calibration.""" + +from __future__ import annotations + +import hashlib +import math +from pathlib import Path +from typing import Any + +from ..contracts import ContractError, digest + +VERSION = "baseline-fp64-v1" +SEEDS = (0, 1, 2) +PROFILES = (("recorded", 0), ("random", 1), ("random", 2), ("zeros", 0), ("cancellation", 0)) + + +def implementation_digest() -> str: + """Bind calibration to the trusted policy/oracle source, including across machines.""" + root = Path(__file__).resolve().parents[3] + paths = ("scripts/kernel_tuning/discovery/numerics.py", "benchmarks/kernel_tuning/calibration.py") + return digest( + { + name: hashlib.sha256((root / name).read_text(encoding="utf-8").encode()).hexdigest() + for name in paths + } + ) + + +def policy(identity: dict[str, Any]) -> str | None: + """Return a trusted reference adapter, or retain exact validation for other semantics.""" + if identity.get("mutates") or identity.get("rng"): + return None + dtypes = identity.get("input_dtypes", []) + identity.get("output_dtypes", []) + if not any(dtype in {"float16", "bfloat16", "float32"} for dtype in dtypes): + return None + if any(dtype.startswith(("float8", "float4")) or dtype == "float64" for dtype in dtypes): + return None + if any(dtype not in {"float16", "bfloat16", "float32"} for dtype in identity.get("output_dtypes", [])): + return None + if identity["kind"] == "registry": + # Registry operators define their mathematics by a Torch reference, + # which also evaluates it in FP64 (embodiinfer.backend.torch.ops). + return "registry" + return None + + +def profile_key(profile: str, seed: int) -> str: + """Stable identity of one real, random, or boundary validation input.""" + return f"{profile}:{seed}" + + +def execution_precision(captured: dict[str, Any]) -> dict[str, Any]: + """Extract execution switches without inheriting legacy capture-wide bit exactness.""" + return {key: captured[key] for key in ("matmul_precision", "allow_tf32")} + + +def task_precision(contract: dict[str, Any], captured: dict[str, Any]) -> dict[str, Any]: + """Summarize bounds for reporting; evaluation uses each profile/output's own bound.""" + bounds = [bound for case in contract["cases"].values() for sample in case.values() for bound in sample] + return { + **execution_precision(captured), + "mode": "tolerance", + "atol": max(bound["atol"] for bound in bounds), + "rtol": max(bound["rtol"] for bound in bounds), + "reason": "Frozen per-workload/output FP64 baseline calibration in numerics.json; " + "different accumulation orders are allowed only within those measured bounds.", + } + + +def validate_contract( + contract: dict[str, Any], + identity: dict[str, Any], + cases: dict[str, Any], + precision: dict[str, Any], + flags: dict[str, Any], +) -> None: + """Reject incomplete, misbound, or nonfinite calibration before task generation/evaluation.""" + expected = { + "version": VERSION, + "implementation_digest": implementation_digest(), + "reference": policy(identity), + "operator_identity": digest(identity), + "precision": execution_precision(precision), + "flags": flags, + "profiles": [list(pair) for pair in PROFILES], + "seeds": list(SEEDS), + } + if not expected["reference"] or any(contract.get(k) != v for k, v in expected.items()): + raise ContractError("Numerical calibration policy/operator/execution settings changed; recalibrate") + if not contract.get("environment") or set(contract.get("cases", {})) != set(cases): + raise ContractError("Numerical calibration requires the target environment and every workload") + for key, case in cases.items(): + if contract.get("workloads", {}).get(key) != digest(case): + raise ContractError("Numerical calibration workload/fixture identity changed") + if any("fixture" not in spec for spec in case["inputs"]): + raise ContractError( + "Calibration requires complete real input fixtures; recapture with a larger --fixture-bytes" + ) + samples = contract["cases"][key] + if set(samples) != {profile_key(*pair) for pair in PROFILES}: + raise ContractError("Calibration must cover recorded, random and boundary inputs") + for bounds in samples.values(): + if len(bounds) != len(case["outputs"]): + raise ContractError("Calibration must cover every output") + for bound, spec in zip(bounds, case["outputs"]): + if bound.get("dtype") != spec["dtype"] or bound.get("reference_dtype") != "float64": + raise ContractError("Calibration output/reference dtype changed") + for field in ("atol", "rtol", "baseline_max_abs", "residual_max", "rms", "epsilon", "tiny"): + value = bound.get(field) + if type(value) not in (float, int) or not math.isfinite(value) or value < 0: + raise ContractError(f"Invalid calibration statistic: {field}") + atol = max(bound["epsilon"] * bound["rms"], 2 * bound["residual_max"], bound["tiny"]) + if bound["atol"] != atol or bound["rtol"] != bound["epsilon"]: + raise ContractError("Calibration threshold does not follow the frozen policy") diff --git a/scripts/kernel_tuning/discovery/tasks.py b/scripts/kernel_tuning/discovery/tasks.py new file mode 100644 index 0000000..6e05da2 --- /dev/null +++ b/scripts/kernel_tuning/discovery/tasks.py @@ -0,0 +1,365 @@ +"""Static task construction from verified model captures, with no Torch imports.""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + +from ..artifacts import atomic_json +from ..contracts import ContractError, TaskPackage, read_json +from ..generate import _dump_yaml, tuning_settings, vendor_sources +from ..operators import Operator, Workload +from .contracts import REPLAY_SCHEMA, ModelCapture +from .numerics import SEEDS, execution_precision, policy, task_precision, validate_contract + +# FlashInfer's dtypes plus the explicit local replay-schema float64 extension. +_DTYPES = { + "float64", + "float32", + "float16", + "bfloat16", + "float8_e4m3fn", + "float8_e5m2", + "float4_e2m1", + "int64", + "int32", + "int16", + "int8", + "bool", +} + + +def _expression(recipe: dict[str, Any]) -> str: + kind, value = next(iter(recipe.items())) + if kind == "tensor": + return f"x{value}" + if kind == "constant": + return repr(value) + if kind == "torch_value": + if value.startswith("torch.") and value[6:].isidentifier(): + return value + return f"torch.device({value!r})" + if kind in ("tuple", "list"): + entries = ", ".join(_expression(item) for item in value) + return f"({entries},)" if kind == "tuple" and value else "()" if kind == "tuple" else f"[{entries}]" + if kind == "dict": + return "{" + ", ".join(f"{key!r}: {_expression(item)}" for key, item in value.items()) + "}" + raise ContractError(f"Unrecognized argument recipe {kind}") + + +def _return_expressions( + recipe: dict[str, Any], prefix: str = "result", found: dict[int, str] | None = None +) -> list[str]: + found = {} if found is None else found + if "tensor" in recipe: + found.setdefault(recipe["tensor"], prefix) + for key in ("tuple", "list"): + for index, item in enumerate(recipe.get(key, [])): + _return_expressions(item, f"{prefix}[{index}]", found) + for key, item in recipe.get("dict", {}).items(): + _return_expressions(item, f"{prefix}[{key!r}]", found) + return [found[index] for index in sorted(found)] + + +def _sources(op: dict[str, Any]) -> tuple[str, dict[str, str]]: + identity = op["identity"] + files = {} + if identity["kind"] != "registry": + raise ContractError("Only registry tasks are supported; recapture with capture-ops") + module, entry = identity["module"], identity["entry"] + if module != "embodiinfer.backend.torch.ops" or not entry.isidentifier(): + raise ContractError("Registry tasks take their baseline from embodiinfer.backend.torch.ops") + for name, key in ((module, "source_digests"), (identity["start_module"], "start_digests")): + vendored = vendor_sources(name) + if {path: source_digest for path, (_, source_digest) in vendored.items()} != identity[key]: + raise ContractError(f"Registry operator source changed since capture: {name}") + files.update({path: text for path, (text, _) in vendored.items()}) + imports, callable_name = f"import torch\nfrom .vendor.{module} import {entry}\n", entry + args, kwargs = identity["arguments"]["tuple"] + names = [f"x{i}" for i in range(len(identity["input_dtypes"]))] + returns = _return_expressions(identity["returns"]) + [f"x{i}" for i in identity["mutates"]] + layouts = _production_layouts(op) + if layouts: + imports += _RELAYOUT.format(layouts=repr(layouts)) + code = imports + f"\ndef run({', '.join(names)}):\n" + code += f" result = {callable_name}(*{_expression(args)}, **{_expression(kwargs)})\n" + if layouts: + code += f" outputs = _relayout(({', '.join(names)},), ({', '.join(returns)},))\n" + returns = [f"outputs[{i}]" for i in range(len(returns))] + code += " return " + (returns[0] if len(returns) == 1 else "(" + ", ".join(returns) + ")") + "\n" + return code, files + + +# Registry captures take output layouts from the production call. Where the Torch +# reference lays a result out differently, the baseline copies it into that layout: +# values are unchanged, and the reference meets the same contract as candidates. +_RELAYOUT = """ +# Production output layout per recorded input layout: (shape, stride) per output, or +# None where the Torch reference already matches. +_LAYOUTS = {layouts} + + +def _relayout(inputs, outputs): + key = tuple((tuple(x.shape), tuple(x.stride()), x.storage_offset()) for x in inputs) + targets = _LAYOUTS.get(key) + if targets is None: + return outputs + result = [] + for value, target in zip(outputs, targets): + if target is not None and tuple(value.stride()) != target[1]: + copy = torch.empty_strided(target[0], target[1], dtype=value.dtype, device=value.device) + copy.copy_(value) + value = copy + result.append(value) + return tuple(result) +""" + + +def _production_layouts(op: dict[str, Any]) -> dict[tuple[Any, ...], tuple[Any, ...]]: + """Recorded input layout -> production output layouts, for cases whose reference differs.""" + layouts = {} + for case in op["workloads"]: + reference = case.get("reference_outputs") + if reference is None: + continue + if len(reference) != len(case["outputs"]): + raise ContractError(f"Reference/production output counts differ in {op['name']}") + key = tuple((tuple(s["shape"]), tuple(s["stride"]), s["offset"]) for s in case["inputs"]) + layouts[key] = tuple( + None if ref == spec else (tuple(spec["shape"]), tuple(spec["stride"])) + for spec, ref in zip(case["outputs"], reference) + ) + return layouts + + +def _plan( + capture: ModelCapture, op: dict[str, Any], overrides: dict[str, Any], hardware_notes: Path | None +) -> dict[str, str | bytes]: + name = f"registry_{op['identity']['op']}_{op['id']}" + unsupported = { + spec["dtype"] for case in op["workloads"] for spec in case["inputs"] + case["outputs"] + } - _DTYPES + if unsupported: + raise ContractError(f"Replay definition needs a custom dtype adapter for {sorted(unsupported)}") + code, sources = _sources(op) + # The reference can import snapshotted helper sources through this task's + # benchmark adapter; its standalone mathematical body remains frozen. + identity = op["identity"] + axes, inputs, outputs = {}, {}, {} + for kind, specs in (("inputs", inputs), ("outputs", outputs)): + for index, tensor in enumerate(op["workloads"][0][kind]): + shape = [f"{kind[0]}{index}d{d}" for d in range(len(tensor["shape"]))] + for axis in shape: + axes[axis] = {"type": "var"} + specs[f"{'x' if kind == 'inputs' else 'y'}{index}"] = {"shape": shape, "dtype": tensor["dtype"]} + definition = { + "name": name, + "op_type": "custom", + "description": f"Captured single operator {op['name']}", + "tags": [f"model:{capture.manifest['target']['model']}"], + "axes": axes, + "inputs": inputs, + "outputs": outputs, + "constraints": [], + "reference": code, + } + workloads = [] + cases = {} + for case in op["workloads"]: + values = {} + for kind in ("inputs", "outputs"): + if len(case[kind]) != len(definition[kind]): + raise ContractError(f"Output/input count changes in {op['name']}") + for index, tensor in enumerate(case[kind]): + spec = list(definition[kind].values())[index] + if len(spec["shape"]) != len(tensor["shape"]) or spec["dtype"] != tensor["dtype"]: + raise ContractError(f"Rank/dtype changes require distinct task signatures: {op['name']}") + values.update(zip(spec["shape"], tensor["shape"])) + workloads.append( + { + "definition": name, + "workload": { + "uuid": case["id"], + "axes": values, + "inputs": {key: {"type": "random"} for key in inputs}, + }, + "solution": None, + "evaluation": None, + } + ) + cases[case["id"]] = case + for tensor in case["inputs"]: + if "fixture" in tensor: + sources[tensor["fixture"]] = (capture.root / tensor["fixture"]).read_bytes() + numerical = None + precision = {"mode": "bit_exact", **execution_precision(capture.manifest["precision"])} + if policy(identity) is not None: + path = capture.root / "calibration.json" + numerical = read_json(path).get(op["id"]) if path.is_file() else None + if numerical is None: + raise ContractError( + "Missing numerical calibration; run calibrate --captured CAPTURE --output NEW_CAPTURE in the model environment" + ) + validate_contract( + numerical, identity, cases, capture.manifest["precision"], capture.manifest["flags"] + ) + precision = task_precision(numerical, capture.manifest["precision"]) + sources["numerics.json"] = json.dumps(numerical, indent=2) + "\n" + descriptor = Operator( + name=name, + summary=op["name"], + models=(capture.manifest["target"]["model"],), + op_type="custom", + source="", + entry="run", + call="", + axes={}, + inputs={}, + outputs={}, + reference=code, + workloads=tuple(Workload(case["id"], {}, weight=case["count"]) for case in op["workloads"]), + notes="Registry operator; the Torch reference is the baseline.", + precision=precision, + # Registry operators run inside the model's CUDA graphs in production. + timing="cuda_graph", + ) + settings = tuning_settings(descriptor, overrides) + if settings["precision"] != precision: + raise ContractError("Model task precision cannot override its frozen numerical contract") + if numerical is not None and settings["seeds"] != list(SEEDS): + raise ContractError("Calibrated validation seeds cannot be overridden") + if 0 not in settings["seeds"]: + raise ContractError("Model replay must include seed 0 for the unchanged captured fixtures") + device = next( + ( + spec["device"] + for spec in op["workloads"][0]["inputs"] + op["workloads"][0]["outputs"] + if spec["device"].startswith(capture.manifest["device_type"]) + ), + None, + ) + if "device" in overrides and overrides["device"] != device: + raise ContractError("Capture and evaluator devices differ; prepare on the target device") + settings["device"] = device + replay = { + "schema_version": 1, + "definition_schema": REPLAY_SCHEMA, + "operator": op["name"], + "kind": identity["kind"], + "mutates": identity["mutates"], + "cases": cases, + "torch_version": capture.manifest["torch_version"], + "flags": capture.manifest["flags"], + "capture_identity": capture.identity, + "target": capture.manifest["target"], + } + if numerical is not None: + replay["numerical_identity"] = identity + sources.update( + { + "definition.json": json.dumps(definition, indent=2) + "\n", + "baseline.py": code, + "workloads.jsonl": "".join(json.dumps(work) + "\n" for work in workloads), + "tuning.yaml": _dump_yaml(settings), + "replay.json": json.dumps(replay, indent=2) + "\n", + "benchmark.py": "from benchmarks.kernel_tuning.replay_adapter import ReplayAdapter\n\ndef create_adapter(task):\n return ReplayAdapter(task)\n", + } + ) + start = identity["start_module"].replace(".", "/") + sources["README.md"] = ( + f"# {identity['op']}\n\nTarget: `{capture.manifest['target']['model']}`.\n" + f"Capture: `{capture.identity}`.\n\n" + "Write a faster implementation of this operator-registry entry. The baseline to " + "beat is the Torch reference in baseline.py " + f"(`vendor/embodiinfer/backend/torch/ops.py::{identity['entry']}`); the repository's " + f"current implementation, `vendor/{start}.py::{identity['start_entry']}`, is a valid " + "starting point. Keep the call signature, return outputs with exactly the recorded " + "shapes, strides and dtypes in replay.json, use return-value style " + "(destination_passing_style=false), and do not change dtypes, masks or numerical " + "settings. Every recorded workload is mandatory and weighted by its call count in " + "the model. The installed kernel is used only for these recorded argument layouts, " + "so you may specialize on them.\n" + ) + if _production_layouts(op): + sources["README.md"] += ( + "\nOutput layouts are those of the implementation the model ran. Where the Torch " + "reference lays a result out differently (`reference_outputs` in replay.json), " + "baseline.py copies it into the production layout; the copy is part of the " + "baseline's time and leaves values unchanged.\n" + ) + sources["generated.json"] = json.dumps( + { + "generator": "scripts.kernel_tuning.regions", + "capture_identity": capture.identity, + "operator_id": op["id"], + "registry_op": identity["op"], + } + ) + if hardware_notes is not None: + sources["HARDWARE.md"] = hardware_notes.read_text(encoding="utf-8") + if numerical is not None: + sources["README.md"] += ( + "\nRead numerics.json: all floating elements are compared to the trusted FP64 " + "reference using frozen per-workload/profile/output bounds. tuning.yaml contains " + "only their reporting maxima. Recorded inputs, fresh random inputs, zeros and " + "cancellation inputs must all pass. Reassociation is permitted within these " + "bounds; changing dtype, masks, accumulation settings or the contract is not.\n" + ) + if any(spec["dtype"] == "float64" for kind in (inputs, outputs) for spec in kind.values()): + sources["README.md"] += ( + "\nThis task uses the embodiinfer-replay-v1 float64 definition extension and its " + "benchmark.py adapter. Keep all float64 inputs, intermediates and outputs at their " + "declared precision. The upstream FlashInfer 0.1.2 Definition parser alone cannot " + "load this extension.\n" + ) + return sources + + +def render_model( + capture: ModelCapture, + destination: Path, + *, + overrides: dict[str, Any] | None = None, + hardware_notes: Path | None = None, +) -> list[TaskPackage]: + """Build all eligible tasks or reject the entire model before writing any task.""" + capture.require_ready() + # Recheck snapshots immediately before construction; do not trust a previously loaded object. + checked = ModelCapture.load(capture.root, capture.manifest["target"]["model"]) + if checked.identity != capture.identity: + raise ContractError("Capture changed before task generation") + destination = destination.resolve() + if destination.exists() and any(destination.iterdir()): + raise ContractError("Model tasks require a new or empty destination") + plans = {} + failures = [] + for op in capture.operators: + if op["status"] != "ready": + continue + try: + plans[op["id"]] = _plan(capture, op, overrides or {}, hardware_notes) + except (ContractError, KeyError, ValueError) as exc: + failures.append(f"{op['name']}: {exc}") + if failures: + raise ContractError("Task coverage gaps; no model batch may start:\n" + "\n".join(failures)) + tasks = [] + for op_id, files in plans.items(): + root = destination / op_id + for name, content in files.items(): + path = root / name + path.parent.mkdir(parents=True, exist_ok=True) + if isinstance(content, bytes): + path.write_bytes(content) + else: + path.write_text(content, encoding="utf-8", newline="\n") + tasks.append(TaskPackage.load(root)) + atomic_json( + destination / "model.json", + { + "capture_identity": capture.identity, + "target": capture.manifest["target"], + "tasks": {task.root.name: task.identity for task in tasks}, + }, + ) + return tasks diff --git a/scripts/kernel_tuning/discovery/tensors.py b/scripts/kernel_tuning/discovery/tensors.py new file mode 100644 index 0000000..5c39b09 --- /dev/null +++ b/scripts/kernel_tuning/discovery/tensors.py @@ -0,0 +1,154 @@ +"""Tensor argument recipes and storage-preserving replay in the model runtime.""" + +from __future__ import annotations + +import math +from pathlib import Path +from typing import Any + +import torch + +from ..contracts import ContractError + + +def tensor_spec(value: torch.Tensor) -> dict[str, Any]: + """Describe layout without copying data or synchronizing a device.""" + if value.layout != torch.strided or value.is_quantized or value.is_conj() or value.is_neg(): + raise ContractError("Only ordinary dense strided tensors have a replay recipe") + if any(stride < 0 for stride in value.stride()): + raise ContractError("Negative strides require a dedicated replay adapter") + return { + "shape": list(value.shape), + "stride": list(value.stride()), + "offset": value.storage_offset(), + "dtype": str(value.dtype).removeprefix("torch."), + "device": str(value.device), + } + + +def flatten(value: Any, tensors: list[torch.Tensor]) -> Any: + """Replace tensor leaves with positions, preserving supported argument containers.""" + if isinstance(value, torch.Tensor): + index = next((i for i, tensor in enumerate(tensors) if tensor is value), None) + if index is None: + index = len(tensors) + tensors.append(value) + return {"tensor": index} + if isinstance(value, (tuple, list)): + return {"tuple" if isinstance(value, tuple) else "list": [flatten(x, tensors) for x in value]} + if isinstance(value, dict) and all(isinstance(key, str) for key in value): + return {"dict": {key: flatten(item, tensors) for key, item in value.items()}} + if value is None or isinstance(value, (bool, str, int)): + return {"constant": value} + if isinstance(value, float) and math.isfinite(value): + return {"constant": value} + if isinstance(value, (torch.dtype, torch.device, torch.layout, torch.memory_format)): + return {"torch_value": str(value)} + raise ContractError( + f"No replay recipe for argument type {type(value).__module__}.{type(value).__qualname__}" + ) + + +def restore(recipe: Any, tensors: list[torch.Tensor]) -> Any: + """Rebuild a captured argument tree with current input tensors.""" + kind, value = next(iter(recipe.items())) + if kind == "tensor": + return tensors[value] + if kind == "constant": + return value + if kind == "torch_value": + return getattr(torch, value[6:]) if value.startswith("torch.") else torch.device(value) + if kind in ("tuple", "list"): + result = [restore(x, tensors) for x in value] + return tuple(result) if kind == "tuple" else result + if kind == "dict": + return {key: restore(item, tensors) for key, item in value.items()} + raise ContractError(f"Unknown argument recipe: {kind}") + + +def storage_key(tensor: torch.Tensor) -> tuple[Any, ...]: + """Identify shared storage within one call, including zero-sized storage objects.""" + return str(tensor.device), tensor.untyped_storage()._cdata + + +def describe_inputs(tensors: list[torch.Tensor]) -> list[dict[str, Any]]: + """Assign stable per-call alias groups to tensor inputs.""" + groups: dict[tuple[Any, ...], int] = {} + result = [] + for tensor in tensors: + key = storage_key(tensor) + result.append({**tensor_spec(tensor), "storage": groups.setdefault(key, len(groups))}) + for group in groups.values(): + if len({spec["dtype"] for spec in result if spec["storage"] == group}) != 1: + raise ContractError("Reinterpreted storage with different dtypes needs a dedicated adapter") + return result + + +def extent(spec: dict[str, Any]) -> int: + """Number of storage elements needed for a layout, including its offset.""" + if not all(spec["shape"]): + return spec["offset"] + return spec["offset"] + 1 + sum((n - 1) * stride for n, stride in zip(spec["shape"], spec["stride"])) + + +def materialize(specs: list[dict[str, Any]], root: Path, seed: int) -> list[torch.Tensor]: + """Restore captured storage and aliasing, or generate unconstrained floating storage.""" + buffers = {} + for group in {spec["storage"] for spec in specs}: + members = [spec for spec in specs if spec["storage"] == group] + spec = members[0] + dtype = getattr(torch, spec["dtype"]) + size = max(extent(member) for member in members) + generator = torch.Generator(device=spec["device"]).manual_seed(seed + group) + if "fixture" in spec: + data = bytearray((root / spec["fixture"]).read_bytes()) + buffer = ( + torch.frombuffer(data, dtype=dtype).clone().to(spec["device"]) + if data + else torch.empty(0, dtype=dtype, device=spec["device"]) + ) + if buffer.numel() != size: + raise ContractError("Captured fixture size differs from its storage layout") + # Different seeds exercise changed floating values while respecting integer/index fixtures. + if buffer.is_floating_point() and seed: + buffer = buffer * (1.0 + (seed % 7) / 128.0) + elif dtype.is_floating_point: + buffer = torch.randn(size, dtype=dtype, device=spec["device"], generator=generator) + else: + raise ContractError("Non-floating inputs require a valid captured fixture") + buffers[group] = buffer + return [ + buffers[spec["storage"]].as_strided(spec["shape"], spec["stride"], spec["offset"]) for spec in specs + ] + + +def clone_inputs(inputs: list[torch.Tensor]) -> list[torch.Tensor]: + """Clone backing storage once per alias group instead of breaking views apart.""" + specs = describe_inputs(inputs) + buffers = {} + for tensor, spec in zip(inputs, specs): + group = spec["storage"] + if group not in buffers: + size = max(extent(item) for item in specs if item["storage"] == group) + buffers[group] = tensor.as_strided((size,), (1,), 0).clone() + return [ + buffers[spec["storage"]].as_strided(spec["shape"], spec["stride"], spec["offset"]) for spec in specs + ] + + +def tensor_outputs(result: Any) -> list[torch.Tensor]: + """Flatten tensor-only outputs without silently dropping data-dependent scalar results.""" + tensors: list[torch.Tensor] = [] + recipe = flatten(result, tensors) + + def check(node: Any) -> None: + if "constant" in node and node["constant"] is not None: + raise ContractError("Non-tensor results require a dedicated output contract") + for key in ("tuple", "list"): + for child in node.get(key, []): + check(child) + for child in node.get("dict", {}).values(): + check(child) + + check(recipe) + return tensors diff --git a/scripts/kernel_tuning/flow.py b/scripts/kernel_tuning/flow.py new file mode 100644 index 0000000..a2eb9d1 --- /dev/null +++ b/scripts/kernel_tuning/flow.py @@ -0,0 +1,223 @@ +"""The KDA minimal loop as a resumable Humanize2 Flow. + +Only this module needs hmz. Importing it defines a flow; executing the developer +CLI's run/resume command is what opens coding-agent sessions. +""" + +from __future__ import annotations + +import asyncio +import contextlib +import copy +import time +from datetime import timedelta +from pathlib import Path +from typing import Any + +from hmz.flows import ( + AgentCollection, + Budget, + BudgetExceeded, + EnvCollection, + FlowContext, + FlowParams, + LocalEnv, + flow, +) + +from scripts.kernel_tuning.artifacts import RunStore, atomic_json +from scripts.kernel_tuning.contracts import ( + ContractError, + promotion, + read_json, + tree_hashes, + validate_measurement, + validate_solution, +) +from scripts.kernel_tuning.permissions import Coder, session_agent +from scripts.kernel_tuning.prompts import implementation_prompt, plan_prompt +from scripts.kernel_tuning.runner import evaluate + + +class Agents(AgentCollection): + """An existing coding agent, selected by an explicit hmz agent specification.""" + + coder: Coder + + +class Envs(EnvCollection): + """Humanize2's local workspace, containing only this run's scratch attempts.""" + + workspace: LocalEnv + + +class Params(FlowParams): + """All tunable conditions are in the snapshotted task; only its run path varies.""" + + run_dir: str + + +@flow(agents=Agents, envs=Envs, params=Params, resumable=True, name="kernel_tuning") +async def kernel_tuning( + task: str, *, agents: Agents, envs: Envs, params: Params, ctx: FlowContext +) -> dict[str, Any]: + """Plan, generate, verify, measure, and retain one candidate at a time.""" + del task # The immutable task package supplies the complete objective. + store = RunStore(Path(params.run_dir)) + state = ctx.state + if state is None: + raise ContractError("This flow requires Humanize2 resumable state") + if "ledger" in state: + ledger = state["ledger"] + if ledger["run_id"] != store.manifest["run_id"] or ledger["task_digest"] != store.task.identity: + raise ContractError("Humanize2 journal belongs to a different task/run") + store.manifest = ledger + store.verify() + store.recover() + coder = session_agent(agents["coder"], unfenced=bool(store.manifest.get("unfenced"))) + cfg = store.task.settings + base_elapsed = store.manifest["elapsed_seconds"] + started = time.monotonic() + + def checkpoint() -> None: + store.manifest["elapsed_seconds"] = base_elapsed + time.monotonic() - started + # The framework journal is authoritative after a crash between these writes. + state["ledger"] = copy.deepcopy(store.manifest) + store.save() + + def remaining() -> float: + return max(0.0, cfg.search.max_seconds - (base_elapsed + time.monotonic() - started)) + + async def heartbeat() -> None: + while True: + await asyncio.sleep(1) + checkpoint() + + async def measure( + output: Path, *, operation: str = "evaluate", candidate: Path | None = None + ) -> dict[str, Any]: + if remaining() <= 0: + raise TimeoutError("Search duration exhausted") + return await evaluate( + store, + output, + operation=operation, + candidate=candidate, + timeout=min(cfg.search.evaluation_timeout_seconds, remaining()), + ) + + store.manifest.update(status="running", stop_reason=None) + checkpoint() + ticker = asyncio.create_task(heartbeat()) + try: + # Re-probe on every resume before any agent or kernel is allowed to run. + report = await measure(store.root / "preflight", operation="probe") + store.manifest["environment"] = report["environment"] + checkpoint() + if not store.manifest.get("baseline_verified"): + report = await measure(store.root / "baseline/evaluation") + validate_measurement(store.task, report["measurement"]) + store.manifest["baseline_verified"] = True + checkpoint() + while True: + reason = None + if len(store.manifest["attempts"]) >= cfg.search.max_candidates: + reason = "max_candidates" + elif store.manifest["streak"] >= cfg.search.patience: + reason = "patience" + elif remaining() <= 0: + reason = "max_seconds" + if reason: + store.manifest.update(status="completed", stop_reason=reason) + break + candidate, archive, scratch = store.begin_attempt(persist=checkpoint) + checkpoint() + solution = None + session = None + try: + work_env = await envs["workspace"].derive_subdir(subdir=candidate) + session = await coder.spawn(env=work_env) + for stage, prompt in ( + ("plan", plan_prompt(store.task)), + ("implement", implementation_prompt(store.task)), + ): + store.verify() + if remaining() <= 0: + raise TimeoutError("Search duration exhausted") + seconds = min(cfg.search.agent_timeout_seconds, remaining()) + (archive / f"{stage}-prompt.txt").write_text(prompt, encoding="utf-8") + reply = await asyncio.wait_for( + coder.run( + prompt, + session=session, + budget=Budget(duration=timedelta(seconds=seconds), graceful=False), + ), + timeout=seconds, + ) + (archive / f"{stage}-reply.txt").write_text(reply, encoding="utf-8") + plan = scratch / "PLAN.md" + if ( + not plan.is_file() + or plan.is_symlink() + or not plan.read_text(encoding="utf-8").strip() + ): + raise ContractError("Agent did not write a nonempty PLAN.md") + (archive / "PLAN.md").write_bytes(plan.read_bytes()) + del session + store.verify() + if tree_hashes(scratch / "task") != store.task.hashes: + raise ContractError("Agent modified its frozen task copy") + solution_file = scratch / "solution.json" + if solution_file.is_symlink(): + raise ContractError("Candidate solution must be a regular file") + solution = validate_solution(read_json(solution_file), store.task) + atomic_json(archive / "solution.json", solution) + report = await measure(archive / "evaluation", candidate=archive / "solution.json") + decision = promotion(store.task, report["measurement"]) + status = "promoted" if decision["promoted"] else "rejected" + store.finish_attempt(candidate, status=status, decision=decision, solution=solution) + checkpoint() + if cfg.profile: + # Profiling is evidence only, outside promotion and performance timing. + try: + await measure( + archive / "profile", operation="profile", candidate=archive / "solution.json" + ) + except Exception as exc: + atomic_json(archive / "profile/error.json", {"error": str(exc)}) + except (asyncio.CancelledError, KeyboardInterrupt): + raise + except Exception as exc: + # Per-turn budget exhaustion is an attempt failure. The whole-run + # duration is checked separately, so resume cannot reset budgets. + store.finish_attempt( + candidate, + status="failed", + decision={"promoted": False, "reasons": [f"{type(exc).__name__}: {exc}"]}, + solution=solution, + ) + checkpoint() + # A changed reference/tool/baseline is a run error, not another candidate. + store.verify() + finally: + session = None # hmz closes sessions when the last holder lets go. + checkpoint() + return copy.deepcopy(store.manifest) + except (BudgetExceeded, TimeoutError): + store.recover() + store.manifest.update(status="completed", stop_reason="max_seconds") + raise + except BaseException as exc: + store.recover() + store.manifest.update( + status="interrupted" + if isinstance(exc, (asyncio.CancelledError, KeyboardInterrupt)) + else "failed", + stop_reason=f"{type(exc).__name__}: {exc}", + ) + raise + finally: + ticker.cancel() + with contextlib.suppress(asyncio.CancelledError): + await ticker + checkpoint() diff --git a/scripts/kernel_tuning/generate.py b/scripts/kernel_tuning/generate.py new file mode 100644 index 0000000..89ed760 --- /dev/null +++ b/scripts/kernel_tuning/generate.py @@ -0,0 +1,461 @@ +"""Render catalog operators into self-contained KDA task packages. + +The baseline is EmbodiInfer's production kernel, copied byte-for-byte from the +repository together with the repository modules it imports. Generation is +static: it reads source files and writes the task; it never imports Torch, +Triton, or the copied kernels. +""" + +from __future__ import annotations + +import ast +import hashlib +import json +import re +import shutil +import struct +import subprocess +from dataclasses import asdict +from pathlib import Path, PurePosixPath +from typing import Any + +from .artifacts import REPOSITORY +from .contracts import ContractError, Settings, TaskPackage, read_json +from .operators import Operator + +VENDOR = "vendor" +MARKER = "generated.json" +PACKAGE = "embodiinfer" +SAFETENSORS_DTYPES = {"int32": ("I32", " str: + """Digest source text with normalized newlines, so Windows checkouts agree with Linux.""" + return hashlib.sha256(path.read_text(encoding="utf-8").encode()).hexdigest() + + +def _module_file(module: str) -> Path: + """Resolve a dotted repository module to its file, refusing package initializers.""" + base = REPOSITORY.joinpath(*module.split(".")) + if base.with_suffix(".py").is_file(): + return base.with_suffix(".py") + if (base / "__init__.py").is_file(): + raise ContractError( + f"Cannot vendor package {module}: import the defining submodule instead of the package" + ) + raise ContractError(f"Repository module {module} does not exist") + + +def _imports(path: Path) -> list[tuple[ast.ImportFrom | ast.Import, list[str]]]: + """List each repository import in ``path`` with the dotted modules it loads.""" + relative = path.relative_to(REPOSITORY).with_suffix("") + package = list(relative.parts[:-1]) + found: list[tuple[ast.ImportFrom | ast.Import, list[str]]] = [] + for node in ast.walk(ast.parse(path.read_text(encoding="utf-8"))): + if isinstance(node, ast.Import): + if any(alias.name.split(".")[0] == PACKAGE for alias in node.names): + raise ContractError(f"{path}: use 'from ... import' for repository modules") + continue + if not isinstance(node, ast.ImportFrom): + continue + if node.level: + if node.level > len(package): + raise ContractError(f"{path}: relative import escapes the repository package") + anchor = package[: len(package) - node.level + 1] + elif node.module and node.module.split(".")[0] == PACKAGE: + anchor = [] + else: + continue + target = [*anchor, *(node.module or "").split(".")] if node.module else anchor + target = [part for part in target if part] + if node.module: + found.append((node, [".".join(target)])) + else: # ``from . import module`` loads sibling modules. + found.append((node, [".".join([*target, alias.name]) for alias in node.names])) + return found + + +def _relative_import(source: str, target: str) -> str: + """Spell ``target`` relative to the package containing module ``source``.""" + origin = source.split(".")[:-1] + goal = target.split(".") + common = 0 + while common < min(len(origin), len(goal)) and origin[common] == goal[common]: + common += 1 + return "." * (len(origin) - common + 1) + ".".join(goal[common:]) + + +def vendor_sources(entry: str) -> dict[str, tuple[str, str]]: + """Copy a module and its repository imports; return task paths to (text, source digest). + + Relative imports keep working because the package layout is preserved. + Absolute ``embodiinfer.`` imports are rewritten to relative form, the only + edit made to a copied file. Package ``__init__`` modules are never copied, so + importing one kernel does not import the whole backend. + """ + pending = [entry] + result: dict[str, tuple[str, str]] = {} + while pending: + module = pending.pop() + name = f"{VENDOR}/{module.replace('.', '/')}.py" + if name in result: + continue + path = _module_file(module) + text = path.read_text(encoding="utf-8") + lines = text.splitlines(keepends=True) + for node, targets in _imports(path): + for target in targets: + _module_file(target) + pending.append(target) + if not node.level: + line = node.lineno - 1 + spelled = _relative_import(module, node.module) + lines[line], count = re.subn( + rf"\bfrom\s+{re.escape(node.module)}\s+import\b", + f"from {spelled} import", + lines[line], + count=1, + ) + if not count: + raise ContractError(f"{path}:{node.lineno}: cannot rewrite multi-line import") + result[name] = ("".join(lines), _text_digest(path)) + return dict(sorted(result.items())) + + +def safetensors_bytes(tensors: dict[str, tuple[str, tuple[int, ...]]]) -> bytes: + """Encode one-dimensional integer tensors in the safetensors format without Torch.""" + header: dict[str, Any] = {} + payload = b"" + for key, (dtype, values) in sorted(tensors.items()): + code, fmt = SAFETENSORS_DTYPES[dtype] + data = struct.pack(f"<{len(values)}{fmt[1]}", *values) + header[key] = { + "dtype": code, + "shape": [len(values)], + "data_offsets": [len(payload), len(payload) + len(data)], + } + payload += data + encoded = json.dumps(header, separators=(",", ":")).encode() + encoded += b" " * (-len(encoded) % 8) + return struct.pack(" str: + try: + result = subprocess.run( + ["git", "-C", str(REPOSITORY), "describe", "--always", "--dirty", "--abbrev=40"], + capture_output=True, + text=True, + check=True, + timeout=10, + ) + except (OSError, subprocess.SubprocessError): + return "unknown" + return result.stdout.strip() + + +def _yaml_value(text: str) -> Any: + import yaml + + return yaml.safe_load(text) + + +def parse_overrides(items: list[str]) -> dict[str, Any]: + """Parse ``section.key=value`` tuning.yaml overrides; values use YAML syntax.""" + overrides: dict[str, Any] = {} + for item in items: + key, separator, value = item.partition("=") + if not separator or not key: + raise ContractError(f"Override must be key=value: {item!r}") + overrides[key.strip()] = _yaml_value(value) + return overrides + + +def tuning_settings(operator: Operator, overrides: dict[str, Any] | None = None) -> dict[str, Any]: + """Default contract for a catalog operator, with validated dotted overrides.""" + settings: dict[str, Any] = { + "schema_version": 1, + "language": "triton", + "agent": "", + "evaluator_python": "python", + "device": "cuda:0", + "target_hardware": ["cuda"], + "seeds": [0, 1, 2], + "precision": {"matmul_precision": "highest", "allow_tf32": False, **operator.precision}, + "timing": { + "mode": operator.timing, + "warmup": 10, + "iterations": 1000, + "trials": 3, + "paired_rounds": 2, + }, + "acceptance": {"min_improvement": 0.03, "max_regression": 0.05}, + "search": { + "max_candidates": 20, + "max_seconds": 7200, + "patience": 5, + "agent_timeout_seconds": 600, + "evaluation_timeout_seconds": 600, + }, + "profile": False, + } + weights = {work.uuid: work.weight for work in operator.workloads} + if len(set(weights.values())) > 1: + settings["weights"] = weights + for key, value in (overrides or {}).items(): + *parents, leaf = key.split(".") + node = settings + for parent in parents: + if not isinstance(node.get(parent), dict): + raise ContractError(f"Unknown tuning.yaml section in override: {key}") + node = node[parent] + node[leaf] = value + Settings.load_mapping(settings) + return settings + + +def _definition(operator: Operator) -> dict[str, Any]: + axes = { + name: {"type": "var"} if value is None else {"type": "const", "value": value} + for name, value in operator.axes.items() + } + + def tensors(specs: dict[str, Any]) -> dict[str, Any]: + return { + name: {"shape": None if spec.shape is None else list(spec.shape), "dtype": spec.dtype} + for name, spec in specs.items() + } + + return { + "name": f"embodiinfer_{operator.name}", + "op_type": operator.op_type, + "description": operator.summary, + "tags": [f"model:{model}" for model in operator.models], + "axes": axes, + "inputs": tensors(operator.inputs), + "outputs": tensors(operator.outputs), + "constraints": list(operator.constraints), + "reference": operator.reference, + } + + +def _workload_inputs(operator: Operator, work: Any) -> dict[str, Any]: + inputs: dict[str, Any] = {} + for name, spec in operator.inputs.items(): + if name in work.fixed: + inputs[name] = { + "type": "safetensors", + "path": f"data/{work.uuid}.safetensors", + "tensor_key": name, + } + elif name in work.scalars: + inputs[name] = {"type": "scalar", "value": work.scalars[name]} + elif spec.shape is None: + raise ContractError(f"{operator.name}/{work.uuid}: scalar input {name} needs a value") + else: + inputs[name] = {"type": "random"} + return inputs + + +def _baseline(operator: Operator) -> str: + module = operator.source.removesuffix(".py").replace("/", ".") + arguments = ", ".join(operator.inputs) + return f'''"""Generated baseline: EmbodiInfer's production kernel, called as in the model. + +The kernel source under {VENDOR}/ is copied from the repository; see README.md. +""" + +from .{VENDOR}.{module} import {operator.entry} + + +def run({arguments}): + return {operator.call} +''' + + +def _readme(operator: Operator, vendored: dict[str, tuple[str, str]], revision: str, hardware: bool) -> str: + def shape(spec: Any) -> str: + return "scalar" if spec.shape is None else "[" + ", ".join(spec.shape) + "]" + + tensors = "\n".join( + f"| {kind} | `{name}` | {shape(spec)} | {spec.dtype} |" + for kind, specs in (("input", operator.inputs), ("output", operator.outputs)) + for name, spec in specs.items() + ) + cases = "\n".join( + f"| `{work.uuid}` | " + + ", ".join(f"{k}={v}" for k, v in work.axes.items()) + + "".join(f", {k}={v}" for k, v in work.scalars.items()) + + f" | {work.weight:g} | {work.note} |" + for work in operator.workloads + ) + sources = "\n".join( + f"- `{path}` (source sha256 `{digest[:16]}`)" for path, (_, digest) in vendored.items() + ) + if operator.precision["mode"] == "bit_exact": + precision = "Outputs must be **bit-identical** to the reference on every seed." + else: + precision = ( + f"Outputs must satisfy `|actual - reference| <= {operator.precision['atol']} + " + f"{operator.precision['rtol']} * |reference|` elementwise. Reason: {operator.precision['reason']}" + ) + graph_note = ( + "\nTiming captures each call in a CUDA Graph, so a candidate must not synchronize with the host." + if operator.timing == "cuda_graph" + else "" + ) + origin = ( + "Shapes and weights come from captured production calls (`capture`); weights are call counts." + if all(work.uuid.startswith("captured") for work in operator.workloads) + else "Shapes are representative of the listed models, not traced traffic." + ) + hardware_note = ( + "\nTarget-hardware notes supplied for this machine are in `HARDWARE.md`.\n" if hardware else "" + ) + return f"""# {operator.name} + +> Generated by `python -m scripts.kernel_tuning generate` from repository revision +> `{revision}`. Do not edit by hand; change `scripts/kernel_tuning/operators.py` +> or the kernel source, then regenerate. + +{operator.summary} Used by: {", ".join(operator.models)}. + +| role | name | shape | dtype | +|---|---|---|---| +{tensors} + +Constraints: {", ".join(f"`{c}`" for c in operator.constraints) or "none"}. + +{operator.notes} + +## Baseline + +`baseline.py` calls the production kernel `{operator.entry}` from `{operator.source}` +exactly as the inference engine does. It is the latency to beat; the +mathematical reference in `definition.json` defines correctness. The production +kernel and its repository imports are copied verbatim under `{VENDOR}/`: + +{sources} + +Start from that source. Its input checks show which layouts the engine passes; +a candidate only has to handle the declared definition and workloads.{graph_note} + +## Precision + +{precision} + +## Workloads + +{origin} + +| uuid | axes and scalars | weight | where | +|---|---|---|---| +{cases} +{hardware_note}""" + + +def render( + operator: Operator, + destination: Path, + *, + overrides: dict[str, Any] | None = None, + hardware_notes: Path | None = None, + force: bool = False, +) -> TaskPackage: + """Write one task package, replacing only a directory this generator created.""" + if destination.exists(): + if not (destination / MARKER).is_file() and any(destination.iterdir()): + raise ContractError(f"{destination} exists and was not generated; choose another output") + if not force: + raise ContractError(f"{destination} exists; pass --force to regenerate it") + shutil.rmtree(destination) + module = operator.source.removesuffix(".py").replace("/", ".") + vendored = vendor_sources(module) + revision = _git_revision() + settings = tuning_settings(operator, overrides) + staging = destination.with_name(f".{destination.name}.tmp") + if staging.exists(): + shutil.rmtree(staging) + try: + files: dict[str, str | bytes] = { + "definition.json": json.dumps(_definition(operator), indent=2) + "\n", + "workloads.jsonl": "".join( + json.dumps( + { + "definition": f"embodiinfer_{operator.name}", + "workload": { + "uuid": work.uuid, + "axes": work.axes, + "inputs": _workload_inputs(operator, work), + }, + "solution": None, + "evaluation": None, + } + ) + + "\n" + for work in operator.workloads + ), + "baseline.py": _baseline(operator), + "tuning.yaml": _dump_yaml(settings), + "README.md": _readme(operator, vendored, revision, hardware_notes is not None), + } + files.update({path: text for path, (text, _) in vendored.items()}) + for work in operator.workloads: + if work.fixed: + files[f"data/{work.uuid}.safetensors"] = safetensors_bytes( + {name: (operator.inputs[name].dtype, values) for name, values in work.fixed.items()} + ) + if hardware_notes is not None: + files["HARDWARE.md"] = hardware_notes.read_text(encoding="utf-8") + files[MARKER] = ( + json.dumps( + { + "generator": "scripts.kernel_tuning.generate", + "operator": operator.name, + "repository_revision": revision, + "sources": {path: digest for path, (_, digest) in vendored.items()}, + "catalog_entry": _catalog_entry(operator), + }, + indent=2, + ) + + "\n" + ) + for name, content in files.items(): + path = staging / PurePosixPath(name) + path.parent.mkdir(parents=True, exist_ok=True) + if isinstance(content, bytes): + path.write_bytes(content) + else: + path.write_text(content, encoding="utf-8", newline="\n") + TaskPackage.load(staging) + staging.replace(destination) + finally: + if staging.exists(): + shutil.rmtree(staging) + return TaskPackage.load(destination) + + +def _catalog_entry(operator: Operator) -> dict[str, Any]: + entry = asdict(operator) + entry["workloads"] = [asdict(work) for work in operator.workloads] + return json.loads(json.dumps(entry)) + + +def _dump_yaml(settings: dict[str, Any]) -> str: + import yaml + + return ( + "# Generated; see README.md. Machine-specific values may be overridden at generation.\n" + + yaml.safe_dump(settings, sort_keys=False) + ) + + +def stale_sources(task: Path) -> list[str]: + """Copied sources whose repository file has changed since generation.""" + marker = read_json(task / MARKER) + changed = [] + for path, digest in marker["sources"].items(): + source = REPOSITORY / PurePosixPath(path).relative_to(VENDOR) + if not source.is_file() or _text_digest(source) != digest: + changed.append(path) + return changed diff --git a/scripts/kernel_tuning/hotspots.py b/scripts/kernel_tuning/hotspots.py new file mode 100644 index 0000000..0788155 --- /dev/null +++ b/scripts/kernel_tuning/hotspots.py @@ -0,0 +1,661 @@ +"""Profile end-to-end inference and rank where GPU time goes. + +``python -m scripts.kernel_tuning profile --output DIR -- DRIVER`` runs an +unmodified inference driver (for example ``-m scripts.kernel_tuning.discovery.drive``) +in the model runtime. Each top-level call of a model-agnostic inference entry +(``EngineCore.execute``, ``GenerationBackend`` rollouts, serving ``infer``) is one +step; the first ``warmup`` steps are skipped and the next ``steps`` are profiled. + +Every GPU kernel is attributed through its launch (CUDA correlation id) to the +profiling ranges enclosing that launch: + +* ``stage:``: policy and decoder protocol methods (``encode_prefix``, + ``denoise_step``, ``integrate``, ...), +* ``op:``: operator-registry calls (:mod:`embodiinfer.layers.ops`), +* ``attention:``: attention-backend calls, +* ``module:``: ``nn.Module`` calls, with layer indices folded to ``*``. + +A kernel's region is its stage ranges plus the innermost operator, attention or +module range enclosing its launch. + +Kernels replayed from a CUDA graph are labelled from the graph's capture: a +short profile around each capture outside the measured window records the +ranges of every node launch, and replayed kernels take those labels in launch +order when the counts match (otherwise they keep their replay-time stage). +Ranges are skipped inside ``torch.compile`` tracing, so compiled regions report +at the level of the range that called them. +""" + +from __future__ import annotations + +import contextlib +import functools +import importlib +import json +import re +import runpy +import sys +import threading +from collections import defaultdict +from collections.abc import Callable, Iterable, Iterator +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from .contracts import ContractError + +RANGE = "kt/" +STEP = f"{RANGE}step:" +GRAPH = f"{RANGE}graph:" +_ENTRIES = ( + ("embodiinfer.engine.core", "EngineCore", ("execute", "execute_pipelined")), + ( + "embodiinfer.engine.rollout.generation_backend", + "GenerationBackend", + ("generate", "generate_with_logprob", "sample_group", "best_of_n"), + ), +) +_POLICY_STAGES = ("collate", "encode_prefix", "denoise_step", "finalize_actions") +_DECODER_STAGES = ("decode", "produce_chunk", "integrate", "sample_with_logprob") +_LAYER_INDEX = re.compile(r"(?<=\.)\d+(?=\.|$)") +_KERNEL_ACTIVITIES = {"kernel", "gpu_memset", "gpu_memcpy"} +_LAUNCH_ACTIVITIES = {"cuda_runtime", "cuda_driver"} + + +@dataclass(frozen=True) +class Range: + """A profiling range on one host thread, in nanoseconds.""" + + name: str + thread: int + start: int + end: int + + +@dataclass(frozen=True) +class Launch: + """A host-side launch (kernel, memset/memcpy or graph) and its correlation id.""" + + name: str + thread: int + start: int + correlation: int + + +@dataclass(frozen=True) +class DeviceKernel: + """A device activity with the correlation id of its launch.""" + + name: str + start: int + duration_ns: int + correlation: int + + +@dataclass(frozen=True) +class Kernel: + """One profiled kernel attributed to a step and its enclosing ranges.""" + + step: int | None + path: tuple[str, ...] + name: str + duration_us: float + graph: bool + labelled: bool = False + + +def fold_layers(name: str) -> str: + """Fold numeric path components so repeated layers aggregate (``layers.3`` -> ``layers.*``).""" + return _LAYER_INDEX.sub("*", name) + + +def _enclosing(ranges: Iterable[Range], launches: Iterable[Launch]) -> dict[int, tuple[Range, ...]]: + """Map each launch's correlation id to its enclosing ranges, outermost first. + + Ranges on one thread nest (``record_function`` scopes), so one sweep per + thread with a stack of open ranges finds them in O(n log n). + """ + range_lists: dict[int, list[Range]] = defaultdict(list) + launch_lists: dict[int, list[Launch]] = defaultdict(list) + for item in ranges: + range_lists[item.thread].append(item) + for launch in launches: + launch_lists[launch.thread].append(launch) + result: dict[int, tuple[Range, ...]] = {} + for thread, thread_launches in launch_lists.items(): + thread_ranges = sorted(range_lists.get(thread, []), key=lambda item: (item.start, -item.end)) + stack: list[Range] = [] + index = 0 + for launch in sorted(thread_launches, key=lambda item: item.start): + while index < len(thread_ranges) and thread_ranges[index].start <= launch.start: + while stack and stack[-1].end < thread_ranges[index].start: + stack.pop() + stack.append(thread_ranges[index]) + index += 1 + while stack and stack[-1].end < launch.start: + stack.pop() + result[launch.correlation] = tuple(stack) + return result + + +def _labels(enclosing: Iterable[Range]) -> tuple[int | None, int | None, tuple[str, ...]]: + """Step index, replayed graph id and region path of one launch.""" + step = graph = None + stages: list[str] = [] + innermost = None + for item in enclosing: + if item.name.startswith(STEP): + step = int(item.name[len(STEP) :]) + elif item.name.startswith(GRAPH): + graph = int(item.name[len(GRAPH) :]) + elif item.name.startswith(RANGE + "stage:"): + name = item.name[len(RANGE) :] + if not stages or stages[-1] != name: + stages.append(name) + else: + innermost = fold_layers(item.name[len(RANGE) :]) + return step, graph, (*stages, innermost) if innermost else tuple(stages) + + +_KERNEL_CLASSES = ( + ("attention", re.compile(r"attention|flash|fmha|_attn", re.IGNORECASE)), + ("gemm", re.compile(r"gemm|nvjet|cutlass|xmma|_mm_|matmul|cublas", re.IGNORECASE)), + ("inductor", re.compile(r"^triton_")), + ("copy", re.compile(r"memcpy|memset|copy|cat", re.IGNORECASE)), + ("elementwise", re.compile(r"elementwise|reduce|softmax|norm", re.IGNORECASE)), +) + + +def kernel_class(name: str) -> str: + """Coarse kernel category from its name, for regions without finer ranges.""" + for label, pattern in _KERNEL_CLASSES: + if pattern.search(name): + return label + return "other" + + +def _with_class(path: tuple[str, ...], name: str) -> tuple[str, ...]: + """Append the kernel category when no operator, attention or module range encloses it.""" + if all(part.startswith("stage:") for part in path): + return (*path, f"kernels:{kernel_class(name)}") + return path + + +def is_node_launch(name: str) -> bool: + """Whether a host API call adds one node (kernel, memset or memcpy) to a captured graph.""" + return "LaunchKernel" in name or name.startswith( + ("cudaMemsetAsync", "cudaMemcpyAsync", "cuMemsetD", "cuMemcpy") + ) + + +def capture_labels(ranges: Iterable[Range], launches: Iterable[Launch]) -> list[tuple[str, ...]]: + """Region paths of a graph's node launches, in capture order.""" + nodes = sorted( + (launch for launch in launches if is_node_launch(launch.name)), key=lambda item: item.start + ) + enclosing = _enclosing(ranges, nodes) + return [_labels(enclosing[launch.correlation])[2] for launch in nodes] + + +def attribute( + ranges: Iterable[Range], + launches: Iterable[Launch], + kernels: Iterable[DeviceKernel], + graph_labels: dict[int, list[tuple[str, ...]]] | None = None, +) -> list[Kernel]: + """Attribute each device kernel to the ranges enclosing its host launch. + + Kernels of one graph replay share the replay's launch; when ``graph_labels`` + holds as many capture labels as the replay ran kernels, they are matched in + device start order. + """ + kernels = list(kernels) + wanted = {kernel.correlation for kernel in kernels} + launch_by_id = {launch.correlation: launch for launch in launches if launch.correlation in wanted} + enclosing_by_id = _enclosing(ranges, launch_by_id.values()) + replays: dict[int, list[int]] = defaultdict(list) + result: list[Kernel] = [] + for kernel in kernels: + launch = launch_by_id.get(kernel.correlation) + if launch is None: + result.append(Kernel(None, ("",), kernel.name, kernel.duration_ns / 1000, False)) + continue + step, _, path = _labels(enclosing_by_id[kernel.correlation]) + graph = "Graph" in launch.name + if graph: + replays[kernel.correlation].append(len(result)) + result.append( + Kernel(step, _with_class(path, kernel.name), kernel.name, kernel.duration_ns / 1000, graph) + ) + for correlation, indices in replays.items(): + graph_id = _labels(enclosing_by_id[correlation])[1] + labels = (graph_labels or {}).get(graph_id) if graph_id is not None else None + if labels is None or len(labels) != len(indices): + continue + ordered = sorted(indices, key=lambda index: kernels[index].start) + replay_stages = tuple( + part for part in _labels(enclosing_by_id[correlation])[2] if part.startswith("stage:") + ) + for index, label in zip(ordered, labels, strict=True): + item = result[index] + # The capture profiler starts inside the stage that captures, so its + # labels lack the outer stages; the replay supplies them. + merged = (*replay_stages, *(part for part in label if part not in replay_stages)) + path = _with_class(merged, item.name) + result[index] = Kernel(item.step, path, item.name, item.duration_us, True, True) + return result + + +def summarize(kernels: Iterable[Kernel], steps: int, *, small_us: float = 5.0) -> dict[str, Any]: + """Aggregate attributed kernels per step into regions ranked by GPU time.""" + measured = [kernel for kernel in kernels if kernel.step is not None] + if steps <= 0 or not measured: + raise ContractError("No kernels were profiled inside the measured steps") + total = sum(kernel.duration_us for kernel in measured) + regions: dict[tuple[str, ...], dict[str, Any]] = {} + for kernel in measured: + region = regions.setdefault( + kernel.path, + {"us": 0.0, "kernels": 0, "small": 0, "graph": 0, "names": defaultdict(lambda: [0, 0.0])}, + ) + region["us"] += kernel.duration_us + region["kernels"] += 1 + region["small"] += kernel.duration_us < small_us + region["graph"] += kernel.graph + entry = region["names"][kernel.name] + entry[0] += 1 + entry[1] += kernel.duration_us + ranked = [] + for path, region in sorted(regions.items(), key=lambda item: -item[1]["us"]): + names = sorted(region["names"].items(), key=lambda item: -item[1][1]) + ranked.append( + { + "region": "/".join(path) or "", + "us_per_step": region["us"] / steps, + "share": region["us"] / total, + "kernels_per_step": region["kernels"] / steps, + "mean_kernel_us": region["us"] / region["kernels"], + "small_kernel_fraction": region["small"] / region["kernels"], + "graph_replayed_fraction": region["graph"] / region["kernels"], + "top_kernels": [ + {"name": name, "calls_per_step": calls / steps, "us_per_step": us / steps} + for name, (calls, us) in names[:5] + ], + } + ) + graph_us = sum(kernel.duration_us for kernel in measured if kernel.graph) + labelled_us = sum(kernel.duration_us for kernel in measured if kernel.labelled) + return { + "steps": steps, + "graph_us_per_step": graph_us / steps, + "graph_labelled_fraction": labelled_us / graph_us if graph_us else 1.0, + "kernel_us_per_step": total / steps, + "kernels_per_step": len(measured) / steps, + "small_kernel_us": small_us, + "regions": ranked, + } + + +def render(summary: dict[str, Any], *, limit: int = 30) -> str: + """Markdown table of the highest-cost regions.""" + lines = [ + f"# Hotspots ({summary['steps']} steps)", + "", + f"GPU kernel time per step: {summary['kernel_us_per_step'] / 1000:.2f} ms in " + f"{summary['kernels_per_step']:.0f} kernels. Small kernels are under {summary['small_kernel_us']:g} us.", + f"CUDA graph replays: {summary['graph_us_per_step'] / 1000:.2f} ms per step, " + f"{summary['graph_labelled_fraction']:.0%} labelled from their capture.", + "", + "| Region | ms/step | Share | Kernels/step | Mean us | Small | Graph | Top kernel |", + "| --- | ---: | ---: | ---: | ---: | ---: | ---: | --- |", + ] + for region in summary["regions"][:limit]: + top = region["top_kernels"][0]["name"] if region["top_kernels"] else "" + lines.append( + f"| `{region['region']}` | {region['us_per_step'] / 1000:.3f} | {region['share']:.1%} " + f"| {region['kernels_per_step']:.1f} | {region['mean_kernel_us']:.1f} " + f"| {region['small_kernel_fraction']:.0%} | {region['graph_replayed_fraction']:.0%} | `{top[:60]}` |" + ) + return "\n".join(lines) + "\n" + + +def _compiling() -> bool: + import torch + + return torch.compiler.is_compiling() + + +def _ranged(function: Callable[..., Any], name: str) -> Callable[..., Any]: + from torch.profiler import record_function + + @functools.wraps(function) + def wrapper(*args: Any, **kwargs: Any) -> Any: + if _compiling(): + return function(*args, **kwargs) + with record_function(RANGE + name): + return function(*args, **kwargs) + + return wrapper + + +class ProfileSession: + """Count inference steps, profile a window of them and label their GPU work.""" + + def __init__(self, warmup: int, steps: int, *, modules: bool = True) -> None: + if warmup < 0 or steps <= 0: + raise ContractError("profile needs warmup >= 0 and steps > 0") + self.warmup, self.steps, self.modules = warmup, steps, modules + self.started = 0 + self.depth = threading.local() + self.lock = threading.Lock() + self.profiler: Any = None + self.trace: Any = None + self._restore: list[Callable[[], None]] = [] + self._module_names: dict[int, str] = {} + self._module_ranges = threading.local() + self._capture: tuple[Any, Any] | None = None + self.graph_labels: dict[int, list[tuple[str, ...]]] = {} + + @property + def recording(self) -> bool: + """Whether a profiler (measured window or graph capture) is collecting ranges.""" + return self._capture is not None or (self.profiler is not None and self.trace is None) + + def _replace(self, owner: Any, name: str, value: Any) -> None: + """Set ``owner.name`` (class, module or instance) and remember how to undo it.""" + had = name in vars(owner) + original = vars(owner).get(name) + setattr(owner, name, value) + + def restore() -> None: + if had: + setattr(owner, name, original) + else: + delattr(owner, name) + + self._restore.append(restore) + + def _step(self, function: Callable[..., Any]) -> Callable[..., Any]: + session = self + + @functools.wraps(function) + def wrapper(*args: Any, **kwargs: Any) -> Any: + depth = getattr(session.depth, "value", 0) + if depth: + return function(*args, **kwargs) + with session.lock: + index = session.started + session.started += 1 + if index == session.warmup: + session._start() + measured = session.warmup <= index < session.warmup + session.steps + session.depth.value = 1 + try: + if not measured: + return function(*args, **kwargs) + from torch.profiler import record_function + + with record_function(f"{STEP}{index - session.warmup}"): + return function(*args, **kwargs) + finally: + session.depth.value = 0 + if index == session.warmup + session.steps - 1: + session._stop() + + return wrapper + + def _start(self) -> None: + import torch + from torch.profiler import ProfilerActivity, profile + + if torch.cuda.is_available(): + torch.cuda.synchronize() + activities = [ProfilerActivity.CPU] + if torch.cuda.is_available(): + activities.append(ProfilerActivity.CUDA) + self.profiler = profile(activities=activities) + self.profiler.start() + + def _stop(self) -> None: + import torch + + if self.profiler is None or self.trace is not None: + return + if torch.cuda.is_available(): + torch.cuda.synchronize() + self.profiler.stop() + self.trace = self.profiler + + def _bind_policy(self, policy: Any) -> None: + import torch + + if isinstance(policy, torch.nn.Module): + self._module_names.update({id(module): name for name, module in policy.named_modules() if name}) + for name in _POLICY_STAGES: + method = getattr(policy, name, None) + if callable(method) and not isinstance(method, torch.nn.Module): + self._replace(policy, name, _ranged(method, f"stage:{name}")) + decoder = getattr(policy, "decoder", None) + for name in _DECODER_STAGES: + method = getattr(decoder, name, None) + if callable(method): + self._replace(decoder, name, _ranged(method, f"stage:{name}")) + + def install(self) -> None: + """Patch inference entries, protocol stages, operator and attention resolution.""" + for module_name, owner, methods in _ENTRIES: + holder = getattr(importlib.import_module(module_name), owner) + for method in methods: + if hasattr(holder, method): + self._replace(holder, method, self._step(getattr(holder, method))) + core = importlib.import_module("embodiinfer.engine.core").EngineCore + original_init = core.__init__ + session = self + + @functools.wraps(original_init) + def init(engine: Any, policy: Any, *args: Any, **kwargs: Any) -> None: + session._bind_policy(policy) + original_init(engine, policy, *args, **kwargs) + + self._replace(core, "__init__", init) + + from embodiinfer.layers.ops import KernelSelection + + original_resolve = KernelSelection.resolve + + @functools.wraps(original_resolve) + def resolve(selection: Any, op: str, *args: Any, **kwargs: Any) -> Any: + return _ranged(original_resolve(selection, op, *args, **kwargs), f"op:{op}") + + self._replace(KernelSelection, "resolve", resolve) + + attention = importlib.import_module("embodiinfer.layers.attention") + original_backend = attention._resolve_attention_backend + + def resolve_backend(name: str) -> Any: + cls = original_backend(name) + for method in ("attend", "attend_split"): + function = cls.__dict__.get(method) + if function is not None and not getattr(function, "_kt_ranged", False): + ranged = _ranged(function, f"attention:{cls.__name__}") + ranged._kt_ranged = True # type: ignore[attr-defined] + self._replace(cls, method, ranged) + return cls + + self._replace(attention, "_resolve_attention_backend", resolve_backend) + self._install_graph_hooks() + if self.modules: + self._install_module_hooks() + + def _install_graph_hooks(self) -> None: + """Label CUDA graph nodes at capture and tag each measured replay with its graph.""" + import torch + from torch.profiler import ProfilerActivity, profile, record_function + + graph_class = torch.cuda.CUDAGraph + begin, end, replay = graph_class.capture_begin, graph_class.capture_end, graph_class.replay + session = self + + @functools.wraps(begin) + def capture_begin(graph: Any, *args: Any, **kwargs: Any) -> Any: + if session.profiler is None and session._capture is None: + labeller = profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) + labeller.start() + session._capture = (graph, labeller) + return begin(graph, *args, **kwargs) + + @functools.wraps(end) + def capture_end(graph: Any, *args: Any, **kwargs: Any) -> Any: + try: + return end(graph, *args, **kwargs) + finally: + if session._capture is not None and session._capture[0] is graph: + labeller = session._capture[1] + session._capture = None + labeller.stop() + ranges, launches, _ = _events(labeller) + session.graph_labels[id(graph)] = capture_labels(ranges, launches) + + @functools.wraps(replay) + def tagged_replay(graph: Any, *args: Any, **kwargs: Any) -> Any: + if session.profiler is None or session.trace is not None: + return replay(graph, *args, **kwargs) + with record_function(f"{GRAPH}{id(graph)}"): + return replay(graph, *args, **kwargs) + + self._replace(graph_class, "capture_begin", capture_begin) + self._replace(graph_class, "capture_end", capture_end) + self._replace(graph_class, "replay", tagged_replay) + + def _install_module_hooks(self) -> None: + import torch + from torch.nn.modules import module as module_hooks + + def before(module: Any, _inputs: Any) -> None: + if _compiling() or not self.recording: + return + name = self._module_names.get(id(module)) + if name is None: + return + handle = torch.ops.profiler._record_function_enter_new(f"{RANGE}module:{name}", None) + stack = getattr(self._module_ranges, "stack", None) + if stack is None: + stack = self._module_ranges.stack = [] + stack.append((id(module), handle)) + + def after(module: Any, _inputs: Any, _output: Any) -> None: + if _compiling(): + return + stack = getattr(self._module_ranges, "stack", None) + if stack and stack[-1][0] == id(module): + torch.ops.profiler._record_function_exit._RecordFunction(stack.pop()[1]) + + pre = module_hooks.register_module_forward_pre_hook(before) + post = module_hooks.register_module_forward_hook(after) + self._restore.append(pre.remove) + self._restore.append(post.remove) + + def close(self) -> None: + """Undo every patch, stopping the profiler if the driver ended early.""" + self._stop() + while self._restore: + self._restore.pop()() + + def kernels(self) -> list[Kernel]: + """Attributed kernels of the profiled window.""" + if self.trace is None or self.started < self.warmup + self.steps: + raise ContractError( + f"Only {self.started} inference steps ran; profiling needs warmup {self.warmup} + {self.steps}" + ) + return attribute(*_events(self.trace), self.graph_labels) + + +def _events(profiler: Any) -> tuple[list[Range], list[Launch], list[DeviceKernel]]: + """Ranges, host launches and device kernels recorded by a stopped profiler.""" + ranges, launches, devices = [], [], [] + events = profiler.profiler.kineto_results.events() + if events and not hasattr(events[0], "activity_type"): + import torch + + raise RuntimeError( + f"Hotspot attribution needs Kineto events with activity_type(); Torch {torch.__version__} lacks it" + ) + for event in events: + activity = event.activity_type() + name = event.name() + if activity == "user_annotation" and name.startswith(RANGE): + start = event.start_ns() + ranges.append(Range(name, event.start_thread_id(), start, start + event.duration_ns())) + elif activity in _LAUNCH_ACTIVITIES: + launches.append(Launch(name, event.start_thread_id(), event.start_ns(), event.correlation_id())) + elif activity in _KERNEL_ACTIVITIES: + devices.append(DeviceKernel(name, event.start_ns(), event.duration_ns(), event.correlation_id())) + return ranges, launches, devices + + +def _environment() -> dict[str, Any]: + import torch + + environment: dict[str, Any] = {"torch": torch.__version__, "python": sys.version.split()[0]} + try: + import triton + + environment["triton"] = triton.__version__ + except ImportError: + environment["triton"] = None + if torch.cuda.is_available(): + environment["device_name"] = torch.cuda.get_device_name() + environment["capability"] = list(torch.cuda.get_device_capability()) + return environment + + +@contextlib.contextmanager +def _driver_argv(command: list[str]) -> Iterator[None]: + argv, path = sys.argv, sys.path[:] + try: + yield + finally: + sys.argv, sys.path = argv, path + + +def run( + command: list[str], + output: Path, + *, + warmup: int = 10, + steps: int = 20, + modules: bool = True, + trace: bool = False, +) -> dict[str, Any]: + """Run an inference driver, profile ``steps`` steps after ``warmup`` and write the report.""" + if not command: + raise ContractError("Supply the inference driver after --") + if output.exists(): + raise ContractError(f"{output} already exists; choose a new profile directory") + session = ProfileSession(warmup, steps, modules=modules) + session.install() + try: + with _driver_argv(command): + if command[0] == "-m": + sys.argv = command[1:] + runpy.run_module(command[1], run_name="__main__", alter_sys=True) + else: + script = Path(command[0]).resolve() + sys.argv = [str(script), *command[1:]] + sys.path.insert(0, str(script.parent)) + runpy.run_path(str(script), run_name="__main__") + except SystemExit as exc: + if exc.code not in (None, 0): + raise + finally: + session.close() + summary = summarize(session.kernels(), steps) + summary.update( + {"command": command, "warmup": warmup, "environment": _environment(), "module_ranges": modules} + ) + output.mkdir(parents=True) + (output / "hotspots.json").write_text(json.dumps(summary, indent=2, ensure_ascii=False), encoding="utf-8") + (output / "hotspots.md").write_text(render(summary), encoding="utf-8") + if trace: + session.trace.export_chrome_trace(str(output / "trace.json")) + return summary diff --git a/scripts/kernel_tuning/install.py b/scripts/kernel_tuning/install.py new file mode 100644 index 0000000..dcae27f --- /dev/null +++ b/scripts/kernel_tuning/install.py @@ -0,0 +1,219 @@ +"""Install an exported registry-operator solution into a hardware kernel profile. + +``python -m scripts.kernel_tuning install EXPORT --profile NAME`` runs in the model +runtime on the target machine. It copies the winning sources of a registry task +(see :mod:`scripts.kernel_tuning.regions`) into +``embodiinfer/backend/triton/tuned/NAME//`` and records the operator in that +profile's ``profile.json``, matched to the current device and Torch/Triton +versions. The generated package only applies to argument layouts the evaluator +verified (every recorded workload); other calls fall back to the policy's +preferred implementation. Installing changes no default: select the profile +with ``kernel_profile``. +""" + +from __future__ import annotations + +import json +import shutil +import tempfile +from pathlib import Path +from typing import Any + +from .contracts import ContractError, TaskPackage, read_json, relative_path + +_MATCH = ("device_name", "capability", "torch", "triton") +PROFILE_ROOT = Path(__file__).resolve().parents[2] / "embodiinfer" / "backend" / "triton" / "tuned" +_PACKAGE = '''"""Tuned ``{op}`` kernel for the ``{profile}`` profile. + +Generated by ``scripts.kernel_tuning install`` from run {run_id}; do not edit, +reinstall instead. The kernel applies only to the argument layouts verified by +the evaluator. Strides of size-one dimensions and of empty tensors address no +other element, so they are compared as zero. Like the evaluator, the kernel +receives only the tensor arguments; constants are part of the verified layout. +""" + +import torch + +from .{module} import {symbol} as run + +_VERIFIED = frozenset({cases}) +_ARGUMENTS = {arguments} +_TENSORS = {tensors} + + +def _strides(shape, stride): + if 0 in shape: + return (0,) * len(shape) + return tuple(0 if size == 1 else step for size, step in zip(shape, stride)) + + +def _leaves(value, out): + if isinstance(value, torch.Tensor): + dtype = str(value.dtype).removeprefix("torch.") + shape = tuple(value.shape) + out.append(("tensor", shape, _strides(shape, value.stride()), dtype)) + elif isinstance(value, (tuple, list)): + for item in value: + _leaves(item, out) + elif isinstance(value, dict): + for key in sorted(value): + _leaves(value[key], out) + else: + out.append(("constant", repr(value))) + return out + + +def applies(*args, **kwargs): + """Whether these arguments match a layout the evaluator verified.""" + return tuple(_leaves((args, kwargs), [])) in _VERIFIED + + +def _bind(recipe, value, tensors): + kind, item = next(iter(recipe.items())) + if kind == "tensor": + tensors[item] = value + elif kind in ("tuple", "list"): + for entry, element in zip(item, value): + _bind(entry, element, tensors) + elif kind == "dict": + for key, entry in item.items(): + _bind(entry, value[key], tensors) + + +def call(*args, **kwargs): + """Call the kernel as the evaluator did: with the captured tensor arguments x0..xN.""" + tensors = {{}} + _bind(_ARGUMENTS, (args, kwargs), tensors) + return run(*(tensors[index] for index in range(_TENSORS))) +''' + + +def canonical_strides(shape: tuple[int, ...], stride: tuple[int, ...]) -> tuple[int, ...]: + """Strides with those that address no other element (size-one or empty tensors) as zero. + + Eager and compiled runs of the same model can differ only there, e.g. a + ``(1, 1, L, D)`` cache with stride ``(L*D, L*D, D, 1)`` or ``(L*D, D, D, 1)``. + Mirrors ``_strides`` in the generated package. + """ + if 0 in shape: + return (0,) * len(shape) + return tuple(0 if size == 1 else step for size, step in zip(shape, stride)) + + +def recorded_signature(recipe: dict[str, Any], inputs: list[dict[str, Any]]) -> tuple[Any, ...]: + """Leaf signature of a captured call, matching the generated ``applies`` traversal.""" + leaves: list[tuple[Any, ...]] = [] + + def walk(node: dict[str, Any]) -> None: + kind, value = next(iter(node.items())) + if kind == "tensor": + spec = inputs[value] + shape = tuple(spec["shape"]) + leaves.append(("tensor", shape, canonical_strides(shape, tuple(spec["stride"])), spec["dtype"])) + elif kind in ("tuple", "list"): + for item in value: + walk(item) + elif kind == "dict": + for key in sorted(value): + walk(value[key]) + elif kind == "constant": + leaves.append(("constant", repr(value))) + else: + raise ContractError(f"Registry kernels cannot gate arguments of kind {kind}") + + walk(recipe) + return tuple(leaves) + + +def install( + export: Path, profile: str, *, root: Path = PROFILE_ROOT, replace: bool = False +) -> dict[str, Any]: + """Copy one exported registry solution into ``root/profile`` and update its manifest.""" + from embodiinfer.backend.triton.tuned import current_environment, load_profile, source_digest + from embodiinfer.layers.ops import available_ops + + if not profile.isidentifier(): + raise ContractError("Profile names must be Python identifiers, e.g. thor or rtx4090") + export = export.resolve() + task = TaskPackage.load(export / "task") + generated = read_json(task.root / "generated.json") + op = generated.get("registry_op") + if generated.get("generator") != "scripts.kernel_tuning.regions" or not op: + raise ContractError("Only registry-operator tasks (capture-ops) can be installed") + if op not in available_ops(): + raise ContractError(f"{op!r} is not registered in this source tree; merge its registration first") + solution = read_json(export / "solution.json") + spec = solution["spec"] + if spec.get("language") != "triton": + raise ContractError("Only Python/Triton solutions can be installed into a kernel profile") + entry_file, symbol = spec["entry_point"].split("::") + module = Path(relative_path(entry_file)).with_suffix("").as_posix().replace("/", ".") + replay = read_json(task.root / "replay.json") + identity = replay.get("numerical_identity") or {} + if identity.get("op") != op: + raise ContractError("Task replay identity does not match its registry operator") + environment = current_environment() + if replay["torch_version"] != environment["torch"]: + raise ContractError("The task was captured with another Torch version; tune on this runtime") + cases = sorted( + {recorded_signature(identity["arguments"], case["inputs"]) for case in replay["cases"].values()} + ) + run = read_json(export / "manifest.json") + attempt = next(item for item in run["attempts"] if item["id"] == run["best"]) + directory = root / profile + manifest_path = directory / "profile.json" + manifest = ( + read_json(manifest_path) + if manifest_path.is_file() + else {"schema": 1, "match": {key: environment.get(key) for key in _MATCH}, "ops": {}} + ) + if any(manifest["match"].get(key) != environment.get(key) for key in _MATCH): + raise ContractError(f"Profile {profile!r} was verified for {manifest['match']}, not {environment}") + if op in manifest["ops"] and not replace: + raise ContractError(f"{profile}/{op} is already installed; pass --replace to overwrite it") + directory.mkdir(parents=True, exist_ok=True) + staging = Path(tempfile.mkdtemp(prefix=f".{op}-", dir=directory)) + try: + package = staging / op + shutil.copytree(export / "sources", package) + if (package / "__init__.py").exists(): + raise ContractError("Solution sources may not define the package __init__.py") + (package / "__init__.py").write_text( + _PACKAGE.format( + op=op, + profile=profile, + run_id=run["run_id"], + module=module, + symbol=symbol, + cases=repr(cases), + arguments=repr(identity["arguments"]), + tensors=len(identity["input_dtypes"]), + ), + encoding="utf-8", + newline="\n", + ) + target = directory / op + if target.exists(): + shutil.rmtree(target) + package.rename(target) + finally: + shutil.rmtree(staging, ignore_errors=True) + manifest["ops"][op] = { + "module": op, + "function": "call", + "applies": "applies", + "sha256": source_digest(directory / op), + "verified_layouts": len(cases), + "evidence": { + "run_id": run["run_id"], + "task_digest": task.identity, + "capture_identity": generated["capture_identity"], + "solution_digest": attempt.get("solution_digest"), + "decision": attempt.get("decision"), + "export_checksums": read_json(export / "checksums.json").get("manifest.json"), + }, + } + manifest_path.write_text(json.dumps(manifest, indent=2, ensure_ascii=False) + "\n", encoding="utf-8") + loaded = load_profile(directory) + loaded.kernel(op) # Import once so a broken package fails here, not in the model. + return {"profile": profile, "op": op, "path": str(directory / op), "verified_layouts": len(cases)} diff --git a/scripts/kernel_tuning/operators.py b/scripts/kernel_tuning/operators.py new file mode 100644 index 0000000..6590e00 --- /dev/null +++ b/scripts/kernel_tuning/operators.py @@ -0,0 +1,515 @@ +"""Catalog of EmbodiInfer's core Triton operators, rendered into tuning task packages. + +Each entry states only what a person must decide: the mathematical reference, the +production kernel call used as the baseline, model-derived workload shapes, and +the numerical contract. ``generate.py`` turns an entry into a complete KDA task. + +The ``workloads`` here are estimates from model configurations. Shapes captured +from real runs (``benchmarks/kernel_tuning/captures/*.jsonl``, see +``capture.py``) replace them by default; the estimates remain the fallback for +operators no capture covers. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +# Tolerances count BF16 roundings that the production kernel and the eager +# reference do not share. Each contributes at most the unit roundoff 2**-8 +# relative error; the two final output roundings alone account for two steps. +# FP32 reordering is negligible except where outputs cancel toward zero, which +# atol covers. ``test_production_baseline_meets_generated_contract`` reports how +# much of each bound the production kernel uses on the target GPU. +BF16_ATOL = 2**-10 + + +@dataclass(frozen=True) +class Tensor: + """A FlashInfer TensorSpec: ``shape=None`` declares a scalar argument.""" + + shape: tuple[str, ...] | None + dtype: str + + +@dataclass(frozen=True) +class Workload: + """One declared case; tensors not listed in ``scalars``/``fixed`` are random.""" + + uuid: str + axes: dict[str, int] + scalars: dict[str, int | float] = field(default_factory=dict) + # Valid structured integer inputs (e.g. segment offsets) cannot be random. + # They are written into the task as safetensors data. + fixed: dict[str, tuple[int, ...]] = field(default_factory=dict) + weight: float = 1.0 + note: str = "" + + +@dataclass(frozen=True) +class Operator: + """A production operator plus the frozen contract used to tune it.""" + + name: str + summary: str + models: tuple[str, ...] + op_type: str + source: str # Repository module holding the production kernel. + entry: str # Function imported from ``source`` and called by the baseline. + call: str # Baseline return expression over the definition inputs. + axes: dict[str, int | None] # None declares a variable axis. + inputs: dict[str, Tensor] + outputs: dict[str, Tensor] + reference: str + workloads: tuple[Workload, ...] + notes: str + constraints: tuple[str, ...] = () + precision: dict[str, Any] = field(default_factory=lambda: {"mode": "bit_exact"}) + timing: str = "cuda_graph" + + +def _tolerance(reason: str, steps: int = 2) -> dict[str, Any]: + """Tolerance for ``steps`` unshared BF16 roundings (two: only the final casts differ).""" + return {"mode": "tolerance", "atol": BF16_ATOL, "rtol": steps * 2**-8, "reason": reason} + + +BF16 = "bfloat16" + +# pi0.5 action expert (Gemma 300M): width 1024, MLP 4096, 8 query heads, 1 KV +# head, head_dim 256, 50 action tokens. Prefix = 3 x 256 image + 200 text tokens. +# B=1 is serving; B=8 is a grouped RL rollout. +PI05_BATCHES = (1, 8) +PI05_TOKENS = 50 +PI05_PREFIX = 968 + +# Qwen R2R low-level / panoramic policies (Qwen2.5-VL-3B ViT through +# EmbodiInfer's Triton vision path): 16 heads, head_dim 80, 64-patch windows. +# ActiveVLN runs the Transformers vision tower and uses none of these kernels. +# StreamVLN language model (Qwen2-7B): hidden 3584, MLP 18944, decode rows. + +OPERATORS: tuple[Operator, ...] = ( + Operator( + name="ada_rms_norm", + summary="Adaptive RMSNorm with per-batch scale/shift and a materialized residual gate.", + models=("pi05",), + op_type="norm", + source="embodiinfer/backend/triton/norm.py", + entry="ada_rms_norm", + call="ada_rms_norm(x, modulation, eps)", + axes={"B": None, "T": None, "H": None, "H3": None, "ONE": 1}, + inputs={ + "x": Tensor(("B", "T", "H"), BF16), + "modulation": Tensor(("B", "H3"), BF16), + "eps": Tensor(None, "float32"), + }, + outputs={"output": Tensor(("B", "T", "H"), BF16), "gate": Tensor(("B", "ONE", "H"), BF16)}, + constraints=("H3 == 3 * H",), + reference="""import torch + + +def run(x, modulation, eps): + values = x.float() + variance = torch.mean(torch.square(values), dim=-1, keepdim=True) + scale, shift, gate = modulation.unsqueeze(1).chunk(3, dim=-1) + output = values * torch.rsqrt(variance + eps) + output = output * (1.0 + scale.float()) + shift.float() + return output.to(x.dtype), gate.to(x.dtype) +""", + workloads=tuple( + Workload( + f"pi05-expert-b{b}", + {"B": b, "T": PI05_TOKENS, "H": 1024, "H3": 3072}, + scalars={"eps": 1e-6}, + note="every action-expert layer, twice per layer", + ) + for b in PI05_BATCHES + ), + precision=_tolerance( + "The production kernel reduces the FP32 variance in a different order than torch.mean." + ), + notes="`modulation` is the AdaRMS projection in scale/shift/gate order, in the checkpoint " + "dtype (BF16 in pi05_libero_finetuned_v044, as captured). The gate output is cast to the " + "hidden dtype, one row per batch item.", + ), + Operator( + name="gated_residual", + summary="Residual plus gated update with BF16 rounding of the product before addition.", + models=("pi05",), + op_type="elementwise", + source="embodiinfer/backend/triton/norm.py", + entry="gated_residual", + call="gated_residual(residual, update, gate)", + axes={"B": None, "T": None, "H": None, "G": None}, + inputs={ + "residual": Tensor(("B", "T", "H"), BF16), + "update": Tensor(("B", "T", "H"), BF16), + "gate": Tensor(("B", "G", "H"), BF16), + }, + outputs={"output": Tensor(("B", "T", "H"), BF16)}, + constraints=("G == 1 or G == T",), + reference="""import torch + + +def run(residual, update, gate): + return residual + update * gate +""", + workloads=tuple( + Workload( + f"pi05-expert-b{b}", + {"B": b, "T": PI05_TOKENS, "H": 1024, "G": 1}, + note="after attention and MLP in every action-expert layer", + ) + for b in PI05_BATCHES + ), + notes="The gate broadcasts over tokens. Contracting multiply-add into an FMA " + "violates the bit-exact contract; the production kernel disables FP fusion.", + ), + Operator( + name="gated_gelu", + summary="Tanh-approximate GELU of the gate projection times the up projection.", + models=("pi05",), + op_type="activation", + source="embodiinfer/backend/triton/activation.py", + entry="gated_gelu", + call="gated_gelu(gate, up)", + axes={"B": None, "T": None, "I": None}, + inputs={"gate": Tensor(("B", "T", "I"), BF16), "up": Tensor(("B", "T", "I"), BF16)}, + outputs={"output": Tensor(("B", "T", "I"), BF16)}, + reference="""import torch.nn.functional as F + + +def run(gate, up): + return F.gelu(gate, approximate="tanh") * up +""", + workloads=tuple( + Workload(f"pi05-expert-b{b}", {"B": b, "T": PI05_TOKENS, "I": 4096}, note="every expert MLP") + for b in PI05_BATCHES + ), + precision=_tolerance( + "The production kernel evaluates tanh-GELU through its sigmoid identity in FP32, " + "which differs from Torch's tanh formulation; both round GELU(gate) to BF16.", + steps=4, + ), + notes="Rounding GELU(gate) to BF16 before multiplying by `up` matches the eager path.", + ), + Operator( + name="rotate_qk", + summary="Rotate-half RoPE on BHSD query and key with BF16 rounding of each product.", + models=("pi05",), + op_type="rope", + source="embodiinfer/backend/triton/rotary.py", + entry="rotate_qk", + call="rotate_qk(query, key, cos, sin)", + axes={"B": None, "HQ": None, "HK": None, "S": None, "D": None}, + inputs={ + "query": Tensor(("B", "HQ", "S", "D"), BF16), + "key": Tensor(("B", "HK", "S", "D"), BF16), + "cos": Tensor(("B", "S", "D"), BF16), + "sin": Tensor(("B", "S", "D"), BF16), + }, + outputs={ + "query_out": Tensor(("B", "HQ", "S", "D"), BF16), + "key_out": Tensor(("B", "HK", "S", "D"), BF16), + }, + constraints=("D == 64 or D == 128 or D == 256",), + reference="""import torch + + +def _rotate_half(x): + half = x.shape[-1] // 2 + return torch.cat((-x[..., half:], x[..., :half]), dim=-1) + + +def run(query, key, cos, sin): + cos = cos.unsqueeze(1) + sin = sin.unsqueeze(1) + return query * cos + _rotate_half(query) * sin, key * cos + _rotate_half(key) * sin +""", + workloads=tuple( + Workload( + f"pi05-expert-b{b}", + {"B": b, "HQ": 8, "HK": 1, "S": PI05_TOKENS, "D": 256}, + note="every action-expert attention layer", + ) + for b in PI05_BATCHES + ), + notes="In the model, query/key are transposed views; this task uses contiguous BHSD tensors.", + ), + Operator( + name="split_kv_attention", + summary="Non-causal GQA attention over a cached prefix K/V block and the current K/V block.", + models=("pi05",), + op_type="attention", + source="embodiinfer/backend/triton/split_kv_attention.py", + entry="split_kv_attention", + call="split_kv_attention(query, prefix_key, prefix_value, suffix_key, suffix_value, None, scaling)", + axes={"B": None, "HQ": None, "HK": None, "P": None, "S": None, "D": None}, + inputs={ + "query": Tensor(("B", "HQ", "S", "D"), BF16), + "prefix_key": Tensor(("B", "HK", "P", "D"), BF16), + "prefix_value": Tensor(("B", "HK", "P", "D"), BF16), + "suffix_key": Tensor(("B", "HK", "S", "D"), BF16), + "suffix_value": Tensor(("B", "HK", "S", "D"), BF16), + "scaling": Tensor(None, "float32"), + }, + outputs={"output": Tensor(("B", "HQ", "S", "D"), BF16)}, + constraints=("HQ % HK == 0", "D == 64 or D == 128 or D == 256"), + reference="""import torch + + +def run(query, prefix_key, prefix_value, suffix_key, suffix_value, scaling): + repeats = query.shape[1] // prefix_key.shape[1] + key = torch.cat((prefix_key, suffix_key), dim=2).repeat_interleave(repeats, dim=1) + value = torch.cat((prefix_value, suffix_value), dim=2).repeat_interleave(repeats, dim=1) + scores = torch.matmul(query.float(), key.float().transpose(-1, -2)) * scaling + return torch.matmul(torch.softmax(scores, dim=-1), value.float()).to(query.dtype) +""", + workloads=( + *( + Workload( + f"pi05-denoise-b{b}-prefix{PI05_PREFIX}", + {"B": b, "HQ": 8, "HK": 1, "P": PI05_PREFIX, "S": PI05_TOKENS, "D": 256}, + scalars={"scaling": 0.0625}, + note="every denoise step, every expert layer; three cameras", + ) + for b in PI05_BATCHES + ), + Workload( + "pi05-denoise-b1-prefix712", + {"B": 1, "HQ": 8, "HK": 1, "P": 712, "S": PI05_TOKENS, "D": 256}, + scalars={"scaling": 0.0625}, + note="two cameras", + ), + ), + precision=_tolerance( + "The production kernel uses split-KV online softmax with exp2; FP32 accumulation " + "order differs from the materialized Torch softmax." + ), + notes="The padding mask is omitted (all keys valid). The production call also " + "accepts a shared [B,1,1,P+S] padding mask; masked tuning needs a custom benchmark.py.", + ), + Operator( + name="rotate_half_rope", + summary="Rotate-half RoPE for packed vision Q/K, computed in FP32.", + models=("qwen25-vln",), + op_type="rope", + source="embodiinfer/backend/triton/packed_rope.py", + entry="rotate_half_rope", + call="rotate_half_rope(q, k, cos, sin)", + axes={"T": None, "HQ": None, "HK": None, "D": None}, + inputs={ + "q": Tensor(("T", "HQ", "D"), BF16), + "k": Tensor(("T", "HK", "D"), BF16), + "cos": Tensor(("T", "D"), BF16), + "sin": Tensor(("T", "D"), BF16), + }, + outputs={"q_out": Tensor(("T", "HQ", "D"), BF16), "k_out": Tensor(("T", "HK", "D"), BF16)}, + constraints=("D % 2 == 0", "D <= 256"), + reference="""import torch + + +def _rotate_half(x): + half = x.shape[-1] // 2 + return torch.cat((-x[..., half:], x[..., :half]), dim=-1) + + +def run(q, k, cos, sin): + cosine = cos.unsqueeze(-2).float() + sine = sin.unsqueeze(-2).float() + q_float, k_float = q.float(), k.float() + q_out = (q_float * cosine + _rotate_half(q_float) * sine).to(q.dtype) + k_out = (k_float * cosine + _rotate_half(k_float) * sine).to(k.dtype) + return q_out, k_out +""", + workloads=( + Workload("vit-1frame", {"T": 1024, "HQ": 16, "HK": 16, "D": 80}, note="one 448x448 frame"), + Workload("vit-4frames", {"T": 4096, "HQ": 16, "HK": 16, "D": 80}, note="four frames"), + ), + notes="Every vision block applies this before attention. head_dim 80 is not a power of two. " + "Products of BF16 values are exact in FP32, so the production kernel is bit-exact.", + ), + Operator( + name="segmented_attention", + summary="Non-causal attention independently inside packed variable-length segments.", + models=("qwen25-vln",), + op_type="attention", + source="embodiinfer/backend/triton/segmented_attention.py", + entry="segmented_attention", + call="segmented_attention(q, k, v, segment_offsets, max_query_length=max_length, " + "max_key_length=max_length)", + axes={"T": None, "HQ": None, "HK": None, "D": None, "N": None}, + inputs={ + "q": Tensor(("T", "HQ", "D"), BF16), + "k": Tensor(("T", "HK", "D"), BF16), + "v": Tensor(("T", "HK", "D"), BF16), + "segment_offsets": Tensor(("N",), "int32"), + "max_length": Tensor(None, "int32"), + }, + outputs={"output": Tensor(("T", "HQ", "D"), BF16)}, + constraints=("HQ % HK == 0", "D <= 256", "N >= 2"), + reference="""import torch + + +def run(q, k, v, segment_offsets, max_length): + del max_length # A bound for graph-safe kernels; the reference uses the offsets. + output = torch.empty_like(q) + repeats = q.shape[1] // k.shape[1] + scale = q.shape[-1] ** -0.5 + bounds = segment_offsets.tolist() + for start, end in zip(bounds[:-1], bounds[1:]): + query = q[start:end].float().transpose(0, 1) + key = k[start:end].float().repeat_interleave(repeats, dim=1).transpose(0, 1) + value = v[start:end].float().repeat_interleave(repeats, dim=1).transpose(0, 1) + probs = torch.softmax(torch.matmul(query, key.transpose(-1, -2)) * scale, dim=-1) + output[start:end] = torch.matmul(probs, value).transpose(0, 1).to(q.dtype) + return output +""", + workloads=( + Workload( + "vit-window-1frame", + {"T": 1024, "HQ": 16, "HK": 16, "D": 80, "N": 17}, + scalars={"max_length": 64}, + fixed={"segment_offsets": tuple(range(0, 1025, 64))}, + note="window-attention blocks, 16 windows of 64 patches", + ), + Workload( + "vit-window-ragged", + {"T": 1156, "HQ": 16, "HK": 16, "D": 80, "N": 21}, + scalars={"max_length": 64}, + fixed={"segment_offsets": (*range(0, 1025, 64), 1057, 1089, 1121, 1156)}, + note="partial edge windows of a non-multiple-of-window frame", + ), + Workload( + "vit-full-1frame", + {"T": 1024, "HQ": 16, "HK": 16, "D": 80, "N": 2}, + scalars={"max_length": 1024}, + fixed={"segment_offsets": (0, 1024)}, + note="the full-attention blocks, one segment per frame", + ), + ), + precision=_tolerance( + "The production kernel uses blockwise online softmax with exp2; FP32 accumulation " + "order differs from the materialized Torch softmax." + ), + notes="`segment_offsets` holds cumulative token positions; queries and keys share them. " + "`max_length` bounds every segment so the kernel can launch without host synchronization.", + ), + Operator( + name="rms_norm", + summary="RMSNorm of one decode row with FP32 variance.", + models=("streamvln",), + op_type="norm", + source="embodiinfer/backend/triton/norm.py", + entry="rms_norm", + call="rms_norm(x, weight, eps)", + axes={"M": None, "H": None}, + inputs={ + "x": Tensor(("M", "H"), BF16), + "weight": Tensor(("H",), BF16), + "eps": Tensor(None, "float32"), + }, + outputs={"output": Tensor(("M", "H"), BF16)}, + constraints=("M == 1",), + reference="""import torch + + +def run(x, weight, eps): + values = x.float() + normalized = values * torch.rsqrt(values.pow(2).mean(dim=-1, keepdim=True) + eps) + return normalized.to(x.dtype) * weight +""", + workloads=(Workload("qwen2-7b-decode", {"M": 1, "H": 3584}, scalars={"eps": 1e-6}),), + precision=_tolerance( + "The production kernel multiplies by the weight in FP32 and rounds once; the eager " + "path rounds the normalized row to BF16 first.", + steps=3, + ), + notes="The production dispatcher only routes single-row (decode) inputs to this kernel.", + ), + Operator( + name="add_rms_norm", + summary="Residual addition fused with RMSNorm of the sum for one decode row.", + models=("streamvln",), + op_type="norm", + source="embodiinfer/backend/triton/norm.py", + entry="add_rms_norm", + call="add_rms_norm(residual, update, weight, eps)", + axes={"M": None, "H": None}, + inputs={ + "residual": Tensor(("M", "H"), BF16), + "update": Tensor(("M", "H"), BF16), + "weight": Tensor(("H",), BF16), + "eps": Tensor(None, "float32"), + }, + outputs={"summed": Tensor(("M", "H"), BF16), "normalized": Tensor(("M", "H"), BF16)}, + constraints=("M == 1",), + reference="""import torch + + +def run(residual, update, weight, eps): + summed = residual + update + values = summed.float() + normalized = values * torch.rsqrt(values.pow(2).mean(dim=-1, keepdim=True) + eps) + return summed, normalized.to(summed.dtype) * weight +""", + workloads=(Workload("qwen2-7b-decode", {"M": 1, "H": 3584}, scalars={"eps": 1e-6}),), + precision=_tolerance( + "The production kernel normalizes the unrounded FP32 sum and applies the weight in " + "FP32; the eager path rounds the sum and the normalized row to BF16 first.", + steps=4, + ), + notes="Both outputs are compared. The production dispatcher only routes single-row inputs.", + ), + Operator( + name="swiglu", + summary="SiLU(gate) times up over a packed [gate | up] projection.", + models=("streamvln",), + op_type="activation", + source="embodiinfer/backend/triton/activation.py", + entry="swiglu", + # [M, 2, I] is the contiguous [M, 2 * I] projection output, viewed so that + # every axis is inferable from the inputs. + call="swiglu(packed.view(packed.shape[0], -1))", + axes={"M": None, "TWO": 2, "I": None}, + inputs={"packed": Tensor(("M", "TWO", "I"), BF16)}, + outputs={"output": Tensor(("M", "I"), BF16)}, + reference="""import torch.nn.functional as F + + +def run(packed): + gate, up = packed.unbind(1) + return F.silu(gate) * up +""", + workloads=( + Workload("qwen2-7b-decode", {"M": 1, "I": 18944}, weight=3.0, note="every decode step"), + Workload("qwen2-7b-chunk256", {"M": 256, "I": 18944}, note="prefill chunk"), + ), + precision=_tolerance( + "The production kernel keeps SiLU(gate) in FP32 and rounds once; the eager path " + "rounds SiLU(gate) to BF16 before the multiplication.", + steps=3, + ), + notes="In the model the input is the contiguous [M, 2 * I] gate_up projection: the first " + "half of each row is the gate, the second half is up. `packed[:, 0]` is the gate.", + ), +) + + +def catalog() -> dict[str, Operator]: + """Operators by task name, in declaration order.""" + return {operator.name: operator for operator in OPERATORS} + + +def select(names: list[str] | None = None, models: list[str] | None = None) -> list[Operator]: + """Choose operators by name and/or model; no filter selects every core operator.""" + known = catalog() + if unknown := set(names or ()) - known.keys(): + raise ValueError(f"Unknown operators: {sorted(unknown)}; known: {sorted(known)}") + known_models = {model for operator in OPERATORS for model in operator.models} + if unknown := set(models or ()) - known_models: + raise ValueError(f"Unknown models: {sorted(unknown)}; known: {sorted(known_models)}") + return [ + operator + for operator in OPERATORS + if (not names or operator.name in names) and (not models or set(models) & set(operator.models)) + ] diff --git a/scripts/kernel_tuning/permissions.py b/scripts/kernel_tuning/permissions.py new file mode 100644 index 0000000..b945ee1 --- /dev/null +++ b/scripts/kernel_tuning/permissions.py @@ -0,0 +1,36 @@ +"""What the coding agent of a tuning or registration flow may touch. + +Humanize2 holds a session to anything narrower than every scope ``ALL`` with an +OS fence that needs Linux Landlock. Roles are therefore declared with every +scope, so an explicitly unfenced run can start where Landlock is unavailable; +every other run narrows the agent back to Humanize2's default fence before its +first session opens. Only modules that already require hmz import this one. +""" + +from __future__ import annotations + +from hmz.flows import Agent, Permission, PermissionKind + +#: Humanize2's default grant: write the session workdir, read the rest, stay online. +FENCED = Permission() +#: Nothing fenced. Isolation is then the operator's account or container. +UNFENCED = Permission( + local=PermissionKind.ALL, + user=PermissionKind.ALL, + system=PermissionKind.ALL, + online=PermissionKind.ALL, +) + + +class Coder(Agent): + """A coding-agent role that may run unfenced; ``session_agent`` narrows it otherwise.""" + + _permission = UNFENCED + + +def session_agent(agent: Agent, *, unfenced: bool) -> Agent: + """Return the agent sessions must be opened with for the run's recorded fence choice. + + Fenced runs keep the exact permission the flows used before this option existed. + """ + return agent if unfenced else agent.derive(permission=FENCED) diff --git a/scripts/kernel_tuning/prompts.py b/scripts/kernel_tuning/prompts.py new file mode 100644 index 0000000..16b057e --- /dev/null +++ b/scripts/kernel_tuning/prompts.py @@ -0,0 +1,94 @@ +"""Task-specific coding instructions; the evaluator owns acceptance decisions.""" + +from __future__ import annotations + +import json +from dataclasses import asdict + +from .contracts import TaskPackage + + +def plan_prompt(task: TaskPackage) -> str: + """Ask for a reviewable hypothesis and executable plan before source generation.""" + replay = ( + "Read task/replay.json for the frozen model, tensor layouts, aliases and mutations. Optimize this one observed call; no neighboring-operator fusion or quantization.\n" + if "replay.json" in task.hashes + else "" + ) + if "numerics.json" in task.hashes: + replay += "Read task/numerics.json. It freezes per-workload/profile/output bounds against FP64, not byte equality with the production baseline. All elements and all profiles must pass; tuning.yaml thresholds are reporting maxima only.\n" + return f"""Optimize the operator in task/ for this target machine. +Read task/README.md, task/definition.json, task/workloads.jsonl, task/baseline.py, +task/tuning.yaml, environment.json, incumbent.json, and feedback.json before choosing one change. +Read task/HARDWARE.md when present for target-specific observations and prior experiments. +Feedback includes complete source for recent evaluated candidates, including rejected ones. +Compare their implementations with their measured failures before selecting a materially +different configuration or algorithm; do not repeat a rejected implementation by guessing its source. +{replay}\ +Write PLAN.md with: the current bottleneck; one falsifiable optimization hypothesis; +the source/launch configuration changes; correctness hazards; and validation steps. +Describe the validation criteria for the controller; do not compile, run GPU code, +benchmark, or inspect generated assembly during either agent turn. The agent's +command environment may restrict GPU access and compiler cache operations even +when environment.json confirms that the separate evaluator has a working GPU. +Finish this turn when PLAN.md contains that executable plan. The next turn implements it. +The mathematical reference and baseline have different roles: preserve the former's +outputs and beat the latter's latency over every declared workload. +Frozen conditions: {json.dumps(asdict(task.settings), ensure_ascii=False)} +Use only this attempt directory for edits. Task files, environment.json, incumbent.json, feedback.json, +the evaluator, run archives, and precision/timing settings are fixed inputs. +Create candidate source beside PLAN.md; do not add files inside task/. +Do not commit, install dependencies, alter global settings, or start another search. +""" + + +def implementation_prompt(task: TaskPackage) -> str: + """Specify a complete official Solution rather than interpreting PASS in prose.""" + schema = { + "name": "descriptive_candidate_name", + "definition": task.definition["name"], + "author": "coding-agent", + "spec": { + "language": task.settings.language, + "target_hardware": list(task.settings.target_hardware), + "entry_point": "kernel.py::run" if task.settings.language == "triton" else "kernel.cu::run", + "destination_passing_style": False, + **({"binding": "torch"} if task.settings.language == "cuda" else {}), + }, + "sources": [ + { + "path": "kernel.py" if task.settings.language == "triton" else "kernel.cu", + "content": "complete source text", + } + ], + } + replay = ( + "This captured registry task requires destination_passing_style=false. Preserve replay.json input/output layouts, aliases and mutations, returning mutated buffers in definition order. No new operator fusion or quantization.\n" + if "replay.json" in task.hashes + else "" + ) + return f"""Implement PLAN.md and write solution.json using the official FlashInfer +Solution structure below. Include every source/helper in sources; relative POSIX +paths and multiple files are supported. Include actual source text, not placeholders. +You may write source files here and serialize them into solution.json. +{json.dumps(schema, indent=2)} +The entry point matches definition.inputs in order. If destination_passing_style +is true, append all definition.outputs as output buffers; otherwise return outputs +in definition order. For CUDA binding=torch, export the symbol with +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m). tvm-ffi binding is also supported when installed. +{replay}\ +Preserve dtype and the declared numerical contract; bit_exact includes intermediate rounding, signed zero and NaN +bit patterns. Write the configured operator in {task.settings.language}; target +the declared hardware. Use only the permitted input metadata for specialization. +Return values must depend on the supplied tensors on every call and graph replay. +Launch on the caller's current stream; any helper-stream work must join it before +returning. The caller's device, stream, and numerical settings stay unchanged. +Keep evaluation logic, expected answers, precision, workloads, and metrics untouched. +The controller will compile, check all seeds, measure, and promote from structured +evidence after this turn returns. Do not compile, run GPU code, benchmark, inspect +assembly, or work around agent-environment permissions/cache failures. Static +source and JSON checks are sufficient for this handoff. A failed controller check +will be recorded as feedback for the next candidate. +Finish when solution.json contains the complete implementation and PLAN.md explains +the implemented change. The controller retains this attempt even when it fails. +""" diff --git a/scripts/kernel_tuning/pyproject.toml b/scripts/kernel_tuning/pyproject.toml new file mode 100644 index 0000000..e56e5ce --- /dev/null +++ b/scripts/kernel_tuning/pyproject.toml @@ -0,0 +1,15 @@ +[project] +name = "embodiinfer-kernel-tuning-tools" +version = "0.1.0" +description = "Optional repository-only Humanize2 kernel tuning environment" +requires-python = ">=3.12" +dependencies = [ + "hmz @ https://github.com/humanfia/humanize/archive/b70d4427f11308442aef20f9b2233dfbbc34a39e.zip", + "pyyaml>=6.0.3,<7", +] + +[dependency-groups] +dev = ["pytest>=8,<10", "ruff>=0.10,<1"] + +[tool.uv] +package = false diff --git a/scripts/kernel_tuning/regions.py b/scripts/kernel_tuning/regions.py new file mode 100644 index 0000000..a59341b --- /dev/null +++ b/scripts/kernel_tuning/regions.py @@ -0,0 +1,449 @@ +"""Capture operator-registry calls of an end-to-end run as a tunable model capture. + +``python -m scripts.kernel_tuning capture-ops --model MODEL --output DIR -- DRIVER`` +runs an unmodified inference driver in the model runtime. Calls that the policy +resolves through :mod:`embodiinfer.layers.ops` during the measured steps (after +``--warmup``) are recorded with their real inputs. The result is a model capture +in the shared capture format, so ``calibrate``, ``generate`` and ``tune-all`` apply +unchanged. Each operator's identity has ``kind: registry``: its baseline and +its FP64 oracle are the Torch reference in ``embodiinfer.backend.torch.ops``, +and the implementation the model ran is shipped as the agent's starting point. +Output layouts and aliases come from the call the model made, not from the +reference: a tuned kernel must hand its consumers what production hands them. +Where the reference lays its outputs out differently, the generated baseline +copies them into the production layout (values unchanged). + +Record in eager mode (no CUDA graphs, no ``torch.compile``): graph replays and +compiled regions do not call the registry from Python. +""" + +from __future__ import annotations + +import functools +import hashlib +import importlib +import runpy +import sys +import threading +from collections import Counter +from collections.abc import Callable +from pathlib import Path +from typing import Any + +from .artifacts import atomic_json +from .contracts import ContractError, digest, tree_hashes +from .hotspots import _ENTRIES + +REFERENCE_MODULE = "embodiinfer.backend.torch.ops" + + +def _aliases(sources: list[Any], outputs: list[Any]) -> list[list[int]]: + """For each output, the positions in ``sources`` that share its storage.""" + from .discovery.tensors import storage_key + + return [ + [i for i, value in enumerate(sources) if storage_key(value) == storage_key(output)] + for output in outputs + ] + + +def _dense(shape: list[int], stride: list[int]) -> bool: + """Whether a layout covers exactly its elements once, as ``torch.empty_strided`` can allocate.""" + if 0 in shape: + return True + expected = 1 + for step, size in sorted((step, size) for size, step in zip(shape, stride) if size != 1): + if step != expected: + return False + expected *= size + return True + + +def _portable(value: Any) -> Any: + if value is None or isinstance(value, (str, int, float, bool)): + return value + if isinstance(value, dict): + return {str(k): _portable(v) for k, v in value.items()} + if isinstance(value, (tuple, list)): + return [_portable(v) for v in value] + return str(value) + + +class OpCapture: + """Record registry-operator calls of the measured inference steps.""" + + def __init__( + self, + model: str, + output: Path, + *, + ops: set[str] | None = None, + warmup: int = 2, + steps: int = 8, + fixture_bytes: int = 1 << 30, + ) -> None: + if warmup < 0 or steps <= 0: + raise ContractError("capture-ops needs warmup >= 0 and steps > 0") + self.model, self.root, self.ops = model, output.resolve(), ops + self.warmup, self.steps, self.fixture_bytes = warmup, steps, fixture_bytes + if self.root.exists(): + raise ContractError(f"{self.root} already exists; captures are immutable") + self.root.mkdir(parents=True) + self.started = 0 + self.lock = threading.Lock() + self.local = threading.local() + self.operators: dict[str, dict[str, Any]] = {} + self.errors: list[str] = [] + self.target: dict[str, Any] | None = None + self.used_bytes = 0 + self._digests: dict[str, dict[str, str]] = {} + self._fixtures: dict[str, list[dict[str, Any]]] = {} + self._restore: list[Callable[[], None]] = [] + import torch + + self.precision = { + "matmul_precision": torch.get_float32_matmul_precision(), + "allow_tf32": torch.backends.cuda.matmul.allow_tf32, + } + self.flags = { + "cudnn_allow_tf32": torch.backends.cudnn.allow_tf32, + "fp16_reduced_reduction": torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction, + "bf16_reduced_reduction": torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction, + "deterministic_algorithms": torch.are_deterministic_algorithms_enabled(), + } + + @property + def measuring(self) -> bool: + """Whether the current thread is inside a measured inference step.""" + return getattr(self.local, "measured", False) + + def _replace(self, owner: Any, name: str, value: Any) -> None: + had = name in vars(owner) + original = vars(owner).get(name) + setattr(owner, name, value) + self._restore.append(lambda: setattr(owner, name, original) if had else delattr(owner, name)) + + def _replace_everywhere(self, original: Any, value: Any) -> None: + for module in list(sys.modules.values()): + if module is None: + continue + for name, item in list(vars(module).items()): + if item is original: + self._replace(module, name, value) + + def _step(self, function: Callable[..., Any]) -> Callable[..., Any]: + @functools.wraps(function) + def wrapper(*args: Any, **kwargs: Any) -> Any: + if getattr(self.local, "depth", 0): + return function(*args, **kwargs) + with self.lock: + index = self.started + self.started += 1 + self.local.depth = 1 + self.local.measured = self.warmup <= index < self.warmup + self.steps + try: + return function(*args, **kwargs) + finally: + self.local.depth = 0 + self.local.measured = False + + return wrapper + + def install(self) -> None: + """Patch policy construction, inference entries and registry resolution.""" + from embodiinfer.layers.ops import KernelSelection, get_op + from embodiinfer.policies import factory + + for module_name, owner, methods in _ENTRIES: + holder = getattr(importlib.import_module(module_name), owner) + for method in methods: + if hasattr(holder, method): + self._replace(holder, method, self._step(getattr(holder, method))) + + original_factory = factory.make_policy + + @functools.wraps(original_factory) + def make_policy(name: str, **kwargs: Any) -> Any: + if name != self.model: + raise ContractError(f"The driver constructed {name!r}, expected {self.model!r}") + policy = original_factory(name, **kwargs) + self.target = { + "model": name, + "type": f"{type(policy).__module__}.{type(policy).__qualname__}", + "config": _portable(kwargs), + } + return policy + + self._replace_everywhere(original_factory, make_policy) + original_resolve = KernelSelection.resolve + capture = self + + @functools.wraps(original_resolve) + def resolve(selection: Any, op: str, prefer: str = "torch", *args: Any, **kwargs: Any) -> Any: + function = original_resolve(selection, op, prefer, *args, **kwargs) + if capture.ops is not None and op not in capture.ops: + return function + implementation = get_op(op, "torch" if selection.mode == "torch" else prefer) + return capture._recording(op, implementation, function) + + self._replace(KernelSelection, "resolve", resolve) + + def _recording(self, op: str, implementation: Any, function: Any) -> Any: + @functools.wraps(function) + def wrapper(*args: Any, **kwargs: Any) -> Any: + import torch + + if ( + not self.measuring + or torch.compiler.is_compiling() + or (torch.cuda.is_available() and torch.cuda.is_current_stream_capturing()) + ): + return function(*args, **kwargs) + return self.record(op, implementation, function, args, kwargs) + + return wrapper + + def _source_digests(self, module: str) -> dict[str, str]: + if module not in self._digests: + from .generate import vendor_sources + + self._digests[module] = {path: value for path, (_, value) in vendor_sources(module).items()} + return self._digests[module] + + def record(self, op: str, implementation: Any, function: Any, args: tuple, kwargs: dict) -> Any: + """Snapshot one call, run it once, and take its output contract from that call. + + The output layout, aliases and dtypes recorded for the case are those of the + implementation the model actually ran (``function``), so downstream operators + and the tuned kernel see the production layout. The Torch reference runs on + cloned inputs and must agree on everything except strides and offsets; when + its layout differs, the case keeps it as ``reference_outputs`` and the + generated baseline returns the production layout. + """ + from embodiinfer.layers.ops import get_op + + from .discovery.tensors import ( + clone_inputs, + describe_inputs, + flatten, + restore, + tensor_outputs, + tensor_spec, + ) + + try: + tensors: list[Any] = [] + recipe = flatten((args, kwargs), tensors) + specs = describe_inputs(tensors) + reference = get_op(op, "torch") + local = clone_inputs(tensors) + ref_args, ref_kwargs = restore(recipe, local) + result = reference(*ref_args, **ref_kwargs) + reference_outputs = tensor_outputs(result) + if not reference_outputs: + raise ContractError("registry operators must return tensors") + identity = { + "kind": "registry", + "name": f"registry.{op}", + "op": op, + "module": REFERENCE_MODULE, + "entry": op, + "source_digests": self._source_digests(REFERENCE_MODULE), + "start_module": implementation.__module__, + "start_entry": implementation.__name__, + "start_digests": self._source_digests(implementation.__module__), + "arguments": recipe, + "input_dtypes": [spec["dtype"] for spec in specs], + "input_ranks": [len(spec["shape"]) for spec in specs], + "mutates": [], + "returns": flatten(result, []), + } + # The first call of each structural case supplies its real inputs; + # later calls with the same layout only add to the case's count. + # Registry operators are pure, so the inputs are unchanged until the call below. + structure = digest([identity, specs]) + if structure not in self._fixtures: + self._fixtures[structure] = self._snapshot(tensors, specs) + except (ContractError, RuntimeError, TypeError, ValueError) as exc: + self._error(op, exc) + return function(*args, **kwargs) + # The call the model makes, exactly once; its result is what the model uses. + actual = function(*args, **kwargs) + try: + outputs = tensor_outputs(actual) + if flatten(actual, []) != identity["returns"]: + raise ContractError( + "the production implementation returns another structure than the Torch reference" + ) + case = { + "inputs": self._fixtures[structure], + "outputs": [tensor_spec(tensor) for tensor in outputs], + "output_aliases": _aliases(tensors, outputs), + "output_groups": _aliases(outputs, outputs), + } + expected = [tensor_spec(tensor) for tensor in reference_outputs] + for spec, ref in zip(case["outputs"], expected): + if any(spec[key] != ref[key] for key in ("shape", "dtype", "device")): + raise ContractError( + "production and Torch reference outputs differ in shape, dtype or device" + ) + if ( + _aliases(local, reference_outputs) != case["output_aliases"] + or _aliases(reference_outputs, reference_outputs) != case["output_groups"] + ): + raise ContractError("production and Torch reference outputs alias different storage") + if expected != case["outputs"]: + for spec, ref, aliases in zip(case["outputs"], expected, case["output_aliases"]): + if spec != ref and ( + aliases or spec["offset"] or not _dense(spec["shape"], spec["stride"]) + ): + raise ContractError( + "the production output layout cannot be rebuilt from the Torch reference " + "(it aliases an input, has an offset, or is not dense)" + ) + case["reference_outputs"] = expected + identity["output_dtypes"] = [spec["dtype"] for spec in case["outputs"]] + identity["output_ranks"] = [len(spec["shape"]) for spec in case["outputs"]] + self._add(identity, case) + except (ContractError, RuntimeError, TypeError, ValueError) as exc: + self._error(op, exc) + return actual + + def _error(self, op: str, exc: Exception) -> None: + message = f"registry.{op}: {exc}" + if message not in self.errors: + self.errors.append(message) + + def _snapshot(self, tensors: list[Any], specs: list[dict[str, Any]]) -> list[dict[str, Any]]: + import torch + + from .discovery.tensors import extent + + specs = [dict(spec) for spec in specs] + for group in {spec["storage"] for spec in specs}: + indices = [i for i, spec in enumerate(specs) if spec["storage"] == group] + tensor = tensors[indices[0]] + size = max(extent(specs[i]) for i in indices) + data = tensor.detach().as_strided((size,), (1,), 0).contiguous().view(torch.uint8).cpu() + data = data.numpy().tobytes() + path = f"fixtures/{hashlib.sha256(data).hexdigest()}.bin" + file = self.root / path + if not file.exists(): + if self.used_bytes + len(data) > self.fixture_bytes: + raise ContractError("input fixtures exceed --fixture-bytes; raise the budget") + file.parent.mkdir(exist_ok=True) + file.write_bytes(data) + self.used_bytes += len(data) + for i in indices: + specs[i]["fixture"] = path + return specs + + def _add(self, identity: dict[str, Any], case: dict[str, Any]) -> None: + op_id = digest(identity)[:20] + op = self.operators.setdefault( + op_id, + {"id": op_id, "name": identity["name"], "identity": identity, "status": "ready", "workloads": []}, + ) + case_id = digest(case)[:20] + existing = next((item for item in op["workloads"] if item["id"] == case_id), None) + if existing is None: + existing = {**case, "id": case_id, "count": 0} + op["workloads"].append(existing) + existing["count"] += 1 + + def close(self, *, complete: bool) -> None: + """Restore every patch and publish the hashed capture.""" + import torch + + from .discovery.contracts import SCHEMA + from .generate import _git_revision + + while self._restore: + self._restore.pop()() + if self.target is None: + self.errors.append("No policy was built through make_policy") + if self.started < self.warmup + self.steps: + self.errors.append( + f"Only {self.started} inference steps ran; capture needs warmup {self.warmup} + {self.steps}" + ) + recorded = {op["identity"]["op"] for op in self.operators.values()} + for op in sorted((self.ops or set()) - recorded): + self.errors.append( + f"registry.{op} was not called during the measured steps; run the driver without " + "CUDA graphs or torch.compile so registry calls happen in Python" + ) + if not self.operators and not self.errors: + self.errors.append("No registry operator calls were recorded") + operators = list(self.operators.values()) + atomic_json(self.root / "operators.json", operators) + atomic_json( + self.root / "coverage.json", + {"operators": dict(Counter(op["status"] for op in operators)), "errors": self.errors}, + ) + config = (self.target or {}).get("config", {}) + atomic_json( + self.root / "manifest.json", + { + "schema": SCHEMA, + "capture": "registry", + "target": self.target or {}, + "complete": complete, + "errors": self.errors, + "window": {"warmup": self.warmup, "steps": self.steps}, + "repository_revision": _git_revision(), + "torch_version": str(torch.__version__), + "precision": self.precision, + "flags": self.flags, + "device_type": "cuda", + "compile_backend": config.get("compile_backend", "none"), + "fixture_budget_bytes": self.fixture_bytes, + "fixture_bytes": self.used_bytes, + "files": tree_hashes(self.root), + }, + ) + + +def run( + command: list[str], + output: Path, + model: str, + *, + ops: set[str] | None = None, + warmup: int = 2, + steps: int = 8, + fixture_bytes: int = 1 << 30, +) -> dict[str, Any]: + """Run an inference driver and write a registry-operator model capture.""" + from .discovery.contracts import ModelCapture + + if not command: + raise ContractError("Supply the inference driver after --") + capture = OpCapture(model, output, ops=ops, warmup=warmup, steps=steps, fixture_bytes=fixture_bytes) + argv, path = sys.argv, sys.path[:] + complete = False + try: + capture.install() + if command[0] == "-m": + sys.argv = command[1:] + runpy.run_module(command[1], run_name="__main__", alter_sys=True) + else: + script = Path(command[0]).resolve() + sys.argv = [str(script), *command[1:]] + sys.path.insert(0, str(script.parent)) + runpy.run_path(str(script), run_name="__main__") + complete = True + except SystemExit as exc: + complete = exc.code in (None, 0) + if not complete: + raise + finally: + sys.argv, sys.path = argv, path + capture.close(complete=complete) + result = ModelCapture.load(output, model) + return { + "model": model, + "identity": result.identity, + "operators": { + op["identity"]["op"]: sum(case["count"] for case in op["workloads"]) for op in result.operators + }, + } diff --git a/scripts/kernel_tuning/registration/README.md b/scripts/kernel_tuning/registration/README.md new file mode 100644 index 0000000..9faabc2 --- /dev/null +++ b/scripts/kernel_tuning/registration/README.md @@ -0,0 +1,359 @@ +# 算子注册工具（`register`） + +`python -m scripts.kernel_tuning register …` 把一个尚未注册的数学算子接入算子注册表 +（`embodiinfer/layers/ops.py`），并把选定的调用点改为经过注册表调用。设计归属见 +[0008 §4.2](../../../docs/proposals/0008-kernel-tuning.md#42-select-and-if-necessary-register-an-operator)。 + +工具由一个控制器驱动编码智能体，循环执行"计划 → 实现 → CPU 检查 → 反馈修复"，直到全部检查通过或预算用尽。 +最终产物是一个待人工审阅的分支和 `patch.diff`。 + +注册只改调用路径，不改数值：参考实现必须逐字保留原表达式的算术。以下工作不在本工具范围内： + +- 选择要注册的算子和调用点（由调用方指定）； +- 融合多个相邻算子； +- 调优、安装 kernel（`tune-all`/`run`、`install`、`verify` 等）； +- GPU 上的整模型逐位一致验证（合并前必须单独完成，见[合并前](#合并前)）； +- 合并分支（由人工完成）。 + +## 1. 前提条件 + +- **工具环境**：Python ≥ 3.12，并已安装 Humanize2 和 Ruff，与 `tune-all` 共用： + `uv sync --project scripts/kernel_tuning --python 3.12 --frozen`。 + `run` 和 `resume` 只能在 POSIX 系统上执行；`new`、`status`、`finalize` 不启动智能体。 + Humanize2 默认用 Linux Landlock 给智能体会话加 fence。如果内核没有 Landlock（例如 Jetson L4T）， + 需要在 `register run` 时加 `--unfenced`，并在没有 `sudo`、不在 `docker` 组的专用账号或容器中运行。 + 启动后工具环境的包版本（`hmz`、`pydantic`、`pyyaml`、`psutil`）被记录下来，变化后需要新建运行。 +- **测试解释器 `test_python`**：一个已准备好的模型运行环境，需包含 Torch、pytest，以及规格中 CPU + 测试需要的模型依赖。所有需要导入 `embodiinfer` 的检查都在这个解释器的子进程中执行。工具环境自带 + Torch 时可以省略，否则必须通过 `--set test_python=/absolute/python` 指定。 +- **干净的仓库**：`git status --porcelain` 必须为空（包括未跟踪文件）。输出目录应放在仓库外或被 + `.gitignore` 忽略的位置，默认的 `results/kernel_tuning/register` 满足要求。 +- **策略已接入注册表**：调用点所在文件必须已有模块级 `_op` helper，并能拿到 `KernelSelection`。 + 目前只有 pi0.5（`embodiinfer/policies/pi05/modeling_pi05.py`）满足。其他策略需要先人工接入一次。 + +## 2. 快速开始 + +```bash +TOOL="uv run --project scripts/kernel_tuning python -m scripts.kernel_tuning" + +# 1. 建立运行：校验规格、建立 worktree、定位调用点，打印 sites.md 后停下 +$TOOL register new --spec linear.yaml --set test_python=/abs/runtime/python +# 输出最后一行为 "Run: " + +# 2. 审阅 RUN/sites.md（需要时补充 RUN/sites.json），然后启动智能体循环 +$TOOL register run RUN --agent 'HARNESS/MODEL:EFFORT' + +# 中断后继续；查看状态 +$TOOL register resume RUN +$TOOL register status RUN + +# 3. 导出已验证的补丁；提交和推送都需要显式参数 +$TOOL register finalize RUN [--commit] [--push] +``` + +所有命令都在仓库根目录执行。`run`/`resume` 的退出码在状态为 `passed` 或 `finalized` 时为 0， +否则为 1。 + +## 3. 规格 + +规格决定本次运行注册哪个算子、修改哪些调用点、用哪些测试回归，以及预算。 + +```yaml +model: pi05 # 调用点所属的策略目录名 +op: linear # 新算子名 +description: > # 算子的数学含义，写入智能体上下文和 SUMMARY.md + F.linear(x, weight, bias)：weight 为 [out, in]，bias 可为 None；不改变 dtype。 +sites: # 要改的调用点：文件、所在函数、原表达式 + - source: embodiinfer/policies/pi05/modeling_pi05.py + function: _fused_qkv + expression: "F.linear(x, weight)" + - source: embodiinfer/policies/pi05/modeling_pi05.py + function: _mlp + expression: "mlp.down_proj(activated)" +preserve: # 必须保持原样的路径，作为审阅说明交给智能体 + - "QuantizedLinear 继续走量化注册表" +tests: # 每次检查都要通过的已有 CPU 测试 + - tests/test_kernel_ops.py + - tests/test_pi05_cpu_routes.py +budget: + max_attempts: 5 # 最多尝试次数 + agent_timeout_seconds: 1800 # 单次智能体调用的超时 + max_seconds: 7200 # 整次运行的累计时长上限，中断后继续累计 +``` + +校验规则（`spec.py`），不满足时在建立 worktree 之前就报错： + +- 不允许未知字段。 +- `model` 必须是标识符。 +- `op` 必须是公开的 Python 标识符：不能是关键字，不能以 `_` 开头。 + 它按数学含义命名，不能包含模型名。下游调优工具要求 `embodiinfer.backend.torch.ops` 中的参考函数与算子同名， + `install` 也用它作目录名。 +- `sites` 至少一项，且不能重复。 + - `source` 必须是 `embodiinfer/policies//` 下的 `.py` 相对路径。 + - `function` 是函数名，或 `Class.method` 形式的限定名。 + - `expression` 必须是一个调用表达式，按 `ast.unparse` 规范化后比较。 +- `tests` 是 `tests/` 下的 pytest 节点，默认 `tests/test_kernel_ops.py`。这些文件在 BASE 中必须存在。 +- `budget` 的值必须是正整数。 + +也可以不写规格文件，直接用命令行参数。此时 `description`、`preserve`、`tests` 和 `budget` 取默认值： + +```bash +$TOOL register new --model pi05 --op linear \ + --site 'embodiinfer/policies/pi05/modeling_pi05.py::_fused_qkv::F.linear(x, weight)' \ + --site 'embodiinfer/policies/pi05/modeling_pi05.py::_mlp::mlp.down_proj(activated)' +``` + +## 4. 流程 + +### 4.1 `register new`：建立运行 + +`new` 的参数： + +| 参数 | 含义 | +|---|---| +| `--spec FILE` 或 `--model/--op/--site` | 规格，二选一 | +| `--repository PATH` | 要注册进去的干净仓库，默认是控制器所在仓库。控制器代码有未提交改动时，可以指向另一个干净检出 | +| `--output DIR` | 运行目录的父目录，默认 `results/kernel_tuning/register` | +| `--set test_python=/abs/python` | 测试解释器，必须是已存在的绝对路径 | +| `--yes` | 在记录中标记已确认调用点；`new` 仍会停下，不会启动智能体 | + +控制器依次执行以下步骤，任何一步失败都会报错，并撤销已经建立的 worktree 和分支： + +1. 确认仓库干净，记录 `BASE = HEAD`。 +2. 用 `ast` 定位每个调用点（见 4.2），并确认规格中的测试文件存在。 +3. 生成 `run_id`（UTC 时间加随机后缀），建立 `RUN = /`。 +4. **预检**：`backend/torch/ops.py` 中不能已有同名函数。然后在 `test_python` 中以 `PYTHONPATH=<仓库>` + 导入 `embodiinfer.layers.available_ops`，算子不能已经注册，日志写入 `preflight.log`。 + 预检后再次确认仓库仍是干净的 `BASE`。 +5. 建立两个 worktree，主检出在整个过程中不被修改： + - `RUN/base`：detached 检出 `BASE`，作为只读基线； + - `RUN/worktree`：新分支 `registry/-`，智能体只在这里修改。 +6. 写入冻结的 `spec.json`（规格和 BASE）、`sites.json`、`sites.md` 和 `manifest.json`。 + 此时状态为 `awaiting_confirmation`。 + +最后打印 `sites.md` 和 `Run: `。 + +### 4.2 审阅调用点 + +`sites.json` 的每个调用点记录以下内容： + +- **位置**：表达式在函数体内必须恰好匹配一处，不进入嵌套函数，记录行号和列号。 +- **guard**：外层 `if` 的条件和分支。包含调用点的分支要改；不包含它的兄弟分支、以及调用点之前 + 会提前 `return` 的分支，标记为 `preserve_body`，这些分支的代码必须原样保留。 +- **`KernelSelection` 的来源**： + - 函数已有 `kernels` 参数时用 `kernels`； + - 方法所在文件用到 `self._kernels` 时用 `self._kernels`； + - 否则需要新增 `kernels` 参数，并由同一文件中列出的调用方逐层传递。 +- **调用方**：同一文件内直接或间接调用该函数的位置。 +- **源码片段**：调用点前后几行。 + +`sites.md` 是同一份内容的可读版本。审阅时可以： + +- 在 `sites.json` 的 `preserve` 中追加说明； +- 给调用点追加额外字段作为备注。 + +以下内容不能删改：调用点的 `id`、`source`、`function`、`expression`、`guards`，以及规格中已有的 +`preserve` 条目。需要换调用点时，新建一次运行。 + +### 4.3 `register run`：启动智能体 + +`run RUN --agent HARNESS/MODEL:EFFORT [--set test_python=…] [--yes] [--unfenced]` 执行一次即视为确认调用点，然后： + +1. 恢复并校验运行目录（见 4.6）。 +2. 用当前或新指定的 `test_python` 在 `RUN/base` 上重新做预检。 +3. 将 `sites.json` 与 BASE 上重新定位的结果比对，然后冻结 agent、`test_python`、`unfenced` 和 `sites.json` 哈希，状态变为 `ready`。 +4. 建立只读上下文 `RUN/context/`，内容如下： + - `AGENTS.md`； + - `embodiinfer/backend/torch/ops.py`、`embodiinfer/layers/ops.py`； + - `embodiinfer/policies/pi05/modeling_pi05.py`：已经通过 `_op` 和 `KernelSelection` + 走注册表的调用点，作为改法示例，取自 BASE； + - `tests/test_kernel_ops.py`：测试写法示例； + - 各调用点所在的源文件、`sites.json`、`RULES.md`。 +5. 启动 Humanize2 flow `operator_registration`（`flow.py`）。Humanize2 的 workspace 设为 `RUN`， + 智能体的 `LocalEnv` 是从中派生出的 `RUN/worktree`；框架日志归档到 `RUN/workspace/`， + 实际的 epic 路径记录在 `humanize.json` 中。 + +`run` 只能执行一次。之后使用 `resume`，它不接受更换 agent、`test_python` 或 `unfenced`。 + +**每次尝试的阶段：** + +| 尝试 | 阶段 | 会话 | +|---|---|---| +| 第一次 | `plan` → `implement` → 检查 | 新会话 | +| 之后 | `repair` → 检查 | 沿用上一次的会话 | +| 同一项检查连续失败两次后 | `plan` → `implement` → 检查 | 新会话，worktree 保留 | +| 智能体调用出错之后（超时、没写 `PLAN.md` 等） | `repair` → 检查 | 新会话 | + +- `plan`：只写 `PLAN.md`，内容包括数学定义及其与每个原表达式的逐项对应、参数中哪些是张量哪些是常量、 + `KernelSelection` 如何到达每个调用点、每个 preserve 分支如何保留、要新增哪些测试。 +- `implement` / `repair`：修改代码，并写出 `REGISTRATION.json`： + + ```json + { + "op": "linear", + "reference": "embodiinfer/backend/torch/ops.py::linear", + "sites": [{"id": "S1", "file": "embodiinfer/policies/pi05/modeling_pi05.py", + "function": "_fused_qkv", "call": "_op(kernels, \"linear\", False)(x, weight)"}], + "preserved": [{"site": "S1", "branch": "exact original guard"}], + "tests": ["tests/test_registry_linear.py::test_reference"] + } + ``` + +每轮智能体调用结束后，`PLAN.md` 必须存在且不为空。单轮超时取 `agent_timeout_seconds` 与剩余总时长中较小的一个。 +每个 prompt 都包含冻结的规格、规则全文和上一次的反馈，新开的会话也能独立接手。 + +**规则摘要**（全文见 `prompts.py` 中的 `RULES`，由检查强制执行）： + +1. 参考实现写在 `embodiinfer/backend/torch/ops.py`，函数名等于算子名。带类型注解和说明契约的 + docstring，只导入 `torch` 和标准库。参数只能是张量或常量，返回张量或张量元组，不原地修改输入。 +2. 逐字保留原算术：ATen 调用、顺序、形状、中间 dtype 转换和 `contiguous` 调用都不变。 + 例如不能把 `F.linear` 改写成 `x @ w.T`，不能合并或拆分 GEMM。 +3. 在 `layers/ops.py` 中单独写一行顶层的 + `register_lazy_op("", "torch", "embodiinfer.backend.torch.ops", "")`， + 不放进同时注册 Torch 和 Triton 的循环。 +4. 调用点改为 `_op(kernels, "", False)(...)`，即偏好 Torch 实现。`KernelSelection` 通过参数显式传递，不使用全局变量。 +5. preserve 分支原样保留。区分普通 `nn.Linear` 时用 `type(module) is nn.Linear`，不按模型名分支。 +6. 新增测试覆盖三类行为： + - 参考实现与原表达式在 CPU FP32 和 BF16 随机输入上 `torch.equal`； + - 调用点经过一个会记录调用的假调优档案； + - preserve 分支不经过注册表。 +7. 允许修改的范围见 C1。禁止删改已有测试，禁止安装依赖、运行 GPU 代码、执行 `git commit`/`push`/merge、 + 修改 Git 元数据或运行目录中控制器拥有的文件。 + +### 4.4 CPU 检查（C1–C5） + +控制器按顺序执行检查，前一项失败时后面的不再执行。所有检查子进程都满足以下条件： + +- 以 `PYTHONPATH=<检出目录>` 运行，并断言导入的 `embodiinfer` 确实来自该检出目录（开发环境中 + EmbodiInfer 通常是可编辑安装）； +- 隐藏 CUDA/ROCm 设备，禁止自动加载 pytest 插件； +- 超时或取消时终止整个进程组。 + +| 检查 | 内容 | +|---|---| +| C1 范围 | 改动文件取自 `git diff ` 加上所有未跟踪文件（包括被忽略的，缓存目录除外）。只允许出现：`embodiinfer/backend/torch/ops.py`、`embodiinfer/layers/ops.py`、调用点源文件、`tests/test_kernel_ops.py`、`tests/test_registry_.py`，以及 `PLAN.md`、`REGISTRATION.json`。不允许删除文件。已有测试文件只能新增内容，原有行和顶层语句不变。`KernelSelection` 类不变 | +| C2 结构 | ① `REGISTRATION.json` 指向正确的算子和参考实现；覆盖每个调用点恰好一次；列出的测试是 BASE 中不存在的新测试函数，位于上述两个测试文件中，不能是文件或单个参数化实例。② 参考函数是模块级同步函数，带 docstring 和合法注解，没有 `*args`/`**kwargs`，所在文件只导入 `torch` 和标准库，函数内没有原地修改。③ `layers/ops.py` 中恰好有一行对应的顶层 `register_lazy_op`。④ 每个调用点都包含报告的 `_op(..., False)` 调用；`kernels` 是函数参数，同一文件内的调用链逐层传递它；guard 及 preserve 分支与 BASE 相同；原表达式只出现在 preserve 分支中，或出现在 `type(x) is nn.Linear` 不成立的那一支。⑤ 在 `test_python` 中确认：`available_ops()` 包含该算子的 `torch` 实现，`get_op` 返回参考实现，`KernelSelection("none").prepare()` 成功，`KernelSelection("torch").resolve` 返回参考实现 | +| C3 格式 | 改动文件通过 `ruff check` 和 `ruff format --check`（在工具环境中运行） | +| C4 测试 | 在 `RUN/worktree` 中以 `-m "not gpu"` 运行：规格中的测试一起跑；每个新测试单独跑。新测试必须实际执行，不能被跳过 | +| C5 有效性 | 把新测试文件复制到 `RUN/base`，逐个运行新测试。每个新测试的所有用例都必须执行，并以断言失败（退出码 1）。通过、跳过、收集错误或依赖缺失都不能作为有效证据 | + +检查结果写入 `attempts//checks.json`，并与当时源码的指纹绑定。中断后恢复时，只有源码字节未变， +已通过的检查才会复用。检查过程中源码发生变化会被判为失败。 + +### 4.5 反馈、重试和停止 + +检查失败时，控制器写出 `attempts//feedback.json`，下一次尝试的 prompt 会附上它： + +```json +{"check": "C2", "summary": "S2: original expression remains outside a preserve branch at line 315", "details": {}} +``` + +智能体调用本身出错时，`check` 为 `agent`。每次尝试结束时，当时的补丁都会保存为 `attempts//patch.diff`。 + +| 结果 | 状态 | +|---|---| +| C1–C5 全部通过 | `passed`，记录已验证的源码指纹，并写出 `RUN/patch.diff` 和 `RUN/SUMMARY.md` | +| 达到 `max_attempts` 或 `max_seconds` | `exhausted`，worktree 保留，供人工接手 | +| 冻结输入被改动（规格、BASE、`RUN/base`、分支、提交历史、已审阅的调用点、只读上下文），或控制器、运行环境出错 | `failed`，不会转为反馈交给智能体重试 | +| Ctrl-C 或取消 | `interrupted`，可用 `resume` 继续 | + +### 4.6 `resume` 和 `status` + +- `status RUN` 打印 `manifest.json`，不加锁，也不启动智能体。 +- `resume RUN` 根据 `manifest.json` 和 Humanize2 的状态台账继续： + - 已完成的阶段不重跑； + - 已用时长每秒落盘一次，中断不会重置预算； + - 状态已是 `passed`、`finalized`、`failed` 或 `exhausted` 时直接返回。 + +所有修改运行目录的命令都持有 `RUN/.run.lock`。C5 往 `RUN/base` 复制测试之前，会先把这些文件记入 +`manifest.json` 的 `base_overlay`；任何失败、取消或下一次启动时，都会先用 BASE 还原这些文件。 + +### 4.7 `register finalize`：导出和提交 + +`finalize` 只接受状态为 `passed` 或 `finalized`、且源码指纹与通过检查时一致的运行。 +它会重新写出 `RUN/patch.diff` 和 `RUN/SUMMARY.md`。补丁不包含 `PLAN.md` 和 `REGISTRATION.json`。 + +| 参数 | 行为 | +|---|---| +| 无 | 打印 worktree 路径、分支名、建议的提交说明和产物路径，不修改 Git | +| `--commit` | 只暂存通过 C1 的文件（暂存区里有其他文件时报错），在 worktree 中以 `feat: route calls through the registry operator` 提交。提交使用当前 Git 身份，会运行仓库的 hook。状态变为 `finalized` | +| `--push` | 需要同时传 `--commit`，或此前已经提交。推送前再次确认文件字节与已验证的一致，然后执行 `git push -u origin registry/-` | + +本工具不提供合并功能。 + +#### 合并前 + +CPU 检查不能证明 GPU 上逐位一致。例如 pi0.5 的 native 路径只在 CUDA 上启用，CPU 测试执行不到其中的调用点。 +合并前必须在目标 GPU 上按 [0008 §6](../../../docs/proposals/0008-kernel-tuning.md#6-losslessness-and-precision-criterion) +确认：`kernel_profile: none` 时，整模型输出与 BASE 逐位一致。两边要用相同的 checkpoint、dtype、 +attention 语义、随机种子和噪声、matmul precision、batch 形状和 decode/denoise 步数，并设置 +`TORCHINDUCTOR_DETERMINISTIC=1`。`SUMMARY.md` 的末尾也会写明这一项尚未完成。 + +## 5. Git 的使用和清理 + +智能体执行命令时没有交互确认，控制器因此用 Git 做隔离和取证： + +- **隔离**：两个 worktree 共享主仓库的对象库，主检出始终不被修改。 +- **基线**：`RUN/base` 固定在 BASE。C5 需要在未修改的代码上运行新测试；还原 overlay 时用 + `git restore --source=`。 +- **防篡改**：每轮都核对 `RUN/base` 的 HEAD 和文件哈希、`RUN/worktree` 的分支和 HEAD。 + 智能体自行提交会导致运行失败。 +- **范围和证据**：C1 用 `git diff` 和 `git ls-files --others` 列出改动；补丁用 `git diff --binary` + 生成，再补上未跟踪的新文件。 +- **交付**：只有在显式传入 `--commit`/`--push` 时才提交或推送。 + +运行结束后，worktree 和分支会留在仓库里，需要手动清理： + +```bash +git worktree remove --force RUN/worktree +git worktree remove RUN/base +git branch -D registry/- # 已推送或已合并后再删 +rm -rf RUN +# 如果先删除了 RUN 目录：git worktree prune +``` + +## 6. 运行目录 + +``` +// + manifest.json 状态、尝试记录、已用时长、冻结的 agent/test_python、各类哈希 + spec.json 冻结的规格和 BASE + sites.json sites.md 调用点及审阅说明 + preflight.log new 的注册表预检；run 的预检写入 test-runtime.log + context/ 智能体的只读上下文（校验哈希） + base/ BASE 的 detached worktree + worktree/ 分支 registry/- 的 worktree + workspace/ 归档的 Humanize2 日志；humanize.json 记录 epic 路径 + attempts// -prompt.txt、-reply.txt、PLAN.md、REGISTRATION.json、 + checks.json、feedback.json、patch.diff，以及 registry/ruff/tests/baseline 的日志和 JUnit 报告 + patch.diff SUMMARY.md 通过后的导出 + .run.lock +``` + +状态的变化顺序为：`awaiting_confirmation` → `ready` → `running` → `passed` → `finalized`。 +中途可能变为 `exhausted`、`failed` 或 `interrupted`。 + +## 7. 代码布局 + +| 文件 | 内容 | +|---|---| +| `spec.py` | 规格的解析（YAML 或命令行）和校验 | +| `sites.py` | 用 `ast` 定位调用点、提取 guard 和调用方，审阅前后比对，生成 `sites.md` | +| `workspace.py` | `RegistrationRun`：worktree、冻结输入的校验、只读上下文、子进程执行、改动列表、补丁、C1 | +| `checks.py` | C2–C5 以及检查调度 | +| `prompts.py` | 规则和 `plan`/`implement`/`repair` prompt | +| `flow.py` | Humanize2 flow `operator_registration` | +| `finalize.py` | `SUMMARY.md`、补丁导出、提交和推送 | +| `cli.py` | `register` 子命令、环境检查、Humanize2 启动 | + +CPU 测试位于 `tests/kernel_tuning/test_registration.py`，使用仓库的临时克隆、真实的 CPU 子进程和 +Humanize2 自带的假智能体，不需要 checkpoint 或 GPU。 + +## 8. 限制 + +- CPU 检查和有限的测试不能证明 GPU 上逐位一致，也不能保证智能体写的测试覆盖了全部行为。 +- 每次运行只注册一个算子，不支持融合。调用点只能位于 `embodiinfer/policies//` 下， + 且所在文件必须已有 `_op` 集成。`kernels` 传递链的检查只覆盖同一文件内的调用方。 +- 检查全部通过的补丁仍可能结构不佳（例如传参链过长），需要人工审阅。 +- 目录边界、哈希校验和隐藏加速设备都不构成操作系统级的沙箱。智能体和测试以当前用户的权限运行， + 必要时请在隔离的账户或容器中运行。 diff --git a/scripts/kernel_tuning/registration/__init__.py b/scripts/kernel_tuning/registration/__init__.py new file mode 100644 index 0000000..1b8b973 --- /dev/null +++ b/scripts/kernel_tuning/registration/__init__.py @@ -0,0 +1 @@ +"""Agent-assisted, controller-verified operator registration (developer tools only).""" diff --git a/scripts/kernel_tuning/registration/checks.py b/scripts/kernel_tuning/registration/checks.py new file mode 100644 index 0000000..4b522f5 --- /dev/null +++ b/scripts/kernel_tuning/registration/checks.py @@ -0,0 +1,417 @@ +"""Ordered controller checks; candidate claims never substitute for execution.""" + +from __future__ import annotations + +import ast +import shutil +import sys +import time +import xml.etree.ElementTree as ET +from pathlib import Path +from typing import Any + +from scripts.kernel_tuning.artifacts import atomic_json +from scripts.kernel_tuning.contracts import ContractError, read_json + +from .sites import body_nodes, branch_dump, function_node, functions, source_path +from .spec import test_node +from .workspace import REFERENCE, REGISTRY, RegistrationRun, check_scope, command, file_hash, import_guard + + +def registration(run: RegistrationRun) -> dict[str, Any]: + """Validate agent output and require newly introduced executable pytest nodes.""" + document = read_json(source_path(run.worktree, "REGISTRATION.json")) + if not isinstance(document, dict): + raise ContractError("REGISTRATION.json must be an object") + if document.get("op") != run.spec.op or document.get("reference") != f"{REFERENCE}::{run.spec.op}": + raise ContractError("REGISTRATION.json names the wrong operator/reference") + tests = document.get("tests") + if ( + not isinstance(tests, list) + or not tests + or not all(isinstance(t, str) for t in tests) + or len(tests) != len(set(tests)) + ): + raise ContractError("REGISTRATION.json must list unique new test node IDs") + allowed = {"tests/test_kernel_ops.py", f"tests/test_registry_{run.spec.op}.py"} + for item in tests: + parts = test_node(item).split("::") + if len(parts) not in (2, 3) or parts[0] not in allowed or "[" in item: + raise ContractError("List whole new test functions (all parametrizations), not files or cases") + name = ".".join(parts[1:]) + if not parts[-1].startswith("test_"): + raise ContractError(f"Not an executable test node: {item}") + tree = ast.parse(source_path(run.worktree, parts[0]).read_text(encoding="utf-8")) + function_node(tree, name) + if (run.base / parts[0]).exists(): + old = ast.parse((run.base / parts[0]).read_text(encoding="utf-8")) + if any(qualified == name for qualified, _ in functions(old)): + raise ContractError(f"Test already existed at BASE: {item}") + sites = document.get("sites") + expected = read_json(run.root / "sites.json")["sites"] + if ( + not isinstance(sites, list) + or len(sites) != len(expected) + or not all(isinstance(s, dict) and isinstance(s.get("id"), str) for s in sites) + or {s.get("id") for s in sites} != {s["id"] for s in expected} + ): + raise ContractError("REGISTRATION.json must describe every selected site exactly once") + if not isinstance(document.get("preserved"), list): + raise ContractError("REGISTRATION.json must include preserved branch evidence") + for item in document["preserved"]: + if ( + not isinstance(item, dict) + or item.get("site") not in {s["id"] for s in expected} + or not isinstance(item.get("branch"), str) + ): + raise ContractError("Malformed preserved branch evidence") + return document + + +def _annotation(node: ast.expr | None, *, result: bool = False) -> bool: + if node is None: + return False + if isinstance(node, ast.Constant): + if isinstance(node.value, str): + return _annotation(ast.parse(node.value, mode="eval").body, result=result) + return node.value is Ellipsis if result else node.value is None + if isinstance(node, ast.Name): + return node.id in ({"Tensor"} if result else {"Tensor", "int", "float", "bool", "str"}) + if isinstance(node, ast.Attribute): + return ast.unparse(node) == "torch.Tensor" + if isinstance(node, ast.BinOp) and isinstance(node.op, ast.BitOr): + return _annotation(node.left, result=result) and _annotation(node.right, result=result) + if isinstance(node, ast.Subscript) and ast.unparse(node.value) in ( + {"tuple", "Tuple"} if result else {"Optional", "Union"} + ): + items = node.slice.elts if isinstance(node.slice, ast.Tuple) else [node.slice] + return all(_annotation(item, result=result) for item in items) + return False + + +def _torch_route(call: ast.AST, op: str) -> bool: + if not isinstance(call, ast.Call) or not isinstance(call.func, ast.Call): + return False + helper = call.func + return ( + isinstance(helper.func, ast.Name) + and helper.func.id == "_op" + and len(helper.args) == 3 + and ast.unparse(helper.args[0]) in {"kernels", "self._kernels"} + and isinstance(helper.args[1], ast.Constant) + and helper.args[1].value == op + and isinstance(helper.args[2], ast.Constant) + and helper.args[2].value is False + and not helper.keywords + ) + + +def _preserved_expression(node: ast.AST, function: ast.AST, guards: list[dict[str, Any]]) -> bool: + for parent in body_nodes(function): + if not isinstance(parent, ast.If): + continue + for arm in ("body", "orelse"): + if not any(node in ast.walk(stmt) for stmt in getattr(parent, arm)): + continue + condition = ast.unparse(parent.test) + if any(g["preserve_body"] and g["condition"] == condition and g["arm"] == arm for g in guards): + return True + # The original module call may survive on the non-plain-Linear arm. + test = parent.test + if isinstance(test, ast.Compare) and len(test.ops) == 1 and len(test.comparators) == 1: + left = test.left + is_type = ( + isinstance(left, ast.Call) and isinstance(left.func, ast.Name) and left.func.id == "type" + ) + plain_linear = ast.unparse(test.comparators[0]) in {"nn.Linear", "torch.nn.Linear"} + if ( + is_type + and plain_linear + and ( + (isinstance(test.ops[0], ast.Is) and arm == "orelse") + or (isinstance(test.ops[0], ast.IsNot) and arm == "body") + ) + ): + return True + return False + + +def _check_parameter_chain(tree: ast.AST, function: ast.FunctionDef | ast.AsyncFunctionDef) -> None: + """Reject dropped selections in same-file callers, including newly added parameters.""" + pending = [function] + seen = set() + while pending: + callee = pending.pop() + if callee.name in seen: + continue + seen.add(callee.name) + positional = [a.arg for a in (*callee.args.posonlyargs, *callee.args.args)] + if "kernels" not in positional + [a.arg for a in callee.args.kwonlyargs]: + continue + for qualified, caller in functions(tree): + for call in body_nodes(caller): + if ( + not isinstance(call, ast.Call) + or not isinstance(call.func, ast.Name) + or call.func.id != callee.name + ): + continue + value = next((kw.value for kw in call.keywords if kw.arg == "kernels"), None) + if value is None and "kernels" in positional and len(call.args) > positional.index("kernels"): + value = call.args[positional.index("kernels")] + if value is None or ast.unparse(value) not in {"kernels", "self._kernels"}: + raise ContractError( + f"KernelSelection not forwarded by {qualified}:{call.lineno} to {callee.name}" + ) + if ast.unparse(value) == "kernels": + parameters = { + a.arg for a in (*caller.args.posonlyargs, *caller.args.args, *caller.args.kwonlyargs) + } + if "kernels" not in parameters: + raise ContractError(f"KernelSelection is not a parameter of {qualified}") + pending.append(caller) + + +def check_structure(run: RegistrationRun) -> dict[str, Any]: + """C2 static gate: reference signature, explicit registration and routed sites.""" + document = registration(run) + reference_tree = ast.parse((run.worktree / REFERENCE).read_text(encoding="utf-8")) + reference = function_node(reference_tree, run.spec.op) + if reference not in reference_tree.body or isinstance(reference, ast.AsyncFunctionDef): + raise ContractError("Reference must be a module-level synchronous function") + args = (*reference.args.posonlyargs, *reference.args.args, *reference.args.kwonlyargs) + if ( + not ast.get_docstring(reference) + or not _annotation(reference.returns, result=True) + or not all(_annotation(a.annotation) for a in args) + or reference.args.vararg + or reference.args.kwarg + ): + raise ContractError("Reference requires a contract docstring and tensor/constant annotations") + for node in ast.walk(reference_tree): + if isinstance(node, (ast.Import, ast.ImportFrom)): + modules = ( + [alias.name for alias in node.names] if isinstance(node, ast.Import) else [node.module or ""] + ) + if (isinstance(node, ast.ImportFrom) and node.level) or any( + m.split(".")[0] not in sys.stdlib_module_names | {"torch"} for m in modules + ): + raise ContractError("Torch reference may import only torch and the standard library") + for node in ast.walk(reference): + if isinstance(node, (ast.AugAssign, ast.Delete)) or ( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and node.func.attr.endswith("_") + and not node.func.attr.endswith("__") + ): + raise ContractError("Reference must not mutate inputs in place") + if isinstance(node, ast.Assign) and any( + isinstance(t, (ast.Subscript, ast.Attribute)) for t in node.targets + ): + raise ContractError("Reference must not assign into tensor inputs") + registry = ast.parse((run.worktree / REGISTRY).read_text(encoding="utf-8")) + expected = ast.parse( + f"register_lazy_op({run.spec.op!r}, 'torch', 'embodiinfer.backend.torch.ops', {run.spec.op!r})" + ).body[0] + if sum(ast.dump(n) == ast.dump(expected) for n in registry.body) != 1: + raise ContractError( + "Operator requires its own top-level register_lazy_op line, outside the Triton loop" + ) + reported = {item["id"]: item for item in document["sites"]} + for site in read_json(run.root / "sites.json")["sites"]: + item = reported[site["id"]] + if item.get("file") != site["source"] or item.get("function") != site["function"]: + raise ContractError(f"{site['id']}: reported site does not match reviewed source") + tree = ast.parse(source_path(run.worktree, site["source"]).read_text(encoding="utf-8")) + function = function_node(tree, site["function"]) + call = ast.parse(item.get("call", ""), mode="eval").body + if not _torch_route(call, run.spec.op): + raise ContractError(f"{site['id']}: use _op(kernels, op, False)(...) with Torch preference") + nodes = body_nodes(function) + if not any(isinstance(n, ast.Call) and ast.dump(n) == ast.dump(call) for n in nodes): + raise ContractError(f"{site['id']}: reported registry call is absent") + selection = ast.unparse(call.func.args[0]) + args = {a.arg for a in (*function.args.posonlyargs, *function.args.args, *function.args.kwonlyargs)} + if (selection == "kernels" and "kernels" not in args) or ( + selection == "self._kernels" and "self" not in args + ): + raise ContractError(f"{site['id']}: KernelSelection must be passed, not a global") + if selection == "kernels": + _check_parameter_chain(tree, function) + for guard in site["guards"]: + matches = [ + n for n in nodes if isinstance(n, ast.If) and ast.unparse(n.test) == guard["condition"] + ] + if not matches or ( + guard["preserve_body"] + and not any(branch_dump(getattr(n, guard["arm"])) == guard["body"] for n in matches) + ): + raise ContractError(f"{site['id']}: preserve branch changed: {guard['condition']}") + for node in nodes: + if ( + isinstance(node, ast.expr) + and ast.unparse(node) == site["expression"] + and not _preserved_expression(node, function, site["guards"]) + ): + raise ContractError( + f"{site['id']}: original expression remains outside a preserve branch at line {node.lineno}" + ) + return {"reference": document["reference"], "sites": list(reported), "tests": document["tests"]} + + +async def registry_probe(run: RegistrationRun, output: Path, timeout: float) -> dict[str, Any]: + """Check actual public registry resolution in the selected model interpreter.""" + code = import_guard(run.worktree) + ( + "from embodiinfer.layers import available_ops, get_op, KernelSelection\n" + f"from embodiinfer.backend.torch.ops import {run.spec.op} as reference\n" + f"assert 'torch' in available_ops()[{run.spec.op!r}]\n" + f"assert get_op({run.spec.op!r}) is reference\n" + "KernelSelection('none').prepare()\n" + f"assert KernelSelection('torch').resolve({run.spec.op!r}) is reference\n" + ) + result = await command( + [run.manifest["test_python"], "-c", code], root=run.worktree, output=output, timeout=timeout + ) + if result["returncode"]: + raise ContractError(f"Public registry probe failed: {result['output']}") + return result + + +async def pytest_nodes( + run: RegistrationRun, root: Path, nodes: list[str], output: Path, timeout: float +) -> dict[str, Any]: + """Run only CPU tests with an origin check and machine-readable outcome evidence.""" + report = output.with_suffix(".xml") + report.unlink(missing_ok=True) + args = [*nodes, "-q", "-m", "not gpu", "-p", "no:cacheprovider", f"--junitxml={report}"] + code = import_guard(root) + f"import pytest\nraise SystemExit(pytest.main({args!r}))\n" + result = await command( + [run.manifest["test_python"], "-c", code], root=root, output=output, timeout=timeout + ) + result["tests"] = nodes + cases = [] + if report.is_file(): + for case in ET.parse(report).iter("testcase"): + failure = case.find("failure") + cases.append( + { + "name": case.get("name"), + "class": case.get("classname"), + "failed": case.find("failure") is not None, + "error": case.find("error") is not None, + "skipped": case.find("skipped") is not None, + "assertion": failure is not None + and failure.get("message", "").startswith(("assert ", "AssertionError")), + } + ) + result["cases"] = cases + return result + + +async def baseline_tests(run: RegistrationRun, output: Path, deadline: float) -> dict[str, Any]: + """C5: every new test must fail on BASE; restore even on worker cancellation.""" + document = registration(run) + files = {node.split("::")[0] for node in document["tests"]} + run.verify() + run.manifest["base_overlay"] = {name: file_hash(source_path(run.worktree, name)) for name in files} + run.save() # Record expected overlay bytes before touching the detached checkout. + results = [] + try: + for name in files: + target = source_path(run.base, name) + target.parent.mkdir(parents=True, exist_ok=True) + shutil.copyfile(source_path(run.worktree, name), target) + for index, node in enumerate(document["tests"]): + result = await pytest_nodes( + run, run.base, [node], output / f"baseline-{index}.log", deadline - time.monotonic() + ) + results.append(result) + # Collection/import/environment errors are not evidence of an effective test. + if ( + result["returncode"] != 1 + or not result["cases"] + or not all( + c["failed"] and c["assertion"] and not c["skipped"] and not c["error"] + for c in result["cases"] + ) + ): + raise ContractError( + f"C5: {node} did not fail as an executed test on BASE (exit {result['returncode']}): {result['output']}" + ) + finally: + run.recover_base() + atomic_json(output / "baseline-results.json", results) + run.verify() + return {"commands": results} + + +async def run_checks( + run: RegistrationRun, archive: Path, *, seconds: float +) -> tuple[list[dict[str, Any]], dict[str, Any] | None]: + """Execute C1–C5 in order; reuse passed gates only for unchanged source bytes.""" + deadline = time.monotonic() + seconds + run.verify() + identity = run.fingerprint() + path = archive / "checks.json" + previous = read_json(path) if path.exists() else {} + results = previous.get("checks", []) if previous.get("fingerprint") == identity else [] + results = [item for item in results if item["passed"]] + for check in ("C1", "C2", "C3", "C4", "C5"): + if any(item["check"] == check for item in results): + continue + details: dict[str, Any] = {} + try: + if time.monotonic() >= deadline: + raise TimeoutError("Registration duration exhausted") + if check == "C1": + details = check_scope(run) + elif check == "C2": + details = check_structure(run) + details["command"] = await registry_probe( + run, archive / "registry.log", deadline - time.monotonic() + ) + elif check == "C3": + details["commands"] = [] + files = check_scope(run)["files"] + for label, args in (("lint", ["check"]), ("format", ["format", "--check"])): + result = await command( + [sys.executable, "-m", "ruff", *args, *files], + root=run.worktree, + output=archive / f"ruff-{label}.log", + timeout=deadline - time.monotonic(), + ) + details["commands"].append(result) + if result["returncode"]: + raise ContractError(f"Ruff {label} failed: {result['output']}") + elif check == "C4": + details["commands"] = [] + # Run declared regression tests together, new nodes separately so that + # a deselected/skipped parametrization cannot hide among passing tests. + selections = [list(run.spec.tests), *[[node] for node in registration(run)["tests"]]] + for index, nodes in enumerate(selections): + result = await pytest_nodes( + run, run.worktree, nodes, archive / f"tests-{index}.log", deadline - time.monotonic() + ) + details["commands"].append(result) + if ( + result["returncode"] + or not result["cases"] + or (index and any(c["skipped"] for c in result["cases"])) + ): + raise ContractError( + f"CPU tests failed or new tests were skipped: {nodes}: {result['output']}" + ) + else: + details = await baseline_tests(run, archive, deadline) + run.verify() + if run.fingerprint() != identity: + raise ContractError("Source changed while checks were running") + except (ContractError, SyntaxError, ValueError, KeyError, TypeError, OSError) as exc: + results.append({"check": check, "passed": False, "details": details, "summary": str(exc)[:12000]}) + feedback = {"check": check, "summary": str(exc)[:12000], "details": details} + atomic_json(path, {"fingerprint": identity, "checks": results}) + atomic_json(archive / "feedback.json", feedback) + return results, feedback + results.append({"check": check, "passed": True, "details": details}) + atomic_json(path, {"fingerprint": identity, "checks": results}) + return results, None diff --git a/scripts/kernel_tuning/registration/cli.py b/scripts/kernel_tuning/registration/cli.py new file mode 100644 index 0000000..637254b --- /dev/null +++ b/scripts/kernel_tuning/registration/cli.py @@ -0,0 +1,178 @@ +"""The register command group; inspection never starts a coding agent.""" + +from __future__ import annotations + +import argparse +import asyncio +import json +import os +import shutil +import sys +from pathlib import Path +from typing import Any + +from scripts.kernel_tuning.artifacts import ( + REPOSITORY, + UNFENCED_HELP, + atomic_json, + run_lock, + runtime_identity, + warn_unfenced, +) +from scripts.kernel_tuning.contracts import ContractError + +from .finalize import finalize +from .spec import load_spec +from .workspace import RegistrationRun, check_unregistered + + +def add_parser(commands: Any) -> None: + """Add a distinct namespace without changing existing tuning commands.""" + parser = commands.add_parser("register", help="Register one operator through a CPU-checked agent loop") + actions = parser.add_subparsers(dest="registration_command", required=True) + new = actions.add_parser("new", help="Create worktrees and stop for call-site review") + new.add_argument("--spec", type=Path) + new.add_argument("--model") + new.add_argument("--op") + new.add_argument("--site", action="append", default=[]) + new.add_argument( + "--repository", type=Path, default=REPOSITORY, help="Clean model checkout to register into" + ) + new.add_argument("--output", type=Path, default=REPOSITORY / "results/kernel_tuning/register") + new.add_argument("--set", action="append", default=[], metavar="test_python=/abs/python") + new.add_argument( + "--yes", action="store_true", help="Record call-site confirmation without an interactive prompt" + ) + run = actions.add_parser("run", help="Confirm reviewed sites and start plan/implement/check/repair") + run.add_argument("run", type=Path) + run.add_argument("--agent", required=True) + run.add_argument("--set", action="append", default=[], metavar="test_python=/abs/python") + run.add_argument("--yes", action="store_true", help="Explicitly confirm the reviewed sites") + run.add_argument("--unfenced", action="store_true", help=UNFENCED_HELP) + for name in ("resume", "status"): + actions.add_parser(name).add_argument("run", type=Path) + finish = actions.add_parser("finalize", help="Print verified branch/patch; optionally commit or push") + finish.add_argument("run", type=Path) + finish.add_argument("--commit", action="store_true") + finish.add_argument("--push", action="store_true") + + +def test_python(settings: list[str]) -> str | None: + """Accept only an existing absolute interpreter path, never shell fragments.""" + if not settings: + return None + if len(settings) != 1 or not settings[0].startswith("test_python="): + raise ContractError("The only register --set override is test_python=/absolute/python") + value = settings[0].split("=", 1)[1] + path = Path(value) + if not path.is_absolute() or not path.is_file(): + raise ContractError("test_python must be an existing absolute executable path") + return str(path) # Do not resolve a venv's python symlink into its base interpreter. + + +def execute(run: RegistrationRun, *, resume: bool) -> dict[str, Any]: + """Start the fixed Humanize2 flow in its POSIX tooling environment.""" + if os.name != "posix": + raise ContractError("Humanize2 registration run/resume requires POSIX; use the Linux target host") + if sys.version_info < (3, 12): + raise ContractError("Humanize2 requires a separate Python 3.12+ tooling environment") + from hmz.flows import BudgetExceeded + from hmz.sdk import Hmz + + run.recover_base() + run.verify() + if runtime_identity() != run.manifest["tool_runtime"]: + raise ContractError("Tool runtime changed; start a new registration run") + remaining = run.spec.budget.max_seconds - run.manifest["elapsed_seconds"] + if run.manifest["status"] in {"passed", "finalized", "failed", "exhausted"}: + return run.manifest + if remaining <= 0: + run.manifest.update(status="exhausted", stop_reason="max_seconds") + run.save() + return run.manifest + warn_unfenced(run.manifest) + # LocalEnv cannot derive an absolute path or '..'. Root it at RUN, then give + # only RUN/worktree to the agent. Archive framework journals in workspace/. + running = Hmz(workspace=run.root).run( + flow=Path(__file__).with_name("flow.py"), + task=f"Register {run.spec.op} and route its call sites through the registry", + agents={"coder": run.manifest["agent"]}, + params={"run_dir": str(run.root)}, + budget={"duration": remaining, "graceful": False}, + resume=resume, + ) + try: + return running.run() + except BudgetExceeded: + stopped = RegistrationRun(run.root) + stopped.recover_base() + stopped.manifest.update(status="exhausted", stop_reason="max_seconds") + stopped.save() + return stopped.manifest + finally: + archived = None + if running.epic: + archived = run.root / "workspace" / running.epic.name + archived.mkdir(parents=True, exist_ok=True) + for path in running.epic.glob("*.jsonl"): + if path.is_file() and not path.is_symlink(): + shutil.copyfile(path, archived / path.name) + atomic_json( + run.root / "humanize.json", + { + "epic": str(running.epic) if running.epic else None, + "archived": str(archived) if archived else None, + }, + ) + + +def dispatch(args: argparse.Namespace) -> int: + """Perform one register action, locking every mutation of an existing run.""" + action = args.registration_command + if action == "new": + spec = load_spec(args.spec, model=args.model, op=args.op, sites=args.site) + python = test_python(args.set) or sys.executable + run = asyncio.run(RegistrationRun.create(args.repository, args.output, spec, python, yes=args.yes)) + print((run.root / "sites.md").read_text(encoding="utf-8")) + print(f"Run: {run.root}") + return 0 + run = RegistrationRun(args.run) + if action == "status": + print(json.dumps(run.manifest, indent=2, ensure_ascii=False)) + return 0 + with run_lock(run.root): + run = RegistrationRun(args.run) # Reload after acquiring the writer lock. + if action == "finalize": + result = finalize(run, commit=args.commit, push=args.push) + else: + if action == "run": + if not args.agent or "/" not in args.agent: + raise ContractError("Choose --agent HARNESS/MODEL:EFFORT explicitly") + if run.manifest["sites_hash"]: + raise ContractError("Registration already started; use register resume RUN") + python = test_python(args.set) + # Invoking `run` is the documented confirmation after reviewing sites.md. + run.recover_base() + run.verify() + asyncio.run( + check_unregistered( + run.base, + run.spec.op, + python or run.manifest["test_python"], + run.root / "test-runtime.log", + ) + ) + run.prepare(args.agent, python, unfenced=args.unfenced) + elif not run.manifest["sites_hash"]: + raise ContractError("Use register run RUN --agent ... to confirm the call sites first") + try: + result = execute(run, resume=action == "resume") + except ContractError as exc: + # Platform/dependency issues leave a prepared run available on its host. + if os.name == "posix": + stopped = RegistrationRun(run.root) + stopped.manifest.update(status="failed", stop_reason=str(exc)) + stopped.save() + raise + print(json.dumps(result, indent=2, ensure_ascii=False)) + return 0 if action == "finalize" or result["status"] in {"passed", "finalized"} else 1 diff --git a/scripts/kernel_tuning/registration/finalize.py b/scripts/kernel_tuning/registration/finalize.py new file mode 100644 index 0000000..900d565 --- /dev/null +++ b/scripts/kernel_tuning/registration/finalize.py @@ -0,0 +1,99 @@ +"""Review artifacts and explicitly requested commit/push operations.""" + +from __future__ import annotations + +import json +from typing import Any + +from scripts.kernel_tuning.contracts import ContractError, read_json + +from .workspace import RegistrationRun, check_scope, git + + +def write_outputs(run: RegistrationRun) -> None: + """Export the complete verified patch and CPU evidence without Git mutations.""" + run.verify() + if run.manifest["status"] not in {"passed", "finalized"} or run.fingerprint() != run.manifest.get( + "verified_fingerprint" + ): + raise ContractError("Only the unchanged, fully checked patch can be finalized") + (run.root / "patch.diff").write_text(run.patch(), encoding="utf-8") + review = read_json(run.root / "sites.json") + lines = [ + f"# Register {run.spec.op}", + "", + run.spec.description or f"Register {run.spec.op} for {run.spec.model}.", + "", + f"BASE: `{run.base_commit}`", + f"Branch: `{run.manifest['branch']}`", + f"Attempts: {len(run.manifest['attempts'])}", + "", + "## Call sites", + "", + ] + lines.extend( + f"- {s['id']}: `{s['source']}::{s['function']}` — `{s['expression']}`" for s in review["sites"] + ) + lines.extend(["", "## Preserved branches", "", *[f"- {item}" for item in review["preserve"]]]) + lines.extend( + f"- {s['id']}: `{g['condition']}` ({g['arm']})" for s in review["sites"] for g in s["guards"] + ) + lines.extend(["", "## Controller checks", ""]) + for attempt in run.manifest["attempts"]: + lines.append(f"### Attempt {attempt['id']}: {attempt['status']}") + for check in attempt["checks"]: + lines.extend( + [ + "", + f"{check['check']}: {'passed' if check['passed'] else 'failed'}", + "", + "```json", + json.dumps(check, indent=2, ensure_ascii=False), + "```", + ] + ) + lines.extend( + [ + "", + "## Verification still required", + "", + "GPU whole-model bit-exact parity has NOT been run. Before merging, compare against BASE on the target GPU under proposal 0009 §6 with kernel_profile: none, identical checkpoint, dtype, attention semantics, seeds/noise, matmul precision, batch shape and decode/denoise schedule. CPU gates do not prove GPU parity.", + "", + ] + ) + (run.root / "SUMMARY.md").write_text("\n".join(lines), encoding="utf-8") + + +def finalize(run: RegistrationRun, *, commit: bool = False, push: bool = False) -> dict[str, Any]: + """Commit/push only with explicit flags, never merge or touch the main checkout.""" + run.recover_base() + write_outputs(run) + message = f"feat: route {run.spec.model} calls through the {run.spec.op} registry operator" + if push and not commit and not run.manifest.get("commit"): + raise ContractError("--push requires --commit or an already finalized commit") + if commit and not run.manifest.get("commit"): + files = check_scope(run)["files"] + # Agent staging must not smuggle PLAN.md or other artifacts into the commit. + staged = set(filter(None, git(run.worktree, "diff", "--cached", "--name-only", "-z").split("\0"))) + if not staged <= set(files): + raise ContractError("Index contains files outside the verified patch") + git(run.worktree, "add", "--", *files) + git(run.worktree, "commit", "-m", message) + run.manifest.update(commit=git(run.worktree, "rev-parse", "HEAD").strip(), status="finalized") + run.save() + if push: + # Commit hooks may change files or the index. Never push bytes that were + # not checked, even when commit and push were requested together. + write_outputs(run) + git(run.worktree, "push", "-u", "origin", run.manifest["branch"]) + run.manifest["pushed"] = True + run.save() + return { + "worktree": str(run.worktree), + "branch": run.manifest["branch"], + "suggested_commit": message, + "commit": run.manifest.get("commit"), + "pushed": run.manifest.get("pushed", False), + "patch": str(run.root / "patch.diff"), + "summary": str(run.root / "SUMMARY.md"), + } diff --git a/scripts/kernel_tuning/registration/flow.py b/scripts/kernel_tuning/registration/flow.py new file mode 100644 index 0000000..a460af8 --- /dev/null +++ b/scripts/kernel_tuning/registration/flow.py @@ -0,0 +1,214 @@ +"""Resumable Humanize2 plan/implementation/check/repair loop for registration.""" + +from __future__ import annotations + +import asyncio +import contextlib +import copy +import time +from datetime import timedelta +from pathlib import Path +from typing import Any + +from hmz.flows import ( + AgentCollection, + Budget, + BudgetExceeded, + EnvCollection, + FlowContext, + FlowParams, + LocalEnv, + flow, +) + +from scripts.kernel_tuning.artifacts import atomic_json +from scripts.kernel_tuning.contracts import ContractError +from scripts.kernel_tuning.permissions import Coder, session_agent +from scripts.kernel_tuning.registration.checks import run_checks +from scripts.kernel_tuning.registration.finalize import write_outputs +from scripts.kernel_tuning.registration.prompts import prompt +from scripts.kernel_tuning.registration.workspace import RegistrationRun, source_path + + +class Agents(AgentCollection): + """One explicitly selected coding-agent harness.""" + + coder: Coder + + +class Envs(EnvCollection): + """Run directory; only its worktree child is given to the coding agent.""" + + workspace: LocalEnv + + +class Params(FlowParams): + """All mutable progress is in the run; the flow accepts only its identity.""" + + run_dir: str + + +@flow(agents=Agents, envs=Envs, params=Params, resumable=True, name="operator_registration") +async def operator_registration( + task: str, *, agents: Agents, envs: Envs, params: Params, ctx: FlowContext +) -> dict[str, Any]: + """Keep a session across repairs, restarting planning after two equal failures.""" + del task + run = RegistrationRun(Path(params.run_dir)) + run.recover_base() # An interrupted C5 overlay lives in the external ledger. + state = ctx.state + if state is None: + raise ContractError("Registration requires resumable Humanize2 state") + if "ledger" in state: + ledger = state["ledger"] + if any(ledger[k] != run.manifest[k] for k in ("run_id", "spec_hash", "base_commit", "sites_hash")): + raise ContractError("Humanize2 journal belongs to a different registration") + ledger["elapsed_seconds"] = max(ledger["elapsed_seconds"], run.manifest["elapsed_seconds"]) + ledger["base_overlay"] = {} + run.manifest = ledger + coder = session_agent(agents["coder"], unfenced=bool(run.manifest.get("unfenced"))) + cfg = run.spec.budget + elapsed, started = run.manifest["elapsed_seconds"], time.monotonic() + + def checkpoint() -> None: + run.manifest["elapsed_seconds"] = elapsed + time.monotonic() - started + state["ledger"] = copy.deepcopy(run.manifest) + run.save() + + def remaining() -> float: + return max(0.0, cfg.max_seconds - elapsed - (time.monotonic() - started)) + + async def heartbeat() -> None: + while True: + await asyncio.sleep(1) + checkpoint() + + ticker = None + session = None + try: + run.verify() + if run.manifest["status"] in {"passed", "finalized", "exhausted", "failed"}: + return copy.deepcopy(run.manifest) + if not run.manifest["sites_hash"] or not run.manifest["confirmed"]: + raise ContractError("Call sites must be confirmed before starting an agent") + work_env = await envs["workspace"].derive_subdir(subdir="worktree") + if Path(str(work_env.workdir)).resolve() != run.worktree.resolve(): + raise ContractError("Humanize2 LocalEnv must be rooted at RUN") + run.manifest["status"] = "running" + checkpoint() + ticker = asyncio.create_task(heartbeat()) + while True: + run.verify() + if remaining() <= 0 or ( + not run.manifest["pending"] and len(run.manifest["attempts"]) >= cfg.max_attempts + ): + run.manifest.update( + status="exhausted", stop_reason="max_seconds" if remaining() <= 0 else "max_attempts" + ) + break + if run.manifest["pending"]: + attempt = run.manifest["attempts"][-1] + else: + previous = run.manifest["attempts"] + restart = run.manifest.get("failure_streak", 0) >= 2 + stage = "plan" if not previous or restart else "repair" + if restart: + session = None + run.manifest.update(failure_streak=0, last_failure=None) + attempt = { + "id": f"{len(previous) + 1:04d}", + "status": "running", + "stage": stage, + "checks": [], + } + previous.append(attempt) + run.manifest["pending"] = attempt["id"] + checkpoint() + archive = run.root / "attempts" / attempt["id"] + archive.mkdir(parents=True, exist_ok=True) + feedback = None + try: + while attempt["stage"] != "checks": + stage = attempt["stage"] + run.verify() + if remaining() <= 0: + raise TimeoutError("Registration duration exhausted") + if session is None: + session = await coder.spawn(env=work_env) + previous_feedback = run.manifest.get("feedback") + instruction = prompt(run, stage, previous_feedback) + (archive / f"{stage}-prompt.txt").write_text(instruction, encoding="utf-8") + seconds = min(cfg.agent_timeout_seconds, remaining()) + reply = await asyncio.wait_for( + coder.run( + instruction, + session=session, + budget=Budget(duration=timedelta(seconds=seconds), graceful=False), + ), + timeout=seconds, + ) + (archive / f"{stage}-reply.txt").write_text(reply, encoding="utf-8") + run.verify() + plan = source_path(run.worktree, "PLAN.md") + if not plan.is_file() or not plan.read_text(encoding="utf-8").strip(): + raise ContractError("Agent did not write a nonempty PLAN.md") + (archive / "PLAN.md").write_bytes(plan.read_bytes()) + attempt["stage"] = "implement" if stage == "plan" else "checks" + checkpoint() + agent_record = source_path(run.worktree, "REGISTRATION.json") + if agent_record.is_file(): + (archive / "REGISTRATION.json").write_bytes(agent_record.read_bytes()) + attempt["checks"], feedback = await run_checks(run, archive, seconds=remaining()) + except (asyncio.CancelledError, KeyboardInterrupt): + raise + except Exception as exc: + run.recover_base() + run.verify() # Immutable-input damage is fatal, never agent repair feedback. + feedback = { + "check": "agent", + "summary": f"{type(exc).__name__}: {exc}"[:12000], + "details": {}, + } + atomic_json(archive / "feedback.json", feedback) + session = None + finally: + (archive / "patch.diff").write_text(run.patch(), encoding="utf-8") + attempt.update( + status="failed" if feedback else "passed", + failure=feedback["check"] if feedback else None, + stage="done", + ) + run.manifest.update(pending=None, feedback=feedback) + if feedback: + same = feedback["check"] == run.manifest.get("last_failure") + run.manifest.update( + failure_streak=run.manifest.get("failure_streak", 0) + 1 if same else 1, + last_failure=feedback["check"], + ) + if feedback is None: + run.verify() + run.manifest.update(status="passed", verified_fingerprint=run.fingerprint()) + checkpoint() + write_outputs(run) + break + checkpoint() + checkpoint() + return copy.deepcopy(run.manifest) + except (BudgetExceeded, TimeoutError): + run.manifest.update(status="exhausted", stop_reason="max_seconds") + raise + except BaseException as exc: + run.manifest.update( + status="interrupted" + if isinstance(exc, (asyncio.CancelledError, KeyboardInterrupt)) + else "failed", + stop_reason=f"{type(exc).__name__}: {exc}", + ) + raise + finally: + if ticker: + ticker.cancel() + with contextlib.suppress(asyncio.CancelledError): + await ticker + session = None + checkpoint() diff --git a/scripts/kernel_tuning/registration/prompts.py b/scripts/kernel_tuning/registration/prompts.py new file mode 100644 index 0000000..62cdeff --- /dev/null +++ b/scripts/kernel_tuning/registration/prompts.py @@ -0,0 +1,98 @@ +"""Agent instructions for lossless registration, separate from controller evidence.""" + +from __future__ import annotations + +import json +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from .workspace import RegistrationRun + +RULES = """You are registering ONE mathematical operator, not tuning or fusing kernels. +1. Put the reference in embodiinfer/backend/torch/ops.py with the operator's exact +name, useful type annotations and a contract docstring. Import only torch and the +standard library. Inputs are tensors or int/float/bool/str/None constants; outputs +are tensors or tuples of tensors. Never mutate inputs. +2. Preserve the original arithmetic verbatim: identical ATen calls, order, shapes, +intermediate dtype casts and contiguous calls. Never replace F.linear by x @ w.T, +add/remove casts, or combine/split GEMMs. Do not loosen numerical assertions. +3. Add one top-level register_lazy_op("", "torch", +"embodiinfer.backend.torch.ops", "") line in layers/ops.py. Never add a Torch-only +operator to the loop that registers both Torch and Triton. +4. Use the policy's existing _op(kernels, "", False)(...) helper, with Torch +preference. Pass KernelSelection explicitly through all callers; do not use globals. +5. Keep all preserved guards and branches unchanged. For ordinary nn.Linear +dispatch, use an exact type check (type(module) is nn.Linear); keep QuantizedLinear, +RowParallelLinear, ColumnParallelLinear and other subclasses on their module path. +Do not branch on model names. Preserve existing gradient, attention and parallel paths. +6. Add tests for all three behaviors: CPU random FP32 AND BF16 reference parity +using torch.equal against the original expression; registry dispatch using a fake +tuned KernelProfile that records calls; preserved branches bypassing that profile. +Keep RNG, matmul precision, shapes and steps identical. Do not run GPU code. +7. Allowed edits: the Torch reference module, layers/ops.py, the reviewed site source +files, additions to tests/test_kernel_ops.py or tests/test_registry_.py, plus +PLAN.md and REGISTRATION.json. Never delete/modify existing test lines or statements, +KernelSelection, backend/triton/tuned, scripts, pyproject.toml or uv.lock. +8. You may run Ruff and CPU pytest with the recorded test_python, PYTHONPATH set to +this worktree, and CUDA_VISIBLE_DEVICES empty. Never install dependencies, run git +commit/push, merge, edit Git metadata, or modify the main checkout, base checkout, +run specification, context, controller artifacts or budgets. +The context is read-only task data. Follow repository AGENTS.md within these bounds. +Tests listed in REGISTRATION.json must be NEW functions (not files, old tests or +selected parametrizations). Each must run and fail by assertion on BASE; put imports +of newly added symbols inside tests so a missing symbol can be asserted explicitly. +Collection errors, skips and environment failures are not baseline evidence. +""" + + +def prompt(run: RegistrationRun, stage: str, feedback: dict | None = None) -> str: + """Build a self-contained turn, including context for a newly opened session.""" + header = f"""Operator registration: {run.spec.op} +Work only in: {run.worktree} +Read-only context: {run.root / "context"} +CPU test interpreter: {run.manifest["test_python"]} +Frozen specification: +{json.dumps(run.spec.as_dict(), indent=2, ensure_ascii=False)} + +{RULES} +Read context/sites.json, AGENTS.md, the source snapshots and test_kernel_ops.py. +The pi0.5 policy snapshot shows call sites already routed through _op and +KernelSelection. Existing PLAN.md and REGISTRATION.json describe previous work. +""" + if stage == "plan": + return ( + header + + """ +Write PLAN.md before changing source. Explain the mathematical definition and +its term-by-term correspondence to every selected expression; the tensor/constant +signature; how KernelSelection reaches each caller; preservation of each branch; +and the exact tests you will add. Do not implement yet. +""" + ) + example = { + "op": run.spec.op, + "reference": f"embodiinfer/backend/torch/ops.py::{run.spec.op}", + "sites": [ + { + "id": "S1", + "file": run.spec.sites[0].source, + "function": run.spec.sites[0].function, + "call": f'_op(kernels, "{run.spec.op}", False)(x, weight)', + } + ], + "preserved": [{"site": "S1", "branch": "exact original guard"}], + "tests": [f"tests/test_registry_{run.spec.op}.py::test_reference"], + } + instruction = ( + "Implement PLAN.md. " + if stage == "implement" + else "Repair the current worktree using the controller feedback below; update PLAN.md as needed. " + ) + return ( + header + + instruction + + "Write REGISTRATION.json covering EVERY selected site and ALL new tests. Schema example:\n" + + json.dumps(example, indent=2) + + "\nController feedback:\n" + + json.dumps(feedback, indent=2, ensure_ascii=False) + ) diff --git a/scripts/kernel_tuning/registration/sites.py b/scripts/kernel_tuning/registration/sites.py new file mode 100644 index 0000000..243621e --- /dev/null +++ b/scripts/kernel_tuning/registration/sites.py @@ -0,0 +1,210 @@ +"""AST call-site discovery and reviewable preservation constraints.""" + +from __future__ import annotations + +import ast +from pathlib import Path +from typing import Any + +from scripts.kernel_tuning.contracts import ContractError + +from .spec import RegistrationSpec + + +def source_path(root: Path, name: str) -> Path: + """Resolve a regular file inside a checkout, refusing symlink components.""" + candidate = root / name + if not candidate.resolve().is_relative_to(root.resolve()): + raise ContractError(f"File escapes checkout: {name}") + if any(p.is_symlink() for p in (candidate, *candidate.parents) if p != root.parent): + raise ContractError(f"Symlink is not a registration input: {name}") + return candidate + + +def functions(tree: ast.AST) -> list[tuple[str, ast.FunctionDef | ast.AsyncFunctionDef]]: + """Enumerate functions with lexical qualification, including class methods.""" + result = [] + + def visit(node: ast.AST, parents: tuple[str, ...]) -> None: + for child in ast.iter_child_nodes(node): + if isinstance(child, (ast.ClassDef, ast.FunctionDef, ast.AsyncFunctionDef)): + names = (*parents, child.name) + if not isinstance(child, ast.ClassDef): + result.append((".".join(names), child)) + visit(child, names) + else: + visit(child, parents) + + visit(tree, ()) + return result + + +def function_node(tree: ast.AST, name: str) -> ast.FunctionDef | ast.AsyncFunctionDef: + """Require exactly one matching function; qualified names disambiguate methods.""" + matches = [ + node + for qualified, node in functions(tree) + if qualified == name or ("." not in name and node.name == name) + ] + if len(matches) != 1: + raise ContractError(f"Function {name!r} matches {len(matches)} definitions; qualify it explicitly") + return matches[0] + + +def body_nodes(function: ast.AST) -> list[ast.AST]: + """Walk one function without treating nested definitions as its call sites.""" + result = [] + + def visit(node: ast.AST) -> None: + result.append(node) + for child in ast.iter_child_nodes(node): + if not isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef, ast.Lambda)): + visit(child) + + visit(function) + return result + + +def branch_dump(statements: list[ast.stmt]) -> str: + """Normalize a preserved branch while ignoring source positions/formatting.""" + return ast.dump(ast.Module(body=statements, type_ignores=[]), include_attributes=False) + + +def locate(root: Path, spec: RegistrationSpec) -> dict[str, Any]: + """Locate each exact call, surrounding guards and same-file parameter callers.""" + records = [] + for index, site in enumerate(spec.sites, 1): + source = source_path(root, site.source).read_text(encoding="utf-8") + tree = ast.parse(source) + function = function_node(tree, site.function) + nodes = body_nodes(function) + matches = [n for n in nodes if isinstance(n, ast.expr) and ast.unparse(n) == site.expression] + if len(matches) != 1: + raise ContractError(f"{site.source}::{site.function}: expression matches {len(matches)} sites") + selected = matches[0] + guards = [] + for node in nodes: + if not isinstance(node, ast.If): + continue + for arm in ("body", "orelse"): + statements = getattr(node, arm) + contains = any(selected in ast.walk(stmt) for stmt in statements) + sibling_contains = any( + selected in ast.walk(stmt) + for stmt in getattr(node, "orelse" if arm == "body" else "body") + ) + early_return = node.lineno < selected.lineno and any( + isinstance(n, ast.Return) for stmt in statements for n in ast.walk(stmt) + ) + if contains or (statements and (sibling_contains or early_return)): + guards.append( + { + "condition": ast.unparse(node.test), + "arm": arm, + "line": node.lineno, + "preserve_body": not contains, + "body": branch_dump(statements), + } + ) + args = [a.arg for a in (*function.args.posonlyargs, *function.args.args, *function.args.kwonlyargs)] + if "kernels" in args: + access = "kernels" + elif "self" in args and any( + isinstance(n, ast.Attribute) and ast.unparse(n) == "self._kernels" for n in ast.walk(tree) + ): + access = "self._kernels" + else: + access = "add kernels parameter and propagate through the callers below" + if not any(isinstance(n, ast.FunctionDef) and n.name == "_op" for n in tree.body): + raise ContractError(f"{site.source} needs an existing _op/KernelSelection integration") + callers = [] + frontier, seen = {function.name}, set() + while frontier: + target = frontier.pop() + if target in seen: + continue + seen.add(target) + for qualified, caller in functions(tree): + for call in body_nodes(caller): + if ( + isinstance(call, ast.Call) + and isinstance(call.func, ast.Name) + and call.func.id == target + ): + callers.append( + { + "function": qualified, + "callee": target, + "line": call.lineno, + "call": ast.unparse(call), + } + ) + frontier.add(caller.name) + records.append( + { + "id": f"S{index}", + "source": site.source, + "function": site.function, + "expression": site.expression, + "line": selected.lineno, + "column": selected.col_offset, + "end_line": selected.end_lineno, + "kernels": access, + "guards": guards, + "callers": callers, + "snippet": "\n".join( + source.splitlines()[max(0, selected.lineno - 4) : selected.end_lineno + 3] + ), + } + ) + return {"schema_version": 1, "preserve": list(spec.preserve), "sites": records} + + +def validate_review(review: dict[str, Any], baseline: dict[str, Any]) -> None: + """Permit review notes/extra guards, but never silently drop selected sites.""" + if not isinstance(review, dict) or not isinstance(review.get("sites"), list): + raise ContractError("sites.json must contain the reviewed sites") + if len(review["sites"]) != len(baseline["sites"]): + raise ContractError("Changing the selected sites requires a new specification") + for new, old in zip(review["sites"], baseline["sites"], strict=True): + if any(new.get(k) != old[k] for k in ("id", "source", "function", "expression", "guards")): + raise ContractError("Cannot remove source identities or baseline guards from sites.json") + if not isinstance(review.get("preserve"), list) or any( + not isinstance(x, str) for x in review["preserve"] + ): + raise ContractError("sites.json preserve must be a list of review notes") + if not set(baseline["preserve"]) <= set(review["preserve"]): + raise ContractError("Cannot remove specification preservation requirements") + + +def render(review: dict[str, Any]) -> str: + """Render the source and branch evidence shown before any agent starts.""" + lines = [ + "# Operator registration call sites", + "", + "Review sites.json, then run `register run RUN --agent HARNESS/MODEL:EFFORT`.", + "", + ] + for site in review["sites"]: + lines.extend( + [ + f"## {site['id']}: {site['source']}:{site['line']} ({site['function']})", + "", + "```python", + site["snippet"], + "```", + "", + f"KernelSelection: {site['kernels']}", + "", + ] + ) + lines.extend( + f"- Guard ({g['arm']}): `{g['condition']}`; preserve body: {g['preserve_body']}" + for g in site["guards"] + ) + lines.extend(f"- Caller: {c['function']}:{c['line']} — `{c['call']}`" for c in site["callers"]) + lines.append("") + lines.extend( + ["## Additional preservation requirements", "", *[f"- {item}" for item in review["preserve"]], ""] + ) + return "\n".join(lines) diff --git a/scripts/kernel_tuning/registration/spec.py b/scripts/kernel_tuning/registration/spec.py new file mode 100644 index 0000000..27b70eb --- /dev/null +++ b/scripts/kernel_tuning/registration/spec.py @@ -0,0 +1,141 @@ +"""Frozen input contracts for one operator-registration run.""" + +from __future__ import annotations + +import ast +import keyword +from dataclasses import asdict, dataclass, field +from pathlib import Path, PurePosixPath +from typing import Any + +from scripts.kernel_tuning.contracts import ContractError + + +def relative_file(value: Any, *, prefix: str, suffix: str = ".py") -> str: + """Accept only portable repository-relative files in the owning directory.""" + if not isinstance(value, str) or "\\" in value or ":" in value: + raise ContractError(f"Expected a relative {prefix} file: {value!r}") + path = PurePosixPath(value) + if path.is_absolute() or ".." in path.parts or path.as_posix() != value: + raise ContractError(f"Unsafe repository path: {value!r}") + if not value.startswith(prefix + "/") or path.suffix != suffix: + raise ContractError(f"Expected a {suffix} file under {prefix}: {value!r}") + return value + + +def test_node(value: Any) -> str: + """Validate a pytest node ID without accepting options or external paths.""" + if not isinstance(value, str): + raise ContractError("Test selectors must be strings") + parts = value.split("::") + relative_file(parts[0], prefix="tests") + if any(not part or "\n" in part for part in parts): + raise ContractError(f"Invalid pytest node: {value!r}") + return value + + +@dataclass(frozen=True) +class Site: + """An unambiguous expression within a named (optionally qualified) function.""" + + source: str + function: str + expression: str + + +@dataclass(frozen=True) +class RegistrationBudget: + """Hard attempt and wall-clock limits that survive interruption.""" + + max_attempts: int = 5 + agent_timeout_seconds: int = 1800 + max_seconds: int = 7200 + + +@dataclass(frozen=True) +class RegistrationSpec: + """Immutable mathematical objective, source selection, tests and budget.""" + + model: str + op: str + sites: tuple[Site, ...] + description: str = "" + preserve: tuple[str, ...] = () + tests: tuple[str, ...] = ("tests/test_kernel_ops.py",) + budget: RegistrationBudget = field(default_factory=RegistrationBudget) + + def as_dict(self) -> dict[str, Any]: + """Return JSON-compatible data for the frozen specification.""" + return asdict(self) + + @classmethod + def parse(cls, value: Any) -> RegistrationSpec: + """Reject malformed/unknown settings before creating any worktrees.""" + if not isinstance(value, dict) or set(value) - set(cls.__dataclass_fields__): + raise ContractError("Unknown or malformed registration specification") + model, op = value.get("model"), value.get("op") + if not isinstance(model, str) or not model.isidentifier(): + raise ContractError("model must be a canonical policy identifier") + if not isinstance(op, str) or not op.isidentifier() or keyword.iskeyword(op) or op.startswith("_"): + raise ContractError("op must be a public Python identifier, not a keyword") + if model.lower().replace("_", "") in op.lower().replace("_", ""): + raise ContractError("Name the operator by its mathematics, not its model") + description = value.get("description", "") + preserve, tests = value.get("preserve", []), value.get("tests", ["tests/test_kernel_ops.py"]) + if not isinstance(description, str) or not isinstance(preserve, (list, tuple)): + raise ContractError("description must be text; preserve must be a list of text") + if any(not isinstance(item, str) or not item.strip() for item in preserve): + raise ContractError("preserve entries must be nonempty text") + if not isinstance(tests, (list, tuple)) or not tests: + raise ContractError("Specify at least one existing CPU test") + tests = tuple(test_node(item) for item in tests) + raw_sites = value.get("sites") + if not isinstance(raw_sites, (list, tuple)) or not raw_sites: + raise ContractError("Specify at least one call site") + sites = [] + for raw in raw_sites: + if not isinstance(raw, dict) or set(raw) != {"source", "function", "expression"}: + raise ContractError("Each site requires source, function and expression") + source = relative_file(raw["source"], prefix=f"embodiinfer/policies/{model}") + function, expression = raw["function"], raw["expression"] + if not isinstance(function, str) or not all(p.isidentifier() for p in function.split(".")): + raise ContractError("site.function must be a function name or qualified name") + try: + parsed = ast.parse(expression, mode="eval") + except (SyntaxError, TypeError, ValueError) as exc: + raise ContractError("site.expression must be a Python expression") from exc + if not isinstance(parsed.body, ast.Call): + raise ContractError("A registration site must select a call expression") + sites.append(Site(source, function, ast.unparse(parsed.body))) + if len(set(sites)) != len(sites): + raise ContractError("Duplicate call sites") + raw_budget = value.get("budget", {}) + if not isinstance(raw_budget, dict) or set(raw_budget) - set(RegistrationBudget.__dataclass_fields__): + raise ContractError("Unknown or malformed budget") + if any(type(v) is not int or v <= 0 for v in raw_budget.values()): + raise ContractError("Budget values must be positive integers") + return cls( + model, op, tuple(sites), description, tuple(preserve), tests, RegistrationBudget(**raw_budget) + ) + + +def load_spec(path: Path | None, *, model: str | None, op: str | None, sites: list[str]) -> RegistrationSpec: + """Load YAML/JSON or the equivalent repeated --site command-line form.""" + if path is not None: + if model or op or sites: + raise ContractError("Use either --spec or --model/--op/--site") + import yaml + + try: + value = yaml.safe_load(path.read_text(encoding="utf-8")) + except yaml.YAMLError as exc: + raise ContractError(f"Invalid specification: {exc}") from exc + else: + parsed = [] + for item in sites: + parts = item.split("::", 2) + if len(parts) != 3: + raise ContractError("--site expects SOURCE::FUNCTION::EXPRESSION") + parsed.append(dict(zip(("source", "function", "expression"), parts, strict=True))) + value = {"model": model, "op": op, "sites": parsed} + return RegistrationSpec.parse(value) diff --git a/scripts/kernel_tuning/registration/workspace.py b/scripts/kernel_tuning/registration/workspace.py new file mode 100644 index 0000000..04b8632 --- /dev/null +++ b/scripts/kernel_tuning/registration/workspace.py @@ -0,0 +1,444 @@ +"""Isolated Git checkouts, bounded CPU workers and durable registration state.""" + +from __future__ import annotations + +import ast +import asyncio +import contextlib +import difflib +import hashlib +import os +import shutil +import subprocess +import uuid +from datetime import datetime, timezone +from pathlib import Path +from typing import Any + +from scripts.kernel_tuning.artifacts import atomic_json, runtime_identity +from scripts.kernel_tuning.contracts import ContractError, digest, read_json +from scripts.kernel_tuning.runner import _kill_tree + +from .sites import locate, render, source_path, validate_review +from .spec import RegistrationSpec + +REFERENCE = "embodiinfer/backend/torch/ops.py" +REGISTRY = "embodiinfer/layers/ops.py" +# Call sites already routed through the registry; the agent reads them from the +# current BASE as a worked example instead of a historical commit. +ROUTED_EXAMPLE = "embodiinfer/policies/pi05/modeling_pi05.py" +ARTIFACTS = {"PLAN.md", "REGISTRATION.json"} + + +def git(root: Path, *args: str) -> str: + """Execute a bounded Git command without a shell or ambient Git redirection.""" + env = {k: v for k, v in os.environ.items() if not k.startswith("GIT_")} + env["GIT_TERMINAL_PROMPT"] = "0" + try: + result = subprocess.run( + ["git", "-C", str(root), *args], + capture_output=True, + encoding="utf-8", + errors="replace", + timeout=60, + env=env, + ) + except subprocess.TimeoutExpired as exc: + raise ContractError(f"Git command timed out: {args}") from exc + if result.returncode: + raise ContractError(f"git {' '.join(args)}: {result.stderr.strip()[:4000]}") + return result.stdout + + +def file_hash(path: Path) -> str: + """Hash artifact bytes, rejecting links and missing files.""" + if path.is_symlink() or not path.is_file(): + raise ContractError(f"Expected a regular file: {path}") + return hashlib.sha256(path.read_bytes()).hexdigest() + + +def tracked_hashes(root: Path) -> dict[str, str]: + """Hash tracked files and link targets without following unrelated repository links.""" + hashes = {} + for name in git(root, "ls-files", "-z").split("\0"): + if not name: + continue + path = root / name + # Documentation assets may be tracked links on POSIX. Check their parent + # for traversal, then bind the link text itself, never the target bytes. + source_path(root, str(Path(name).parent)) + hashes[name] = ( + "symlink:" + hashlib.sha256(os.fsencode(os.readlink(path))).hexdigest() + if path.is_symlink() + else file_hash(source_path(root, name)) + ) + return hashes + + +def cpu_environment(root: Path) -> dict[str, str]: + """Force subprocess imports into this checkout and hide accelerator devices.""" + env = os.environ.copy() + env.update( + PYTHONPATH=str(root), + PYTHONDONTWRITEBYTECODE="1", + PYTHONNOUSERSITE="1", + PYTHONOPTIMIZE="0", + CUDA_VISIBLE_DEVICES="", + HIP_VISIBLE_DEVICES="", + ROCR_VISIBLE_DEVICES="", + PYTEST_ADDOPTS="", + PYTEST_DISABLE_PLUGIN_AUTOLOAD="1", + ) + return env + + +async def command(argv: list[str], *, root: Path, output: Path, timeout: float) -> dict[str, Any]: + """Capture a CPU subprocess and kill its entire process group on cancellation.""" + if timeout <= 0: + raise TimeoutError("Registration duration exhausted") + output.parent.mkdir(parents=True, exist_ok=True) + record: dict[str, Any] = {"argv": argv, "cwd": str(root), "timeout_seconds": timeout} + with output.open("wb") as stream: + process = await asyncio.create_subprocess_exec( + *argv, + cwd=root, + env=cpu_environment(root), + stdout=stream, + stderr=stream, + start_new_session=os.name != "nt", + ) + try: + await asyncio.wait_for(process.wait(), timeout=timeout) + except BaseException: + await _kill_tree(process) + raise + record["returncode"] = process.returncode + with output.open("rb") as stream: + stream.seek(max(0, output.stat().st_size - 12000)) + record["output"] = stream.read().decode("utf-8", errors="replace") + record["log"] = str(output) + return record + + +def import_guard(root: Path) -> str: + """Check the actual package origin inside the same interpreter as the check.""" + return ( + "import pathlib, embodiinfer\n" + f"assert pathlib.Path(embodiinfer.__file__).resolve().is_relative_to(pathlib.Path({str(root)!r}).resolve()), 'embodiinfer imported outside checkout'\n" + ) + + +async def check_unregistered(root: Path, op: str, python: str, output: Path) -> None: + """Probe the public registry in the selected test environment before mutation.""" + tree = ast.parse(source_path(root, REFERENCE).read_text(encoding="utf-8")) + if any(isinstance(n, (ast.FunctionDef, ast.AsyncFunctionDef)) and n.name == op for n in ast.walk(tree)): + raise ContractError(f"Reference function {op!r} already exists") + code = ( + import_guard(root) + + f"from embodiinfer.layers import available_ops\nassert {op!r} not in available_ops(), 'operator already registered'\n" + ) + result = await command([python, "-c", code], root=root, output=output, timeout=60) + if result["returncode"]: + raise ContractError( + f"Registry preflight failed; select a prepared --set test_python=/abs/python:\n{result['output']}" + ) + + +class RegistrationRun: + """Controller-owned ledger plus two checkouts at an immutable starting commit.""" + + def __init__(self, root: Path) -> None: + self.root = root.resolve(strict=True) + self.manifest = read_json(self.root / "manifest.json") + if file_hash(self.root / "spec.json") != self.manifest["spec_hash"]: + self.manifest.update(status="failed", stop_reason="Frozen specification changed") + self.save() + raise ContractError("Frozen specification changed") + frozen = read_json(self.root / "spec.json") + self.spec = RegistrationSpec.parse(frozen["spec"]) + self.base_commit = frozen["base_commit"] + self.base = self.root / "base" + self.worktree = self.root / "worktree" + + @classmethod + async def create( + cls, repository: Path, output: Path, spec: RegistrationSpec, python: str, *, yes: bool = False + ) -> RegistrationRun: + """Validate first, then create a detached baseline and a review branch.""" + repository, output = repository.resolve(), output.resolve() + if git(repository, "status", "--porcelain", "--untracked-files=all").strip(): + raise ContractError("The repository must be clean before register new") + base_commit = git(repository, "rev-parse", "HEAD").strip() + review = locate(repository, spec) + for name in spec.tests: + if not source_path(repository, name.split("::")[0]).is_file(): + raise ContractError(f"Missing existing CPU test: {name}") + run_id = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ") + "-" + uuid.uuid4().hex[:10] + root = output / run_id + root.mkdir(parents=True) + branch = f"registry/{spec.op}-{run_id}" + created: list[Path] = [] + try: + await check_unregistered(repository, spec.op, python, root / "preflight.log") + if ( + git(repository, "rev-parse", "HEAD").strip() != base_commit + or git(repository, "status", "--porcelain", "--untracked-files=all").strip() + ): + raise ContractError( + "Repository changed during preflight; use an ignored/external output directory" + ) + git(repository, "worktree", "add", "--detach", str(root / "base"), base_commit) + created.append(root / "base") + git(repository, "worktree", "add", "-b", branch, str(root / "worktree"), base_commit) + created.append(root / "worktree") + atomic_json( + root / "spec.json", {"schema_version": 1, "base_commit": base_commit, "spec": spec.as_dict()} + ) + atomic_json(root / "sites.json", review) + (root / "sites.md").write_text(render(review), encoding="utf-8") + manifest = { + "schema_version": 1, + "run_id": run_id, + "repository": str(repository), + "branch": branch, + "base_commit": base_commit, + "spec_hash": file_hash(root / "spec.json"), + "base_hashes": tracked_hashes(root / "base"), + "status": "awaiting_confirmation", + "confirmed": yes, + "elapsed_seconds": 0.0, + "attempts": [], + "pending": None, + "agent": None, + "test_python": python, + "tool_runtime": runtime_identity(), + "sites_hash": None, + "base_overlay": {}, + } + atomic_json(root / "manifest.json", manifest) + return cls(root) + except BaseException: + for checkout in reversed(created): + with contextlib.suppress(ContractError): + git(repository, "worktree", "remove", str(checkout)) + if root / "worktree" in created: + with contextlib.suppress(ContractError): + git(repository, "branch", "-D", branch) + raise + + def save(self) -> None: + """Atomically checkpoint the external ledger for status and crash recovery.""" + atomic_json(self.root / "manifest.json", self.manifest) + + def verify(self) -> None: + """Reject changed specifications, commits, review inputs or baseline files.""" + try: + self._verify() + except (ContractError, OSError, ValueError) as exc: + self.manifest.update(status="failed", stop_reason=str(exc)) + self.save() + raise + + def _verify(self) -> None: + if ( + file_hash(self.root / "spec.json") != self.manifest["spec_hash"] + or self.base_commit != self.manifest["base_commit"] + ): + raise ContractError("Frozen specification or BASE changed") + if git(self.base, "rev-parse", "HEAD").strip() != self.base_commit: + raise ContractError("BASE checkout commit changed") + if git(self.worktree, "symbolic-ref", "--short", "HEAD").strip() != self.manifest["branch"]: + raise ContractError("Registration worktree branch changed") + expected_head = self.manifest.get("commit", self.base_commit) + if git(self.worktree, "rev-parse", "HEAD").strip() != expected_head: + raise ContractError("Agent changed worktree history; git commit is forbidden during registration") + if ( + self.manifest.get("sites_hash") + and file_hash(self.root / "sites.json") != self.manifest["sites_hash"] + ): + raise ContractError("Reviewed sites changed after the run started") + if self.manifest.get("context_hashes"): + actual = { + p.relative_to(self.root / "context").as_posix(): file_hash(p) + for p in (self.root / "context").rglob("*") + if p.is_file() + } + if actual != self.manifest["context_hashes"]: + raise ContractError("Read-only agent context changed") + if self.manifest.get("base_overlay"): + raise ContractError("Baseline test overlay must be recovered before checking integrity") + if ( + tracked_hashes(self.base) != self.manifest["base_hashes"] + or git(self.base, "status", "--porcelain", "--untracked-files=all").strip() + ): + raise ContractError("RUN/base was modified") + + def recover_base(self) -> None: + """Restore only journaled test copies after a killed baseline check.""" + overlay = self.manifest.get("base_overlay", {}) + for name, expected in overlay.items(): + target = source_path(self.base, name) + original = self.manifest["base_hashes"].get(name) + actual = file_hash(target) if target.exists() else None + if actual not in (original, expected): + raise ContractError(f"Baseline overlay was modified unexpectedly: {name}") + if original is None: + target.unlink(missing_ok=True) + else: + git(self.base, "restore", f"--source={self.base_commit}", "--worktree", "--", name) + if overlay: + self.manifest["base_overlay"] = {} + self.save() + + def prepare(self, agent: str, python: str | None = None, *, unfenced: bool = False) -> None: + """Freeze reviewed call sites and build context before the first agent turn. + + The agent, test interpreter and ``unfenced`` choice are frozen together; + resume always reuses them. + """ + self.recover_base() + self.verify() + if self.manifest["sites_hash"]: + if ( + agent != self.manifest["agent"] + or (python and python != self.manifest["test_python"]) + or unfenced != bool(self.manifest.get("unfenced")) + ): + raise ContractError("Agent/test_python/unfenced is frozen; resume without overrides") + return + review = read_json(self.root / "sites.json") + validate_review(review, locate(self.base, self.spec)) + context = self.root / "context" + context.mkdir(exist_ok=True) + names = { + "AGENTS.md", + REFERENCE, + REGISTRY, + ROUTED_EXAMPLE, + "tests/test_kernel_ops.py", + *(s.source for s in self.spec.sites), + } + for name in names: + target = context / name + target.parent.mkdir(parents=True, exist_ok=True) + shutil.copyfile(source_path(self.base, name), target) + shutil.copyfile(self.root / "sites.json", context / "sites.json") + from .prompts import RULES + + (context / "RULES.md").write_text(RULES, encoding="utf-8") + self.manifest.update( + agent=agent, + unfenced=unfenced, + test_python=python or self.manifest["test_python"], + confirmed=True, + sites_hash=file_hash(self.root / "sites.json"), + status="ready", + context_hashes={ + p.relative_to(context).as_posix(): file_hash(p) for p in context.rglob("*") if p.is_file() + }, + ) + self.save() + + def changes(self) -> list[str]: + """Include staged, unstaged, untracked and ignored additions in scope checks.""" + tracked = git(self.worktree, "diff", "--name-only", "-z", self.base_commit, "--").split("\0") + untracked = git(self.worktree, "ls-files", "--others", "-z").split("\0") + # Python/test caches are controller-created and never exported. + return sorted( + { + name + for name in tracked + untracked + if name + and not any( + part in {"__pycache__", ".pytest_cache", ".ruff_cache"} for part in Path(name).parts + ) + } + ) + + def fingerprint(self) -> str: + """Bind check evidence to exact source bytes, including new test files.""" + return digest( + { + name: file_hash(source_path(self.worktree, name)) if (self.worktree / name).exists() else None + for name in self.changes() + if name != "PLAN.md" + } + ) + + def patch(self) -> str: + """Produce a reviewable patch that includes allowed untracked additions.""" + patch = git( + self.worktree, + "diff", + "--binary", + self.base_commit, + "--", + ".", + ":(exclude)PLAN.md", + ":(exclude)REGISTRATION.json", + ) + tracked = set(git(self.worktree, "ls-files", "-z").split("\0")) + for name in self.changes(): + if name in tracked or name in ARTIFACTS: + continue + path = source_path(self.worktree, name) + if path.is_file(): + content = path.read_text(encoding="utf-8") + lines = content.splitlines(keepends=True) + no_newline = bool(content) and not content.endswith("\n") + if no_newline: + lines[-1] += "\n" + patch += f"diff --git a/{name} b/{name}\nnew file mode 100644\n" + patch += "".join( + difflib.unified_diff( + [], + lines, + fromfile="/dev/null", + tofile=f"b/{name}", + ) + ) + if no_newline: + patch += "\\ No newline at end of file\n" + return patch + + +def check_scope(run: RegistrationRun) -> dict[str, Any]: + """C1: restrict every changed path and preserve existing test lines and bodies.""" + allowed = { + REFERENCE, + REGISTRY, + "tests/test_kernel_ops.py", + f"tests/test_registry_{run.spec.op}.py", + *(s.source for s in run.spec.sites), + } + changed = run.changes() + for name in changed: + if name in ARTIFACTS: + source_path(run.worktree, name) + continue + if name not in allowed: + raise ContractError(f"C1: out-of-scope file: {name}") + current = source_path(run.worktree, name) + if not current.is_file(): + raise ContractError(f"C1: deleting files is forbidden: {name}") + if name.startswith("tests/") and (run.base / name).exists(): + before = (run.base / name).read_text(encoding="utf-8") + after = current.read_text(encoding="utf-8") + diff = difflib.ndiff(before.splitlines(), after.splitlines()) + if any(line.startswith("- ") for line in diff): + raise ContractError(f"C1: existing test lines removed or changed: {name}") + old_tree, new_tree = ast.parse(before), ast.parse(after) + for old in old_tree.body: + if ast.dump(old) not in [ast.dump(new) for new in new_tree.body]: + raise ContractError(f"C1: existing test statements changed: {name}:{old.lineno}") + before, after = ( + ast.parse((root / REGISTRY).read_text(encoding="utf-8")) for root in (run.base, run.worktree) + ) + old_classes = [n for n in before.body if isinstance(n, ast.ClassDef) and n.name == "KernelSelection"] + new_classes = [n for n in after.body if isinstance(n, ast.ClassDef) and n.name == "KernelSelection"] + if [ast.dump(n) for n in old_classes] != [ast.dump(n) for n in new_classes]: + raise ContractError("C1: KernelSelection must not be modified") + return { + "files": [name for name in changed if name not in ARTIFACTS], + "command": ["git", "diff", "--name-status", run.base_commit], + } diff --git a/scripts/kernel_tuning/runner.py b/scripts/kernel_tuning/runner.py new file mode 100644 index 0000000..37fac9c --- /dev/null +++ b/scripts/kernel_tuning/runner.py @@ -0,0 +1,136 @@ +"""Bounded evaluator subprocesses, independent of coding-agent transcripts.""" + +from __future__ import annotations + +import asyncio +import contextlib +import os +import signal +import uuid +from pathlib import Path +from typing import Any + +from .artifacts import REPOSITORY, RunStore, atomic_json +from .contracts import ContractError, digest, read_json + + +async def _kill_tree(process: asyncio.subprocess.Process) -> None: + if process.returncode is not None: + return + if os.name == "nt": + # psutil is part of the isolated Humanize2 tooling environment. + import psutil + + with contextlib.suppress(psutil.NoSuchProcess): + parent = psutil.Process(process.pid) + for child in parent.children(recursive=True): + with contextlib.suppress(psutil.NoSuchProcess): + child.kill() + parent.kill() + else: + with contextlib.suppress(ProcessLookupError): + os.killpg(process.pid, signal.SIGKILL) + await process.wait() + + +async def evaluate( + store: RunStore, + output: Path, + *, + operation: str = "evaluate", + candidate: Path | None = None, + timeout: float | None = None, +) -> dict[str, Any]: + """Run a trusted worker and reject stale, incomplete, or mismatched responses. + + Cancellation and timeout terminate the worker and its compiler children. A + fresh response filename/nonce prevents a crashed worker reusing old evidence. + """ + store.verify() + output.mkdir(parents=True, exist_ok=True) + cfg = store.task.settings + baseline = store.solution_path() + paths = { + "baseline": baseline, + "incumbent": store.solution_path(store.manifest["best"]), + "candidate": candidate or baseline, + } + request = { + "schema_version": 1, + "nonce": uuid.uuid4().hex, + "operation": operation, + "task_dir": str(store.task.root), + "task_digest": store.task.identity, + "solutions": {role: str(path.resolve()) for role, path in paths.items()}, + "solution_digests": {role: digest(read_json(path)) for role, path in paths.items()}, + } + request_path = output / "request.json" + response_path = output / f"response-{request['nonce']}.json" + atomic_json(request_path, request) + env = os.environ.copy() + env["PYTHONDONTWRITEBYTECODE"] = "1" + # Both compilers and imported generated Python are confined to this run's cache. + # This is filesystem organization, not a security sandbox. + env["TORCH_EXTENSIONS_DIR"] = str(store.root / "cache/torch") + env["TRITON_CACHE_DIR"] = str(store.root / "cache/triton") + env["FIB_CACHE_PATH"] = str(store.root / "cache/flashinfer") + env["FIB_ENABLE_APPLY"] = "0" + env["FIB_ENABLE_TRACING"] = "0" + argv = [ + cfg.evaluator_python, + str(REPOSITORY / "benchmarks/kernel_tuning/evaluate.py"), + "--request", + str(request_path), + "--output", + str(response_path), + ] + atomic_json( + output / "command.json", + {"argv": argv, "timeout_seconds": timeout or cfg.search.evaluation_timeout_seconds}, + ) + process = None + try: + with (output / "stdout.log").open("wb") as stdout, (output / "stderr.log").open("wb") as stderr: + process = await asyncio.create_subprocess_exec( + *argv, + cwd=output, + env=env, + stdout=stdout, + stderr=stderr, + start_new_session=os.name != "nt", + ) + try: + await asyncio.wait_for( + process.wait(), timeout=timeout or cfg.search.evaluation_timeout_seconds + ) + except BaseException: + await _kill_tree(process) + raise + if process.returncode != 0: + raise ContractError(f"Evaluator failed (exit {process.returncode}); see {output / 'stderr.log'}") + if not response_path.is_file(): + raise ContractError("Evaluator produced no structured response") + response = read_json(response_path) + if ( + response.get("request_digest") != digest(request) + or response.get("task_digest") != store.task.identity + ): + raise ContractError("Evaluator response identity mismatch") + if response.get("solution_digests") != request["solution_digests"]: + raise ContractError("Evaluator measured different solution sources") + if response.get("status") != "ok": + raise ContractError(f"Evaluator rejected request: {response.get('error', 'unknown error')}") + environment = response.get("environment", {}) + if not isinstance(environment.get("identity"), dict) or not environment["identity"]: + raise ContractError("Evaluator must report hardware/software identity") + previous = store.manifest["environment"] + if previous and environment["identity"] != previous["identity"]: + raise ContractError("Evaluator hardware/software environment changed; start a new run") + for role, path in paths.items(): + if digest(read_json(path)) != request["solution_digests"][role]: + raise ContractError("Solution changed during evaluation") + store.verify() + atomic_json(output / "result.json", response) + return response + except TimeoutError as exc: + raise ContractError(f"Evaluator timed out; see {output}") from exc diff --git a/scripts/kernel_tuning/skeleton.py b/scripts/kernel_tuning/skeleton.py new file mode 100644 index 0000000..175376c --- /dev/null +++ b/scripts/kernel_tuning/skeleton.py @@ -0,0 +1,137 @@ +"""Build a randomly initialized copy of a Hugging Face checkpoint for shape capture. + +Kernel shapes depend on the model configuration and its inputs, never on weight +values. This downloads every small file of a repository (configs, tokenizer, +processor) and, for each ``*.safetensors`` file, reads only its header with an +HTTP range request and writes a file with the same tensor names, dtypes, and +shapes filled with small random values. Model loaders then run unchanged. + +The result is for measuring shapes only; its outputs are meaningless. Decode +lengths of generative models follow the random logits, so cap them in the +capture configuration. + +Run in a runtime with Torch and safetensors: +``python -m scripts.kernel_tuning skeleton REPO --revision REV --output DIR`` +""" + +from __future__ import annotations + +import json +import os +import struct +import urllib.request +from pathlib import Path, PurePosixPath +from typing import Any + +from .contracts import ContractError, relative_path + +#: Files larger than this are weights or data, never configuration. +SMALL_FILE_BYTES = 64 * 1024 * 1024 +SAFETENSORS_DTYPES = { + "F32": "float32", + "F16": "float16", + "BF16": "bfloat16", + "F64": "float64", + "I64": "int64", + "I32": "int32", + "I16": "int16", + "I8": "int8", + "U8": "uint8", + "BOOL": "bool", +} + + +def _endpoint() -> str: + return os.environ.get("HF_ENDPOINT", "https://huggingface.co").rstrip("/") + + +def _request(url: str, *, start: int | None = None, end: int | None = None) -> bytes: + headers = {"User-Agent": "embodiinfer-kernel-tuning/1"} + if token := os.environ.get("HF_TOKEN"): + headers["Authorization"] = f"Bearer {token}" + if start is not None: + headers["Range"] = f"bytes={start}-{end}" + with urllib.request.urlopen(urllib.request.Request(url, headers=headers), timeout=60) as response: + return response.read() + + +def list_files(repo: str, revision: str) -> list[dict[str, Any]]: + """Every file of a model repository at a revision, with sizes.""" + url = f"{_endpoint()}/api/models/{repo}/tree/{revision}?recursive=true" + return [entry for entry in json.loads(_request(url)) if entry["type"] == "file"] + + +def file_url(repo: str, revision: str, path: str) -> str: + """Resolve URL of one repository file.""" + return f"{_endpoint()}/{repo}/resolve/{revision}/{path}" + + +def read_header(repo: str, revision: str, path: str) -> dict[str, Any]: + """Read a safetensors header (names, dtypes, shapes) without the tensor data.""" + url = file_url(repo, revision, path) + (size,) = struct.unpack(" SMALL_FILE_BYTES: + raise ContractError(f"{path}: implausible safetensors header of {size} bytes") + return json.loads(_request(url, start=8, end=7 + size)) + + +def _initializer(name: str, shape: list[int], dtype: str, generator: Any) -> Any: + """Small random weights; unit norm scales, zero biases, and index buffers as ranges.""" + import torch + + torch_dtype = getattr(torch, SAFETENSORS_DTYPES[dtype]) + leaf = name.rsplit(".", 1)[-1] + if not torch_dtype.is_floating_point: + if "position_ids" in name and shape: + return torch.arange(shape[-1], dtype=torch_dtype).expand(*shape).contiguous() + return torch.zeros(shape, dtype=torch_dtype) + if leaf == "bias": + return torch.zeros(shape, dtype=torch_dtype) + if "norm" in name.lower() and leaf == "weight": + return torch.ones(shape, dtype=torch_dtype) + values = torch.empty(shape, dtype=torch.float32) + values.normal_(0.0, 0.02, generator=generator) + return values.to(torch_dtype) + + +def synthesize(header: dict[str, Any], destination: Path, *, seed: int = 0) -> int: + """Write a safetensors file matching ``header`` with random values; return tensor count.""" + import torch + from safetensors.torch import save_file + + generator = torch.Generator().manual_seed(seed) + metadata = header.get("__metadata__") + tensors = { + name: _initializer(name, spec["shape"], spec["dtype"], generator) + for name, spec in header.items() + if name != "__metadata__" + } + destination.parent.mkdir(parents=True, exist_ok=True) + save_file(tensors, str(destination), metadata=metadata) + return len(tensors) + + +def build(repo: str, revision: str, output: Path) -> dict[str, Any]: + """Create the skeleton checkpoint directory and return a summary of what was written.""" + if output.exists() and any(output.iterdir()): + raise ContractError(f"{output} is not empty") + files = list_files(repo, revision) + summary: dict[str, Any] = {"repo": repo, "revision": revision, "copied": [], "synthesized": {}} + for index, entry in enumerate(sorted(files, key=lambda e: e["path"])): + path = relative_path(entry["path"]) + target = output / PurePosixPath(path) + target.parent.mkdir(parents=True, exist_ok=True) + if path.endswith(".safetensors"): + summary["synthesized"][path] = synthesize(read_header(repo, revision, path), target, seed=index) + elif entry.get("size", 0) <= SMALL_FILE_BYTES: + target.write_bytes(_request(file_url(repo, revision, path))) + summary["copied"].append(path) + elif path.endswith((".bin", ".pt", ".pth", ".ckpt")): + raise ContractError(f"{path}: only safetensors weights can be synthesized from headers") + if not summary["synthesized"]: + raise ContractError(f"{repo}@{revision} has no safetensors weights") + (output / "SKELETON.json").write_text( + json.dumps({**summary, "note": "random weights for shape capture only"}, indent=2) + "\n", + encoding="utf-8", + ) + return summary diff --git a/scripts/kernel_tuning/uv.lock b/scripts/kernel_tuning/uv.lock new file mode 100644 index 0000000..c249159 --- /dev/null +++ b/scripts/kernel_tuning/uv.lock @@ -0,0 +1,641 @@ +version = 1 +revision = 5 +requires-python = ">=3.12" + +[[package]] +name = "annotated-types" +version = "0.8.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/5f/56/a8120250d128bed162cd73c76d45f6ef9991f3e068f62a8ee060afa3104a/annotated_types-0.8.0.tar.gz", hash = "sha256:13b2beaad985e05e2d6407ee4c4f35590b11f8d693a258a561055cac8f64cab7", size = 15893, upload-time = "2026-07-23T20:16:13.995Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/99/91/8acff4f5e50511b911bbccb72b8628a49c68ce14148cd9f6431094859a90/annotated_types-0.8.0-py3-none-any.whl", hash = "sha256:f072f4d804ea359e4eaf198b1af7a8b0943881a87f31bb764f8bf219bb9419e0", size = 13427, upload-time = "2026-07-23T20:16:12.938Z" }, +] + +[[package]] +name = "certifi" +version = "2026.7.22" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/a3/c2/24167ea9858356b47a87a50d39908bfdb72ceeefe0041586e704e5376b3a/certifi-2026.7.22.tar.gz", hash = "sha256:741e2c3b351ddf169a738da9f2c048608ff7f2c5cc02f1ebc6b118bb090d5d55", size = 138112, upload-time = "2026-07-22T03:35:12.644Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/0b/a7/71ac2cff56fec219ed242bb11b8efb69fcc4bec75db06fb7bfe35de520e6/certifi-2026.7.22-py3-none-any.whl", hash = "sha256:62f22742b58a1a33014a2b6b706588a8d7e2a88ae7bd1a6ebe8c992928483775", size = 136983, upload-time = "2026-07-22T03:35:11.276Z" }, +] + +[[package]] +name = "colorama" +version = "0.4.6" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/d8/53/6f443c9a4a8358a93a6792e2acffb9d9d5cb0a5cfd8802644b7b1c9a02e4/colorama-0.4.6.tar.gz", hash = "sha256:08695f5cb7ed6e0531a20572697297273c47b8cae5a63ffc6d6ed5c201be6e44", size = 27697, upload-time = "2022-10-25T02:36:22.414Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d1/d6/3965ed04c63042e047cb6a3e6ed1a63a35087b6a609aa3a15ed8ac56c221/colorama-0.4.6-py2.py3-none-any.whl", hash = "sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6", size = 25335, upload-time = "2022-10-25T02:36:20.889Z" }, +] + +[[package]] +name = "dateparser" +version = "1.4.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "python-dateutil" }, + { name = "pytz" }, + { name = "regex" }, + { name = "tzlocal" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/c7/5d/bd21ba1519b6b1e222b29878301d2e1fb928e890dc7d085fa4222ac5671b/dateparser-1.4.3.tar.gz", hash = "sha256:bab8c43a746266e68142f4926e69438ce551441aa88e54e78bb6410bf3ee7000", size = 362152, upload-time = "2026-09-03T10:07:54.545Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/30/5f/63f927b5ffd1cc2c030fc2bd362bf920890e8806cf42f701198de5b3f9c7/dateparser-1.4.3-py3-none-any.whl", hash = "sha256:cbce86c64e0cea5c54c84c015d34c903d46c51bfad50632bf2f05a6141d10c05", size = 322384, upload-time = "2026-09-03T10:07:52.913Z" }, +] + +[[package]] +name = "embodiinfer-kernel-tuning-tools" +version = "0.1.0" +source = { virtual = "." } +dependencies = [ + { name = "hmz" }, + { name = "pyyaml" }, +] + +[package.dev-dependencies] +dev = [ + { name = "pytest" }, + { name = "ruff" }, +] + +[package.metadata] +requires-dist = [ + { name = "hmz", url = "https://github.com/humanfia/humanize/archive/b70d4427f11308442aef20f9b2233dfbbc34a39e.zip" }, + { name = "pyyaml", specifier = ">=6.0.3,<7" }, +] + +[package.metadata.requires-dev] +dev = [ + { name = "pytest", specifier = ">=8,<10" }, + { name = "ruff", specifier = ">=0.10,<1" }, +] + +[[package]] +name = "hmz" +version = "0.1.0" +source = { url = "https://github.com/humanfia/humanize/archive/b70d4427f11308442aef20f9b2233dfbbc34a39e.zip" } +dependencies = [ + { name = "dateparser" }, + { name = "psutil" }, + { name = "pydantic" }, + { name = "pyfiglet" }, + { name = "pyyaml" }, + { name = "rich" }, + { name = "sentry-sdk" }, + { name = "textual" }, + { name = "typing-extensions" }, +] +sdist = { hash = "sha256:31706b633e49bd7bec8ac04b3fc14d5aad4617e72fe7e8711a28ff167195fce3" } + +[package.metadata] +requires-dist = [ + { name = "dateparser", specifier = ">=1.4.2" }, + { name = "deepseek-harness-sdk", marker = "extra == 'dsh'", specifier = ">=0.1.1rc1,<0.1.2" }, + { name = "hmz", extras = ["dsh", "kimi"], marker = "extra == 'all'" }, + { name = "psutil", specifier = ">=7.2.2" }, + { name = "pydantic", specifier = ">=2.13.5" }, + { name = "pyfiglet", specifier = ">=1.0" }, + { name = "python-dotenv", marker = "extra == 'dsh'", specifier = ">=1.2.3" }, + { name = "pyyaml", specifier = ">=6.0.3" }, + { name = "rich", specifier = ">=15" }, + { name = "sentry-sdk", specifier = ">=2.68.1" }, + { name = "textual", specifier = ">=8.2.8" }, + { name = "typing-extensions", specifier = ">=4.13" }, + { name = "websockets", marker = "extra == 'kimi'", specifier = ">=15,<18" }, +] +provides-extras = ["dsh", "kimi", "all"] + +[[package]] +name = "iniconfig" +version = "2.3.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/72/34/14ca021ce8e5dfedc35312d08ba8bf51fdd999c576889fc2c24cb97f4f10/iniconfig-2.3.0.tar.gz", hash = "sha256:c76315c77db068650d49c5b56314774a7804df16fee4402c1f19d6d15d8c4730", size = 20503, upload-time = "2025-10-18T21:55:43.219Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/cb/b1/3846dd7f199d53cb17f49cba7e651e9ce294d8497c8c150530ed11865bb8/iniconfig-2.3.0-py3-none-any.whl", hash = "sha256:f631c04d2c48c52b84d0d0549c99ff3859c98df65b3101406327ecc7d53fbf12", size = 7484, upload-time = "2025-10-18T21:55:41.639Z" }, +] + +[[package]] +name = "linkify-it-py" +version = "2.2.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/45/98/7a1a5f31fd5c7ba93e963b168e244b8e3dd705b3d2a718e3c3307583bf57/linkify_it_py-2.2.0.tar.gz", hash = "sha256:907acd2d17ac1fbb9ddb62c8957ccbd6158cac602231a15c3b0cd1e215f03cee", size = 32939, upload-time = "2026-08-29T07:07:08.305Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/13/d4/1152d1c7ab42d8b908be64fd200ddc870dc9d4925e951198702084aa1a7d/linkify_it_py-2.2.0-py3-none-any.whl", hash = "sha256:3adc40eb5af300b2605fcfdb968c24e1d780a90f1f2221af7c15e5111e94d443", size = 21971, upload-time = "2026-08-29T07:07:07.164Z" }, +] + +[[package]] +name = "markdown-it-py" +version = "4.2.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "mdurl" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/06/ff/7841249c247aa650a76b9ee4bbaeae59370dc8bfd2f6c01f3630c35eb134/markdown_it_py-4.2.0.tar.gz", hash = "sha256:04a21681d6fbb623de53f6f364d352309d4094dd4194040a10fd51833e418d49", size = 82454, upload-time = "2026-05-07T12:08:28.36Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b3/81/4da04ced5a082363ecfa159c010d200ecbd959ae410c10c0264a38cac0f5/markdown_it_py-4.2.0-py3-none-any.whl", hash = "sha256:9f7ebbcd14fe59494226453aed97c1070d83f8d24b6fc3a3bcf9a38092641c4a", size = 91687, upload-time = "2026-05-07T12:08:27.182Z" }, +] + +[package.optional-dependencies] +linkify = [ + { name = "linkify-it-py" }, +] + +[[package]] +name = "mdit-py-plugins" +version = "0.6.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "markdown-it-py" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/59/fc/f8d0863f8862f25602c0404d75568e89fb6b4109804645e5cdfb1be5cf56/mdit_py_plugins-0.6.1.tar.gz", hash = "sha256:a2bca0f039f39dbd35fb74ae1b5f998608c437463371f0ff7f49a19a17a114d0", size = 56114, upload-time = "2026-05-13T09:03:38.91Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a5/69/6da5581c6a7fede7dc261bf4e67d6adca4196f176b43288b55b3db395b6e/mdit_py_plugins-0.6.1-py3-none-any.whl", hash = "sha256:214c82fb2ac524472ab6a5bcab1de80f73b50443e187f401bfd77efbc7c6481d", size = 66663, upload-time = "2026-05-13T09:03:37.76Z" }, +] + +[[package]] +name = "mdurl" +version = "0.1.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/d6/54/cfe61301667036ec958cb99bd3efefba235e65cdeb9c84d24a8293ba1d90/mdurl-0.1.2.tar.gz", hash = "sha256:bb413d29f5eea38f31dd4754dd7377d4465116fb207585f97bf925588687c1ba", size = 8729, upload-time = "2022-08-14T12:40:10.846Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b3/38/89ba8ad64ae25be8de66a6d463314cf1eb366222074cfda9ee839c56a4b4/mdurl-0.1.2-py3-none-any.whl", hash = "sha256:84008a41e51615a49fc9966191ff91509e3c40b939176e643fd50a5c2196b8f8", size = 9979, upload-time = "2022-08-14T12:40:09.779Z" }, +] + +[[package]] +name = "packaging" +version = "26.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/7d/fa/3944b40b07da9ce895c0e6303a5ab7d53da063554f534556b134a54d6093/packaging-26.3.tar.gz", hash = "sha256:94edc256424af38762eb31306eed28beb9f0efc50a8837492c9d6fd6004aed79", size = 313412, upload-time = "2026-08-04T18:15:28.737Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/63/34/ba1c580383c9eada3711951fef0795c80b829a078d72188184bcab9dd527/packaging-26.3-py3-none-any.whl", hash = "sha256:d7193f7c8e4e93f444fde0262bf90af30e16fa0ad0ad44cb553c87339b23cd1c", size = 129956, upload-time = "2026-08-04T18:15:27.159Z" }, +] + +[[package]] +name = "platformdirs" +version = "4.12.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/17/c8/721b3855fe457da514fe249247d404b9b39c5d16532278f70ebaa6acf18b/platformdirs-4.12.2.tar.gz", hash = "sha256:eab5f70271a490ef74618bb314fbb86e3c7e82fa3b9c922c2ea0e0a1a155d329", size = 59353, upload-time = "2026-09-29T18:27:56.009Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d0/89/446044f33aba0348d35e433f56d12206d010a5281a1df54054d4cfb82388/platformdirs-4.12.2-py3-none-any.whl", hash = "sha256:29dbf06d96c500bc6bdbce75fb0a14d63279c93b1842f97e72a135b33e856983", size = 32457, upload-time = "2026-09-29T18:27:54.43Z" }, +] + +[[package]] +name = "pluggy" +version = "1.6.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f9/e2/3e91f31a7d2b083fe6ef3fa267035b518369d9511ffab804f839851d2779/pluggy-1.6.0.tar.gz", hash = "sha256:7dcc130b76258d33b90f61b658791dede3486c3e6bfb003ee5c9bfb396dd22f3", size = 69412, upload-time = "2025-05-15T12:30:07.975Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746", size = 20538, upload-time = "2025-05-15T12:30:06.134Z" }, +] + +[[package]] +name = "psutil" +version = "7.2.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/aa/c6/d1ddf4abb55e93cebc4f2ed8b5d6dbad109ecb8d63748dd2b20ab5e57ebe/psutil-7.2.2.tar.gz", hash = "sha256:0746f5f8d406af344fd547f1c8daa5f5c33dbc293bb8d6a16d80b4bb88f59372", size = 493740, upload-time = "2026-01-28T18:14:54.428Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/51/08/510cbdb69c25a96f4ae523f733cdc963ae654904e8db864c07585ef99875/psutil-7.2.2-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:2edccc433cbfa046b980b0df0171cd25bcaeb3a68fe9022db0979e7aa74a826b", size = 130595, upload-time = "2026-01-28T18:14:57.293Z" }, + { url = "https://files.pythonhosted.org/packages/d6/f5/97baea3fe7a5a9af7436301f85490905379b1c6f2dd51fe3ecf24b4c5fbf/psutil-7.2.2-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:e78c8603dcd9a04c7364f1a3e670cea95d51ee865e4efb3556a3a63adef958ea", size = 131082, upload-time = "2026-01-28T18:14:59.732Z" }, + { url = "https://files.pythonhosted.org/packages/37/d6/246513fbf9fa174af531f28412297dd05241d97a75911ac8febefa1a53c6/psutil-7.2.2-cp313-cp313t-manylinux2010_x86_64.manylinux_2_12_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:1a571f2330c966c62aeda00dd24620425d4b0cc86881c89861fbc04549e5dc63", size = 181476, upload-time = "2026-01-28T18:15:01.884Z" }, + { url = "https://files.pythonhosted.org/packages/b8/b5/9182c9af3836cca61696dabe4fd1304e17bc56cb62f17439e1154f225dd3/psutil-7.2.2-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:917e891983ca3c1887b4ef36447b1e0873e70c933afc831c6b6da078ba474312", size = 184062, upload-time = "2026-01-28T18:15:04.436Z" }, + { url = "https://files.pythonhosted.org/packages/16/ba/0756dca669f5a9300d0cbcbfae9a4c30e446dfc7440ffe43ded5724bfd93/psutil-7.2.2-cp313-cp313t-win_amd64.whl", hash = "sha256:ab486563df44c17f5173621c7b198955bd6b613fb87c71c161f827d3fb149a9b", size = 139893, upload-time = "2026-01-28T18:15:06.378Z" }, + { url = "https://files.pythonhosted.org/packages/1c/61/8fa0e26f33623b49949346de05ec1ddaad02ed8ba64af45f40a147dbfa97/psutil-7.2.2-cp313-cp313t-win_arm64.whl", hash = "sha256:ae0aefdd8796a7737eccea863f80f81e468a1e4cf14d926bd9b6f5f2d5f90ca9", size = 135589, upload-time = "2026-01-28T18:15:08.03Z" }, + { url = "https://files.pythonhosted.org/packages/81/69/ef179ab5ca24f32acc1dac0c247fd6a13b501fd5534dbae0e05a1c48b66d/psutil-7.2.2-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:eed63d3b4d62449571547b60578c5b2c4bcccc5387148db46e0c2313dad0ee00", size = 130664, upload-time = "2026-01-28T18:15:09.469Z" }, + { url = "https://files.pythonhosted.org/packages/7b/64/665248b557a236d3fa9efc378d60d95ef56dd0a490c2cd37dafc7660d4a9/psutil-7.2.2-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:7b6d09433a10592ce39b13d7be5a54fbac1d1228ed29abc880fb23df7cb694c9", size = 131087, upload-time = "2026-01-28T18:15:11.724Z" }, + { url = "https://files.pythonhosted.org/packages/d5/2e/e6782744700d6759ebce3043dcfa661fb61e2fb752b91cdeae9af12c2178/psutil-7.2.2-cp314-cp314t-manylinux2010_x86_64.manylinux_2_12_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:1fa4ecf83bcdf6e6c8f4449aff98eefb5d0604bf88cb883d7da3d8d2d909546a", size = 182383, upload-time = "2026-01-28T18:15:13.445Z" }, + { url = "https://files.pythonhosted.org/packages/57/49/0a41cefd10cb7505cdc04dab3eacf24c0c2cb158a998b8c7b1d27ee2c1f5/psutil-7.2.2-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e452c464a02e7dc7822a05d25db4cde564444a67e58539a00f929c51eddda0cf", size = 185210, upload-time = "2026-01-28T18:15:16.002Z" }, + { url = "https://files.pythonhosted.org/packages/dd/2c/ff9bfb544f283ba5f83ba725a3c5fec6d6b10b8f27ac1dc641c473dc390d/psutil-7.2.2-cp314-cp314t-win_amd64.whl", hash = "sha256:c7663d4e37f13e884d13994247449e9f8f574bc4655d509c3b95e9ec9e2b9dc1", size = 141228, upload-time = "2026-01-28T18:15:18.385Z" }, + { url = "https://files.pythonhosted.org/packages/f2/fc/f8d9c31db14fcec13748d373e668bc3bed94d9077dbc17fb0eebc073233c/psutil-7.2.2-cp314-cp314t-win_arm64.whl", hash = "sha256:11fe5a4f613759764e79c65cf11ebdf26e33d6dd34336f8a337aa2996d71c841", size = 136284, upload-time = "2026-01-28T18:15:19.912Z" }, + { url = "https://files.pythonhosted.org/packages/e7/36/5ee6e05c9bd427237b11b3937ad82bb8ad2752d72c6969314590dd0c2f6e/psutil-7.2.2-cp36-abi3-macosx_10_9_x86_64.whl", hash = "sha256:ed0cace939114f62738d808fdcecd4c869222507e266e574799e9c0faa17d486", size = 129090, upload-time = "2026-01-28T18:15:22.168Z" }, + { url = "https://files.pythonhosted.org/packages/80/c4/f5af4c1ca8c1eeb2e92ccca14ce8effdeec651d5ab6053c589b074eda6e1/psutil-7.2.2-cp36-abi3-macosx_11_0_arm64.whl", hash = "sha256:1a7b04c10f32cc88ab39cbf606e117fd74721c831c98a27dc04578deb0c16979", size = 129859, upload-time = "2026-01-28T18:15:23.795Z" }, + { url = "https://files.pythonhosted.org/packages/b5/70/5d8df3b09e25bce090399cf48e452d25c935ab72dad19406c77f4e828045/psutil-7.2.2-cp36-abi3-manylinux2010_x86_64.manylinux_2_12_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:076a2d2f923fd4821644f5ba89f059523da90dc9014e85f8e45a5774ca5bc6f9", size = 155560, upload-time = "2026-01-28T18:15:25.976Z" }, + { url = "https://files.pythonhosted.org/packages/63/65/37648c0c158dc222aba51c089eb3bdfa238e621674dc42d48706e639204f/psutil-7.2.2-cp36-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:b0726cecd84f9474419d67252add4ac0cd9811b04d61123054b9fb6f57df6e9e", size = 156997, upload-time = "2026-01-28T18:15:27.794Z" }, + { url = "https://files.pythonhosted.org/packages/8e/13/125093eadae863ce03c6ffdbae9929430d116a246ef69866dad94da3bfbc/psutil-7.2.2-cp36-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:fd04ef36b4a6d599bbdb225dd1d3f51e00105f6d48a28f006da7f9822f2606d8", size = 148972, upload-time = "2026-01-28T18:15:29.342Z" }, + { url = "https://files.pythonhosted.org/packages/04/78/0acd37ca84ce3ddffaa92ef0f571e073faa6d8ff1f0559ab1272188ea2be/psutil-7.2.2-cp36-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:b58fabe35e80b264a4e3bb23e6b96f9e45a3df7fb7eed419ac0e5947c61e47cc", size = 148266, upload-time = "2026-01-28T18:15:31.597Z" }, + { url = "https://files.pythonhosted.org/packages/b4/90/e2159492b5426be0c1fef7acba807a03511f97c5f86b3caeda6ad92351a7/psutil-7.2.2-cp37-abi3-win_amd64.whl", hash = "sha256:eb7e81434c8d223ec4a219b5fc1c47d0417b12be7ea866e24fb5ad6e84b3d988", size = 137737, upload-time = "2026-01-28T18:15:33.849Z" }, + { url = "https://files.pythonhosted.org/packages/8c/c7/7bb2e321574b10df20cbde462a94e2b71d05f9bbda251ef27d104668306a/psutil-7.2.2-cp37-abi3-win_arm64.whl", hash = "sha256:8c233660f575a5a89e6d4cb65d9f938126312bca76d8fe087b947b3a1aaac9ee", size = 134617, upload-time = "2026-01-28T18:15:36.514Z" }, +] + +[[package]] +name = "pydantic" +version = "2.13.5" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "annotated-types" }, + { name = "pydantic-core" }, + { name = "typing-extensions" }, + { name = "typing-inspection" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/53/ef/fc4f868f4e2cee79f863883abffceff107875f569b848507319842d2a681/pydantic-2.13.5.tar.gz", hash = "sha256:51a9c5f7b2f8e636f04c6cada605d9b6a3bf1348fdf945a3d8869b19bba0ee08", size = 845750, upload-time = "2026-08-28T14:04:00.916Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/eb/47/c95ffc2009878c7aac0c5e08528022dcb885933252a88b5f170058014464/pydantic-2.13.5-py3-none-any.whl", hash = "sha256:346a034f080da3755d8e9cb5e00e8b07de1d39e4f6e2c87d8ab7cafa0b269a73", size = 472589, upload-time = "2026-08-28T14:03:59.136Z" }, +] + +[[package]] +name = "pydantic-core" +version = "2.46.5" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/af/f9/8a06bea35ef8daf588f707784c973a7046e0034c8d8cfb08828eeffb8b75/pydantic_core-2.46.5.tar.gz", hash = "sha256:10416c15b8839ecc4ef4d0885da76da6fd0f67333a0eb8aff6d93c4b8f2910fc", size = 472262, upload-time = "2026-08-28T10:01:31.677Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/82/3f/76358795aa7a8c6d4f36e2cb828ad1c90ee118e1393a9281664f5aade9d4/pydantic_core-2.46.5-cp312-cp312-macosx_10_12_x86_64.whl", hash = "sha256:b9fe6fb92520e3fd61f2e49000b6911b188824f089b75973ea06d6267f0b476d", size = 2076516, upload-time = "2026-08-28T09:58:21.576Z" }, + { url = "https://files.pythonhosted.org/packages/db/50/26b091836076ce4cb2fac264186936acc069e0595772cfd02a563bc4761a/pydantic_core-2.46.5-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:a39ac25a9a2fa4072efdb429833c4a4c8009a51ff9eea3eeae131713cd27991e", size = 1922874, upload-time = "2026-08-28T09:58:23.766Z" }, + { url = "https://files.pythonhosted.org/packages/09/f0/2a8ce3849e299d44e2d2c196b6082643a3235565a735cb51db7a6261f614/pydantic_core-2.46.5-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:4fdc8b93a41521988916eeaa271173fcca7fa0803d62f87675aac8dcec1c8e29", size = 1951772, upload-time = "2026-08-28T09:58:25.435Z" }, + { url = "https://files.pythonhosted.org/packages/87/46/ac0dc8bdd9e6048183a14eb127764e7ad9240021c17513074a4711b0e31e/pydantic_core-2.46.5-cp312-cp312-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:b98134087d9de723658d17a42c7d0da8d6e2ef08015dee7dc93889047315f5e4", size = 2031832, upload-time = "2026-08-28T09:58:27.102Z" }, + { url = "https://files.pythonhosted.org/packages/c4/c2/339de5bef7be36301a2231eaa52e62163742c2281f11b5f4892bc79785cd/pydantic_core-2.46.5-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:e652ab17569c94bff5475520f907b7148b8c24036a8ebbe5cf7cf7493d28579a", size = 2208645, upload-time = "2026-08-28T09:58:28.948Z" }, + { url = "https://files.pythonhosted.org/packages/7b/a0/9ff22b797724262da14427abaed4dd1d864a139693fc5e7809114376a716/pydantic_core-2.46.5-cp312-cp312-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:d925f3d9afd05a8c0fb3a1031463a8d59ebe5e2afad297e29c78be19e13b4e62", size = 2265935, upload-time = "2026-08-28T09:58:30.625Z" }, + { url = "https://files.pythonhosted.org/packages/c0/a4/eb9409ec0736e50aa70a412f16c204ed149516846912f7e6724d4c73ee53/pydantic_core-2.46.5-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:0fc5be0abd4a407e200d844b404e33639a554e7bd0d448e7b9ae181be4789ac2", size = 2066284, upload-time = "2026-08-28T09:58:32.289Z" }, + { url = "https://files.pythonhosted.org/packages/c0/02/7f6156ffc926857f1c37c07d9a388682865a81830ab6a1b637082c25e399/pydantic_core-2.46.5-cp312-cp312-manylinux_2_31_riscv64.whl", hash = "sha256:816ff0a6550ffc06c098ccd2e0698600f9aa7da192a79eaa6f9af504a35db869", size = 2105889, upload-time = "2026-08-28T09:58:33.986Z" }, + { url = "https://files.pythonhosted.org/packages/92/b1/e781d357ebe09fc929f995700f1b3503e8897f1cece183ecb1300d4d67e9/pydantic_core-2.46.5-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:c7ea57fc63aa7da93a1bd2d644e6577befae10c52c4e36377635eea1056a74f5", size = 2158006, upload-time = "2026-08-28T09:58:35.647Z" }, + { url = "https://files.pythonhosted.org/packages/70/0a/644597d84ab400e50609c192120b85c9681c22d3a20461b9060a79be0a7a/pydantic_core-2.46.5-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:efd62a42486f1bda5d24cb4f63d15a3c7768375fe83d36f9417b4ad7a2fb20b3", size = 2158408, upload-time = "2026-08-28T09:58:37.38Z" }, + { url = "https://files.pythonhosted.org/packages/1e/ee/ca3b7b3a4b3769ffe9ce9432a7c9be755de9593a46d3b0d54d0409323e44/pydantic_core-2.46.5-cp312-cp312-musllinux_1_1_armv7l.whl", hash = "sha256:2bc9419666990c06d7397831f2126a1ecc3594aaa3ff7de5bf2d066802f4e07b", size = 2309609, upload-time = "2026-08-28T09:58:39.22Z" }, + { url = "https://files.pythonhosted.org/packages/ce/52/39fa1f451486019524ca685020390e7ca351832fd874530ba30c8628e6dc/pydantic_core-2.46.5-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:18a09e1e1011b462f2e32774f25859ef1223d5c2b0546a633cf56654710721e0", size = 2342618, upload-time = "2026-08-28T09:58:40.89Z" }, + { url = "https://files.pythonhosted.org/packages/81/5e/468fc630568c61dcef3cd47ad32ffbeed9af643f49208d1ea86ab4f890c4/pydantic_core-2.46.5-cp312-cp312-win32.whl", hash = "sha256:5cb482e9e84c851f4e623fe4acc1ced89168cf1fe18f7089db4548c8f5bbb65b", size = 1939475, upload-time = "2026-08-28T09:58:42.591Z" }, + { url = "https://files.pythonhosted.org/packages/cf/c9/4c19f41b84cf6b622a72fbeed7665b25d47a187d68d47d0d430c07f23268/pydantic_core-2.46.5-cp312-cp312-win_amd64.whl", hash = "sha256:5e81740c09e310f5aa5cbd3e434a01c154d4bef93241c7877b39f211d2b78ba8", size = 2043140, upload-time = "2026-08-28T09:58:44.272Z" }, + { url = "https://files.pythonhosted.org/packages/af/dd/0c1a050299147c746e5256db16d645ab5efd4f78c59937d581a0524e74a2/pydantic_core-2.46.5-cp312-cp312-win_arm64.whl", hash = "sha256:f7b0ec93a2893de856652154d73b7ba622f26fa97726487dcac373de5f4c6084", size = 1997729, upload-time = "2026-08-28T09:58:46.13Z" }, + { url = "https://files.pythonhosted.org/packages/f5/37/5abe39a8372a61d3dc3c1338fc504281c01b32fdb3169cd7187153b56d3e/pydantic_core-2.46.5-cp313-cp313-macosx_10_12_x86_64.whl", hash = "sha256:b7ca9034437b6022f941f4857459562ee00a560b97e7cce8a0ec5a74fc6766e0", size = 2075885, upload-time = "2026-08-28T09:58:47.856Z" }, + { url = "https://files.pythonhosted.org/packages/21/43/6323b1f8b217780454c61304bcd2b38ae4762f50754414124603ccc90bb2/pydantic_core-2.46.5-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:f332f0e72a5a0400141f830744e141bf9f97917878dbe968669e8a7fefea78ff", size = 1922768, upload-time = "2026-08-28T09:58:49.58Z" }, + { url = "https://files.pythonhosted.org/packages/0f/a3/c05ca796e1197618a774b01e596aeedfefc2f7d8c01ae3054e910b120e8a/pydantic_core-2.46.5-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:193375f3548919d3f0b60936ca113ada3e38f264f91b9b8e0508efaad57be931", size = 1951241, upload-time = "2026-08-28T09:58:51.511Z" }, + { url = "https://files.pythonhosted.org/packages/68/32/33bc39ac705c52cffc908e8389f9754fdb208aea5c69cceddf4eb3ce99af/pydantic_core-2.46.5-cp313-cp313-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:79bdfa52f843137045b2d081cc05c120ba6665d29b7559c2c47690906f39279f", size = 2031975, upload-time = "2026-08-28T09:58:53.166Z" }, + { url = "https://files.pythonhosted.org/packages/b0/70/2333e885c0f6a67bc105c5916965dac9b57f2718ee20d81d1a06a4ebdc13/pydantic_core-2.46.5-cp313-cp313-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:24922243639cbdac66c75fcb6fd6495a9cb52b213d62f9a0d16f0310b1ff8038", size = 2208542, upload-time = "2026-08-28T09:58:55.017Z" }, + { url = "https://files.pythonhosted.org/packages/f7/ea/296debfb4264207bbda5936133892e027c0a58875ad53ebd512fba8ec3a2/pydantic_core-2.46.5-cp313-cp313-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:c76fe65e607be28c7fd4d56fc3c42b1583aa058ce3408b7ad0fd540171d31f9f", size = 2264692, upload-time = "2026-08-28T09:58:56.767Z" }, + { url = "https://files.pythonhosted.org/packages/d3/f2/9e4de77a6271e07a76d2d58b11c091a979c191ed2939bf80067568b369d2/pydantic_core-2.46.5-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6f7b393a8b3da82f5c1fc0751e6d01ac6c55b93c18226a60bdfba4a724efafd1", size = 2066633, upload-time = "2026-08-28T09:58:58.531Z" }, + { url = "https://files.pythonhosted.org/packages/8d/db/f9e9d0c97445987b2084823d5c240de88087338f04fc2cfaa2df186b8049/pydantic_core-2.46.5-cp313-cp313-manylinux_2_31_riscv64.whl", hash = "sha256:7ac031912d54f3d83ef3b3eb98dfabc1608802e2202263d25957eeed40b94761", size = 2105235, upload-time = "2026-08-28T09:59:00.421Z" }, + { url = "https://files.pythonhosted.org/packages/07/c5/79169b047b3b2c3e99e04bc76372af9637e0bf6db638274fa927df96369e/pydantic_core-2.46.5-cp313-cp313-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:837b396ca3d7b74091ca623f6cbd8351bd42d670a79c2683e79fb089f06a2de5", size = 2157367, upload-time = "2026-08-28T09:59:02.442Z" }, + { url = "https://files.pythonhosted.org/packages/26/b5/ba6057afb7c291bd449f51b867f95aef2072941c4ce4e5c31d6ffd132d3b/pydantic_core-2.46.5-cp313-cp313-musllinux_1_1_aarch64.whl", hash = "sha256:5ee239d575f80b08eca11f6e20f90c4c695de7825c67eefe6091fbf20dda648e", size = 2158420, upload-time = "2026-08-28T09:59:04.2Z" }, + { url = "https://files.pythonhosted.org/packages/6e/28/2057abecaafdc22912afa819603a51f0a62d40643b7c4871c51721fea9be/pydantic_core-2.46.5-cp313-cp313-musllinux_1_1_armv7l.whl", hash = "sha256:e80675d75ae2cd14372cb65cad5400d9347a3d3f6c13000183f22dfd027283ed", size = 2309588, upload-time = "2026-08-28T09:59:06.048Z" }, + { url = "https://files.pythonhosted.org/packages/71/9d/881156dc404e27479c4246128d73538464cab4a239bec61995e227644c30/pydantic_core-2.46.5-cp313-cp313-musllinux_1_1_x86_64.whl", hash = "sha256:9c4b71f10dd532fb7a5cbc8f58707779e64f03a258c2bf8bfbaecfcd9970b519", size = 2341866, upload-time = "2026-08-28T09:59:08.539Z" }, + { url = "https://files.pythonhosted.org/packages/5a/38/d66f443a259f84d13babdceae568e572b0ed26da17ca5d0a649ebb110a67/pydantic_core-2.46.5-cp313-cp313-win32.whl", hash = "sha256:97bf8de4d541598c94a59344eeb988a94c08ff76b5723c41f6567ec18c7892ea", size = 1938580, upload-time = "2026-08-28T09:59:10.402Z" }, + { url = "https://files.pythonhosted.org/packages/2c/1e/1d5371213f4cc9a7ed70c0bfcc7911de22311ee99a662a56077d7292d2ac/pydantic_core-2.46.5-cp313-cp313-win_amd64.whl", hash = "sha256:15f4a94963c95accac15b7b657bb177d3ad82bb90b0d0526d9a9b85079925db5", size = 2041980, upload-time = "2026-08-28T09:59:12.396Z" }, + { url = "https://files.pythonhosted.org/packages/5a/48/4222d90b1c67568bace4dec6dca6271449c66de3595d72b6d098f5fde597/pydantic_core-2.46.5-cp313-cp313-win_arm64.whl", hash = "sha256:d22a945598fb91236b4dd793a6e42e4f3dd7740bb5aace5ebd7d4c08d13bb575", size = 1997213, upload-time = "2026-08-28T09:59:14.245Z" }, + { url = "https://files.pythonhosted.org/packages/8e/8a/14596f2a8367da50cf7cbac48169ee5d9c8e11d486a3b527082384630c72/pydantic_core-2.46.5-cp314-cp314-macosx_10_12_x86_64.whl", hash = "sha256:c1c43ad4339643d70ebb8124e1305a7dab423001eff58bb41a0f731adbc98355", size = 2074081, upload-time = "2026-08-28T09:59:16.141Z" }, + { url = "https://files.pythonhosted.org/packages/ae/d5/d8a4eb6d6c7f66b91dd37c576d76e9e60fba900caf5372c17bcf949febc2/pydantic_core-2.46.5-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:1a353f84de772f423b5ffb11d7ae352fbbef0f446f3c0b0af0f8236d7233606e", size = 1920497, upload-time = "2026-08-28T09:59:18.065Z" }, + { url = "https://files.pythonhosted.org/packages/8e/26/092079428f86e927e030b2c0ced87df69dbb1c875cdeaa67bf42ea2be746/pydantic_core-2.46.5-cp314-cp314-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:5086029a57366b8cf81b130a43908738095c270c21a8d7f0e8bdfdb89718e2f3", size = 1952130, upload-time = "2026-08-28T09:59:20.476Z" }, + { url = "https://files.pythonhosted.org/packages/08/c3/8ec0e290a9ebaebd64047bf5fda94be835c6b1551b02437e4b76778fbcd7/pydantic_core-2.46.5-cp314-cp314-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:46c25dda9d092a06c08db76ffe0a197107904d0dfac653f7d5306bbcd6d6119c", size = 2026371, upload-time = "2026-08-28T09:59:22.227Z" }, + { url = "https://files.pythonhosted.org/packages/01/72/4fd20ad520fb8da0157f95b27a7eb05a72790ef08138e7701ac972c342ea/pydantic_core-2.46.5-cp314-cp314-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:37ea7b83c935e5b0d68c9449b82651accf78a10828b2c02b2f2d9e9496446c21", size = 2202822, upload-time = "2026-08-28T09:59:24.277Z" }, + { url = "https://files.pythonhosted.org/packages/31/b0/d16e0771206b29314f0d52198b720be21e8a99ab2bf11e3bc0d7c9cebdff/pydantic_core-2.46.5-cp314-cp314-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:e64e88d5585bea9ce95861079de72006c7fa6d3df4e3a3b65ba31eb979c15c9f", size = 2262756, upload-time = "2026-08-28T09:59:26.608Z" }, + { url = "https://files.pythonhosted.org/packages/2c/9b/59634b7ac631c63b2a37760eb6943af3e29573d6b59a4abc5e7f019d4cee/pydantic_core-2.46.5-cp314-cp314-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:54d510bac3ee52247af28ed4bb18a1e799f040ac60fd2bf5ccd4c92f1fbe786f", size = 2068352, upload-time = "2026-08-28T09:59:29.044Z" }, + { url = "https://files.pythonhosted.org/packages/08/7c/570abb1ad2155348dc754ea91be22e5aaa18eb6d69a6068f7c6f2679a6ed/pydantic_core-2.46.5-cp314-cp314-manylinux_2_31_riscv64.whl", hash = "sha256:a2a5e1d0ff29adddc9f6d6821a66302e4493f8ca898b715b6b1182c2c201ea0a", size = 2104777, upload-time = "2026-08-28T09:59:30.95Z" }, + { url = "https://files.pythonhosted.org/packages/8e/25/5bf74adc65a1ac5b7be3f6cb0bcb5433615c1598a801c19d830d84c98ded/pydantic_core-2.46.5-cp314-cp314-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:03b9666e41e35d8909852ba191a0607520f81b74eaf12ccf8737005dbb313821", size = 2156312, upload-time = "2026-08-28T09:59:32.604Z" }, + { url = "https://files.pythonhosted.org/packages/90/6a/2ef38830675e050121040618135564ed56b860b45433b02d9b4ebece46f3/pydantic_core-2.46.5-cp314-cp314-musllinux_1_1_aarch64.whl", hash = "sha256:a91c17edf6eea2402cb5457b4c89e99bc5ed1004aa34c4adf1d4258c1a5c22c2", size = 2150067, upload-time = "2026-08-28T09:59:34.453Z" }, + { url = "https://files.pythonhosted.org/packages/90/ef/a7dbb03a14a64c2a4621f989c615ed9a892535a6cad938fc27079f919d80/pydantic_core-2.46.5-cp314-cp314-musllinux_1_1_armv7l.whl", hash = "sha256:b49924c73a235e969511bf2aabdff3beebf9820931f646c80274d5d780010c47", size = 2304516, upload-time = "2026-08-28T09:59:36.194Z" }, + { url = "https://files.pythonhosted.org/packages/68/f8/6bb4c4b80e8a6fde1904c64a51c62a1d04fcdfa3ea521a66b2ddefa1d885/pydantic_core-2.46.5-cp314-cp314-musllinux_1_1_x86_64.whl", hash = "sha256:2cbd9a5eff05e51c447c34dfa4632145b26b09120cf04bd0c871e44c1a5e1c9a", size = 2335223, upload-time = "2026-08-28T09:59:37.931Z" }, + { url = "https://files.pythonhosted.org/packages/2a/80/f46b8c681195190b2c1f1c7c0a81abce60663e987613e09ef64d433dd96b/pydantic_core-2.46.5-cp314-cp314-win32.whl", hash = "sha256:2d5d76654becf5efd62c9e51c3756c67b49498b0c9a40884934c40807adbd074", size = 1934827, upload-time = "2026-08-28T09:59:39.836Z" }, + { url = "https://files.pythonhosted.org/packages/f7/3c/60674207246bc0a4009d2391b7c7251c7159f279c8d2ab8aae8ef46f3dee/pydantic_core-2.46.5-cp314-cp314-win_amd64.whl", hash = "sha256:fa10ef4112775900e7a0661068635eb67b2ab824fbde764de6e0e21982a93db0", size = 2042648, upload-time = "2026-08-28T09:59:41.792Z" }, + { url = "https://files.pythonhosted.org/packages/69/0c/117c562c7c1babdf44576b72a5e496906506c93690387ecfbca7c729ae2e/pydantic_core-2.46.5-cp314-cp314-win_arm64.whl", hash = "sha256:045ab3b6d308439e32b81cc173bba5b9018bc6ed896afd0c65b3b009b1699af5", size = 1989652, upload-time = "2026-08-28T09:59:43.702Z" }, + { url = "https://files.pythonhosted.org/packages/e8/66/9336ae58f9eb68c41d121894e52c4c89eccb07eb8f602a04ee9c3f37736a/pydantic_core-2.46.5-cp314-cp314t-macosx_10_12_x86_64.whl", hash = "sha256:8816f3d218beb4b787de5c9759c259b8fa61f9dec42dc7811f320a33771778b7", size = 2065829, upload-time = "2026-08-28T09:59:45.364Z" }, + { url = "https://files.pythonhosted.org/packages/c5/02/bc19b47a96c2d3109760711acf22369e56bd7e405ca52f7ade164d2ead57/pydantic_core-2.46.5-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:bce57638e08ac148e5778cce7feb968307a727d66f8e2274a543d0cf0c9ad6a3", size = 1905716, upload-time = "2026-08-28T09:59:47.18Z" }, + { url = "https://files.pythonhosted.org/packages/52/a4/70b47c0509923dd98ccfed04fb3e32ea3849c82a0ff2205bb41009b43c00/pydantic_core-2.46.5-cp314-cp314t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:976e1128455aa595ea04c79ccfedff1aaeab96ee013fcc916bed120c4f0ad94f", size = 1934216, upload-time = "2026-08-28T09:59:49.241Z" }, + { url = "https://files.pythonhosted.org/packages/52/ab/aa03b65f7bb198585edf806b906c3223ecf1795543e39e23aec4cce27ad2/pydantic_core-2.46.5-cp314-cp314t-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:e7b891faeedeafba41b2983e5001a81b6a915b69544c7e7570d1989ce1c36ac7", size = 2010635, upload-time = "2026-08-28T09:59:51.692Z" }, + { url = "https://files.pythonhosted.org/packages/3c/8b/0da06343f30b84ec549aafd309c6456223d5dc8bd36af504c573faad561d/pydantic_core-2.46.5-cp314-cp314t-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:5f194189415698233dd1114a093a9b56e61e2c57e11b469be3b0506f46f0771c", size = 2209369, upload-time = "2026-08-28T09:59:53.582Z" }, + { url = "https://files.pythonhosted.org/packages/d6/5b/844c4defaa34a3df66eb9257087d121d70c201298b96abdf9f492fc2f1bf/pydantic_core-2.46.5-cp314-cp314t-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:82a36973cf8a2ef5406f4fe2edbf8ed0c99629535d959e0b100c76a32535a111", size = 2253238, upload-time = "2026-08-28T09:59:55.484Z" }, + { url = "https://files.pythonhosted.org/packages/f4/64/a4e536cb16d7f61a7fd3120b46c577fc7fa7325992f69c4f52bc786d77d8/pydantic_core-2.46.5-cp314-cp314t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:cdbb78909f52b981d3b2d56b97328d71eb0b974c36bd77c920123a7ebb192829", size = 2065740, upload-time = "2026-08-28T09:59:58.038Z" }, + { url = "https://files.pythonhosted.org/packages/5f/75/aaa38c6bc2d085f6605b34eabdc6a8a4e0b2e61fc9c8e6e52b28e97b3125/pydantic_core-2.46.5-cp314-cp314t-manylinux_2_31_riscv64.whl", hash = "sha256:52e24eacdb536cade636aa90fb851835222becff8484b7001fdc78cb0290f2aa", size = 2087425, upload-time = "2026-08-28T09:59:59.898Z" }, + { url = "https://files.pythonhosted.org/packages/55/ae/fcab4cfc39aba3689e1d20c8b5250ad280957022c09af2ed9cd585602a5e/pydantic_core-2.46.5-cp314-cp314t-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:37ae34309d7bd8c0d61ab839668058f2a7962ea1fc51d105d2db228fe0618034", size = 2139306, upload-time = "2026-08-28T10:00:03.057Z" }, + { url = "https://files.pythonhosted.org/packages/2d/f4/f1d03a4bc9d9acbc62f4d742b8a319af52f71885079868b2ff8e48a651ee/pydantic_core-2.46.5-cp314-cp314t-musllinux_1_1_aarch64.whl", hash = "sha256:0cdbada856a1c69a7624a64d3d9aefe79300bd6ef827b43a4f265010b9b55184", size = 2144589, upload-time = "2026-08-28T10:00:05.645Z" }, + { url = "https://files.pythonhosted.org/packages/83/f3/7a53bb1356de514a4cd295f25b6ac39237895620c0462d2592b76c16e114/pydantic_core-2.46.5-cp314-cp314t-musllinux_1_1_armv7l.whl", hash = "sha256:545f26c504b27c3758439a5e6d9349931f0a04f855668d5fe323c89e82300a38", size = 2288882, upload-time = "2026-08-28T10:00:07.931Z" }, + { url = "https://files.pythonhosted.org/packages/cd/94/5a81583660c175c59d49ffb09f4b3a44debeaf86a19fca664ae1cdd9ee32/pydantic_core-2.46.5-cp314-cp314t-musllinux_1_1_x86_64.whl", hash = "sha256:ff218293c9c806138dca139765e3b067621be52bcd93cdc14c7711be7ddc90a9", size = 2335210, upload-time = "2026-08-28T10:00:10.177Z" }, + { url = "https://files.pythonhosted.org/packages/5a/9f/5d685c2693b972d1a59c998586e8823712b66603aeff47ee60a4bdaafd37/pydantic_core-2.46.5-cp314-cp314t-win32.whl", hash = "sha256:97cf3eb53a8cccacf9d46686a0926186c9bfb5574f2ed66d3639d5fe117cd3a9", size = 1921180, upload-time = "2026-08-28T10:00:12.35Z" }, + { url = "https://files.pythonhosted.org/packages/70/12/5c94ee16d65a37a15f9e869f5e6256df111154491173801a4c5e800ab548/pydantic_core-2.46.5-cp314-cp314t-win_amd64.whl", hash = "sha256:d2f9fc07a8042a8f95925b35c4f04f469707c981fc33245b6ca187cf5d2dd290", size = 2020515, upload-time = "2026-08-28T10:00:14.774Z" }, + { url = "https://files.pythonhosted.org/packages/63/19/67830dda664e6bdf9285ee2e40f355d0d7d6b92aa0c42e8d217bb8d33d36/pydantic_core-2.46.5-cp314-cp314t-win_arm64.whl", hash = "sha256:acf8a67ba51f4ca9ddbd0e6b3000a65ac51ab734661778b3e7ba64d99a710f2f", size = 1989276, upload-time = "2026-08-28T10:00:16.984Z" }, + { url = "https://files.pythonhosted.org/packages/df/dd/053c2e4303f791f3b8f8a14ab0b22008e8eb21d868c0c90b4f9be705b76a/pydantic_core-2.46.5-graalpy312-graalpy250_312_native-macosx_10_12_x86_64.whl", hash = "sha256:013d6f3483d81e02e7c328831808f336c8596ee33b4bd4026b9ffb1e960b8942", size = 2062540, upload-time = "2026-08-28T10:01:00.318Z" }, + { url = "https://files.pythonhosted.org/packages/d7/dd/a18df751a5e37dd51bfad7f68e766999125bebe68c9e1d10a493ad01bd63/pydantic_core-2.46.5-graalpy312-graalpy250_312_native-macosx_11_0_arm64.whl", hash = "sha256:e9c134bb666dd54b778b9fc0d2b50cbb7f979b9e3716f26a88c9ab3b6fc1dd0f", size = 1902040, upload-time = "2026-08-28T10:01:02.529Z" }, + { url = "https://files.pythonhosted.org/packages/b7/13/01d40f9d07ce8a779fd6e0bd8ad4fba91309500dd67b869e2e219d261a6d/pydantic_core-2.46.5-graalpy312-graalpy250_312_native-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:347ec774390c87326a2e4929d58d3f7e8763a104d5d35f4cd595a4c952366433", size = 1967479, upload-time = "2026-08-28T10:01:05.004Z" }, + { url = "https://files.pythonhosted.org/packages/fa/04/c81d4841331c2178b6fb09ae225425e110ed72d990c9fe556c4ec03d1013/pydantic_core-2.46.5-graalpy312-graalpy250_312_native-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:8e24d8f05fa2d28513d94e877e9c75ad66175376209b3977f916e240e623193c", size = 2111034, upload-time = "2026-08-28T10:01:07.345Z" }, +] + +[[package]] +name = "pyfiglet" +version = "1.0.4" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/c8/e3/0a86276ad2c383ce08d76110a8eec2fe22e7051c4b8ba3fa163a0b08c428/pyfiglet-1.0.4.tar.gz", hash = "sha256:db9c9940ed1bf3048deff534ed52ff2dafbbc2cd7610b17bb5eca1df6d4278ef", size = 1560615, upload-time = "2025-08-15T18:32:47.302Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/9f/5c/fe9f95abd5eaedfa69f31e450f7e2768bef121dbdf25bcddee2cd3087a16/pyfiglet-1.0.4-py3-none-any.whl", hash = "sha256:65b57b7a8e1dff8a67dc8e940a117238661d5e14c3e49121032bd404d9b2b39f", size = 1806118, upload-time = "2025-08-15T18:32:45.556Z" }, +] + +[[package]] +name = "pygments" +version = "2.21.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/49/2e/ced460408999b33da6b31b0021b0f37d329e202d4169aeb164493778f25b/pygments-2.21.0.tar.gz", hash = "sha256:610ca751c9bc2492b38eb9a38a7fbc93edbbb2d7182edaf34e66ae493dee5c8c", size = 5005329, upload-time = "2026-08-17T08:02:48.824Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/71/46/17f022dd3e953bf20a04a028a21ec746d942f8d2af30fa0f124fa0e6a684/pygments-2.21.0-py3-none-any.whl", hash = "sha256:2363c69b61c4a97c838da3b130dcd6468f4848992b21a82f2a63ec34377137d9", size = 1250147, upload-time = "2026-08-17T08:02:44.912Z" }, +] + +[[package]] +name = "pytest" +version = "9.1.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "iniconfig" }, + { name = "packaging" }, + { name = "pluggy" }, + { name = "pygments" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/e4/47/b9efed96c114afcfa3c9d3fe98a76a1d14c74a9e266d397cf6eb64be5e01/pytest-9.1.1.tar.gz", hash = "sha256:1088fbde8f2b49d95a549a195707afa7a76a3ce9bcadc26b6d71f0ffda5fe313", size = 1636369, upload-time = "2026-06-19T10:58:32.857Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/24/25/1de2678b631f5a49215c6c96fff41ba892b0a34df68d6d80292b1b48aa7f/pytest-9.1.1-py3-none-any.whl", hash = "sha256:37a86b45efb9a47a61a36449063e8e18d0cab3161329fc099eb21783169c4f0c", size = 386536, upload-time = "2026-06-19T10:58:31.347Z" }, +] + +[[package]] +name = "python-dateutil" +version = "2.9.0.post0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "six" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/66/c0/0c8b6ad9f17a802ee498c46e004a0eb49bc148f2fd230864601a86dcf6db/python-dateutil-2.9.0.post0.tar.gz", hash = "sha256:37dd54208da7e1cd875388217d5e00ebd4179249f90fb72437e91a35459a0ad3", size = 342432, upload-time = "2024-03-01T18:36:20.211Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ec/57/56b9bcc3c9c6a792fcbaf139543cee77261f3651ca9da0c93f5c1221264b/python_dateutil-2.9.0.post0-py2.py3-none-any.whl", hash = "sha256:a8b2bc7bffae282281c8140a97d3aa9c14da0b136dfe83f850eea9a5f7470427", size = 229892, upload-time = "2024-03-01T18:36:18.57Z" }, +] + +[[package]] +name = "pytz" +version = "2026.4" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/b0/ed/fa23d28713004418bbf2407127f80f385bfb0bf4e677bef02c25c27b00ee/pytz-2026.4.tar.gz", hash = "sha256:464303645bafafd72418898368b2429458f709cf1eb6a15372fbcc396b64da63", size = 318444, upload-time = "2026-09-24T03:58:40.443Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d5/11/e3df59614bebe0ca50c905442603eda64ae7ac19811d4e4f74aa2d51ca85/pytz-2026.4-py2.py3-none-any.whl", hash = "sha256:9d514388fbc89ca0833203464272ac485b8828568ab73532f5020f17e892a0ff", size = 506747, upload-time = "2026-09-24T03:58:38.629Z" }, +] + +[[package]] +name = "pyyaml" +version = "6.0.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/05/8e/961c0007c59b8dd7729d542c61a4d537767a59645b82a0b521206e1e25c2/pyyaml-6.0.3.tar.gz", hash = "sha256:d76623373421df22fb4cf8817020cbb7ef15c725b9d5e45f17e189bfc384190f", size = 130960, upload-time = "2025-09-25T21:33:16.546Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d1/33/422b98d2195232ca1826284a76852ad5a86fe23e31b009c9886b2d0fb8b2/pyyaml-6.0.3-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:7f047e29dcae44602496db43be01ad42fc6f1cc0d8cd6c83d342306c32270196", size = 182063, upload-time = "2025-09-25T21:32:11.445Z" }, + { url = "https://files.pythonhosted.org/packages/89/a0/6cf41a19a1f2f3feab0e9c0b74134aa2ce6849093d5517a0c550fe37a648/pyyaml-6.0.3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:fc09d0aa354569bc501d4e787133afc08552722d3ab34836a80547331bb5d4a0", size = 173973, upload-time = "2025-09-25T21:32:12.492Z" }, + { url = "https://files.pythonhosted.org/packages/ed/23/7a778b6bd0b9a8039df8b1b1d80e2e2ad78aa04171592c8a5c43a56a6af4/pyyaml-6.0.3-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:9149cad251584d5fb4981be1ecde53a1ca46c891a79788c0df828d2f166bda28", size = 775116, upload-time = "2025-09-25T21:32:13.652Z" }, + { url = "https://files.pythonhosted.org/packages/65/30/d7353c338e12baef4ecc1b09e877c1970bd3382789c159b4f89d6a70dc09/pyyaml-6.0.3-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:5fdec68f91a0c6739b380c83b951e2c72ac0197ace422360e6d5a959d8d97b2c", size = 844011, upload-time = "2025-09-25T21:32:15.21Z" }, + { url = "https://files.pythonhosted.org/packages/8b/9d/b3589d3877982d4f2329302ef98a8026e7f4443c765c46cfecc8858c6b4b/pyyaml-6.0.3-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ba1cc08a7ccde2d2ec775841541641e4548226580ab850948cbfda66a1befcdc", size = 807870, upload-time = "2025-09-25T21:32:16.431Z" }, + { url = "https://files.pythonhosted.org/packages/05/c0/b3be26a015601b822b97d9149ff8cb5ead58c66f981e04fedf4e762f4bd4/pyyaml-6.0.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:8dc52c23056b9ddd46818a57b78404882310fb473d63f17b07d5c40421e47f8e", size = 761089, upload-time = "2025-09-25T21:32:17.56Z" }, + { url = "https://files.pythonhosted.org/packages/be/8e/98435a21d1d4b46590d5459a22d88128103f8da4c2d4cb8f14f2a96504e1/pyyaml-6.0.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:41715c910c881bc081f1e8872880d3c650acf13dfa8214bad49ed4cede7c34ea", size = 790181, upload-time = "2025-09-25T21:32:18.834Z" }, + { url = "https://files.pythonhosted.org/packages/74/93/7baea19427dcfbe1e5a372d81473250b379f04b1bd3c4c5ff825e2327202/pyyaml-6.0.3-cp312-cp312-win32.whl", hash = "sha256:96b533f0e99f6579b3d4d4995707cf36df9100d67e0c8303a0c55b27b5f99bc5", size = 137658, upload-time = "2025-09-25T21:32:20.209Z" }, + { url = "https://files.pythonhosted.org/packages/86/bf/899e81e4cce32febab4fb42bb97dcdf66bc135272882d1987881a4b519e9/pyyaml-6.0.3-cp312-cp312-win_amd64.whl", hash = "sha256:5fcd34e47f6e0b794d17de1b4ff496c00986e1c83f7ab2fb8fcfe9616ff7477b", size = 154003, upload-time = "2025-09-25T21:32:21.167Z" }, + { url = "https://files.pythonhosted.org/packages/1a/08/67bd04656199bbb51dbed1439b7f27601dfb576fb864099c7ef0c3e55531/pyyaml-6.0.3-cp312-cp312-win_arm64.whl", hash = "sha256:64386e5e707d03a7e172c0701abfb7e10f0fb753ee1d773128192742712a98fd", size = 140344, upload-time = "2025-09-25T21:32:22.617Z" }, + { url = "https://files.pythonhosted.org/packages/d1/11/0fd08f8192109f7169db964b5707a2f1e8b745d4e239b784a5a1dd80d1db/pyyaml-6.0.3-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:8da9669d359f02c0b91ccc01cac4a67f16afec0dac22c2ad09f46bee0697eba8", size = 181669, upload-time = "2025-09-25T21:32:23.673Z" }, + { url = "https://files.pythonhosted.org/packages/b1/16/95309993f1d3748cd644e02e38b75d50cbc0d9561d21f390a76242ce073f/pyyaml-6.0.3-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:2283a07e2c21a2aa78d9c4442724ec1eb15f5e42a723b99cb3d822d48f5f7ad1", size = 173252, upload-time = "2025-09-25T21:32:25.149Z" }, + { url = "https://files.pythonhosted.org/packages/50/31/b20f376d3f810b9b2371e72ef5adb33879b25edb7a6d072cb7ca0c486398/pyyaml-6.0.3-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ee2922902c45ae8ccada2c5b501ab86c36525b883eff4255313a253a3160861c", size = 767081, upload-time = "2025-09-25T21:32:26.575Z" }, + { url = "https://files.pythonhosted.org/packages/49/1e/a55ca81e949270d5d4432fbbd19dfea5321eda7c41a849d443dc92fd1ff7/pyyaml-6.0.3-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:a33284e20b78bd4a18c8c2282d549d10bc8408a2a7ff57653c0cf0b9be0afce5", size = 841159, upload-time = "2025-09-25T21:32:27.727Z" }, + { url = "https://files.pythonhosted.org/packages/74/27/e5b8f34d02d9995b80abcef563ea1f8b56d20134d8f4e5e81733b1feceb2/pyyaml-6.0.3-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0f29edc409a6392443abf94b9cf89ce99889a1dd5376d94316ae5145dfedd5d6", size = 801626, upload-time = "2025-09-25T21:32:28.878Z" }, + { url = "https://files.pythonhosted.org/packages/f9/11/ba845c23988798f40e52ba45f34849aa8a1f2d4af4b798588010792ebad6/pyyaml-6.0.3-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:f7057c9a337546edc7973c0d3ba84ddcdf0daa14533c2065749c9075001090e6", size = 753613, upload-time = "2025-09-25T21:32:30.178Z" }, + { url = "https://files.pythonhosted.org/packages/3d/e0/7966e1a7bfc0a45bf0a7fb6b98ea03fc9b8d84fa7f2229e9659680b69ee3/pyyaml-6.0.3-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:eda16858a3cab07b80edaf74336ece1f986ba330fdb8ee0d6c0d68fe82bc96be", size = 794115, upload-time = "2025-09-25T21:32:31.353Z" }, + { url = "https://files.pythonhosted.org/packages/de/94/980b50a6531b3019e45ddeada0626d45fa85cbe22300844a7983285bed3b/pyyaml-6.0.3-cp313-cp313-win32.whl", hash = "sha256:d0eae10f8159e8fdad514efdc92d74fd8d682c933a6dd088030f3834bc8e6b26", size = 137427, upload-time = "2025-09-25T21:32:32.58Z" }, + { url = "https://files.pythonhosted.org/packages/97/c9/39d5b874e8b28845e4ec2202b5da735d0199dbe5b8fb85f91398814a9a46/pyyaml-6.0.3-cp313-cp313-win_amd64.whl", hash = "sha256:79005a0d97d5ddabfeeea4cf676af11e647e41d81c9a7722a193022accdb6b7c", size = 154090, upload-time = "2025-09-25T21:32:33.659Z" }, + { url = "https://files.pythonhosted.org/packages/73/e8/2bdf3ca2090f68bb3d75b44da7bbc71843b19c9f2b9cb9b0f4ab7a5a4329/pyyaml-6.0.3-cp313-cp313-win_arm64.whl", hash = "sha256:5498cd1645aa724a7c71c8f378eb29ebe23da2fc0d7a08071d89469bf1d2defb", size = 140246, upload-time = "2025-09-25T21:32:34.663Z" }, + { url = "https://files.pythonhosted.org/packages/9d/8c/f4bd7f6465179953d3ac9bc44ac1a8a3e6122cf8ada906b4f96c60172d43/pyyaml-6.0.3-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:8d1fab6bb153a416f9aeb4b8763bc0f22a5586065f86f7664fc23339fc1c1fac", size = 181814, upload-time = "2025-09-25T21:32:35.712Z" }, + { url = "https://files.pythonhosted.org/packages/bd/9c/4d95bb87eb2063d20db7b60faa3840c1b18025517ae857371c4dd55a6b3a/pyyaml-6.0.3-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:34d5fcd24b8445fadc33f9cf348c1047101756fd760b4dacb5c3e99755703310", size = 173809, upload-time = "2025-09-25T21:32:36.789Z" }, + { url = "https://files.pythonhosted.org/packages/92/b5/47e807c2623074914e29dabd16cbbdd4bf5e9b2db9f8090fa64411fc5382/pyyaml-6.0.3-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:501a031947e3a9025ed4405a168e6ef5ae3126c59f90ce0cd6f2bfc477be31b7", size = 766454, upload-time = "2025-09-25T21:32:37.966Z" }, + { url = "https://files.pythonhosted.org/packages/02/9e/e5e9b168be58564121efb3de6859c452fccde0ab093d8438905899a3a483/pyyaml-6.0.3-cp314-cp314-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:b3bc83488de33889877a0f2543ade9f70c67d66d9ebb4ac959502e12de895788", size = 836355, upload-time = "2025-09-25T21:32:39.178Z" }, + { url = "https://files.pythonhosted.org/packages/88/f9/16491d7ed2a919954993e48aa941b200f38040928474c9e85ea9e64222c3/pyyaml-6.0.3-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c458b6d084f9b935061bc36216e8a69a7e293a2f1e68bf956dcd9e6cbcd143f5", size = 794175, upload-time = "2025-09-25T21:32:40.865Z" }, + { url = "https://files.pythonhosted.org/packages/dd/3f/5989debef34dc6397317802b527dbbafb2b4760878a53d4166579111411e/pyyaml-6.0.3-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:7c6610def4f163542a622a73fb39f534f8c101d690126992300bf3207eab9764", size = 755228, upload-time = "2025-09-25T21:32:42.084Z" }, + { url = "https://files.pythonhosted.org/packages/d7/ce/af88a49043cd2e265be63d083fc75b27b6ed062f5f9fd6cdc223ad62f03e/pyyaml-6.0.3-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:5190d403f121660ce8d1d2c1bb2ef1bd05b5f68533fc5c2ea899bd15f4399b35", size = 789194, upload-time = "2025-09-25T21:32:43.362Z" }, + { url = "https://files.pythonhosted.org/packages/23/20/bb6982b26a40bb43951265ba29d4c246ef0ff59c9fdcdf0ed04e0687de4d/pyyaml-6.0.3-cp314-cp314-win_amd64.whl", hash = "sha256:4a2e8cebe2ff6ab7d1050ecd59c25d4c8bd7e6f400f5f82b96557ac0abafd0ac", size = 156429, upload-time = "2025-09-25T21:32:57.844Z" }, + { url = "https://files.pythonhosted.org/packages/f4/f4/a4541072bb9422c8a883ab55255f918fa378ecf083f5b85e87fc2b4eda1b/pyyaml-6.0.3-cp314-cp314-win_arm64.whl", hash = "sha256:93dda82c9c22deb0a405ea4dc5f2d0cda384168e466364dec6255b293923b2f3", size = 143912, upload-time = "2025-09-25T21:32:59.247Z" }, + { url = "https://files.pythonhosted.org/packages/7c/f9/07dd09ae774e4616edf6cda684ee78f97777bdd15847253637a6f052a62f/pyyaml-6.0.3-cp314-cp314t-macosx_10_13_x86_64.whl", hash = "sha256:02893d100e99e03eda1c8fd5c441d8c60103fd175728e23e431db1b589cf5ab3", size = 189108, upload-time = "2025-09-25T21:32:44.377Z" }, + { url = "https://files.pythonhosted.org/packages/4e/78/8d08c9fb7ce09ad8c38ad533c1191cf27f7ae1effe5bb9400a46d9437fcf/pyyaml-6.0.3-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:c1ff362665ae507275af2853520967820d9124984e0f7466736aea23d8611fba", size = 183641, upload-time = "2025-09-25T21:32:45.407Z" }, + { url = "https://files.pythonhosted.org/packages/7b/5b/3babb19104a46945cf816d047db2788bcaf8c94527a805610b0289a01c6b/pyyaml-6.0.3-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6adc77889b628398debc7b65c073bcb99c4a0237b248cacaf3fe8a557563ef6c", size = 831901, upload-time = "2025-09-25T21:32:48.83Z" }, + { url = "https://files.pythonhosted.org/packages/8b/cc/dff0684d8dc44da4d22a13f35f073d558c268780ce3c6ba1b87055bb0b87/pyyaml-6.0.3-cp314-cp314t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:a80cb027f6b349846a3bf6d73b5e95e782175e52f22108cfa17876aaeff93702", size = 861132, upload-time = "2025-09-25T21:32:50.149Z" }, + { url = "https://files.pythonhosted.org/packages/b1/5e/f77dc6b9036943e285ba76b49e118d9ea929885becb0a29ba8a7c75e29fe/pyyaml-6.0.3-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:00c4bdeba853cc34e7dd471f16b4114f4162dc03e6b7afcc2128711f0eca823c", size = 839261, upload-time = "2025-09-25T21:32:51.808Z" }, + { url = "https://files.pythonhosted.org/packages/ce/88/a9db1376aa2a228197c58b37302f284b5617f56a5d959fd1763fb1675ce6/pyyaml-6.0.3-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:66e1674c3ef6f541c35191caae2d429b967b99e02040f5ba928632d9a7f0f065", size = 805272, upload-time = "2025-09-25T21:32:52.941Z" }, + { url = "https://files.pythonhosted.org/packages/da/92/1446574745d74df0c92e6aa4a7b0b3130706a4142b2d1a5869f2eaa423c6/pyyaml-6.0.3-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:16249ee61e95f858e83976573de0f5b2893b3677ba71c9dd36b9cf8be9ac6d65", size = 829923, upload-time = "2025-09-25T21:32:54.537Z" }, + { url = "https://files.pythonhosted.org/packages/f0/7a/1c7270340330e575b92f397352af856a8c06f230aa3e76f86b39d01b416a/pyyaml-6.0.3-cp314-cp314t-win_amd64.whl", hash = "sha256:4ad1906908f2f5ae4e5a8ddfce73c320c2a1429ec52eafd27138b7f1cbe341c9", size = 174062, upload-time = "2025-09-25T21:32:55.767Z" }, + { url = "https://files.pythonhosted.org/packages/f1/12/de94a39c2ef588c7e6455cfbe7343d3b2dc9d6b6b2f40c4c6565744c873d/pyyaml-6.0.3-cp314-cp314t-win_arm64.whl", hash = "sha256:ebc55a14a21cb14062aa4162f906cd962b28e2e9ea38f9b4391244cd8de4ae0b", size = 149341, upload-time = "2025-09-25T21:32:56.828Z" }, +] + +[[package]] +name = "regex" +version = "2026.9.29" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/fc/f2/af1da9d3ceed77bfcdce40427d49ba0be94e4fe84245e3bfef68c10e75b6/regex-2026.9.29.tar.gz", hash = "sha256:8b5fcc4771732191b2b7d1dd68d8f0353f47f8d90b6150f6dce58bf1112442cb", size = 419199, upload-time = "2026-09-29T00:49:58.298Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/84/48/3fdcde9a0baa84d7d25571223265d6e434e114763b438601d54a8028bf3e/regex-2026.9.29-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:dc79d36d0618752265f0d575915bdc5c5130ecb9c9f6b3bcefeae32e4bdfafcf", size = 497903, upload-time = "2026-09-29T00:46:38.938Z" }, + { url = "https://files.pythonhosted.org/packages/2e/1c/4ee3e97c76f53940488dfe7a7e18705e78daac8cd7fb161d246b9e328449/regex-2026.9.29-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:3a21a9509d0ee88e7a70e1ad228cd2f0e0fd1e187458db132e8a8d18c97daf9d", size = 296416, upload-time = "2026-09-29T00:46:40.406Z" }, + { url = "https://files.pythonhosted.org/packages/37/14/f3f0ba083d2094392d5eabf56db5ea6ba469fd6e927afd187042054ea68a/regex-2026.9.29-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:f57dc6b8fef170f105d2cf5cdce254f47b137d7755086cf7050f47e16582abba", size = 293633, upload-time = "2026-09-29T00:46:41.959Z" }, + { url = "https://files.pythonhosted.org/packages/c9/72/67e7a8ce17f1aea49df215564048efb49cc8c2b31a0e0fc30f36838f8516/regex-2026.9.29-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f93bc1c3486ef3747e07c9d7c1d0a147b8fbaab975f80e348aed6f71309dfaca", size = 805885, upload-time = "2026-09-29T00:46:43.373Z" }, + { url = "https://files.pythonhosted.org/packages/f6/78/25436bcfd4d2260b4b4090094d55d7ab53ec8a1ab4865a0b8bcb33c7d5c0/regex-2026.9.29-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:9e1d3a4cb7993b708f0ada8d0c84590efd853f169e7147d2202c9da503180242", size = 878344, upload-time = "2026-09-29T00:46:45.328Z" }, + { url = "https://files.pythonhosted.org/packages/97/e6/a09ec3a23ae41d6179880e67f0aace9284b2d95f2d7b326eff203f8eec5e/regex-2026.9.29-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:dabee8f4935e731fb46b2a3091bdda0d3d94b3bbfb907d2b4f12eefce4009619", size = 919181, upload-time = "2026-09-29T00:46:47.041Z" }, + { url = "https://files.pythonhosted.org/packages/26/83/d2fbd2e4e3afb1167daa825187d196f313cbaa1a4768f311fb041bb0e3d2/regex-2026.9.29-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:39ab5894d971f9ac68baa6eca5c50387db579cfcacf36ae8df3feceb1815e6d0", size = 807783, upload-time = "2026-09-29T00:46:48.894Z" }, + { url = "https://files.pythonhosted.org/packages/46/0b/eb429a7016610d44fc89a597163f8c9127505f0d7dc724dc9effbb6a3ac0/regex-2026.9.29-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:c1a9a6651197fbed6f0212591418b9def774fc3f8324f78d1bf0e6a63e5f8aa1", size = 783465, upload-time = "2026-09-29T00:46:50.64Z" }, + { url = "https://files.pythonhosted.org/packages/1b/07/58a3c0153c7476898430f6a7cf3d9062a1d17fbea4f43399ecaf411c7b4c/regex-2026.9.29-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:87fb80cbe3557e27e7b28b995c2b2eedf689b8886f941ab93e0e288f0976518a", size = 793519, upload-time = "2026-09-29T00:46:52.396Z" }, + { url = "https://files.pythonhosted.org/packages/2a/e8/161b94d39164520e21a7befe0245569bf7fda4c7cf1fc4e2df2b5def49da/regex-2026.9.29-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:3c5c2ef13797466aa64170cbb66ad98a32351dd4127694cea7199f80f213750d", size = 869293, upload-time = "2026-09-29T00:46:54.128Z" }, + { url = "https://files.pythonhosted.org/packages/8f/07/3b02ed829aa2decdc1955d222bd1e2f99d1c8bb4873bbb9a66b2f0a36bff/regex-2026.9.29-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:59b49507f47479e299a9e1bc41b5cb83a7afda0540625f1dbae886615978acbf", size = 770239, upload-time = "2026-09-29T00:46:56.106Z" }, + { url = "https://files.pythonhosted.org/packages/42/5b/ba61f6fe062eb8562e742367d177bb75370434138ef6c9d2a27114f8d613/regex-2026.9.29-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:0dd8af32e9f7b56b7f95cc1fd79b23054c3bdc172392ae560acc24d57b7ffe71", size = 861973, upload-time = "2026-09-29T00:46:57.665Z" }, + { url = "https://files.pythonhosted.org/packages/cc/27/767259b20e8a842948990f5e99138d6c077248fd42f8b5468b1d9ca4b814/regex-2026.9.29-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:db5e82ba15c142425b8406690032df89e39cca4a2e8afbbb9a3d84edc2373ac3", size = 796470, upload-time = "2026-09-29T00:46:59.236Z" }, + { url = "https://files.pythonhosted.org/packages/a0/05/2566c4ba849b68a8ab81a6bf428fa79d20aae7ddee83979103c0381df254/regex-2026.9.29-cp312-cp312-win32.whl", hash = "sha256:d0c3082bf79bcd6a614d55916590ad4b8f93200e10b97f463ea5d9d07c9b5f23", size = 269327, upload-time = "2026-09-29T00:47:01.135Z" }, + { url = "https://files.pythonhosted.org/packages/93/19/489bc8db91196381c935752df01ba3f607140daece33b78d88573f028e64/regex-2026.9.29-cp312-cp312-win_amd64.whl", hash = "sha256:fdd88ed5e20b1bcdd234421e454962c971aa44b653bdb7f1ea9ef683e90fb649", size = 280334, upload-time = "2026-09-29T00:47:04.436Z" }, + { url = "https://files.pythonhosted.org/packages/0b/47/fb88ba779d0e5e7d4b0ec1aceeb13845948a2cb876bd572a2d1dfdba090b/regex-2026.9.29-cp312-cp312-win_arm64.whl", hash = "sha256:4fe97894d1b306c919b4e50def1e6f6c522f4d03a7283811f4d108f1ce5d3ac2", size = 279614, upload-time = "2026-09-29T00:47:06.541Z" }, + { url = "https://files.pythonhosted.org/packages/79/d5/6080f7d1a6e7e36aa720f806ac93c035ba39c209ae6cc510e8ef4c0279c6/regex-2026.9.29-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:f1a0d5117230dd46b399a30a38afa44f79c99f3168988fdc4f425c3f928b39df", size = 497622, upload-time = "2026-09-29T00:47:08.251Z" }, + { url = "https://files.pythonhosted.org/packages/00/71/c87fc7a2e21a42f9d57489db32951c37eef56d153840459a80d464f0321d/regex-2026.9.29-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:f0fe9834e5aeccaf19a0d8feb296d66a24be1a7c9922002f842a682cd5abb787", size = 296281, upload-time = "2026-09-29T00:47:09.764Z" }, + { url = "https://files.pythonhosted.org/packages/11/9e/aa0f4cde3bc4688c1d58b0cd8415edd708339bc0bc401a195b0b1e8c8f0c/regex-2026.9.29-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:c90fcf7804ea0a54b896ce0f2b9565350220b8d4890fd0db461a476a4c687963", size = 293472, upload-time = "2026-09-29T00:47:11.723Z" }, + { url = "https://files.pythonhosted.org/packages/90/d4/e835c487850ed922a8d6074f953b888c8ea99775c76b9ed5f8a4d72eab92/regex-2026.9.29-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e11edba5bc344a32b029a7af9d4b3173982dd79eeafa0b9dbd787364414b0509", size = 805966, upload-time = "2026-09-29T00:47:13.235Z" }, + { url = "https://files.pythonhosted.org/packages/2c/57/ba8809847fbae8d2cbc71367c6ded510a7ec88bf52493c65efc1acf4effb/regex-2026.9.29-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:bb90e7177944b6684738c1fc36aabd2dd00d1de3be7dbe09f91e196f1bc0dc81", size = 878346, upload-time = "2026-09-29T00:47:14.877Z" }, + { url = "https://files.pythonhosted.org/packages/1a/52/e3da19fc3cc15ef67ab67e121e87887c3bccfdb683a7a9ec557c460ca5b7/regex-2026.9.29-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:d06fcdecc10fc7954d7c8f27a03c96055fe525274dc84a7b0dbdc3d6b9e03dab", size = 919250, upload-time = "2026-09-29T00:47:16.622Z" }, + { url = "https://files.pythonhosted.org/packages/9a/8e/c1ed81f55f992f6aa0b699a592a50c1ce9e6d44ff1aee2c14c0537dcef9c/regex-2026.9.29-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d49c18f1ea294cf4adde2e5ac256e98c82ea9d708462ce4bf799dffa7cfe8a2c", size = 807902, upload-time = "2026-09-29T00:47:18.268Z" }, + { url = "https://files.pythonhosted.org/packages/ad/bc/5a6886eb470e41040e21e05b75024a18b6ebfe7ea400b72094a60f949101/regex-2026.9.29-cp313-cp313-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:3e778bfccd63075167709136afbc251c1f683758d5bf49c803c60ac3f894ce6b", size = 783514, upload-time = "2026-09-29T00:47:19.921Z" }, + { url = "https://files.pythonhosted.org/packages/cb/52/6d951d453b023c6edb880f1ba474291b53b8ce1cc438b96a9db6d791d991/regex-2026.9.29-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:686ac5350fceae63830bb98805fcb8039325bf4c06d9f6f048ff65229d5bffa5", size = 793466, upload-time = "2026-09-29T00:47:21.552Z" }, + { url = "https://files.pythonhosted.org/packages/99/b9/d5a41adc08360f5eee0dc4846c578f002366947211fc8af5a69a64ee7b9f/regex-2026.9.29-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:26ec4ccce55aa533fbd603d08911b01101a8fcfec987845ac3ae2c7087b2bde3", size = 869464, upload-time = "2026-09-29T00:47:23.276Z" }, + { url = "https://files.pythonhosted.org/packages/4b/32/d76c9d91f5d798e2e9e67f6f85ec4ae35445ac425f7454797311cecb80ca/regex-2026.9.29-cp313-cp313-musllinux_1_2_riscv64.whl", hash = "sha256:a655d34b2a6943af32401f3d94f72e9d731f6ad16285815550bf2b4ee69d420a", size = 770278, upload-time = "2026-09-29T00:47:25.193Z" }, + { url = "https://files.pythonhosted.org/packages/24/00/aeebdb540c620a0f7317f6d6fad80a47729ecf0599a24b5c34ec155351f5/regex-2026.9.29-cp313-cp313-musllinux_1_2_s390x.whl", hash = "sha256:0c992c19cd45058a4b92f68f139c93db168b48fb1f322c9a7cd620806afb6b51", size = 861949, upload-time = "2026-09-29T00:47:27.005Z" }, + { url = "https://files.pythonhosted.org/packages/12/62/d0314bcedfd3586197e4596931fa220260eb2385bf53184e5b9ae67db24b/regex-2026.9.29-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:ebb8912f565b8cdbbf27debfe00df04202c20e2f651b9e32767930c5eace3621", size = 796602, upload-time = "2026-09-29T00:47:29.233Z" }, + { url = "https://files.pythonhosted.org/packages/ae/c7/d5a8c13a613facb03e0fb55c1ebaaf7bb35d8e2c1abe8bef8dca809fc1d9/regex-2026.9.29-cp313-cp313-win32.whl", hash = "sha256:4d7d93613b01b0199961330e49cfc52d479b3d5776c56c691db31130c0a07d91", size = 269306, upload-time = "2026-09-29T00:47:31.14Z" }, + { url = "https://files.pythonhosted.org/packages/80/a7/bf93a3a6afa5f7bc16b7afb94ae581b01cae620b8ad56bd8f9572a985959/regex-2026.9.29-cp313-cp313-win_amd64.whl", hash = "sha256:61956f074ecd123f55adca68ee3eab46e6a07ad3f8e64e6db95dfacb444f55c4", size = 280307, upload-time = "2026-09-29T00:47:32.709Z" }, + { url = "https://files.pythonhosted.org/packages/b2/7d/388274e53605a86297f433a08102a7bbdcf9379d47683d307ccaefd88e2c/regex-2026.9.29-cp313-cp313-win_arm64.whl", hash = "sha256:bfc71e6d970419c1309b3640305298643e2a734cad3f7cfb6d2ddee4175ab53d", size = 279599, upload-time = "2026-09-29T00:47:34.674Z" }, + { url = "https://files.pythonhosted.org/packages/93/1f/d9dc6f02f569625faf67a4daec926cd5023472dcd69bb44286dccd5a5ab3/regex-2026.9.29-cp314-cp314-macosx_10_15_universal2.whl", hash = "sha256:957bb708e8057ab1649ba566456429d691ec9b90d1c9ad1af1ba7ffbbeaf05f2", size = 497598, upload-time = "2026-09-29T00:47:36.541Z" }, + { url = "https://files.pythonhosted.org/packages/9c/83/9b693a3fd1451381e812031a8961ec5b3b8f0c8cc6871f14c5223642804d/regex-2026.9.29-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:c9b602fae1e00b7c035d661ce85575365719192a7b46784bd71cf64c68053aa0", size = 296250, upload-time = "2026-09-29T00:47:38.233Z" }, + { url = "https://files.pythonhosted.org/packages/dd/5f/52bc2abc3fef040cd9de76ab29c918d6a717a454ae2b9dd7938b0c95656d/regex-2026.9.29-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:0166844493626c5015c6088ee15c9ca2fd060ca15b7641d1657da6a58432ae33", size = 293518, upload-time = "2026-09-29T00:47:39.957Z" }, + { url = "https://files.pythonhosted.org/packages/dc/fc/cf50671215ee0057046980b4571ef8646a005819bb67f0957e779ed107a5/regex-2026.9.29-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:b97a38fb4c732b6832db6bf108963adbcd82ef1268ba2025dce390f45af75efa", size = 806037, upload-time = "2026-09-29T00:47:41.676Z" }, + { url = "https://files.pythonhosted.org/packages/14/4b/dddef8fc15c63e4347cc9efb138d0cd306f30e6c98acbcc81a8f780083b9/regex-2026.9.29-cp314-cp314-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:a540abfab208e1b7ef2df231c40ef3b6cbb30a0aad6204e9b6a81c10a6794628", size = 878881, upload-time = "2026-09-29T00:47:43.755Z" }, + { url = "https://files.pythonhosted.org/packages/9f/cb/38daabed32d28f7e58a06e9344ce00dc67952e9996bc578ed6a29fe1240e/regex-2026.9.29-cp314-cp314-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:ddfa987262763c3c22a8367d2a49c244b018a74c3a8e3ab1a864119ad45c5633", size = 918684, upload-time = "2026-09-29T00:47:45.594Z" }, + { url = "https://files.pythonhosted.org/packages/a9/4d/041d9458a645fee4fce4d642a89d27271a3cfcd91095104f6dde44da70bf/regex-2026.9.29-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:2f7f7aa47b229f2b39a2ae2596d2ad5625d77b5eb9856fac2dab3eb506cdd0a0", size = 807176, upload-time = "2026-09-29T00:47:47.372Z" }, + { url = "https://files.pythonhosted.org/packages/bf/c4/4383eed7aa5aef67616cb1b3f3ad06b7c624c4e6cced48630cd5ce133d85/regex-2026.9.29-cp314-cp314-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:d9b77b25b4f395f92de6099ab08e8ae2bc7e51dfe157f22900902243a5cc90c7", size = 784315, upload-time = "2026-09-29T00:47:49.518Z" }, + { url = "https://files.pythonhosted.org/packages/5c/a6/0086ad31cebb183c637d3198547075aa493afde308e1ff61fccccb29ba6e/regex-2026.9.29-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:34b6925af9853bf461950e6508910f179fd6e9b1a7ec8548e069606b7e51a26b", size = 793748, upload-time = "2026-09-29T00:47:51.279Z" }, + { url = "https://files.pythonhosted.org/packages/d5/a0/f9005cba3f629a859573fc5d1224ea4e1f97919ec8581d018e03a351a604/regex-2026.9.29-cp314-cp314-musllinux_1_2_ppc64le.whl", hash = "sha256:addd736a0547d553283adaf4e05d7104e7f2c7b0b092e9b4d28756825f14531f", size = 870302, upload-time = "2026-09-29T00:47:53.368Z" }, + { url = "https://files.pythonhosted.org/packages/01/4f/e1a3e46bb5315a4e18b01a990e7a28e2a16595609d50c442baf2815a3c65/regex-2026.9.29-cp314-cp314-musllinux_1_2_riscv64.whl", hash = "sha256:fe3fa1dd453ed5c7f5ea23a26218329790ed7197a99b90e94330e313959a7f52", size = 770299, upload-time = "2026-09-29T00:47:55.606Z" }, + { url = "https://files.pythonhosted.org/packages/2c/fe/f303b4acfda44e1ff1379368748c1ef2dad04a6a8e9c0ecbc970b19d97ca/regex-2026.9.29-cp314-cp314-musllinux_1_2_s390x.whl", hash = "sha256:0cc63b5e47c12a48d90c7e9d7de6a035dd14f62868aaedbb4e0ff8ba2b8bfe7b", size = 861570, upload-time = "2026-09-29T00:47:57.617Z" }, + { url = "https://files.pythonhosted.org/packages/60/b6/b4f7e99249f596017c60ccad5faf9310fc8e3e59bb2244940a90a1b0bdff/regex-2026.9.29-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:724184b4aafed865e4f13ca313fdcb43024300c028ec67319cfa16847d84685e", size = 795967, upload-time = "2026-09-29T00:47:59.922Z" }, + { url = "https://files.pythonhosted.org/packages/fb/d3/fc865a4638d9f6762192b6bab5b7aa1f33a90e9e99578c2e111e2a63c8c3/regex-2026.9.29-cp314-cp314-win32.whl", hash = "sha256:c6c8fabf1dafc1f1ddcbb67896d3f93efb092e8c4b6322d7389b944e76a484e5", size = 274758, upload-time = "2026-09-29T00:48:01.8Z" }, + { url = "https://files.pythonhosted.org/packages/31/e2/c2b466924ccbeb874862968ca638051b15a8fd29d994a0e99004a5cbf78e/regex-2026.9.29-cp314-cp314-win_amd64.whl", hash = "sha256:1c2a0026062abcc321a53db4a185ceba0b59a66b5d37b0808917a88b55a5257f", size = 283817, upload-time = "2026-09-29T00:48:03.614Z" }, + { url = "https://files.pythonhosted.org/packages/c6/42/ea0f8dbaa924fa75c6338935eaee2f44dab369b27f02db1e03d74344b049/regex-2026.9.29-cp314-cp314-win_arm64.whl", hash = "sha256:121a76a0985db80ceae9e171c337f8c927868e37d01b54e3ce87bc87f9c6a208", size = 283663, upload-time = "2026-09-29T00:48:05.624Z" }, + { url = "https://files.pythonhosted.org/packages/44/48/d58e5081119f5c223bbb37d2340acde3d069e1df8e8cd166c37502eee4da/regex-2026.9.29-cp314-cp314t-macosx_10_15_universal2.whl", hash = "sha256:e31f72490b7c12f7790e1e25c3afffd20503ee1bfb43461d7838b871ff244b19", size = 501393, upload-time = "2026-09-29T00:48:07.833Z" }, + { url = "https://files.pythonhosted.org/packages/72/3c/c49945287d4f9efee7d41f98072f8ad880efb8f430595a612fbdea996a4e/regex-2026.9.29-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:80ea96f5c1a30bf09007d48466521d9c294bebe197c708c3359096e3e3691632", size = 298237, upload-time = "2026-09-29T00:48:09.684Z" }, + { url = "https://files.pythonhosted.org/packages/f9/1f/688cb61c3d4cf7bcc1ed444b5cc49399eba3e51c469ae285cf87fea3022e/regex-2026.9.29-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:554bffadcbcb6d5f4e5fb10a61cc52084b9a63d1dab5f10bcd2c4343972e8e2c", size = 295936, upload-time = "2026-09-29T00:48:11.454Z" }, + { url = "https://files.pythonhosted.org/packages/26/a3/de43ac6b877b7d09c19a3a426b1bd5acdd209eaaf68f406466f80439ccf6/regex-2026.9.29-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:864e9b87ac33c3fb9fb4ad48166d4fdb579c351d5c77deb0d34bccb36a775cd9", size = 816905, upload-time = "2026-09-29T00:48:13.321Z" }, + { url = "https://files.pythonhosted.org/packages/62/14/9940763201c51d537786304984c67d0fc3d2ed18837ffb6f09a869f6b6c9/regex-2026.9.29-cp314-cp314t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:044265d77d94f5e3cb2fd72c76723807c429cb8c533e9d4672d0334a6f14f588", size = 881527, upload-time = "2026-09-29T00:48:15.313Z" }, + { url = "https://files.pythonhosted.org/packages/d3/e1/c842d8df0b23245ebf202f8ab9c39fd48e2db39959454ec39a41c8c72082/regex-2026.9.29-cp314-cp314t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:2089fe39c406784d90101c726755ffa1497bb74638fd434300d2b88006186de8", size = 923115, upload-time = "2026-09-29T00:48:17.328Z" }, + { url = "https://files.pythonhosted.org/packages/d8/c1/98622479e3c354a446a75232e522d747d2b3df23092dcd8a5309380a2020/regex-2026.9.29-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0def9fb6abac55492d6d51cddb7225d07d6f279e774e0adc08569a54a5fc8d46", size = 820674, upload-time = "2026-09-29T00:48:19.32Z" }, + { url = "https://files.pythonhosted.org/packages/6c/d0/5808c95f9c79ed27b5eedaafc3df6239ec56a49f2e23ea8f831b18427c82/regex-2026.9.29-cp314-cp314t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:888d60953908dcf761aa320c3e390ab8556efbdb551ace63921de90f6ae0848d", size = 793306, upload-time = "2026-09-29T00:48:21.615Z" }, + { url = "https://files.pythonhosted.org/packages/bf/d3/021ca2638671ad20603bcd9b4d5bfa35d2610cd216a043ea7f0b44ea39f6/regex-2026.9.29-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:ed511a0708e2297e1d6431e7fb217e3402791e491e02da800658ace4973df1bb", size = 803103, upload-time = "2026-09-29T00:48:23.871Z" }, + { url = "https://files.pythonhosted.org/packages/6b/2d/755c6d13ef9c657378013676c391c7a402166b3f419a464a3e058dcbe533/regex-2026.9.29-cp314-cp314t-musllinux_1_2_ppc64le.whl", hash = "sha256:e1172147d28d8fbcf8cb8d26c41506169f5ad8fe9ec969cb116835a19d4d8eca", size = 872175, upload-time = "2026-09-29T00:48:26.255Z" }, + { url = "https://files.pythonhosted.org/packages/6c/fc/e1cab183b9dafe8597f58c1c766da9bf96204d3b2f232bcf3eeb75ff7b6c/regex-2026.9.29-cp314-cp314t-musllinux_1_2_riscv64.whl", hash = "sha256:92f05c9c42bde5785dc48770bc2194d9f7442544156f951e19cd31b096cec562", size = 777362, upload-time = "2026-09-29T00:48:28.389Z" }, + { url = "https://files.pythonhosted.org/packages/06/7c/e10ea17fba31fb4a1f9d13ed53a2d2a9066a2aea58d7557e263f6d99e7b0/regex-2026.9.29-cp314-cp314t-musllinux_1_2_s390x.whl", hash = "sha256:f37964e4a5e993d2fd45147741e9dff7f34a2d8c00ab94c4ea0514a4677f959e", size = 865606, upload-time = "2026-09-29T00:48:30.4Z" }, + { url = "https://files.pythonhosted.org/packages/8e/6e/69824d9aee1fd41c54ea7264654a47c8d9d84d8a228e11c2bcf4c201ed81/regex-2026.9.29-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:951733b1bbdb71e377cec567b409f1a7881b47cfcad84121aa74cb575fa425ea", size = 807945, upload-time = "2026-09-29T00:48:32.375Z" }, + { url = "https://files.pythonhosted.org/packages/89/22/857050a86e21ce60193e02a8ef662521f2e263a645c8b1b905fc136b61a7/regex-2026.9.29-cp314-cp314t-win32.whl", hash = "sha256:65b408d8fcb273e3499e7ef2ce796810da1becd208c7fb4373692a242d79d461", size = 276758, upload-time = "2026-09-29T00:48:34.72Z" }, + { url = "https://files.pythonhosted.org/packages/4d/96/56808fe029553d7d4c703414f2a527faad2ea2bfa9ca094a2e7f8762b530/regex-2026.9.29-cp314-cp314t-win_amd64.whl", hash = "sha256:bf48516e35cf848390ea68850aba53e7c333720d2945b4d2c25b69fc5171723f", size = 286527, upload-time = "2026-09-29T00:48:36.864Z" }, + { url = "https://files.pythonhosted.org/packages/01/aa/074e2cfb3d8101a6a764aba5f7c5d1e21de087483e35bdc0c4ce2eb60364/regex-2026.9.29-cp314-cp314t-win_arm64.whl", hash = "sha256:9173db3be74a35cb6731701094b98120f7ee4876a287882a59cdea1fa7da342f", size = 285954, upload-time = "2026-09-29T00:48:38.901Z" }, + { url = "https://files.pythonhosted.org/packages/a7/dc/d84990386c9dfdf8c377f00f371b241fdc9a2c8aea0e3d66941b2e51be0b/regex-2026.9.29-cp315-cp315-macosx_10_15_universal2.whl", hash = "sha256:c3589f40749acce747510bf5d589d54e376cb0930ea58b35effac97e5312b0c1", size = 497836, upload-time = "2026-09-29T00:48:40.858Z" }, + { url = "https://files.pythonhosted.org/packages/c2/ab/a569ebde875fa12ff8c6c9a30e07503620f195e4be4d54c3d3ee8eecc283/regex-2026.9.29-cp315-cp315-macosx_10_15_x86_64.whl", hash = "sha256:32ab11df9677ca80bcbb5fe4eb1da9109a5019239a054836efc6fa1c64e683cf", size = 296244, upload-time = "2026-09-29T00:48:42.952Z" }, + { url = "https://files.pythonhosted.org/packages/f3/3e/7d548e82a108e7c8b2d5246650e397a2f8db599f9b2e975466939c5b4e70/regex-2026.9.29-cp315-cp315-macosx_11_0_arm64.whl", hash = "sha256:7c03031610e3e6ed1768a2b7a8fc84637c1257b50c5eacaf094c6e17a84fc563", size = 293748, upload-time = "2026-09-29T00:48:44.985Z" }, + { url = "https://files.pythonhosted.org/packages/40/34/a8e19a52f452bbb07b32a2bef70dcdf90c2737049749f74cc12d7486fb4f/regex-2026.9.29-cp315-cp315-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:42e82e578c904445d4c8a35b8f28052cf567593215fa5db06266fbc6f77aaa2e", size = 807840, upload-time = "2026-09-29T00:48:46.948Z" }, + { url = "https://files.pythonhosted.org/packages/88/7b/11fbd4640b3bb82b72822a63c20ade4013d562d291703a9debeedc24e682/regex-2026.9.29-cp315-cp315-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:0b65c72739f981377c9c22e0c5c3cd7f42da7bd8a3c9209330fac772c7d893ed", size = 879330, upload-time = "2026-09-29T00:48:49.168Z" }, + { url = "https://files.pythonhosted.org/packages/f3/55/de58c74f1f4e31586d83eb39c56872d686c4e0d0966d151884c833b94ced/regex-2026.9.29-cp315-cp315-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:4408b2b27a95ca8cc48b7411945753773353b5c93b307754781086c99d3a576f", size = 919251, upload-time = "2026-09-29T00:48:51.322Z" }, + { url = "https://files.pythonhosted.org/packages/81/42/a8c480f6dd5ac59fa28ddae79afd9d7ac7e596fdb61813adc65bb6e674b8/regex-2026.9.29-cp315-cp315-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a714befaacbd10092ffe4cea0d3c5f008fb9efe9bc322c715bcdfdee414b9a3d", size = 808808, upload-time = "2026-09-29T00:48:53.529Z" }, + { url = "https://files.pythonhosted.org/packages/68/60/0bc0d1ec8b37ad64be6fa30e035251f11de9667a0fac9e82ee74517d81be/regex-2026.9.29-cp315-cp315-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:33026515aebc0e70d1c89978e53e8d695d35d9e472f8d5b34465ba3c74028650", size = 789907, upload-time = "2026-09-29T00:48:56.036Z" }, + { url = "https://files.pythonhosted.org/packages/da/84/116a3ef19b3acfe81077f0bf2cbc7714a5e94bc8935b7243ab61cb0f1c3c/regex-2026.9.29-cp315-cp315-musllinux_1_2_aarch64.whl", hash = "sha256:31b003f9a070335e2a8233ee9b14a3ca8e6d792012ae011f741bf0aaf11744c5", size = 795770, upload-time = "2026-09-29T00:48:58.284Z" }, + { url = "https://files.pythonhosted.org/packages/96/ba/e38c3f203e7e7e18c957d48e6cb6dbf96c11e95a44efa4a480522afc5d6d/regex-2026.9.29-cp315-cp315-musllinux_1_2_ppc64le.whl", hash = "sha256:c03c6eb6ece86dfdcbb34799efaa339b093132e1aceed491ba5e08fe06cdf699", size = 870671, upload-time = "2026-09-29T00:49:00.506Z" }, + { url = "https://files.pythonhosted.org/packages/2f/0f/9ee0b0cb76c55f63684bd7fff554978e8773b4fc86e2bcb2d50772dc1086/regex-2026.9.29-cp315-cp315-musllinux_1_2_riscv64.whl", hash = "sha256:a5300757f8a68f5b6cc33f57338d72a0e3589c5cc9ad5f8504ea06f028be582a", size = 778631, upload-time = "2026-09-29T00:49:02.984Z" }, + { url = "https://files.pythonhosted.org/packages/b6/19/e6e3eeb226af5872c4958002f6edef4e4f40ea4cc5f5665023f2019eb045/regex-2026.9.29-cp315-cp315-musllinux_1_2_s390x.whl", hash = "sha256:80c7cadd3fd2bfde5df8aa0787e315812cad0c313a753095d02f4c2b6c01677b", size = 862187, upload-time = "2026-09-29T00:49:05.264Z" }, + { url = "https://files.pythonhosted.org/packages/5b/62/823c102e106bb2711d6b7dfe5981552fe4467b2969c46a20c5c383cf498c/regex-2026.9.29-cp315-cp315-musllinux_1_2_x86_64.whl", hash = "sha256:3f1e6cb402a89457582cd696f982559217d13484a193202c394015297968c86d", size = 798423, upload-time = "2026-09-29T00:49:07.644Z" }, + { url = "https://files.pythonhosted.org/packages/37/e0/e927776258fa70b2f6feffc3be584ffc85ba4c1e20a320f0aee9a632fc7d/regex-2026.9.29-cp315-cp315-win32.whl", hash = "sha256:a64b85a4760337cfefdb27d42da6ed8b58e8cde3f2d57b6ef43e76ef6ea9ef47", size = 274757, upload-time = "2026-09-29T00:49:10.513Z" }, + { url = "https://files.pythonhosted.org/packages/77/04/358de85d1860238e1b4fa98fc2c80c990124a25d2e14739e28cc02c25562/regex-2026.9.29-cp315-cp315-win_amd64.whl", hash = "sha256:b3e445b66c80b4eb4234e855ce94d9adc183eedbd632816228d89930b91b2c5b", size = 283824, upload-time = "2026-09-29T00:49:12.849Z" }, + { url = "https://files.pythonhosted.org/packages/92/d3/d5c5b264784a5ab2b0f8cf620c1eeb4dbf3440d306761905e7d99345bef5/regex-2026.9.29-cp315-cp315-win_arm64.whl", hash = "sha256:8f39588af4731c8923c26810eb3b33f76f17633985e40f59c3cd45a33805a895", size = 283664, upload-time = "2026-09-29T00:49:15.331Z" }, + { url = "https://files.pythonhosted.org/packages/02/dc/f63ec2c201445ce1150fe780f5c56f16a10124d9a9da3a93161dbb0d8892/regex-2026.9.29-cp315-cp315t-macosx_10_15_universal2.whl", hash = "sha256:fb99cc9d45f48895d9d67f6a0b8a57f08d39c174d9f25ad97a313e0470267b1c", size = 501549, upload-time = "2026-09-29T00:49:17.705Z" }, + { url = "https://files.pythonhosted.org/packages/9a/5c/d2a698dc6bfc11fbce03f1cb0249c13284e93b79ed11f893edf6fac431c9/regex-2026.9.29-cp315-cp315t-macosx_10_15_x86_64.whl", hash = "sha256:720537c7ea6f80dc61913184edb0ce2497a306b39ef19f28505b322553d52bdb", size = 298187, upload-time = "2026-09-29T00:49:20.171Z" }, + { url = "https://files.pythonhosted.org/packages/85/b7/88dcdb38cd3935d4ee9e9ce9b8e56cb3b3518d1f020acfa7dd62ad289bf8/regex-2026.9.29-cp315-cp315t-macosx_11_0_arm64.whl", hash = "sha256:0fd2c901cc307a745ad4bc87f20060d7a0825a3371d1e93488af22e7a387f78f", size = 296157, upload-time = "2026-09-29T00:49:22.342Z" }, + { url = "https://files.pythonhosted.org/packages/d3/8e/ba6c01dde33a69fc294b38b43f6677baaa5735a6248f39708031a738158a/regex-2026.9.29-cp315-cp315t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:b11b589e00095ec69cf79841a76360f9b079e95b0368a25b5ebb951ab0c157ff", size = 818468, upload-time = "2026-09-29T00:49:24.612Z" }, + { url = "https://files.pythonhosted.org/packages/2a/f1/2586693e3a2d6b1247852593d37a6c17b42a92ee44f7cdcb9a0c1494e64a/regex-2026.9.29-cp315-cp315t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:d7cab119d0df0b9413f106b4d7fc34f2872d3574ed3806fb48959c830b1537da", size = 882825, upload-time = "2026-09-29T00:49:26.996Z" }, + { url = "https://files.pythonhosted.org/packages/30/51/084f3e7bdcd0e9c33665c938cf5d134dc3548cbb4a75f0197ec7bfd754b1/regex-2026.9.29-cp315-cp315t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:b89efc38431793d28b7cd91227e2f952ad7c48df19132b17f43a5fec3c14143b", size = 923314, upload-time = "2026-09-29T00:49:29.822Z" }, + { url = "https://files.pythonhosted.org/packages/5a/f1/066c6fc23b7dc229789c21c880b5ba5ad689fb95fed12e078266f55a1f9b/regex-2026.9.29-cp315-cp315t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:80a5ea3b4fd9d6a5b9a44f7976a9acaaab35aa3c1f6b29e5bd857dfabaded223", size = 822222, upload-time = "2026-09-29T00:49:32.404Z" }, + { url = "https://files.pythonhosted.org/packages/0a/56/592cd46fdb8f2f8682a1d7fd1310e4d0bcb93fbd0e6bbe4141ac28240227/regex-2026.9.29-cp315-cp315t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:19959129885356df0e97556856f77eb2888380dac18bed075a7c05c5128c618d", size = 794882, upload-time = "2026-09-29T00:49:35.076Z" }, + { url = "https://files.pythonhosted.org/packages/ee/4d/d65384bb071c864b01aa8314e3a6a687845ebd57588390976edc960c218b/regex-2026.9.29-cp315-cp315t-musllinux_1_2_aarch64.whl", hash = "sha256:6a1a824fbed817e0a891103886b68f063b1e83cc51bc97192a90a60195a9291f", size = 805887, upload-time = "2026-09-29T00:49:37.395Z" }, + { url = "https://files.pythonhosted.org/packages/65/b6/358de0d8f40d5178e4f7e7e121cfd5b961c812b77a055d11f5079e3f8fd7/regex-2026.9.29-cp315-cp315t-musllinux_1_2_ppc64le.whl", hash = "sha256:1ba8c6a416569ce0d37e83e28a254a61dc99a419084dfb6476cea02d997f74fa", size = 872901, upload-time = "2026-09-29T00:49:39.927Z" }, + { url = "https://files.pythonhosted.org/packages/00/06/6bfded72d043240c6b52bbb5e16f639d81affbf7484b4fe2ec45f3d4afc9/regex-2026.9.29-cp315-cp315t-musllinux_1_2_riscv64.whl", hash = "sha256:446654b29bfaa30500d80947eda42cef1449dc8a87f4e3cf061cc8485d3a1f0b", size = 782970, upload-time = "2026-09-29T00:49:42.581Z" }, + { url = "https://files.pythonhosted.org/packages/5a/20/9f418a50baa78b3ed8308fcb0cc49e472dd000b7ef935a7295af202ea744/regex-2026.9.29-cp315-cp315t-musllinux_1_2_s390x.whl", hash = "sha256:bf3c49863c23a1ad6da9c30351aed6cff8d5ddbeb63c5c8420ae54e98c7d0138", size = 865441, upload-time = "2026-09-29T00:49:45.238Z" }, + { url = "https://files.pythonhosted.org/packages/2c/29/817c7eacdeaf8463123e949bd394c39ad024eea1ec38ddf5ad141da2f3bd/regex-2026.9.29-cp315-cp315t-musllinux_1_2_x86_64.whl", hash = "sha256:01000ddf0e3ffef97f2413ceb514f6313040106b6d18a03ee00a4fe35c1eb1db", size = 809431, upload-time = "2026-09-29T00:49:47.878Z" }, + { url = "https://files.pythonhosted.org/packages/63/0b/83aab3b5b739947f744135a7a3a446e25433ebc92b05e01aae197ccbfdda/regex-2026.9.29-cp315-cp315t-win32.whl", hash = "sha256:c4e38dd8f39c43a91d2410ad2b85610701b0979342c3df1d69eaf8e838c757d8", size = 276964, upload-time = "2026-09-29T00:49:50.524Z" }, + { url = "https://files.pythonhosted.org/packages/72/f2/6314b5fc68789b5dcc38885bc6e3d6986b34fb3372b7231088ee5cecaa05/regex-2026.9.29-cp315-cp315t-win_amd64.whl", hash = "sha256:e2c89e9b762c57f59d5e99ee8b20202adb892e35f8d3485741340999ca55058e", size = 286487, upload-time = "2026-09-29T00:49:53.224Z" }, + { url = "https://files.pythonhosted.org/packages/56/bc/97b2245c8c7b2dd01f2db74f2bea003cd33c15009b4996a2447f46b5325c/regex-2026.9.29-cp315-cp315t-win_arm64.whl", hash = "sha256:e8c65ef3862a8ad6e86492b6ed9327805dd66904c012bd3649dc67d822ed6c34", size = 285955, upload-time = "2026-09-29T00:49:55.655Z" }, +] + +[[package]] +name = "rich" +version = "15.0.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "markdown-it-py" }, + { name = "pygments" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/c0/8f/0722ca900cc807c13a6a0c696dacf35430f72e0ec571c4275d2371fca3e9/rich-15.0.0.tar.gz", hash = "sha256:edd07a4824c6b40189fb7ac9bc4c52536e9780fbbfbddf6f1e2502c31b068c36", size = 230680, upload-time = "2026-04-12T08:24:00.75Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/82/3b/64d4899d73f91ba49a8c18a8ff3f0ea8f1c1d75481760df8c68ef5235bf5/rich-15.0.0-py3-none-any.whl", hash = "sha256:33bd4ef74232fb73fe9279a257718407f169c09b78a87ad3d296f548e27de0bb", size = 310654, upload-time = "2026-04-12T08:24:02.83Z" }, +] + +[[package]] +name = "ruff" +version = "0.16.10" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/c4/49/23802c45f093eb14bde54b141d2b2f058edfa63a06db7beed047308cc08f/ruff-0.16.10.tar.gz", hash = "sha256:eff4728c4eaae93f0955cd264d24b2ab348e74bf59986ccf282ba6dc16b3b017", size = 4958724, upload-time = "2026-10-01T18:03:21.697Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2f/21/ebce22e1d90cdb2cd691026b9c6e9bec6499f0b396089481755b5d49efff/ruff-0.16.10-py3-none-linux_armv6l.whl", hash = "sha256:488b0fe3f3574210e5cf80d9f59b9e3ab17a127a8155de3f307b392589cfb511", size = 10095557, upload-time = "2026-10-01T18:02:36.072Z" }, + { url = "https://files.pythonhosted.org/packages/cb/98/a54de85876a8b2612bfa0d84c7b9abfb39c6a3354aee7800b09c1649b9e3/ruff-0.16.10-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:e748ff95c934c4e978783b8e687bc174e7bd84e8ad24e3243e1ecfcda5e0282d", size = 10340577, upload-time = "2026-10-01T18:02:39.258Z" }, + { url = "https://files.pythonhosted.org/packages/9f/16/1a5a4a2657effe4806110f2b907313802f1367fcdbb29e8122766e407fab/ruff-0.16.10-py3-none-macosx_11_0_arm64.whl", hash = "sha256:3031a4a2e8e7b8a46f70be45f198c35a11ece509a94b80334d8d397a33c67550", size = 9774282, upload-time = "2026-10-01T18:02:41.79Z" }, + { url = "https://files.pythonhosted.org/packages/6e/fb/470085af734da396e68cd80fb0a3e7c109a459716ae59588c0f6fab8a17d/ruff-0.16.10-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:494401c86df4c4c25f69b9419605d944467ee98c42fb6ad405ef4fa40b8fb67d", size = 9920895, upload-time = "2026-10-01T18:02:44.458Z" }, + { url = "https://files.pythonhosted.org/packages/57/de/f10cffe4f88a37ea6615ec460f01bf76f0bc1477737a35d9ee61a630a78b/ruff-0.16.10-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:d203abc0ff2b773ee33d00ab8df0bb67046fbc7c07b119332f08b9b345cf8221", size = 9902707, upload-time = "2026-10-01T18:02:47.041Z" }, + { url = "https://files.pythonhosted.org/packages/ff/44/3fdcedf83ae60ef837dd239170e480dce606a9142cfa2db911afdf855fd9/ruff-0.16.10-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:bd83d1235a5258d318477bc5b576303974cbdef5df0c01a1bff14efcc23bd12a", size = 10619287, upload-time = "2026-10-01T18:02:49.485Z" }, + { url = "https://files.pythonhosted.org/packages/c1/62/02e76a5574002153618eafbb70e159468addc72a5d2c488e65cbb4639d2d/ruff-0.16.10-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:bc2610fb269fa56dd8a68669ae470fa6272902668c0fc2ebc3aa112b2633d5b8", size = 11339942, upload-time = "2026-10-01T18:02:52.008Z" }, + { url = "https://files.pythonhosted.org/packages/6b/c4/cde27d47ad8d4126c587e608c67c46e7feac11763f606d261a1bc8a489e6/ruff-0.16.10-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:3e70175e29cc94c26ea296c80e470180b744b7419026898e58f520c6ab32578e", size = 10934316, upload-time = "2026-10-01T18:02:54.811Z" }, + { url = "https://files.pythonhosted.org/packages/e4/03/17234145f302a645a123e8c3bb2411ecf4669fbc350de3b0d803230a1729/ruff-0.16.10-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f33f43a864a8483eebd160e713336c8bab02c934feaff0a33cf5ccb41546d09a", size = 10387968, upload-time = "2026-10-01T18:02:57.497Z" }, + { url = "https://files.pythonhosted.org/packages/71/29/2493af60240ee7644b38a4b821f5f9c3b5a4fa3178770fe1d0c685217395/ruff-0.16.10-py3-none-manylinux_2_31_riscv64.whl", hash = "sha256:1dfc6f0088149fb6a362c1c446bcbb3fd2157b3852fe2fa68409276eab9ad9b3", size = 10537906, upload-time = "2026-10-01T18:03:00.006Z" }, + { url = "https://files.pythonhosted.org/packages/2e/9a/f56b28f3b143bb9e518c34fb89ab191b626bca8b70dbf8af0d2e2d473572/ruff-0.16.10-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:6553498afc35f580f036030795810b9e6bcea31604b0fd9e8d352795473042e3", size = 10012938, upload-time = "2026-10-01T18:03:03.006Z" }, + { url = "https://files.pythonhosted.org/packages/c2/c4/fca37362848ea4d4d80712df13632e645e7c7cfb1bdedf140699ac7a090b/ruff-0.16.10-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:a3b8471dea115d37f123882be852bed13403746d3a76c11de4a19ec5f9ff5a03", size = 9897945, upload-time = "2026-10-01T18:03:05.818Z" }, + { url = "https://files.pythonhosted.org/packages/7d/c7/e0bc57664d6e0c61fd22f620af260fe9662d7e1ddb320ea7f160e76ef165/ruff-0.16.10-py3-none-musllinux_1_2_i686.whl", hash = "sha256:92e59a70bcbd9d3a5483656da906ec28edfdacfce00afd99edb8b4e9d15644be", size = 10333134, upload-time = "2026-10-01T18:03:08.25Z" }, + { url = "https://files.pythonhosted.org/packages/21/aa/5c9f3b68737c0e4a33a1db5dfab44f7b786d96ca91a7233112ddab1e6dc2/ruff-0.16.10-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:7ae7375f803b5520dc9f546bed7e3a0acb70b91812e9bb4b19927de22f25b77d", size = 10741702, upload-time = "2026-10-01T18:03:10.775Z" }, + { url = "https://files.pythonhosted.org/packages/78/fa/0f9c2020dc316c53d983be011720b8157cc052b97e3f4dfa7db6d0f880a6/ruff-0.16.10-py3-none-win32.whl", hash = "sha256:2a12e01cb9156c10c466f63b46eaae5ecea28dfbd21b5836353ae498e7d1349a", size = 10139176, upload-time = "2026-10-01T18:03:13.267Z" }, + { url = "https://files.pythonhosted.org/packages/99/29/cfb0df9448d4d4ad48c2de029ada9ebd71baa6da983a6c77ee6c6cd0fe82/ruff-0.16.10-py3-none-win_amd64.whl", hash = "sha256:97f2015c92aa97105b0eab19eb5d224884399281cfc5da86a92db4ab5e7fb2ca", size = 10584734, upload-time = "2026-10-01T18:03:16.006Z" }, + { url = "https://files.pythonhosted.org/packages/fc/05/c16957eb287c3fc062e032619a25868d93d408a844b725bcf514f7a378ff/ruff-0.16.10-py3-none-win_arm64.whl", hash = "sha256:25a65fe998c4e6861ec079ada5826a2fc605e6cbccbe9dcd7fac1f54e791621b", size = 10366440, upload-time = "2026-10-01T18:03:19.04Z" }, +] + +[[package]] +name = "sentry-sdk" +version = "2.71.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "certifi" }, + { name = "urllib3" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/67/93/9ef70bfb7346c778caf09b7cd10c17c33f601f2d17019fa1aedcda738b5f/sentry_sdk-2.71.0.tar.gz", hash = "sha256:7beb27a22f06396f3a05510c3da6ea10e1c79e02e859993c27ce4ff074e296d6", size = 1052358, upload-time = "2026-09-28T14:04:00.645Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a4/f6/b3e68b2e08e9a52e6c8a279ca2fc0d3a535f3886222375f54ef9efbd694d/sentry_sdk-2.71.0-py3-none-any.whl", hash = "sha256:9ed36b4c206048aa67b5ddba0a6e28d59fc55cf10d60f5f7aea2c3917184bf4a", size = 534254, upload-time = "2026-09-28T14:03:58.668Z" }, +] + +[[package]] +name = "six" +version = "1.17.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/94/e7/b2c673351809dca68a0e064b6af791aa332cf192da575fd474ed7d6f16a2/six-1.17.0.tar.gz", hash = "sha256:ff70335d468e7eb6ec65b95b99d3a2836546063f63acc5171de367e834932a81", size = 34031, upload-time = "2024-12-04T17:35:28.174Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b7/ce/149a00dd41f10bc29e5921b496af8b574d8413afcd5e30dfa0ed46c2cc5e/six-1.17.0-py2.py3-none-any.whl", hash = "sha256:4721f391ed90541fddacab5acf947aa0d3dc7d27b2e1e8eda2be8970586c3274", size = 11050, upload-time = "2024-12-04T17:35:26.475Z" }, +] + +[[package]] +name = "textual" +version = "8.2.8" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "markdown-it-py", extra = ["linkify"] }, + { name = "mdit-py-plugins" }, + { name = "platformdirs" }, + { name = "pygments" }, + { name = "rich" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/00/21/39a76b01bd5eea82a04baaca7580e105d8c59450df03998345bb2cfb307b/textual-8.2.8.tar.gz", hash = "sha256:3f106a9fbc73e39dd266c9712432087de78a6d644084c7c241d6a25c3169115b", size = 1860502, upload-time = "2026-06-30T06:51:24.495Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fb/be/35261223d9416a0751cdff1c7b4a6f881387218a12d439fe22fefebc8c04/textual-8.2.8-py3-none-any.whl", hash = "sha256:267375fd402dc8d981457212efa71f0e3365fd17bba144ba9bb3ed7563cb374a", size = 731418, upload-time = "2026-06-30T06:51:26.364Z" }, +] + +[[package]] +name = "typing-extensions" +version = "4.16.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f6/cc/6253133b5bb138fc3306cebfbda2c520f545d36b5be2c7255cc528bb45d6/typing_extensions-4.16.0.tar.gz", hash = "sha256:dc983d19a509c94dba722ee6abd33940f7c05a89e243c47e907eb4db6f1a43e5", size = 113555, upload-time = "2026-07-02T08:40:05.92Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/49/d3/b8441a820a491ddfc024b0b0cf0393375b75ea13866d9c66727e54c2fc80/typing_extensions-4.16.0-py3-none-any.whl", hash = "sha256:481caa481374e813c1b176ada14e97f1f67a4539ce9cfeb3f350d78d6370c2e8", size = 45571, upload-time = "2026-07-02T08:40:04.659Z" }, +] + +[[package]] +name = "typing-inspection" +version = "0.4.4" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/a3/26/b09b8010994eccc3c09092e6b34058f36a460eea2d4c3e8b910c695975a0/typing_inspection-0.4.4.tar.gz", hash = "sha256:547274fa6b0a561ccf549cc9524b999a578e737d015d8709d021f9d0d13bea47", size = 76928, upload-time = "2026-08-12T12:37:25.997Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/67/81/4add07e5172b7ac40d8ed5ff580409a7801a4fe26d529bdd915401dabfbe/typing_inspection-0.4.4-py3-none-any.whl", hash = "sha256:65b8397ba37ccbce054456aaccddfc91e6e3083c92824df348d96ca832f3f147", size = 14750, upload-time = "2026-08-12T12:37:24.648Z" }, +] + +[[package]] +name = "tzdata" +version = "2026.4" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/e4/31/3d74fa778a63b98b7374323befcc0be5ab3bd94afd4096a0124e7379152c/tzdata-2026.4.tar.gz", hash = "sha256:f1b8bd365d8d210c55353f4d7f8d6d8561c0ba50d704b700d195a9424bba0d79", size = 199350, upload-time = "2026-09-12T12:56:03.251Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f9/bc/8737e8d54cf51106118039b83f485a4783112fab49ea9d044b234978a46e/tzdata-2026.4-py2.py3-none-any.whl", hash = "sha256:c2169a8b0a7a5e9674da5a135ccdfb2b3e671b333ed9fed17b41f73c34476e81", size = 347494, upload-time = "2026-09-12T12:56:01.67Z" }, +] + +[[package]] +name = "tzlocal" +version = "5.4.4" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "tzdata", marker = "sys_platform == 'win32'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/81/5b/879b2f932adfa7a053c360d50bc896c977fa6426109185f7c12ebdd0cb9d/tzlocal-5.4.4.tar.gz", hash = "sha256:8dbb8660838688a7b6ba4fed31d18dedf842afb4d47ca050d6d891c2c15f3be4", size = 31170, upload-time = "2026-06-29T08:03:40.026Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/9e/a4/017a7a6cbe387d961a688ec31364ae60a5c4e22c96ae9921b79a947c855d/tzlocal-5.4.4-py3-none-any.whl", hash = "sha256:aae09f0126a8a86fa736be266eb4a471380d26a0de3bc14844e7821fee3e2a15", size = 18115, upload-time = "2026-06-29T08:03:38.666Z" }, +] + +[[package]] +name = "urllib3" +version = "2.8.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/e3/05/b17359e1cefb4f909b5e40b1b90a496d987258916dbbf88e842c729f510e/urllib3-2.8.0.tar.gz", hash = "sha256:63bf2ead4c879426ebf22ef2a781eeb4aa3b4ae798a0435506f8687fd5bb9b63", size = 458972, upload-time = "2026-09-15T19:29:36.253Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/92/9d/c4e665119135114480843e7ab388fa94d8480650450e6f8e26b70d323a4c/urllib3-2.8.0-py3-none-any.whl", hash = "sha256:0cf3cae568d36aa9576b28dfb35f11328f1cb974ca7647d9475ebb86c75ac6e3", size = 135717, upload-time = "2026-09-15T19:29:34.577Z" }, +] diff --git a/scripts/kernel_tuning/verify.py b/scripts/kernel_tuning/verify.py new file mode 100644 index 0000000..ffe876d --- /dev/null +++ b/scripts/kernel_tuning/verify.py @@ -0,0 +1,288 @@ +"""Model-level numerical gate for an installed kernel profile. + +``python -m scripts.kernel_tuning verify --profile NAME --output REPORT -- DRIVER_ARGS`` +runs in the model runtime. ``DRIVER_ARGS`` are the serving arguments and recorded +inputs of :mod:`scripts.kernel_tuning.discovery.drive`. The same inputs run three +times, each with a fixed per-input seed: + +* reference: FP32 execution with every registry operator on its Torch reference; +* baseline: the requested dtype with ``kernel_profile: torch``; +* candidate: the requested dtype with ``kernel_profile: NAME``. + +Continuous outputs are judged over all inputs together, per output name: the +candidate's mean max-abs and mean RMSE error against the reference may exceed the +baseline's by at most 25%, and its worst max-abs error may be at most twice the +baseline's worst (floors 1e-4 and 1e-5). A single input's baseline error is not a +threshold: Torch itself varies between runs by more than that on inputs where it +happens to land close to the reference. Discrete outputs (integers, strings, +booleans) must equal the baseline exactly. See docs/proposals/0009, Section 6. +""" + +from __future__ import annotations + +import argparse +import gc +import json +import math +import tempfile +from pathlib import Path +from typing import Any + +from .contracts import ContractError + +MAX_ABS_FLOOR = 1e-4 +RMSE_FLOOR = 1e-5 +# Mean errors over all inputs may grow by this factor; the worst input by WORST_FACTOR. +MEAN_FACTOR = 1.25 +WORST_FACTOR = 2.0 +# A baseline this far from the FP32 reference means the two runs did not see the +# same inputs or noise; the reference cannot calibrate a threshold then. +INCOMPARABLE_FRACTION = 0.5 + + +def flatten_outputs(value: Any, prefix: str = "") -> dict[str, tuple[str, Any]]: + """Map output leaves to ``("continuous", floats)`` or ``("discrete", values)``.""" + import torch + + result: dict[str, tuple[str, Any]] = {} + if isinstance(value, torch.Tensor): + tensor = value.detach().cpu() + if tensor.is_floating_point(): + result[prefix] = ("continuous", tensor.double().flatten().tolist()) + else: + result[prefix] = ("discrete", tensor.flatten().tolist()) + elif hasattr(value, "dtype") and hasattr(value, "tolist"): # NumPy arrays and scalars + values = value.tolist() + values = values if isinstance(values, list) else [values] + flat = _flat(values) + kind = "continuous" if value.dtype.kind in "fc" else "discrete" + result[prefix] = (kind, [float(x) for x in flat] if kind == "continuous" else flat) + elif hasattr(value, "__dataclass_fields__"): + for name in value.__dataclass_fields__: + if name in ("timing", "policy_revision"): + continue # Wall-clock timing and version strings are not model outputs. + result.update(flatten_outputs(getattr(value, name), f"{prefix}.{name}")) + elif isinstance(value, dict): + for key in sorted(value, key=str): + result.update(flatten_outputs(value[key], f"{prefix}.{key}")) + elif isinstance(value, (list, tuple)): + if value and all(isinstance(item, float) for item in value): + result[prefix] = ("continuous", [float(item) for item in value]) + elif value and all(isinstance(item, (int, float)) and not isinstance(item, bool) for item in value): + if any(isinstance(item, float) for item in value): + result[prefix] = ("continuous", [float(item) for item in value]) + else: + result[prefix] = ("discrete", list(value)) + else: + for index, item in enumerate(value): + result.update(flatten_outputs(item, f"{prefix}[{index}]")) + elif isinstance(value, float): + result[prefix] = ("continuous", [value]) + else: + result[prefix] = ("discrete", [value]) + return result + + +def _flat(values: list[Any]) -> list[Any]: + return [x for item in values for x in (_flat(item) if isinstance(item, list) else [item])] + + +def _error(actual: list[float], expected: list[float]) -> tuple[float, float]: + if len(actual) != len(expected): + raise ContractError("Output sizes differ between runs") + if not actual: + return 0.0, 0.0 + if not all(math.isfinite(x) for x in actual): + return math.inf, math.inf + differences = [abs(a - b) for a, b in zip(actual, expected)] + return max(differences), math.sqrt(sum(d * d for d in differences) / len(differences)) + + +def gate( + reference: list[dict[str, tuple[str, Any]]], + baseline: list[dict[str, tuple[str, Any]]], + candidate: list[dict[str, tuple[str, Any]]], +) -> dict[str, Any]: + """Compare candidate and baseline errors against the reference over all inputs.""" + if not (len(reference) == len(baseline) == len(candidate)) or not reference: + raise ContractError("All runs must produce one output per input") + rows, errors, passed = [], {}, True + for index, (ref, base, cand) in enumerate(zip(reference, baseline, candidate)): + if set(ref) != set(base) or set(base) != set(cand): + raise ContractError(f"input {index}: output structure differs between runs") + for name in sorted(ref): + kind = base[name][0] + if kind == "discrete": + same = cand[name] == base[name] + passed &= same + rows.append({"input": index, "output": name, "kind": kind, "passed": same}) + continue + if ref[name][0] != "continuous" or cand[name][0] != "continuous": + raise ContractError(f"input {index} {name}: output kind differs between runs") + scale = max((abs(x) for x in ref[name][1]), default=0.0) + base_abs, base_rmse = _error(base[name][1], ref[name][1]) + if base_abs > INCOMPARABLE_FRACTION * max(scale, 1e-6): + raise ContractError( + f"input {index} {name}: baseline differs from the FP32 reference by {base_abs:.3g} " + f"(scale {scale:.3g}); the runs did not share inputs/noise, so no threshold is valid" + ) + cand_abs, cand_rmse = _error(cand[name][1], ref[name][1]) + errors.setdefault(name, []).append((base_abs, base_rmse, cand_abs, cand_rmse)) + rows.append( + { + "input": index, + "output": name, + "kind": kind, + "candidate_max_abs": cand_abs, + "candidate_rmse": cand_rmse, + "baseline_max_abs": base_abs, + "baseline_rmse": base_rmse, + } + ) + outputs = {} + for name, values in sorted(errors.items()): + base_abs, base_rmse, cand_abs, cand_rmse = (list(column) for column in zip(*values)) + checks = { + "mean_max_abs": (_mean(cand_abs), max(MAX_ABS_FLOOR, MEAN_FACTOR * _mean(base_abs))), + "mean_rmse": (_mean(cand_rmse), max(RMSE_FLOOR, MEAN_FACTOR * _mean(base_rmse))), + "worst_max_abs": (max(cand_abs), max(MAX_ABS_FLOOR, WORST_FACTOR * max(base_abs))), + } + ok = all(value <= limit for value, limit in checks.values()) + passed &= ok + outputs[name] = { + "passed": ok, + "ratio": max(value / limit for value, limit in checks.values()), + **{key: {"candidate": value, "limit": limit} for key, (value, limit) in checks.items()}, + "baseline": { + "mean_max_abs": _mean(base_abs), + "mean_rmse": _mean(base_rmse), + "worst_max_abs": max(base_abs), + }, + } + return { + "passed": bool(passed), + "inputs": len(reference), + "outputs": outputs, + "failed": [name for name, item in outputs.items() if not item["passed"]] + + [ + f"{row['output']}[{row['input']}]" + for row in rows + if row["kind"] == "discrete" and not row["passed"] + ], + "worst_ratio": max((item["ratio"] for item in outputs.values()), default=None), + "rows": rows, + } + + +def _mean(values: list[float]) -> float: + return math.fsum(values) / len(values) + + +def _run(args: argparse.Namespace, episodes: list[list[Any]], kind: str, profile: str, dtype: str, seed: int): + """Outputs of every input under one kernel profile and dtype.""" + import torch + + from .discovery import drive + + config = drive._parse_adapter_config(args.adapter_config) + config.setdefault("policy_kwargs", {})["kernel_profile"] = profile + outputs = [] + with tempfile.TemporaryDirectory(prefix="kernel-verify-") as scratch: + path = Path(scratch) / "adapter.json" + path.write_text(json.dumps(config), encoding="utf-8") + local = argparse.Namespace(**{**vars(args), "adapter_config": str(path), "dtype": dtype}) + index = 0 + if kind == "requests": + adapter = drive.build_serving_adapter(local) + try: + for episode in episodes: + for request in episode: + torch.manual_seed(seed + index) + outputs.append(flatten_outputs(adapter.infer(request))) + index += 1 + for session in dict.fromkeys(request.session_id for request in episode): + adapter.reset(session) + finally: + if isinstance(adapter, drive.BatchedServingAdapter): + adapter.shutdown() + del adapter + else: + from embodiinfer.engine.config import EngineConfig + from embodiinfer.engine.core import EngineCore + from embodiinfer.policies.factory import make_policy + from embodiinfer.types import SessionKey + + kwargs = dict(config["policy_kwargs"]) + if local.checkpoint is not None: + kwargs["checkpoint"] = local.checkpoint + policy = make_policy(local.policy, **kwargs) + core = EngineCore( + policy, + EngineConfig( + device=local.device, + dtype=dtype, + max_batch_size=local.max_batch, + num_steps=local.num_steps, + use_cuda_graph=not local.no_cuda_graph, + capture_full_loop=local.capture_full_loop, + ), + ) + for number, episode in enumerate(episodes): + key = SessionKey(env_id="verify", episode_id=number) if policy.is_recurrent else None + for step, observation in enumerate(episode): + torch.manual_seed(seed + index) + batch = policy.collate([observation], [f"{number}-{step}"]) + outputs.append( + flatten_outputs(core.execute(batch, session_ids=None if key is None else [key])) + ) + index += 1 + if key is not None: + core.reset_sessions([key]) + del core, policy + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + return outputs + + +def run(profile: str, output: Path, driver_args: list[str], *, samples: int | None = None, seed: int = 1000): + """Run reference, baseline and candidate on the driver inputs and write the gate report.""" + from embodiinfer.engine.serve.factory import add_policy_arguments + + from .discovery.drive import load_inputs + + if output.exists(): + raise ContractError(f"{output} already exists") + parser = argparse.ArgumentParser(prog="verify driver arguments") + add_policy_arguments(parser) + parser.add_argument("--inputs", type=Path, required=True) + args = parser.parse_args(driver_args) + kind, episodes = load_inputs(args.inputs) + if samples is not None: + remaining, kept = samples, [] + for episode in episodes: + if remaining <= 0: + break + kept.append(episode[:remaining]) + remaining -= len(kept[-1]) + episodes = kept + dtype = args.dtype if args.dtype != "auto" else "bfloat16" + runs = { + "reference": _run(args, episodes, kind, "torch", "float32", seed), + "baseline": _run(args, episodes, kind, "torch", dtype, seed), + "candidate": _run(args, episodes, kind, profile, dtype, seed), + } + report = gate(runs["reference"], runs["baseline"], runs["candidate"]) + report.update( + { + "profile": profile, + "dtype": dtype, + "seed": seed, + "driver_args": driver_args, + "floors": {"max_abs": MAX_ABS_FLOOR, "rmse": RMSE_FLOOR}, + "factors": {"mean": MEAN_FACTOR, "worst": WORST_FACTOR}, + } + ) + output.parent.mkdir(parents=True, exist_ok=True) + output.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8") + return report diff --git a/tests/kernel_tuning/conftest.py b/tests/kernel_tuning/conftest.py new file mode 100644 index 0000000..c8776c4 --- /dev/null +++ b/tests/kernel_tuning/conftest.py @@ -0,0 +1,197 @@ +"""CPU fixtures; use --confcutdir here in the isolated, Torch-free tool environment.""" + +from __future__ import annotations + +import json +import sys +from pathlib import Path +from typing import Any + +import pytest +from scripts.kernel_tuning.artifacts import RunStore +from scripts.kernel_tuning.contracts import TaskPackage + +yaml = pytest.importorskip("yaml", reason="requires the optional kernel-tuning tooling environment") + +FAKE_BENCHMARK = '''"""Deterministic CPU fake: performs no kernel compilation or GPU work.""" +from dataclasses import asdict +import time + +class Adapter: + def __init__(self, task): + self.task = task + + def environment(self): + return {"identity": {"hardware": "fake-cpu", "version": "1"}, "conditions": {"fake": True}} + + def evaluate(self, solutions): + if solutions["candidate"].get("description") == "crash": + raise RuntimeError("simulated compilation failure") + if solutions["candidate"].get("description") == "timeout": + time.sleep(30) + cfg = self.task.settings + roles = ("baseline", "incumbent", "candidate") + checks = [{"workload": w, "role": role, "passed": True, + "seeds": list(cfg.seeds), "precision": asdict(cfg.precision), + "max_abs_error": 0.0, "max_rel_error": 0.0, + "graph_replay": cfg.timing.mode == "cuda_graph"} + for w in self.task.workload_ids for role in roles] + rounds = [{w: {r: [float(solutions[r].get("description") or 10)] * cfg.timing.trials + for r in roles} for w in self.task.workload_ids} + for _ in range(cfg.timing.paired_rounds)] + return {"checks": checks, "rounds": rounds} + +def create_adapter(task): + return Adapter(task) +''' + + +GATED_RESIDUAL_REFERENCE = """import torch + +def run(residual, update, gate): + product = (update * gate).to(update.dtype) + return (residual + product).to(residual.dtype) +""" + +GATED_RESIDUAL_DEFINITION = { + "name": "gated_residual_bf16", + "op_type": "elementwise", + "description": "Residual plus a gated update, with BF16 product rounding before addition.", + "axes": {"B": {"type": "var"}, "T": {"type": "var"}, "H": {"type": "var"}, "G": {"type": "var"}}, + "inputs": { + "residual": {"shape": ["B", "T", "H"], "dtype": "bfloat16"}, + "update": {"shape": ["B", "T", "H"], "dtype": "bfloat16"}, + "gate": {"shape": ["B", "G", "H"], "dtype": "bfloat16"}, + }, + "outputs": {"output": {"shape": ["B", "T", "H"], "dtype": "bfloat16"}}, + "constraints": ["G == 1 or G == T", "B > 0", "T > 0", "H > 0"], + "reference": GATED_RESIDUAL_REFERENCE, +} + +GATED_RESIDUAL_WORKLOADS = { + "broadcast-small": {"B": 1, "T": 16, "H": 1024, "G": 1}, + "broadcast-wide": {"B": 2, "T": 32, "H": 2048, "G": 1}, + "token-gate-tail": {"B": 1, "T": 17, "H": 1000, "G": 17}, +} + +GATED_RESIDUAL_TUNING = { + "schema_version": 1, + "language": "triton", + "agent": "", + "evaluator_python": "python", + "device": "cuda:0", + "target_hardware": ["cuda"], + "seeds": [0, 1, 2], + "precision": {"mode": "bit_exact", "matmul_precision": "highest", "allow_tf32": False}, + "timing": {"mode": "eager", "warmup": 10, "iterations": 100, "trials": 3, "paired_rounds": 2}, + "acceptance": {"min_improvement": 0.03, "max_regression": 0.05}, + "search": { + "max_candidates": 20, + "max_seconds": 7200, + "patience": 5, + "agent_timeout_seconds": 600, + "evaluation_timeout_seconds": 300, + }, + "profile": False, +} + + +@pytest.fixture +def example_task(tmp_path: Path) -> Path: + """Write a hand-written BF16 gated residual task package with synthetic workloads.""" + path = tmp_path / "task" + path.mkdir() + (path / "README.md").write_text("# Gated residual BF16 test task\n", encoding="utf-8") + (path / "definition.json").write_text(json.dumps(GATED_RESIDUAL_DEFINITION, indent=2), encoding="utf-8") + (path / "baseline.py").write_text(GATED_RESIDUAL_REFERENCE, encoding="utf-8") + inputs = {name: {"type": "random"} for name in GATED_RESIDUAL_DEFINITION["inputs"]} + (path / "workloads.jsonl").write_text( + "".join( + json.dumps( + { + "definition": GATED_RESIDUAL_DEFINITION["name"], + "workload": {"uuid": uuid, "axes": axes, "inputs": inputs}, + "solution": None, + "evaluation": None, + } + ) + + "\n" + for uuid, axes in GATED_RESIDUAL_WORKLOADS.items() + ), + encoding="utf-8", + ) + (path / "tuning.yaml").write_text(yaml.safe_dump(GATED_RESIDUAL_TUNING), encoding="utf-8") + return path + + +@pytest.fixture +def task_dir(example_task: Path) -> Path: + """Make a portable task using a fake measurement adapter and explicit Python.""" + path = example_task + (path / "benchmark.py").write_text(FAKE_BENCHMARK, encoding="utf-8") + config = yaml.safe_load((path / "tuning.yaml").read_text(encoding="utf-8")) + config["evaluator_python"] = sys.executable + config["search"].update( + max_candidates=3, patience=2, max_seconds=60, agent_timeout_seconds=5, evaluation_timeout_seconds=10 + ) + (path / "tuning.yaml").write_text(yaml.safe_dump(config), encoding="utf-8") + return path + + +@pytest.fixture +def store(task_dir: Path, tmp_path: Path) -> RunStore: + """Create only the task archive; no agent or evaluator starts here.""" + return RunStore.create(TaskPackage.load(task_dir), tmp_path / "runs", "claude/fake:low") + + +@pytest.fixture +def candidate(store: RunStore) -> dict[str, Any]: + """A source-complete Solution; the fake evaluator reads its latency description.""" + return { + "name": "fake_candidate", + "definition": store.task.definition["name"], + "author": "test", + "description": "8", + "spec": { + "language": "triton", + "target_hardware": ["cuda"], + "entry_point": "kernel.py::run", + "destination_passing_style": False, + }, + "sources": [ + {"path": "kernel.py", "content": "def run(residual, update, gate):\n return residual\n"} + ], + } + + +@pytest.fixture +def registry_capture(): + """Build CPU captures using the same recorder as capture-ops, without a checkpoint.""" + pytest.importorskip("torch") + from scripts.kernel_tuning.artifacts import atomic_json + from scripts.kernel_tuning.contracts import read_json + from scripts.kernel_tuning.discovery.contracts import ModelCapture + from scripts.kernel_tuning.regions import OpCapture + + from embodiinfer.layers.ops import get_op + + def record( + root: Path, calls: list[tuple], *, budget: int = 1024 * 1024, device_type: str = "cpu" + ) -> ModelCapture: + capture = OpCapture("mock_flow_vla", root / "raw", warmup=0, steps=1, fixture_bytes=budget) + capture.target = {"model": "mock_flow_vla", "type": "test.Policy", "config": {}} + capture.started = 1 + try: + for name, args, kwargs in calls: + function = get_op(name, "torch") + capture.record(name, function, function, args, kwargs) + finally: + capture.close(complete=True) + # The CLI targets CUDA; the fixture changes only the declared evaluator + # device, after recording real CPU tensors through OpCapture.record. + manifest = read_json(capture.root / "manifest.json") + manifest["device_type"] = device_type + atomic_json(capture.root / "manifest.json", manifest) + return ModelCapture.load(capture.root, require_ready=False) + + return record diff --git a/tests/kernel_tuning/test_batch.py b/tests/kernel_tuning/test_batch.py new file mode 100644 index 0000000..98bbad8 --- /dev/null +++ b/tests/kernel_tuning/test_batch.py @@ -0,0 +1,242 @@ +"""Batch tuning over generated tasks, with fake run commands and synthetic evidence.""" + +from __future__ import annotations + +import json +from dataclasses import asdict +from pathlib import Path +from typing import Any + +import pytest +from scripts.kernel_tuning import batch +from scripts.kernel_tuning.__main__ import main +from scripts.kernel_tuning.artifacts import RunStore +from scripts.kernel_tuning.contracts import ContractError, TaskPackage, promotion, read_json +from scripts.kernel_tuning.operators import select + + +def measurement(task: TaskPackage, latency: float) -> dict[str, Any]: + """Complete synthetic evidence: the candidate at ``latency``, other roles at 10 ms.""" + roles = ("baseline", "incumbent", "candidate") + return { + "checks": [ + { + "workload": work, + "role": role, + "passed": True, + "seeds": list(task.settings.seeds), + "precision": asdict(task.settings.precision), + "max_abs_error": 0.0, + "max_rel_error": 0.0, + "graph_replay": True, + } + for work in task.workload_ids + for role in roles + ], + "rounds": [ + { + work: { + r: [latency if r == "candidate" else 10.0] * task.settings.timing.trials for r in roles + } + for work in task.workload_ids + } + for _ in range(task.settings.timing.paired_rounds) + ], + } + + +class FakeRuns: + """Stands in for ``run``/``resume`` subprocesses, using the real archive code.""" + + def __init__(self, outcomes: dict[str, list[str]]) -> None: + self.outcomes = outcomes + self.commands: list[list[str]] = [] + + def __call__(self, argv: list[str], log: Path) -> int: + self.commands.append(argv) + log.parent.mkdir(parents=True, exist_ok=True) + if argv[3] == "run": + task = TaskPackage.load(Path(argv[4])) + store = RunStore.create( + task, Path(argv[argv.index("--output") + 1]), argv[argv.index("--agent") + 1] + ) + log.write_text(f"Run archive: {store.root}\n", encoding="utf-8") + else: + store = RunStore(Path(argv[4])) + name = store.task.definition["name"].removeprefix("embodiinfer_") + outcome = self.outcomes[name].pop(0) + if outcome == "interrupt": + store.manifest.update(status="interrupted", stop_reason="KeyboardInterrupt") + store.save() + return 130 + if outcome == "fail": + store.manifest.update(status="failed", stop_reason="ContractError: evaluator unavailable") + store.save() + return 2 + if outcome == "promote": + candidate, _, _ = store.begin_attempt() + solution = { + "name": f"{name}_fast", + "definition": store.task.definition["name"], + "author": "fake-agent", + "spec": { + "language": "triton", + "target_hardware": ["cuda"], + "entry_point": "kernel.py::run", + "destination_passing_style": False, + }, + "sources": [{"path": "kernel.py", "content": "def run(*args):\n return args\n"}], + } + decision = promotion(store.task, measurement(store.task, 8.0)) + store.finish_attempt(candidate, status="promoted", decision=decision, solution=solution) + store.manifest.update(status="completed", stop_reason="max_candidates") + store.save() + return 0 + + +def new_batch(tmp_path: Path, names: list[str]) -> Path: + return batch.create( + select(names), tmp_path / "batches", agent="claude/fake:low", overrides={}, hardware_notes=None + ) + + +def test_batch_runs_every_operator_and_exports_promoted_kernels(tmp_path: Path) -> None: + root = new_batch(tmp_path, ["gated_residual", "swiglu", "rms_norm"]) + runs = FakeRuns({"gated_residual": ["promote"], "swiglu": ["fail", "promote"], "rms_norm": ["complete"]}) + record = batch.execute(root, launcher=runs) + status = {name: item["status"] for name, item in record["items"].items()} + assert status == {"gated_residual": "completed", "swiglu": "failed", "rms_norm": "completed"} + assert record["items"]["swiglu"]["error"] == "ContractError: evaluator unavailable" + assert (root / "exports/gated_residual/sources/kernel.py").is_file() + assert not (root / "exports/rms_norm").exists() + summary = {row["operator"]: row for row in read_json(root / "summary.json")} + assert summary["gated_residual"]["improvement_vs_baseline"] == pytest.approx(0.2) + assert summary["gated_residual"]["export"] == "exports/gated_residual" + assert "| gated_residual | completed | 1 | 1 | 20.0% |" in (root / "summary.md").read_text( + encoding="utf-8" + ) + + # Continuing the batch resumes only the failed run, inside its existing archive. + batch.execute(root, launcher=runs) + assert [command[3] for command in runs.commands] == ["run", "run", "run", "resume"] + assert runs.commands[-1][4] == str(root / record["items"]["swiglu"]["run"]) + assert read_json(root / "batch.json")["items"]["swiglu"]["status"] == "completed" + assert (root / "exports/swiglu").is_dir() + + +def test_interrupt_stops_the_batch_and_resume_continues_it(tmp_path: Path) -> None: + root = new_batch(tmp_path, ["gated_residual", "rms_norm"]) + runs = FakeRuns({"gated_residual": ["interrupt", "complete"], "rms_norm": ["complete"]}) + with pytest.raises(KeyboardInterrupt): + batch.execute(root, launcher=runs) + items = read_json(root / "batch.json")["items"] + assert items["gated_residual"]["status"] == "interrupted" + assert items["rms_norm"]["status"] == "pending" + batch.execute(root, launcher=runs) + assert [command[3] for command in runs.commands] == ["run", "resume", "run"] + + +def test_run_that_fails_before_archiving_is_started_again(tmp_path: Path) -> None: + root = new_batch(tmp_path, ["rms_norm"]) + calls = [] + + def broken(argv: list[str], log: Path) -> int: + calls.append(argv[3]) + log.parent.mkdir(parents=True, exist_ok=True) + log.write_text("kernel_tuning: Install the optional tooling environment first\n", encoding="utf-8") + return 2 + + record = batch.execute(root, launcher=broken) + assert record["items"]["rms_norm"]["run"] is None + assert record["items"]["rms_norm"]["error"] == "exit 2; see logs/rms_norm.log" + batch.execute(root, launcher=broken) + assert calls == ["run", "run"] + + +def test_failed_preflight_skips_the_operator(tmp_path: Path, monkeypatch) -> None: + root = new_batch(tmp_path, ["gated_residual", "rms_norm"]) + + async def evaluate(store: RunStore, output: Path, **_: Any) -> dict[str, Any]: + if store.task.definition["name"].endswith("rms_norm"): + raise ContractError("Evaluator failed (exit 1)") + return {"measurement": measurement(store.task, 10.0)} + + monkeypatch.setattr(batch, "evaluate", evaluate) + assert batch.preflight(root) == {"gated_residual": None, "rms_norm": "Evaluator failed (exit 1)"} + items = read_json(root / "batch.json")["items"] + assert items["gated_residual"]["baseline_us"] == {"pi05-expert-b1": 10000.0, "pi05-expert-b8": 10000.0} + assert "baseline_us" not in items["rms_norm"] + runs = FakeRuns({"gated_residual": ["complete"]}) + record = batch.execute(root, launcher=runs) + assert record["items"]["rms_norm"]["status"] == "preflight_failed" + assert len(runs.commands) == 1 + + +def test_cli_validates_agent_and_resume_arguments(tmp_path: Path, capsys) -> None: + output = tmp_path / "batches" + assert main(["tune-all", "rms_norm", "--output", str(output)]) == 2 + assert "explicit Humanize2 agent" in capsys.readouterr().err + assert not output.exists() + assert main(["tune-all", "--catalog", "--model", "streamvln", "--dry-run", "--output", str(output)]) == 0 + (root,) = output.iterdir() + assert set(read_json(root / "batch.json")["items"]) == {"rms_norm", "add_rms_norm", "swiglu"} + assert "| swiglu | pending |" in capsys.readouterr().out + assert main(["tune-all", "--resume", str(root), "--set", "device=cuda:1"]) == 2 + assert "--resume continues the recorded batch" in capsys.readouterr().err + assert main(["tune-all", "--resume", str(root), "--unfenced"]) == 2 + assert "--resume continues the recorded batch" in capsys.readouterr().err + + +def test_unfenced_batch_starts_every_run_unfenced(tmp_path: Path) -> None: + root = batch.create( + select(["rms_norm"]), + tmp_path / "batches", + agent="claude/fake:low", + overrides={}, + hardware_notes=None, + unfenced=True, + ) + commands: list[list[str]] = [] + + def record(argv: list[str], log: Path) -> int: + commands.append(argv) + log.parent.mkdir(parents=True, exist_ok=True) + log.write_text("stopped before archiving\n", encoding="utf-8") + return 1 + + batch.execute(root, launcher=record) + assert commands and commands[0][3] == "run" and "--unfenced" in commands[0] + assert "unfenced" in (root / "summary.md").read_text(encoding="utf-8") + assert read_json(new_batch(tmp_path / "default", ["rms_norm"]) / "batch.json")["unfenced"] is False + + +def test_cli_generate_writes_overridden_tasks(tmp_path: Path, capsys) -> None: + notes = tmp_path / "thor.md" + notes.write_text("Jetson Thor: sm_110, 128 GB unified memory.\n", encoding="utf-8") + argv = [ + "generate", + "--catalog", + "rms_norm", + "--output", + str(tmp_path / "tasks"), + "--hardware-notes", + str(notes), + ] + assert main([*argv, "--set", "evaluator_python=/opt/runtime/bin/python"]) == 0 + task = TaskPackage.load(tmp_path / "tasks/rms_norm") + assert task.settings.evaluator_python == "/opt/runtime/bin/python" + assert "HARDWARE.md" in task.hashes + assert "HARDWARE.md" in (task.root / "README.md").read_text(encoding="utf-8") + assert main(argv) == 2 + assert "--force" in capsys.readouterr().err + + +def test_variable_axes_must_be_inferable_from_inputs(tmp_path: Path) -> None: + root = new_batch(tmp_path, ["rms_norm"]) + task = root / "tasks/rms_norm" + definition = read_json(task / "definition.json") + definition["axes"]["X"] = {"type": "var"} + definition["outputs"]["output"]["shape"] = ["M", "X"] + (task / "definition.json").write_text(json.dumps(definition), encoding="utf-8") + with pytest.raises(ContractError, match="appear in an input shape"): + TaskPackage.load(task) diff --git a/tests/kernel_tuning/test_calibration.py b/tests/kernel_tuning/test_calibration.py new file mode 100644 index 0000000..52c7454 --- /dev/null +++ b/tests/kernel_tuning/test_calibration.py @@ -0,0 +1,279 @@ +"""Numerical policy, independent calibration, and frozen all-element replay gates.""" + +from __future__ import annotations + +from dataclasses import asdict +from pathlib import Path +from types import SimpleNamespace + +import pytest + +torch = pytest.importorskip("torch") + +from benchmarks.kernel_tuning.calibration import ( # noqa: E402 + calibrate, + captured_settings, + compare_calibrated, + derive_bound, + environment, + high_precision, + production, + validation_inputs, +) +from benchmarks.kernel_tuning.replay_adapter import ReplayAdapter # noqa: E402 +from scripts.kernel_tuning.artifacts import atomic_json # noqa: E402 +from scripts.kernel_tuning.contracts import ContractError, TaskPackage, read_json # noqa: E402 +from scripts.kernel_tuning.discovery.contracts import ModelCapture # noqa: E402 +from scripts.kernel_tuning.discovery.numerics import PROFILES, policy # noqa: E402 +from scripts.kernel_tuning.discovery.tasks import render_model # noqa: E402 + + +@pytest.fixture +def calibrated(tmp_path: Path, registry_capture) -> ModelCapture: + generator = torch.Generator().manual_seed(17) + a = torch.randn(3, 16, generator=generator) + b = torch.randn(16, 5, generator=generator) + c = torch.randn(4, 16, generator=generator) * 64 + raw = registry_capture(tmp_path, [("linear", (a, b.T), {}), ("linear", (c, b.T), {})]) + return calibrate(raw, tmp_path / "calibrated") + + +def test_registry_semantics_keep_mutation_integer_and_fp64_exact() -> None: + spec = { + "kind": "registry", + "name": "registry.linear", + "input_dtypes": ["float32"], + "output_dtypes": ["float32"], + "mutates": [], + } + assert policy(spec) == "registry" + for extra in ( + {"kind": "aten", "name": "aten.mm.default"}, + {"kind": "backend"}, + {"mutates": [0]}, + {"rng": {"kind": "torch_default"}}, + {"input_dtypes": ["float64"]}, + {"output_dtypes": ["int64"]}, + ): + assert policy({**spec, **extra}) is None + + +def test_generation_requires_calibration_and_capture_keeps_only_execution_switches( + tmp_path: Path, registry_capture +) -> None: + x = torch.ones(3, 3) + capture = registry_capture(tmp_path, [("linear", (x, x), {})]) + assert "mode" not in capture.manifest["precision"] + with pytest.raises(ContractError, match="Missing numerical calibration"): + render_model(capture, tmp_path / "tasks") + assert not (tmp_path / "tasks").exists() + + +def test_calibration_is_per_workload_and_covers_all_profiles( + calibrated: ModelCapture, tmp_path: Path +) -> None: + contracts = read_json(calibrated.root / "calibration.json") + contract = next(iter(contracts.values())) + assert len(contract["cases"]) == 2 + recorded_bounds = [case["recorded:0"][0]["atol"] for case in contract["cases"].values()] + assert max(recorded_bounds) > 32 * min(recorded_bounds) + for case in contract["cases"].values(): + assert set(case) == {f"{name}:{seed}" for name, seed in PROFILES} + tasks = render_model(calibrated, tmp_path / "tasks") + task = next(task for task in tasks if "numerics.json" in task.hashes) + assert task.settings.precision.mode == "tolerance" + assert read_json(task.root / "numerics.json") == contract + with pytest.raises(ContractError, match="new destination"): + calibrate(calibrated, calibrated.root) + + +@pytest.mark.parametrize("override", [{"precision.atol": 1.0}, {"seeds": [0, 9]}]) +def test_generation_rejects_calibration_overrides( + calibrated: ModelCapture, tmp_path: Path, override: dict +) -> None: + with pytest.raises(ContractError, match="cannot.*override|cannot be overridden"): + render_model(calibrated, tmp_path / "tasks", overrides=override) + + +@pytest.mark.parametrize("edit", ["threshold", "profile", "dtype", "flags"]) +def test_task_loader_rejects_calibration_edits(calibrated: ModelCapture, tmp_path: Path, edit: str) -> None: + task = next(t for t in render_model(calibrated, tmp_path / "tasks") if "numerics.json" in t.hashes) + contract = read_json(task.root / "numerics.json") + samples = next(iter(contract["cases"].values())) + if edit == "threshold": + samples["recorded:0"][0]["atol"] *= 100 + elif edit == "profile": + del samples["zeros:0"] + elif edit == "dtype": + samples["recorded:0"][0]["dtype"] = "bfloat16" + else: + contract["flags"]["fp16_reduced_reduction"] = not contract["flags"]["fp16_reduced_reduction"] + atomic_json(task.root / "numerics.json", contract) + with pytest.raises(ContractError, match="Calibration|calibration"): + TaskPackage.load(task.root) + + +def test_real_inputs_cannot_silently_fall_back_to_random_when_budget_exhausted( + tmp_path: Path, registry_capture +) -> None: + x = torch.ones(8, 8) + capture = registry_capture( + tmp_path, + [("gelu_tanh", (torch.ones(2, dtype=torch.float64),), {}), ("linear", (x, x), {})], + budget=32, + ) + assert any("fixture" in error for error in capture.manifest["errors"]) + with pytest.raises(ContractError, match="incomplete"): + calibrate(capture, tmp_path / "calibrated") + + +def test_every_element_must_pass_and_nan_or_wrong_dtype_fail() -> None: + ref = torch.tensor([1.0, 0.0, -1.0], dtype=torch.float64) + baseline = ref.float() + bound = derive_bound(baseline, ref) + allowed = baseline.clone() + allowed[0] += torch.finfo(torch.float32).eps + assert not torch.equal(allowed, baseline) + compare_calibrated([allowed], [ref], [bound]) + for value in (torch.tensor([1.0, 1e-3, -1.0]), baseline + float("nan"), baseline.double()): + with pytest.raises(ContractError): + compare_calibrated([value], [ref], [bound]) + with pytest.raises(ContractError, match="Nonfinite"): + derive_bound(baseline + float("inf"), ref) + + +def test_calibration_restores_global_precision_on_failure() -> None: + before = environment("cpu") + flags = { + name: before[name] + for name in ( + "cudnn_allow_tf32", + "fp16_reduced_reduction", + "bf16_reduced_reduction", + "deterministic_algorithms", + ) + } + with ( + pytest.raises(RuntimeError), + captured_settings({"matmul_precision": "high", "allow_tf32": True}, flags), + ): + raise RuntimeError("broken baseline") + assert environment("cpu") == before + + +def test_replay_uses_calibrated_bounds_and_all_profiles( + calibrated: ModelCapture, tmp_path: Path, monkeypatch +) -> None: + task = next(t for t in render_model(calibrated, tmp_path / "tasks") if "numerics.json" in t.hashes) + adapter = object.__new__(ReplayAdapter) + adapter.task, adapter.cfg, adapter.device = task, task.settings, "cpu" + adapter.definition = SimpleNamespace( + outputs=task.definition["outputs"], torch_output_dtypes=[torch.float32] + ) + adapter.replay = read_json(task.root / "replay.json") + adapter.numerical = read_json(task.root / "numerics.json") + adapter._check_flags = lambda: None + adapter._seed = lambda seed: None + monkeypatch.setattr(torch.cuda, "synchronize", lambda *args: None) + operator = adapter.replay["numerical_identity"] + + class Runnable: + metadata = SimpleNamespace(destination_passing_style=False) + + def __init__(self, fn): + self.fn = fn + + def __call__(self, *args): + return self.fn(list(args))[0] + + reference = Runnable(lambda args: production(operator, args)) + more_accurate = Runnable(lambda args: [high_precision(operator, args)[0].float()]) + workload = SimpleNamespace(uuid=task.workload_ids[0]) + for seed in task.settings.seeds: + inputs = adapter._inputs(workload, seed) + adapter._verify_call(more_accurate, reference, inputs, seed) + + def capture_graph(runnable, inputs): + output = runnable(*inputs).clone() + + def replay(): + output.copy_(runnable(*inputs)) + + return SimpleNamespace(replay=replay), inputs, [output] + + adapter._capture = capture_graph + for seed in task.settings.seeds: + adapter._verify_graph(more_accurate, reference, workload, seed) + saved = production(operator, adapter._inputs(workload, 0)) + with pytest.raises(ContractError, match="frozen calibrated"): + adapter._verify_call(Runnable(lambda args: saved), reference, adapter._inputs(workload, 0), 0) + with pytest.raises(ContractError, match="frozen calibrated"): + adapter._verify_graph(Runnable(lambda args: saved), reference, workload, 0) + assert adapter._check_metadata()["numerical_contract"] == task.hashes["numerics.json"] + assert adapter._check_metadata()["validation_profiles"] == [list(pair) for pair in PROFILES] + assert asdict(task.settings.precision)["mode"] == "tolerance" + + +def test_random_profiles_preserve_masks_and_change_real_values(tmp_path: Path, registry_capture) -> None: + q = torch.ones(1, 1, 3, 4) + prefix = torch.ones(1, 1, 2, 4) + suffix = torch.ones(1, 1, 1, 4) + mask = torch.tensor([[[[0.0, -1.0, 0.0]]]]) + capture = registry_capture( + tmp_path, [("split_kv_attention", (q, prefix, prefix, suffix, suffix, mask, 0.5), {})] + ) + (op,) = capture.operators + case = op["workloads"][0] + real = validation_inputs(op["identity"], case, capture.root, "recorded", 0) + random = validation_inputs(op["identity"], case, capture.root, "random", 1) + assert not torch.equal(real[0], random[0]) + assert torch.equal(real[-1], random[-1]) + expected = production(op["identity"], [value.double() for value in real]) + torch.testing.assert_close(high_precision(op["identity"], real)[0], expected[0], atol=0, rtol=0) + + +def test_calibrate_cli_writes_a_new_capture_before_generation( + tmp_path: Path, capsys, registry_capture +) -> None: + from scripts.kernel_tuning.__main__ import main + + value = torch.ones(2, 3) + weight = torch.ones(3, 4) + raw = registry_capture(tmp_path, [("linear", (value, weight.T), {})]) + destination = tmp_path / "cli-calibrated" + assert main(["calibrate", "--captured", str(raw.root), "--output", str(destination)]) == 0 + assert "identity" in capsys.readouterr().out + (task,) = render_model(ModelCapture.load(destination), tmp_path / "tasks") + assert task.settings.precision.mode == "tolerance" + + +def test_promotion_requires_all_profile_and_contract_evidence( + calibrated: ModelCapture, tmp_path: Path +) -> None: + from scripts.kernel_tuning.contracts import validate_measurement + from test_batch import measurement + + task = next(t for t in render_model(calibrated, tmp_path / "tasks") if "numerics.json" in t.hashes) + report = measurement(task, 8.0) + with pytest.raises(ContractError, match="calibration/profile"): + validate_measurement(task, report) + for check in report["checks"]: + check.update( + numerical_contract=task.hashes["numerics.json"], + validation_profiles=[list(pair) for pair in PROFILES], + ) + validate_measurement(task, report) + report["checks"][0]["validation_profiles"].pop() + with pytest.raises(ContractError, match="calibration/profile"): + validate_measurement(task, report) + + +def test_changed_oracle_source_invalidates_calibration( + calibrated: ModelCapture, tmp_path: Path, monkeypatch +) -> None: + from scripts.kernel_tuning.discovery import numerics + + task = next(t for t in render_model(calibrated, tmp_path / "tasks") if "numerics.json" in t.hashes) + monkeypatch.setattr(numerics, "implementation_digest", lambda: "changed") + with pytest.raises(ContractError, match="recalibrate"): + TaskPackage.load(task.root) diff --git a/tests/kernel_tuning/test_capture.py b/tests/kernel_tuning/test_capture.py new file mode 100644 index 0000000..4b7ff0d --- /dev/null +++ b/tests/kernel_tuning/test_capture.py @@ -0,0 +1,255 @@ +"""Shape capture with duck-typed tensors and fake kernel modules; no Torch required.""" + +from __future__ import annotations + +import json +import sys +import types +from pathlib import Path + +import pytest +from scripts.kernel_tuning import capture +from scripts.kernel_tuning.__main__ import main +from scripts.kernel_tuning.contracts import ContractError, TaskPackage +from scripts.kernel_tuning.generate import render +from scripts.kernel_tuning.operators import Operator, Tensor, catalog + + +class FakeTensor: + """Just the metadata the mappings read: shape, dtype, numel, and values.""" + + def __init__(self, *shape: int, dtype: str = "torch.bfloat16", values: list[int] | None = None) -> None: + self.shape = shape + self.dtype = dtype + self.is_cuda = False + self._values = values + + def numel(self) -> int: + count = 1 + for size in self.shape: + count *= size + return count + + def tolist(self) -> list[int]: + assert self._values is not None + return self._values + + +def record(name: str, **arguments: object) -> list[dict]: + recorder = capture.Recorder() + recorder.record(name, arguments) + return recorder.rows() + + +def test_every_catalog_operator_has_a_mapping() -> None: + assert capture.MAPPINGS.keys() == catalog().keys() + + +def test_production_calls_become_task_cases() -> None: + (row,) = record( + "gated_residual", + residual=FakeTensor(8, 50, 1024), + update=FakeTensor(8, 50, 1024), + gate=FakeTensor(8, 1, 1024), + ) + assert row == {"operator": "gated_residual", "axes": {"B": 8, "T": 50, "H": 1024, "G": 1}, "count": 1} + (row,) = record("swiglu", packed=FakeTensor(1, 1, 37888)) + assert row["axes"] == {"M": 1, "I": 18944} + (row,) = record( + "split_kv_attention", + query=FakeTensor(1, 8, 50, 256), + prefix_key=FakeTensor(1, 1, 560, 256), + prefix_value=FakeTensor(1, 1, 560, 256), + suffix_key=FakeTensor(1, 1, 50, 256), + suffix_value=FakeTensor(1, 1, 50, 256), + mask=FakeTensor(1, 1, 1, 610, dtype="torch.float32"), + scaling=0.0625, + ) + assert row["axes"] == {"B": 1, "HQ": 8, "HK": 1, "P": 560, "S": 50, "D": 256} + assert row["scalars"] == {"scaling": 0.0625} + + +def test_segment_offsets_are_captured_as_fixed_inputs_and_unsupported_calls_skipped() -> None: + q = FakeTensor(96, 16, 80) + offsets = FakeTensor(3, dtype="torch.int32", values=[0, 64, 96]) + common = {"q": q, "k": q, "v": q, "q_segment_offsets": offsets, "kv_segment_offsets": None} + (row,) = record("segmented_attention", **common, scaling=None, max_query_length=64, max_key_length=64) + assert row["axes"] == {"T": 96, "HQ": 16, "HK": 16, "D": 80, "N": 3} + assert row["scalars"] == {"max_length": 64} and row["fixed"] == {"segment_offsets": [0, 64, 96]} + (row,) = record("segmented_attention", **common, scaling=0.5, max_query_length=64, max_key_length=64) + assert row == {"operator": "segmented_attention", "skipped": "non-default scaling", "count": 1} + (row,) = record( + "gated_gelu", gate=FakeTensor(1, 50, 4096, dtype="torch.float16"), up=FakeTensor(1, 50, 4096) + ) + assert row["skipped"] == "dtype torch.float16" + (row,) = record("gated_gelu", gate=FakeTensor(1, 50, 4096)) + assert row["skipped"].startswith("unmapped call: KeyError") + + +@pytest.fixture +def fake_kernel(monkeypatch: pytest.MonkeyPatch) -> Operator: + """A kernel module plus a second module that imported the function by name.""" + module = types.ModuleType("fake_kernels.norm") + + def scale(x, factor=2.0): + return [factor * v for v in x.values] + + module.scale = scale + importer = types.ModuleType("fake_model") + importer.scale = scale + monkeypatch.setitem(sys.modules, "fake_kernels", types.ModuleType("fake_kernels")) + monkeypatch.setitem(sys.modules, "fake_kernels.norm", module) + monkeypatch.setitem(sys.modules, "fake_model", importer) + monkeypatch.setitem( + capture.MAPPINGS, "fake", lambda a: {"axes": {"N": len(a["x"].values)}, "scalars": {"f": a["factor"]}} + ) + return Operator( + name="fake", + summary="", + models=(), + op_type="elementwise", + source="fake_kernels/norm.py", + entry="scale", + call="scale(x)", + axes={"N": None}, + inputs={"x": Tensor(("N",), "bfloat16")}, + outputs={"y": Tensor(("N",), "bfloat16")}, + reference="def run(x):\n return x\n", + workloads=(), + notes="", + ) + + +def test_run_wraps_every_holder_records_calls_and_restores(fake_kernel: Operator, tmp_path: Path) -> None: + script = tmp_path / "infer.py" + script.write_text( + "import sys, types, fake_model\n" + "from fake_kernels.norm import scale\n" + "x = types.SimpleNamespace(values=[1, 2, 3])\n" + "assert fake_model.scale(x) == [2.0, 4.0, 6.0]\n" + "scale(x, factor=3.0); scale(x, factor=3.0)\n" + "assert sys.argv[1:] == ['--steps', '2']\n", + encoding="utf-8", + ) + rows = capture.run([str(script), "--steps", "2"], tmp_path / "calls.jsonl", [fake_kernel]) + assert rows == [ + {"operator": "fake", "axes": {"N": 3}, "scalars": {"f": 3.0}, "count": 2}, + {"operator": "fake", "axes": {"N": 3}, "scalars": {"f": 2.0}, "count": 1}, + ] + assert capture.load(tmp_path / "calls.jsonl") == rows + original = sys.modules["fake_kernels.norm"].scale + assert not hasattr(original, "__wrapped_kernel__") and sys.modules["fake_model"].scale is original + + +def test_captured_cases_replace_estimated_workloads(tmp_path: Path) -> None: + rows = [ + {"operator": "split_kv_attention", "axes": axes, "scalars": {"scaling": 0.0625}, "count": count} + for axes, count in ( + ({"B": 1, "HQ": 8, "HK": 1, "P": 560, "S": 50, "D": 256}, 1800), + ({"B": 1, "HQ": 8, "HK": 1, "P": 592, "S": 50, "D": 256}, 180), + ) + ] + rows.append({"operator": "split_kv_attention", "skipped": "dtype torch.float16", "count": 3}) + operator, summary = capture.apply(catalog()["split_kv_attention"], rows) + assert [w.uuid for w in operator.workloads] == [ + "captured1-B1-HQ8-HK1-P560-S50-D256", + "captured2-B1-HQ8-HK1-P592-S50-D256", + ] + assert [w.weight for w in operator.workloads] == [1800.0, 180.0] + assert summary.endswith("2 captured shapes covering 100% of 1980 calls; 3 calls outside the task skipped") + task = render(operator, tmp_path / "task") + assert task.settings.weights == {w.uuid: w.weight for w in operator.workloads} + assert "captured production calls" in (task.root / "README.md").read_text(encoding="utf-8") + kept, note = capture.apply(catalog()["swiglu"], rows) + assert kept == catalog()["swiglu"] and "keeping estimated shapes" in note + with pytest.raises(ContractError, match="binds"): + capture.apply(catalog()["swiglu"], [{"operator": "swiglu", "axes": {"M": 1}, "count": 1}]) + + +def test_captured_segment_offsets_render_as_task_data(tmp_path: Path) -> None: + rows = [ + { + "operator": "segmented_attention", + "axes": {"T": 96, "HQ": 16, "HK": 16, "D": 80, "N": 3}, + "scalars": {"max_length": 64}, + "fixed": {"segment_offsets": [0, 64, 96]}, + "count": 5, + } + ] + operator, _ = capture.apply(catalog()["segmented_attention"], rows) + task = render(operator, tmp_path / "task") + assert "data/captured1-T96-HQ16-HK16-D80-N3.safetensors" in task.hashes + + +def test_cli_capture_and_generate_from_capture( + fake_kernel: Operator, tmp_path: Path, monkeypatch, capsys +) -> None: + script = tmp_path / "infer.py" + script.write_text( + "import types\nfrom fake_kernels.norm import scale\nscale(types.SimpleNamespace(values=[1]))\n" + ) + monkeypatch.setattr( + "scripts.kernel_tuning.operators.select", lambda names=None, models=None: [fake_kernel] + ) + output = tmp_path / "calls.jsonl" + assert main(["capture", "--catalog", "--output", str(output), "--", str(script)]) == 0 + assert "Captured 1 calls in 1 cases" in capsys.readouterr().out + monkeypatch.undo() + rows = [{"operator": "rms_norm", "axes": {"M": 1, "H": 3584}, "scalars": {"eps": 1e-6}, "count": 7}] + captured = tmp_path / "rms.jsonl" + captured.write_text("".join(json.dumps(row) + "\n" for row in rows), encoding="utf-8") + argv = [ + "generate", + "--catalog", + "rms_norm", + "--output", + str(tmp_path / "tasks"), + "--captured", + str(captured), + ] + assert main(argv) == 0 + assert "rms_norm: 1 captured shapes covering 100% of 7 calls" in capsys.readouterr().out + assert TaskPackage.load(tmp_path / "tasks/rms_norm").workload_ids == ("captured1-M1-H3584",) + + +def test_calls_captured_in_a_cuda_graph_count_once_per_replay() -> None: + recorder = capture.Recorder() + arguments = {"packed": FakeTensor(1, 1, 37888)} + recorder.record("swiglu", arguments) # eager warmup before capture + recorder.begin_graph() + recorder.record("swiglu", arguments) + recorder.record("swiglu", arguments) + recorder.end_graph(7) + assert recorder.rows()[0]["count"] == 1 + for _ in range(3): + recorder.replay(7) + recorder.replay(8) # an unrelated graph contributes nothing + assert recorder.rows() == [{"operator": "swiglu", "axes": {"M": 1, "I": 18944}, "count": 7}] + + +def test_committed_captures_replace_estimates_by_default() -> None: + operators, notes = capture.apply_file([catalog()["split_kv_attention"], catalog()["segmented_attention"]]) + split_kv, segmented = operators + assert {w.axes["P"] for w in split_kv.workloads} == {576, 560, 0} + assert all(w.uuid.startswith("captured") for w in segmented.workloads) + assert all(w.fixed["segment_offsets"][-1] == w.axes["T"] for w in segmented.workloads) + assert len(notes) == 2 + kept, notes = capture.apply_file([catalog()["swiglu"]], estimated=True) + assert kept == [catalog()["swiglu"]] and notes == ["using the catalog's estimated shapes"] + + +def test_skeleton_reads_only_the_safetensors_header(monkeypatch) -> None: + from scripts.kernel_tuning import skeleton + from scripts.kernel_tuning.generate import safetensors_bytes + + data = safetensors_bytes({"layer.weight": ("int32", (1, 2, 3, 4))}) + requested = [] + + def request(url: str, *, start=None, end=None) -> bytes: + requested.append((start, end)) + return data[start : end + 1] + + monkeypatch.setattr(skeleton, "_request", request) + header = skeleton.read_header("org/model", "abc", "model.safetensors") + assert header == {"layer.weight": {"dtype": "I32", "shape": [4], "data_offsets": [0, 16]}} + assert requested[0] == (0, 7) and requested[1][1] < len(data) - 16 diff --git a/tests/kernel_tuning/test_cli.py b/tests/kernel_tuning/test_cli.py new file mode 100644 index 0000000..e7e9ef7 --- /dev/null +++ b/tests/kernel_tuning/test_cli.py @@ -0,0 +1,59 @@ +"""Read-only developer commands remain usable without Humanize2 or a GPU.""" + +from __future__ import annotations + +import json +import sys +from pathlib import Path + +import pytest +from scripts.kernel_tuning.__main__ import main + + +def test_check_and_status_do_not_import_an_agent(task_dir: Path, store, monkeypatch, capsys) -> None: + monkeypatch.setitem(sys.modules, "hmz.sdk", None) + assert main(["check", str(task_dir)]) == 0 + assert json.loads(capsys.readouterr().out)["validation"].startswith("structural") + assert main(["status", str(store.root)]) == 0 + assert json.loads(capsys.readouterr().out)["status"] == "created" + + +def test_run_requires_an_explicit_agent_before_creating_artifacts( + task_dir: Path, tmp_path: Path, capsys +) -> None: + output = tmp_path / "unused" + assert main(["run", str(task_dir), "--output", str(output)]) == 2 + assert "explicit Humanize2 agent" in capsys.readouterr().err + assert not output.exists() + + +def test_bad_yaml_is_a_clear_static_validation_error(task_dir: Path, capsys) -> None: + (task_dir / "tuning.yaml").write_text("language: [oops\n", encoding="utf-8") + assert main(["check", str(task_dir)]) == 2 + assert "Invalid tuning.yaml" in capsys.readouterr().err + + +def test_removed_full_model_capture_is_rejected_before_running_driver(tmp_path, capsys): + output = tmp_path / "capture" + with pytest.raises(SystemExit) as exc: + main(["capture", "--model", "pi05", "--output", str(output), "--", "unused.py"]) + assert exc.value.code == 2 + assert not output.exists() + assert main(["generate", "--model", "pi05"]) == 2 + assert "run capture-ops --model first" in capsys.readouterr().err + + +def test_catalog_capture_keeps_its_explicit_entrypoint(tmp_path, monkeypatch, capsys): + from scripts.kernel_tuning import capture + + calls = [] + + def record(command, output, operators): + calls.append((command, output, [op.name for op in operators])) + return [{"operator": "gated_residual", "axes": {}, "count": 3}] + + monkeypatch.setattr(capture, "run", record) + output = tmp_path / "calls.jsonl" + assert main(["capture", "--catalog", "--output", str(output), "gated_residual", "--", "driver.py"]) == 0 + assert calls == [(["driver.py"], output, ["gated_residual"])] + assert "Captured 3 calls" in capsys.readouterr().out diff --git a/tests/kernel_tuning/test_contracts.py b/tests/kernel_tuning/test_contracts.py new file mode 100644 index 0000000..7eda879 --- /dev/null +++ b/tests/kernel_tuning/test_contracts.py @@ -0,0 +1,197 @@ +"""Task validation and promotion rules without agents, CUDA, or source execution.""" + +from __future__ import annotations + +import copy +from dataclasses import asdict +from pathlib import Path +from typing import Any + +import pytest +import yaml +from scripts.kernel_tuning.contracts import ( + Acceptance, + ContractError, + Precision, + Settings, + TaskPackage, + Timing, + parse_json, + promotion, + read_json, + relative_path, + validate_measurement, + validate_solution, +) + + +def measurement(task: TaskPackage, latency: float = 8) -> dict[str, Any]: + """Generate complete synthetic evidence across the task's declared cases.""" + roles = ("baseline", "incumbent", "candidate") + return { + "checks": [ + { + "workload": w, + "role": role, + "passed": True, + "seeds": list(task.settings.seeds), + "precision": asdict(task.settings.precision), + "max_abs_error": 0, + "max_rel_error": 0, + "graph_replay": True, + } + for w in task.workload_ids + for role in roles + ], + "rounds": [ + { + w: {r: [latency if r == "candidate" else 10] * task.settings.timing.trials for r in roles} + for w in task.workload_ids + } + for _ in range(task.settings.timing.paired_rounds) + ], + } + + +def test_task_validation_does_not_execute_reference(task_dir: Path) -> None: + definition = read_json(task_dir / "definition.json") + definition["reference"] = "raise RuntimeError('not imported')\n" + definition["reference"] + from scripts.kernel_tuning.artifacts import atomic_json + + atomic_json(task_dir / "definition.json", definition) + task = TaskPackage.load(task_dir) + assert len(task.workload_ids) == 3 + + +@pytest.mark.parametrize( + "path", ["../a", "/tmp/x", "C:/x", "a\\b", "a/../b", "a//b", "NUL", "a/CON.txt", "a."] +) +def test_unsafe_or_nonportable_paths_fail(path: str) -> None: + with pytest.raises(ContractError): + relative_path(path) + + +@pytest.mark.parametrize("source", ['{"axes":{},"axes":{}}', '{"value":NaN}', '{"value":Infinity}']) +def test_all_json_records_reject_duplicate_keys_and_nonfinite_values(source: str) -> None: + with pytest.raises(ContractError): + parse_json(source) + + +def test_task_changes_are_detected(store) -> None: + with (store.task.root / "baseline.py").open("a", encoding="utf-8") as handle: + handle.write("\n# changed\n") + with pytest.raises(ContractError, match="changed"): + store.verify() + + +def test_duplicate_workloads_and_unknown_settings_fail(task_dir: Path) -> None: + workloads = task_dir / "workloads.jsonl" + text = workloads.read_text(encoding="utf-8") + workloads.write_text(text + text.splitlines()[0] + "\n", encoding="utf-8") + with pytest.raises(ContractError, match="distinct"): + TaskPackage.load(task_dir) + workloads.write_text(text, encoding="utf-8") + with (task_dir / "tuning.yaml").open("a", encoding="utf-8") as handle: + handle.write("typo: true\n") + with pytest.raises(ContractError, match="Unknown"): + TaskPackage.load(task_dir) + + +@pytest.mark.parametrize( + "config", + [ + {"language": "python"}, + {"language": "triton", "seeds": ()}, + {"language": "cuda", "weights": {"one": 0}}, + ], +) +def test_invalid_settings_fail(config: dict) -> None: + with pytest.raises(ContractError): + Settings(**config) + + +def test_precision_and_remeasurement_cannot_be_silently_relaxed() -> None: + with pytest.raises(ContractError): + Precision(atol=1e-2) + with pytest.raises(ContractError): + Precision(mode="tolerance", atol=1e-2) + assert Precision(mode="tolerance", atol=1e-4, reason="fixed reduction ordering allowance") + with pytest.raises(ContractError): + Timing(paired_rounds=1) + with pytest.raises(ContractError): + Acceptance(min_improvement=0) + + +def test_solution_multi_file_and_frozen_language(store, candidate: dict) -> None: + candidate["sources"].append({"path": "include/helpers.h", "content": "// helper"}) + assert validate_solution(candidate, store.task) == candidate + candidate["spec"]["language"] = "cuda" + with pytest.raises(ContractError, match="language"): + validate_solution(candidate, store.task) + + +@pytest.mark.parametrize( + "mutation", + [ + lambda r: r["checks"].pop(), + lambda r: r["checks"].append(r["checks"][0]), + lambda r: r["checks"][0].update(passed="true"), + lambda r: r["checks"][0].update(max_abs_error=0.001), + lambda r: r["checks"][0].update(seeds=[99]), + lambda r: r["rounds"].pop(), + lambda r: r["rounds"][0].pop("broadcast-small"), + lambda r: r["rounds"][0]["broadcast-small"]["candidate"].__setitem__(0, float("nan")), + lambda r: r["rounds"][0]["broadcast-small"]["candidate"].__setitem__(0, 0), + ], +) +def test_incomplete_or_malformed_evidence_never_promotes(store, mutation) -> None: + report = measurement(store.task) + mutation(report) + with pytest.raises(ContractError): + promotion(store.task, report) + + +def test_gain_must_repeat_and_each_workload_must_pass(store) -> None: + report = measurement(store.task) + assert promotion(store.task, report)["promoted"] + report["rounds"][1]["broadcast-small"]["candidate"] = [10.6] * 3 + decision = promotion(store.task, report) + assert not decision["promoted"] + assert any("regresses" in r for r in decision["reasons"]) + report = measurement(store.task, 9.8) + assert not promotion(store.task, report)["promoted"] + + +def test_incumbent_and_initial_baseline_both_constrain_promotion(store) -> None: + report = measurement(store.task, 8) + for row in report["rounds"][1].values(): + row["incumbent"] = [8.1] * 3 + assert not promotion(store.task, report)["promoted"] + + +def test_explicit_weights_change_the_mean_not_the_regression_guard(task_dir: Path) -> None: + path = task_dir / "tuning.yaml" + config = yaml.safe_load(path.read_text()) + config["weights"] = {"broadcast-small": 100, "broadcast-wide": 1, "token-gate-tail": 1} + path.write_text(yaml.safe_dump(config)) + task = TaskPackage.load(task_dir) + report = measurement(task) + for round_ in report["rounds"]: + round_["token-gate-tail"]["candidate"] = [10.4] * 3 + assert promotion(task, report)["promoted"] + bad = copy.deepcopy(report) + for round_ in bad["rounds"]: + round_["token-gate-tail"]["candidate"] = [10.6] * 3 + assert not promotion(task, bad)["promoted"] + + +def test_graph_evidence_must_include_changed_inputs(task_dir: Path) -> None: + path = task_dir / "tuning.yaml" + config = yaml.safe_load(path.read_text()) + config["timing"]["mode"] = "cuda_graph" + path.write_text(yaml.safe_dump(config)) + task = TaskPackage.load(task_dir) + report = measurement(task) + report["checks"][0]["graph_replay"] = False + with pytest.raises(ContractError, match="replay"): + validate_measurement(task, report) diff --git a/tests/kernel_tuning/test_drive.py b/tests/kernel_tuning/test_drive.py new file mode 100644 index 0000000..c327d21 --- /dev/null +++ b/tests/kernel_tuning/test_drive.py @@ -0,0 +1,54 @@ +"""The shared inference driver remains usable by profiling, capture-ops and verify.""" + +from __future__ import annotations + +import json + +import pytest + +torch = pytest.importorskip("torch") + + +def test_driver_replays_recorded_observations_through_the_engine(tmp_path, monkeypatch): + from scripts.kernel_tuning.discovery.drive import main, save_inputs + + from embodiinfer.engine import EngineCore + from embodiinfer.policies.mock import preset_config + from embodiinfer.types import Observation + + cfg = preset_config("tiny") + + def observation(): + return Observation( + images=torch.rand(cfg.num_cameras, 3, cfg.image_size, cfg.image_size), + state=torch.rand(cfg.state_dim), + instruction_tokens=torch.randint(0, cfg.vocab_size, (cfg.max_lang_len,)), + ) + + inputs = tmp_path / "inputs.pt" + save_inputs(inputs, "observations", [[observation(), observation()], [observation()]]) + config = tmp_path / "serving.json" + config.write_text(json.dumps({"policy_kwargs": {"preset": "tiny"}}), encoding="utf-8") + original_execute = EngineCore.execute + calls = [] + + def execute(self, *args, **kwargs): + result = original_execute(self, *args, **kwargs) + calls.append(result) + return result + + monkeypatch.setattr(EngineCore, "execute", execute) + main( + [ + "--policy", + "mock_flow_vla", + "--device", + "cpu", + "--no-cuda-graph", + "--adapter-config", + str(config), + "--inputs", + str(inputs), + ] + ) + assert len(calls) == 3 diff --git a/tests/kernel_tuning/test_flow.py b/tests/kernel_tuning/test_flow.py new file mode 100644 index 0000000..aed5405 --- /dev/null +++ b/tests/kernel_tuning/test_flow.py @@ -0,0 +1,243 @@ +"""Run the real Humanize2 engine on its official fake agents and CPU evaluators.""" + +from __future__ import annotations + +import asyncio +import copy +import os +import tempfile +from pathlib import Path +from typing import Any + +import pytest + +pytest.importorskip("hmz", reason="requires the optional Humanize2 tooling environment") +from hmz.sdk import fakes # noqa: E402 +from scripts.kernel_tuning.artifacts import RunStore, atomic_json, run_lock # noqa: E402 +from scripts.kernel_tuning.contracts import ContractError, read_json # noqa: E402 +from scripts.kernel_tuning.flow import kernel_tuning # noqa: E402 +from scripts.kernel_tuning.runner import evaluate # noqa: E402 + + +@pytest.fixture(autouse=True) +def windows_journal_io(monkeypatch) -> None: + """Adapt only hmz's POSIX atomic-file primitive for Windows CPU tests. + + The real framework state, journal reading, sessions, budgets and resume run + unchanged. Production run/resume explicitly requires POSIX; this shim is not + shipped as part of the workflow and does not emulate agent/kernel execution. + """ + if os.name == "nt": + from hmz.coganchor import atomic + + def writes(at: Path, said, *, mode=None) -> None: + fd, name = tempfile.mkstemp(dir=at.parent) + try: + with os.fdopen(fd, "wb") as handle: + if isinstance(said, str): + handle.write(said.encode("utf-8")) + elif isinstance(said, bytes): + handle.write(said) + else: + handle.writelines(said) + handle.flush() + os.fsync(handle.fileno()) + os.replace(name, at) + finally: + Path(name).unlink(missing_ok=True) + + monkeypatch.setattr(atomic, "writes", writes) + + +def fake_coder(store: RunStore, candidate: dict, latencies: list[Any]): + """Write candidate files only when a fake agent is asked to implement its plan.""" + + def answer(prompt: str, **kwargs: Any) -> str: + scratch = Path(str(kwargs["session"].placement.workdir)) + if "Write PLAN.md" in prompt: + (scratch / "PLAN.md").write_text( + "Hypothesis: fewer launches. Validate every workload.", encoding="utf-8" + ) + return "plan written" + value = latencies.pop(0) + if isinstance(value, BaseException): + raise value + solution = copy.deepcopy(candidate) + solution["name"] += "_" + scratch.name + solution["description"] = str(value) + atomic_json(scratch / "solution.json", solution) + return "solution written" + + return fakes.FakeAgentDriver(reply=answer) + + +def run_fake(store: RunStore, coder, *, resume: bool = False) -> dict: + """Exercise hmz journaling/session behavior; no real harness is constructed.""" + return asyncio.run( + fakes.run_fake( + kernel_tuning, + "test only", + agents={"coder": coder}, + local=fakes.FakeEnvDriver(workdir=(store.root / "workspace").as_posix()), + params={"run_dir": str(store.root)}, + journal=store.root / "fake-hmz.jsonl", + resume=resume, + ) + ) + + +def test_full_loop_promotes_then_preserves_failures_and_exports(store, candidate, tmp_path: Path) -> None: + coder = fake_coder(store, candidate, [8, "crash", 9]) + result = run_fake(store, coder) + assert result["status"] == "completed" + assert result["best"] == "0001" + assert [a["status"] for a in result["attempts"]] == ["promoted", "failed", "rejected"] + assert len(coder.prompts) == 6 + assert coder.peak == 1 + assert coder.live == 0 + saved = RunStore(store.root) + assert (saved.root / "candidates/0001/sources/kernel.py").read_bytes() == candidate["sources"][0][ + "content" + ].encode("utf-8") + assert read_json(saved.root / "best.json")["candidate"] == "0001" + exported = saved.export(tmp_path / "export") + assert (exported / "sources/kernel.py").read_text() == candidate["sources"][0]["content"] + assert (exported / "evidence/evaluation/result.json").exists() + assert (exported / "task/definition.json").exists() + assert read_json(exported / "checksums.json") + with pytest.raises(ContractError, match="exists"): + saved.export(exported) + + +def test_rejected_source_is_available_to_the_next_planning_turn(store, candidate) -> None: + """A fresh agent can inspect rejected code without mistaking it for the incumbent.""" + witnessed = [] + + def answer(prompt: str, **kwargs: Any) -> str: + scratch = Path(str(kwargs["session"].placement.workdir)) + if "Write PLAN.md" in prompt: + if scratch.name == "0002": + rejected = read_json(scratch / "feedback.json")[0] + assert rejected["status"] == "rejected" + assert rejected["solution"]["name"] == candidate["name"] + "_0001" + assert rejected["solution"]["sources"] == candidate["sources"] + assert read_json(scratch / "incumbent.json")["name"] != rejected["solution"]["name"] + witnessed.append(rejected["id"]) + (scratch / "PLAN.md").write_text("Revise the previous implementation.", encoding="utf-8") + else: + solution = copy.deepcopy(candidate) + solution["name"] += "_" + scratch.name + solution["description"] = {"0001": "12", "0002": "8", "0003": "9"}[scratch.name] + atomic_json(scratch / "solution.json", solution) + return "ready" + + result = run_fake(store, fakes.FakeAgentDriver(reply=answer)) + assert witnessed == ["0001"] + assert result["best"] == "0002" + + +def test_resume_counts_interrupted_attempt_and_keeps_elapsed_budget(store, candidate) -> None: + coder = fake_coder(store, candidate, [8, asyncio.CancelledError()]) + with pytest.raises(asyncio.CancelledError): + run_fake(store, coder) + interrupted = RunStore(store.root) + assert interrupted.manifest["status"] == "interrupted" + assert interrupted.manifest["best"] == "0001" + assert len(interrupted.manifest["attempts"]) == 2 + elapsed = interrupted.manifest["elapsed_seconds"] + second = fake_coder(interrupted, candidate, [7]) + resumed = run_fake(interrupted, second, resume=True) + assert resumed["best"] == "0003" + assert len(resumed["attempts"]) == 3 + assert resumed["elapsed_seconds"] > elapsed + assert len(second.prompts) == 2 + + +def test_early_agent_failure_counts_toward_patience(store, candidate) -> None: + coder = fake_coder( + store, candidate, [RuntimeError("agent unavailable"), RuntimeError("agent unavailable")] + ) + result = run_fake(store, coder) + assert result["stop_reason"] == "patience" + assert result["best"] is None + assert len(result["attempts"]) == 2 + with pytest.raises(ContractError, match="No candidate"): + RunStore(store.root).export(store.root.parent / "export") + + +def test_worker_timeout_is_failure_and_cleanup_finishes(store, candidate) -> None: + path = store.root / "test-solution.json" + candidate["description"] = "timeout" + atomic_json(path, candidate) + with pytest.raises(ContractError, match="timed out"): + asyncio.run(evaluate(store, store.root / "timeout", candidate=path, timeout=1)) + assert not list((store.root / "timeout").glob("result.json")) + + +def test_environment_change_refuses_measurement_before_agent_turn(store, candidate) -> None: + store.manifest["environment"] = {"identity": {"hardware": "different"}} + store.save() + coder = fake_coder(store, candidate, []) + with pytest.raises(ContractError, match="environment changed"): + run_fake(store, coder) + assert not coder.prompts + + +def test_run_lock_refuses_concurrent_writer(store) -> None: + with run_lock(store.root), pytest.raises(ContractError, match="Another process"), run_lock(store.root): + pass + + +def test_export_rejects_mutated_winning_source(store, candidate, tmp_path: Path) -> None: + run_fake(store, fake_coder(store, candidate, [8, 9, 9])) + saved = RunStore(store.root) + atomic_json(saved.root / "candidates/0001/solution.json", {"changed": True}) + with pytest.raises(ContractError, match="changed"): + saved.export(tmp_path / "export") + + +def test_export_rejects_mutated_winning_evidence(store, candidate, tmp_path: Path) -> None: + run_fake(store, fake_coder(store, candidate, [8, 9, 9])) + saved = RunStore(store.root) + atomic_json(saved.root / "candidates/0001/evaluation/result.json", {"status": "forged"}) + with pytest.raises(ContractError, match="evidence changed"): + saved.export(tmp_path / "export") + + +def test_export_stays_available_after_tool_update(store, candidate, tmp_path: Path, monkeypatch) -> None: + run_fake(store, fake_coder(store, candidate, [8, 9, 9])) + saved = RunStore(store.root) + monkeypatch.setattr("scripts.kernel_tuning.artifacts.tool_identity", lambda: "new tool version") + with pytest.raises(ContractError, match="source changed"): + saved.verify() + exported = saved.export(tmp_path / "export") + assert (exported / "tool/scripts/kernel_tuning/flow.py").is_file() + + +def test_crash_during_attempt_preparation_remains_counted(store, monkeypatch) -> None: + def failed_copy(*args, **kwargs): + raise OSError("interrupted file copy") + + monkeypatch.setattr("scripts.kernel_tuning.artifacts.shutil.copytree", failed_copy) + with pytest.raises(OSError, match="interrupted"): + store.begin_attempt() + resumed = RunStore(store.root) + assert resumed.manifest["pending"] == "0001" + resumed.recover() + resumed.save() + assert resumed.manifest["attempts"][0]["status"] == "interrupted" + assert resumed.manifest["streak"] == 1 + + +@pytest.mark.parametrize("unfenced", [False, True]) +def test_sessions_run_under_the_recorded_fence(store, candidate, unfenced: bool) -> None: + """Fenced runs keep hmz's default grant; only a recorded --unfenced run widens it.""" + from scripts.kernel_tuning.permissions import FENCED, UNFENCED + + store.manifest["unfenced"] = unfenced + store.save() + coder = fake_coder(store, candidate, [8, 9, 9]) + run_fake(store, coder) + assert coder.sessions + expected = UNFENCED if unfenced else FENCED + assert all(session.permission == expected for session in coder.sessions) diff --git a/tests/kernel_tuning/test_generated_tasks.py b/tests/kernel_tuning/test_generated_tasks.py new file mode 100644 index 0000000..6588467 --- /dev/null +++ b/tests/kernel_tuning/test_generated_tasks.py @@ -0,0 +1,234 @@ +"""Generated task packages: static structure on CPU, baseline contracts on a GPU host.""" + +from __future__ import annotations + +import importlib +import importlib.util +import json +import struct +import sys +from pathlib import Path +from typing import Any + +import pytest +from scripts.kernel_tuning.capture import apply_file +from scripts.kernel_tuning.contracts import ContractError, TaskPackage, read_json +from scripts.kernel_tuning.generate import ( + MARKER, + parse_overrides, + render, + safetensors_bytes, + stale_sources, + tuning_settings, + vendor_sources, +) +from scripts.kernel_tuning.operators import OPERATORS, catalog, select + + +@pytest.fixture(scope="module") +def generated(tmp_path_factory: pytest.TempPathFactory) -> dict[str, TaskPackage]: + """Every catalog operator, rendered once with default settings.""" + root = tmp_path_factory.mktemp("generated") + return {operator.name: render(operator, root / operator.name) for operator in OPERATORS} + + +def test_every_catalog_operator_renders_a_valid_task(generated: dict[str, TaskPackage]) -> None: + for name, task in generated.items(): + operator = catalog()[name] + assert task.workload_ids == tuple(work.uuid for work in operator.workloads) + assert task.settings.precision.mode == operator.precision["mode"] + baseline = task.baseline() + paths = {source["path"] for source in baseline["sources"]} + assert "baseline.py" in paths + assert f"vendor/{operator.source}" in paths + assert not stale_sources(task.root) + + +def test_vendored_kernels_are_verbatim_and_import_closed(generated: dict[str, TaskPackage]) -> None: + from scripts.kernel_tuning.artifacts import REPOSITORY + + for task in generated.values(): + for path in read_json(task.root / MARKER)["sources"]: + original = REPOSITORY / Path(path).relative_to("vendor") + assert (task.root / path).read_bytes() == original.read_bytes().replace(b"\r\n", b"\n") + for line in (task.root / path).read_text(encoding="utf-8").splitlines(): + assert not line.lstrip().startswith(("from embodiinfer", "import embodiinfer")) + # packed_rope imports .capability; the closure must carry it. + assert "vendor/embodiinfer/backend/triton/capability.py" in generated["rotate_half_rope"].hashes + + +def test_absolute_repository_imports_are_rewritten(monkeypatch, tmp_path: Path) -> None: + from scripts.kernel_tuning import generate + + package = tmp_path / "embodiinfer/backend/triton" + package.mkdir(parents=True) + (package / "helper.py").write_text("VALUE = 1\n", encoding="utf-8") + (package / "kernel.py").write_text( + "from embodiinfer.backend.triton.helper import VALUE\nfrom . import helper\n", encoding="utf-8" + ) + (package / "__init__.py").write_text("", encoding="utf-8") + (package / "bad.py").write_text("from embodiinfer.backend.triton import helper\n", encoding="utf-8") + monkeypatch.setattr(generate, "REPOSITORY", tmp_path) + vendored = vendor_sources("embodiinfer.backend.triton.kernel") + assert vendored["vendor/embodiinfer/backend/triton/kernel.py"][0].startswith("from .helper import VALUE") + assert "vendor/embodiinfer/backend/triton/helper.py" in vendored + with pytest.raises(ContractError, match="package"): + vendor_sources("embodiinfer.backend.triton.bad") + + +def test_fixed_integer_inputs_round_trip_through_safetensors() -> None: + encoded = safetensors_bytes({"offsets": ("int32", (0, 64, 128))}) + (size,) = struct.unpack(" None: + operator = catalog()["gated_residual"] + overrides = parse_overrides(["evaluator_python=/opt/thor/bin/python", "search.max_candidates=4"]) + task = render(operator, tmp_path / "task", overrides=overrides) + assert task.settings.evaluator_python == "/opt/thor/bin/python" + assert task.settings.search.max_candidates == 4 + with pytest.raises(ContractError, match="--force"): + render(operator, tmp_path / "task") + assert render(operator, tmp_path / "task", overrides=overrides, force=True).identity == task.identity + with pytest.raises(ContractError): + tuning_settings(operator, parse_overrides(["search.unknown=1"])) + with pytest.raises(ContractError, match="Unknown tuning.yaml section"): + tuning_settings(operator, parse_overrides(["nope.key=1"])) + (tmp_path / "manual").mkdir() + (tmp_path / "manual/README.md").write_text("hand-written", encoding="utf-8") + with pytest.raises(ContractError, match="not generated"): + render(operator, tmp_path / "manual", force=True) + + +def test_operator_selection_by_name_and_model() -> None: + assert {op.name for op in select(models=["pi05"])} == { + "ada_rms_norm", + "gated_residual", + "gated_gelu", + "rotate_qk", + "split_kv_attention", + } + assert [op.name for op in select(["swiglu"])] == ["swiglu"] + with pytest.raises(ValueError, match="Unknown operators"): + select(["nope"]) + with pytest.raises(ValueError, match="Unknown models"): + select(models=["nope"]) + + +# ---- Target-host checks: the production kernel must satisfy its own task contract ---- + + +def _load_baseline(task: TaskPackage, root: Path) -> Any: + """Import the baseline the way FlashInfer's Python builder does: as a namespace package.""" + package = f"kernel_task_{task.identity[:12]}" + for source in task.baseline()["sources"]: + path = root / package / source["path"] + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(source["content"], encoding="utf-8") + sys.path.insert(0, str(root)) + return importlib.import_module(f"{package}.baseline").run + + +def _inputs(task: TaskPackage, workload: dict[str, Any], seed: int, device: str) -> list[Any]: + """Mirror FlashInfer's generator: torch.randn floats, scalars, and safetensors data.""" + import torch + + torch.manual_seed(seed) + definition = task.definition + axes = {k: v["value"] for k, v in definition["axes"].items() if v["type"] == "const"} + axes.update(workload["axes"]) + values = [] + for name, spec in definition["inputs"].items(): + source = workload["inputs"][name] + dtype = getattr(torch, spec["dtype"]) + if source["type"] == "scalar": + values.append(source["value"]) + elif source["type"] == "safetensors": + data = (task.root / source["path"]).read_bytes() + (size,) = struct.unpack(" str | None: + try: + import torch + import triton # noqa: F401 + except ImportError: + return "requires Torch and Triton" + return None if torch.cuda.is_available() else "requires CUDA" + + +@pytest.mark.gpu +@pytest.mark.parametrize("name", [op.name for op in OPERATORS]) +def test_production_baseline_meets_generated_contract(name: str, tmp_path: Path) -> None: + """Calibrates each contract: a production kernel that fails its own task is a catalog bug.""" + if reason := _gpu_ready(): + pytest.skip(reason) + import torch + from scripts.kernel_tuning.artifacts import REPOSITORY + + spec = importlib.util.spec_from_file_location( + "generated_task_adapter", REPOSITORY / "benchmarks/kernel_tuning/flashinfer_adapter.py" + ) + adapter = importlib.util.module_from_spec(spec) + spec.loader.exec_module(adapter) + compare = adapter.compare_outputs + + (operator,), _ = apply_file([catalog()[name]]) # calibrate on the committed captured shapes + task = render(operator, tmp_path / "task") + baseline = _load_baseline(task, tmp_path / "build") + reference_scope: dict[str, Any] = {} + exec(compile(task.definition["reference"], "reference", "exec"), reference_scope) + reference = reference_scope["run"] + + def outputs(value: Any) -> list[Any]: + return list(value) if isinstance(value, (tuple, list)) else [value] + + device = "cuda" + precision = task.settings.precision + for trace in task.workloads: + work = trace["workload"] + exact, margin = True, 0.0 + for seed in task.settings.seeds: + inputs = _inputs(task, work, seed, device) + expected = [ + v.clone() + for v in outputs(reference(*[x.clone() if hasattr(x, "clone") else x for x in inputs])) + ] + actual = outputs(baseline(*inputs)) + torch.cuda.synchronize() + compare(actual, expected, precision) + for value, target in zip(actual, expected): + exact = exact and torch.equal(value.view(torch.uint8), target.view(torch.uint8)) + if precision.mode == "tolerance" and value.is_floating_point(): + bound = precision.atol + precision.rtol * target.double().abs() + margin = max(margin, ((value.double() - target.double()).abs() / bound).max().item()) + # Report how much of the bound the production kernel uses, to calibrate the catalog. + print(f"{name}/{work['uuid']}: bit_exact={exact} bound_used={margin:.2f}") + # The evaluator replays a captured graph with changed inputs; the baseline must allow it. + static = _inputs(task, work, task.settings.seeds[0], device) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + baseline(*static) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + captured = outputs(baseline(*static)) + changed = _inputs(task, work, task.settings.seeds[0] + 1, device) + for dst, src in zip(static, changed): + if isinstance(dst, torch.Tensor): + dst.copy_(src) + graph.replay() + torch.cuda.synchronize() + compare(captured, outputs(reference(*changed)), precision) diff --git a/tests/kernel_tuning/test_hotspots.py b/tests/kernel_tuning/test_hotspots.py new file mode 100644 index 0000000..0f427e8 --- /dev/null +++ b/tests/kernel_tuning/test_hotspots.py @@ -0,0 +1,205 @@ +"""CPU tests for end-to-end hotspot attribution and the profiling session.""" + +from __future__ import annotations + +import sys +import types + +import pytest +from scripts.kernel_tuning import hotspots +from scripts.kernel_tuning.contracts import ContractError +from scripts.kernel_tuning.hotspots import ( + DeviceKernel, + Launch, + Range, + attribute, + fold_layers, + render, + summarize, +) + + +def test_fold_layers_merges_repeated_layer_indices(): + assert fold_layers("module:model.layers.12.mlp.gate_proj") == "module:model.layers.*.mlp.gate_proj" + assert fold_layers("module:blocks.3") == "module:blocks.*" + assert fold_layers("op:rotate_qk") == "op:rotate_qk" + + +def test_kernels_follow_their_launch_into_enclosing_ranges(): + ranges = [ + Range("kt/step:0", thread=1, start=0, end=100), + Range("kt/stage:encode_prefix", thread=1, start=5, end=50), + Range("kt/module:layers.0.mlp", thread=1, start=10, end=20), + Range("kt/module:layers.1", thread=1, start=28, end=42), + Range("kt/module:layers.1.mlp", thread=1, start=30, end=40), + Range("kt/stage:integrate", thread=1, start=60, end=95), + Range("kt/step:1", thread=1, start=200, end=300), + Range("kt/stage:other-thread", thread=2, start=0, end=1000), + ] + launches = [ + Launch("cudaLaunchKernel", thread=1, start=12, correlation=1), + Launch("cudaLaunchKernel", thread=1, start=35, correlation=2), + Launch("cudaLaunchKernel", thread=1, start=45, correlation=3), + Launch("cudaGraphLaunch", thread=1, start=70, correlation=4), + Launch("cudaLaunchKernel", thread=1, start=250, correlation=5), + Launch("cudaLaunchKernel", thread=1, start=150, correlation=6), + ] + kernels = [ + DeviceKernel("gemm", 0, 3000, 1), + DeviceKernel("gemm", 0, 3000, 2), + DeviceKernel("norm", 0, 1000, 3), + DeviceKernel("graph_a", 0, 2000, 4), + DeviceKernel("graph_b", 0, 2000, 4), + DeviceKernel("late", 0, 500, 5), + DeviceKernel("between_steps", 0, 500, 6), + DeviceKernel("orphan", 0, 500, 99), + ] + result = attribute(ranges, launches, kernels) + assert [(k.step, k.path, k.graph) for k in result] == [ + (0, ("stage:encode_prefix", "module:layers.*.mlp"), False), + (0, ("stage:encode_prefix", "module:layers.*.mlp"), False), + (0, ("stage:encode_prefix", "kernels:elementwise"), False), + (0, ("stage:integrate", "kernels:other"), True), + (0, ("stage:integrate", "kernels:other"), True), + (1, ("kernels:other",), False), + (None, ("kernels:other",), False), + (None, ("",), False), + ] + assert result[0].duration_us == 3.0 + + +def test_region_keeps_stages_and_innermost_range_only(): + ranges = [ + Range("kt/step:0", 1, 0, 100), + Range("kt/stage:integrate", 1, 1, 99), + Range("kt/stage:integrate", 1, 2, 98), + Range("kt/attention:Split", 1, 3, 90), + Range("kt/attention:Split", 1, 4, 80), + Range("kt/module:layers.4.self_attn", 1, 5, 70), + ] + kernel = attribute(ranges, [Launch("cudaLaunchKernel", 1, 10, 1)], [DeviceKernel("k", 0, 1000, 1)])[0] + assert kernel.path == ("stage:integrate", "module:layers.*.self_attn") + + +def test_summary_ranks_regions_per_step(): + ranges = [Range(f"kt/step:{i}", 1, i * 100, i * 100 + 99) for i in range(2)] + [ + Range(f"kt/op:{name}", 1, i * 100 + offset, i * 100 + offset + 10) + for i in range(2) + for name, offset in (("gated_gelu", 10), ("rotate_qk", 50)) + ] + launches, kernels = [], [] + for i in range(2): + for correlation, (time, name, ns) in enumerate( + ((12, "gelu", 2000), (14, "gelu_tail", 1000), (52, "rope", 30000)), start=i * 10 + ): + launches.append(Launch("cudaLaunchKernel", 1, i * 100 + time, correlation)) + kernels.append(DeviceKernel(name, 0, ns, correlation)) + summary = summarize(attribute(ranges, launches, kernels), steps=2) + assert summary["kernel_us_per_step"] == pytest.approx(33.0) + first, second = summary["regions"] + assert first["region"] == "op:rotate_qk" and first["us_per_step"] == pytest.approx(30.0) + assert second["region"] == "op:gated_gelu" + assert second["kernels_per_step"] == 2 and second["small_kernel_fraction"] == 1.0 + assert second["top_kernels"][0] == {"name": "gelu", "calls_per_step": 1.0, "us_per_step": 2.0} + assert "`op:rotate_qk`" in render(summary) + + +def test_summary_without_measured_kernels_is_an_error(): + with pytest.raises(ContractError, match="No kernels"): + summarize([], steps=3) + + +@pytest.fixture +def fake_engine(monkeypatch): + torch = pytest.importorskip("torch") + pytest.importorskip("embodiinfer") + module = types.ModuleType("kt_fake_engine") + + class Engine: + def __init__(self) -> None: + from embodiinfer.layers import KernelSelection + + self.op = KernelSelection("none") + + def execute(self, nested: bool = False) -> None: + residual = torch.ones(1, 2, 4) + self.op.resolve("gated_residual")(residual, residual, residual) + if not nested: + self.execute(nested=True) + + module.Engine = Engine + monkeypatch.setitem(sys.modules, module.__name__, module) + monkeypatch.setattr(hotspots, "_ENTRIES", ((module.__name__, "Engine", ("execute",)),)) + return Engine + + +def test_session_profiles_the_window_and_restores_patches(fake_engine): + original = fake_engine.execute + session = hotspots.ProfileSession(warmup=1, steps=2, modules=False) + session.install() + try: + engine = fake_engine() + for _ in range(4): + engine.execute() + finally: + session.close() + assert session.started == 4 # nested calls are not steps + assert fake_engine.execute is original + names = [event.name() for event in session.trace.profiler.kineto_results.events()] + assert names.count("kt/step:0") == 1 and names.count("kt/step:1") == 1 + assert "kt/step:2" not in names + assert names.count("kt/op:gated_residual") >= 2 + if hasattr(session.trace.profiler.kineto_results.events()[0], "activity_type"): + assert session.kernels() == [] # CPU-only run: no device kernels + else: + with pytest.raises(RuntimeError, match="activity_type"): + session.kernels() + + +def test_session_requires_enough_steps(fake_engine): + session = hotspots.ProfileSession(warmup=2, steps=3, modules=False) + session.install() + try: + fake_engine().execute() + finally: + session.close() + with pytest.raises(ContractError, match="Only 1 inference steps ran"): + session.kernels() + + +def test_graph_replays_take_their_capture_labels_in_order(): + capture_ranges = [ # the capture profiler starts inside the stage: no stage range + Range("kt/op:ada_rms_norm", 1, 10, 20), + Range("kt/attention:Split", 1, 30, 40), + ] + capture_launches = [ + Launch("cudaLaunchKernel", 1, 35, 2), + Launch("cudaStreamIsCapturing", 1, 12, 9), + Launch("cuLaunchKernel", 1, 15, 1), + Launch("cudaMemsetAsync", 1, 50, 3), + ] + labels = hotspots.capture_labels(capture_ranges, capture_launches) + assert labels == [("op:ada_rms_norm",), ("attention:Split",), ()] + + ranges = [ + Range("kt/step:0", 1, 0, 1000), + Range("kt/stage:integrate", 1, 1, 999), + Range("kt/graph:77", 1, 10, 20), + Range("kt/graph:88", 1, 30, 40), + ] + launches = [Launch("cudaGraphLaunch", 1, 12, 5), Launch("cudaGraphLaunch", 1, 32, 6)] + kernels = [ + DeviceKernel("attn", 200, 4000, 5), + DeviceKernel("norm", 100, 1000, 5), + DeviceKernel("memset", 300, 500, 5), + DeviceKernel("unknown_graph", 400, 500, 6), + ] + result = attribute(ranges, launches, kernels, {77: labels, 88: labels}) + assert [(k.name, k.path, k.labelled) for k in result] == [ + ("attn", ("stage:integrate", "attention:Split"), True), + ("norm", ("stage:integrate", "op:ada_rms_norm"), True), + ("memset", ("stage:integrate", "kernels:copy"), True), + ("unknown_graph", ("stage:integrate", "kernels:other"), False), # 1 kernel vs 3 labels + ] + summary = summarize(result, steps=1) + assert summary["graph_labelled_fraction"] == pytest.approx(5.5 / 6.0) diff --git a/tests/kernel_tuning/test_install.py b/tests/kernel_tuning/test_install.py new file mode 100644 index 0000000..650f16e --- /dev/null +++ b/tests/kernel_tuning/test_install.py @@ -0,0 +1,155 @@ +"""CPU tests for installing a registry solution into a kernel profile.""" + +from __future__ import annotations + +import json +from types import SimpleNamespace + +import pytest + +torch = pytest.importorskip("torch") +pytest.importorskip("embodiinfer") + +from scripts.kernel_tuning import install as installer # noqa: E402 +from scripts.kernel_tuning.contracts import ContractError # noqa: E402 + +import embodiinfer.backend.triton.tuned as tuned # noqa: E402 +from embodiinfer.layers import KernelSelection # noqa: E402 + +KERNEL = """CALLS = [] + + +def fast(residual, update, gate): + CALLS.append(tuple(residual.shape)) + return residual + update * gate +""" + + +def _spec(shape, dtype="float32"): + stride, step = [], 1 + for size in reversed(shape): + stride.insert(0, step) + step *= size + return { + "shape": list(shape), + "stride": stride, + "offset": 0, + "dtype": dtype, + "device": "cpu", + "storage": 0, + } + + +@pytest.fixture +def export(tmp_path, monkeypatch): + root = tmp_path / "export" + (root / "task").mkdir(parents=True) + (root / "sources").mkdir() + (root / "sources" / "kernel.py").write_text(KERNEL, encoding="utf-8") + recipe = {"tuple": [{"tuple": [{"tensor": 0}, {"tensor": 1}, {"tensor": 2}]}, {"dict": {}}]} + shape, gate = [1, 2, 4], [1, 1, 4] + inputs = [_spec(shape), {**_spec(shape), "storage": 1}, {**_spec(gate), "storage": 2}] + files = { + "task/generated.json": { + "generator": "scripts.kernel_tuning.regions", + "registry_op": "gated_residual", + "capture_identity": "capture", + }, + "task/replay.json": { + "torch_version": str(torch.__version__), + "numerical_identity": { + "op": "gated_residual", + "arguments": recipe, + "input_dtypes": ["float32"] * 3, + }, + "cases": {"case": {"inputs": inputs}}, + }, + "solution.json": {"spec": {"language": "triton", "entry_point": "kernel.py::fast"}}, + "manifest.json": { + "run_id": "run-1", + "best": "c1", + "attempts": [{"id": "c1", "solution_digest": "sol", "decision": {"promoted": True}}], + }, + "checksums.json": {"manifest.json": "abc"}, + } + for name, value in files.items(): + (root / name).write_text(json.dumps(value), encoding="utf-8") + environment = { + "device_name": "Test GPU", + "capability": [9, 9], + "torch": str(torch.__version__), + "triton": "x", + } + monkeypatch.setattr( + installer.TaskPackage, "load", lambda path: SimpleNamespace(root=path, identity="task") + ) + monkeypatch.setattr(tuned, "current_environment", lambda: dict(environment)) + return root, environment + + +def test_installed_kernel_applies_only_to_verified_layouts(export, tmp_path): + root, _ = export + profiles = tmp_path / "tuned" + result = installer.install(root, "testgpu", root=profiles) + assert result["verified_layouts"] == 1 + profile = tuned.load_profile(profiles / "testgpu") + assert profile.ops["gated_residual"]["evidence"]["run_id"] == "run-1" + op = KernelSelection(profile).resolve("gated_residual") + calls = profile.kernel("gated_residual").fn.__globals__["run"].__globals__["CALLS"] + x, gate = torch.ones(1, 2, 4), torch.full((1, 1, 4), 2.0) + torch.testing.assert_close(op(x, x, gate), x + x * gate) + assert calls == [(1, 2, 4)] + y = torch.ones(1, 3, 4) + torch.testing.assert_close(op(y, y, gate), y + y * gate) + assert calls == [(1, 2, 4)] # unverified layout fell back + z = torch.ones(2, 4).t().contiguous().t().unsqueeze(0) # (1, 2, 4) with transposed strides + torch.testing.assert_close(op(z, z, gate), z + z * gate) + assert calls == [(1, 2, 4)] # a different layout of the same shape also falls back + odd = torch.ones(5, 2, 4)[2:3] # size-one batch dimension with another stride + torch.testing.assert_close(op(odd, odd, gate), odd + odd * gate) + assert calls == [(1, 2, 4), (1, 2, 4)] # its stride addresses nothing, so it matches + assert op.__closure__ is not None # a dispatching wrapper, not the raw kernel + + +def test_reinstall_requires_replace_and_matching_environment(export, tmp_path, monkeypatch): + root, environment = export + profiles = tmp_path / "tuned" + installer.install(root, "testgpu", root=profiles) + with pytest.raises(ContractError, match="already installed"): + installer.install(root, "testgpu", root=profiles) + installer.install(root, "testgpu", root=profiles, replace=True) + monkeypatch.setattr(tuned, "current_environment", lambda: {**environment, "device_name": "Other"}) + with pytest.raises(ContractError, match="was verified for"): + installer.install(root, "testgpu", root=profiles, replace=True) + + +def test_edited_installed_kernel_is_rejected(export, tmp_path): + root, _ = export + profiles = tmp_path / "tuned" + installer.install(root, "testgpu", root=profiles) + (profiles / "testgpu" / "gated_residual" / "kernel.py").write_text( + KERNEL + "\n# edit\n", encoding="utf-8" + ) + with pytest.raises(ValueError, match="sha256"): + tuned.load_profile(profiles / "testgpu").kernel("gated_residual") + + +def test_only_registry_tasks_install(export, tmp_path): + root, _ = export + (root / "task" / "generated.json").write_text(json.dumps({"generator": "other"}), encoding="utf-8") + with pytest.raises(ContractError, match="Only registry-operator tasks"): + installer.install(root, "testgpu", root=tmp_path / "tuned") + + +def test_installed_kernel_receives_only_tensor_arguments(export, tmp_path): + root, _ = export + replay = json.loads((root / "task" / "replay.json").read_text(encoding="utf-8")) + arguments = replay["numerical_identity"]["arguments"]["tuple"][0]["tuple"] + arguments.append({"constant": 0.5}) # e.g. an attention scale baked into the task + (root / "task" / "replay.json").write_text(json.dumps(replay), encoding="utf-8") + profiles = tmp_path / "tuned" + installer.install(root, "testgpu", root=profiles) + kernel = tuned.load_profile(profiles / "testgpu").kernel("gated_residual") + x, gate = torch.ones(1, 2, 4), torch.full((1, 1, 4), 2.0) + assert kernel.applies(x, x, gate, 0.5) and not kernel.applies(x, x, gate, 0.25) + torch.testing.assert_close(kernel.fn(x, x, gate, 0.5), x + x * gate) diff --git a/tests/kernel_tuning/test_model_capture.py b/tests/kernel_tuning/test_model_capture.py new file mode 100644 index 0000000..ba3facc --- /dev/null +++ b/tests/kernel_tuning/test_model_capture.py @@ -0,0 +1,235 @@ +"""Registry capture artifacts, replay contracts and the all-task preflight barrier.""" + +from __future__ import annotations + +import importlib.util +import sys + +import pytest + +torch = pytest.importorskip("torch") + +from benchmarks.kernel_tuning.replay_adapter import ReplayAdapter # noqa: E402 +from scripts.kernel_tuning import batch # noqa: E402 +from scripts.kernel_tuning.__main__ import main # noqa: E402 +from scripts.kernel_tuning.artifacts import atomic_json # noqa: E402 +from scripts.kernel_tuning.contracts import ContractError, TaskPackage, digest, read_json, tree_hashes # noqa: E402 +from scripts.kernel_tuning.discovery.contracts import ModelCapture # noqa: E402 +from scripts.kernel_tuning.discovery.tasks import render_model # noqa: E402 +from scripts.kernel_tuning.discovery.tensors import clone_inputs, materialize, storage_key # noqa: E402 + + +def _capture(root, registry_capture): + value = torch.ones(4, 4, dtype=torch.float64) + return registry_capture(root, [("gelu_tanh", (value,), {}), ("linear", (value, value), {})]) + + +def test_capture_retains_every_layout_count_and_input_alias(tmp_path, registry_capture): + value = torch.arange(16, dtype=torch.float64).reshape(4, 4).T + other = torch.ones(2, 4, dtype=torch.float64) + weight = value.clone() + capture = registry_capture( + tmp_path, [("linear", (value, weight), {})] * 3 + [("linear", (other, weight), {})] + ) + (op,) = capture.operators + assert sorted(case["count"] for case in op["workloads"]) == [1, 3] + (task,) = render_model(capture, tmp_path / "tasks") + replay = read_json(task.root / "replay.json") + case = next(case for case in replay["cases"].values() if case["count"] == 3) + inputs = materialize(case["inputs"], task.root, 0) + assert len(inputs) == 2 + assert torch.equal(inputs[0], value) and inputs[0].stride() == value.stride() + views = [value, value[:, :2]] + cloned = clone_inputs(views) + assert storage_key(cloned[0]) == storage_key(cloned[1]) + assert storage_key(cloned[0]) != storage_key(value) + assert cloned[1].stride() == views[1].stride() + assert replay["kind"] == "registry" + assert task.settings.timing.mode == "cuda_graph" + + +def test_capture_model_and_fixture_integrity_remain_required(tmp_path, registry_capture): + capture = _capture(tmp_path, registry_capture) + with pytest.raises(ContractError, match="differs from requested model"): + ModelCapture.load(capture.root, "pi05") + fixture = next((capture.root / "fixtures").iterdir()) + fixture.write_bytes(fixture.read_bytes() + b"changed") + with pytest.raises(ContractError, match="content changed"): + render_model(capture, tmp_path / "tasks") + assert not (tmp_path / "tasks").exists() + + +@pytest.mark.parametrize("legacy", ["manifest", "aten", "backend"]) +def test_removed_discovery_artifacts_fail_with_migration_guidance(tmp_path, registry_capture, legacy): + capture = _capture(tmp_path, registry_capture) + manifest = read_json(capture.root / "manifest.json") + if legacy == "manifest": + manifest.pop("capture") + else: + operators = read_json(capture.root / "operators.json") + operators[0]["identity"]["kind"] = legacy + operators[0]["id"] = digest(operators[0]["identity"])[:20] + atomic_json(capture.root / "operators.json", operators) + manifest["files"] = {k: v for k, v in tree_hashes(capture.root).items() if k != "manifest.json"} + atomic_json(capture.root / "manifest.json", manifest) + with pytest.raises(ContractError, match="recapture with capture-ops"): + ModelCapture.load(capture.root, require_ready=False) + + +def test_all_baselines_must_pass_before_any_model_task_runs(tmp_path, registry_capture, monkeypatch): + from test_batch import measurement + + capture = _capture(tmp_path, registry_capture) + root = batch.create_model( + capture, tmp_path / "batches", agent="fake/model", overrides={}, hardware_notes=None + ) + commands = [] + + def launcher(*args): + commands.append(args) + + with pytest.raises(ContractError, match="every task"): + batch.execute(root, launcher=launcher) + + async def evaluator(store, output, **kwargs): + if read_json(store.task.root / "replay.json")["operator"] == "registry.linear": + raise ContractError("missing evaluator support") + return { + "measurement": measurement(store.task, 10.0), + "environment": {"identity": {"hardware": "fake"}}, + } + + monkeypatch.setattr(batch, "evaluate", evaluator) + results = batch.preflight(root) + assert any(error is None for error in results.values()) and any(results.values()) + with pytest.raises(ContractError, match="every task"): + batch.execute(root, launcher=launcher) + assert commands == [] + assert not read_json(root / "batch.json")["preflight_complete"] + + async def success(store, output, **kwargs): + return { + "measurement": measurement(store.task, 10.0), + "environment": {"identity": {"hardware": "fake"}}, + } + + monkeypatch.setattr(batch, "evaluate", success) + assert all(error is None for error in batch.preflight(root).values()) + record = read_json(root / "batch.json") + assert record["preflight_complete"] + for item in record["items"].values(): + assert batch.verify_model_task(TaskPackage.load(root / item["task"])) == root + task_root = root / next(iter(record["items"].values()))["task"] + (task_root / "README.md").write_text("changed", encoding="utf-8") + with pytest.raises(ContractError, match="changed after generation"): + batch.execute(root, launcher=launcher) + assert commands == [] + + +def test_registry_cli_keeps_full_capture_and_preflight_requirements(tmp_path, registry_capture, capsys): + capture = _capture(tmp_path, registry_capture) + args = ["--model", "mock_flow_vla", "--captured", str(capture.root)] + assert main(["generate", *args, "--list"]) == 0 + assert "registry.linear" in capsys.readouterr().out + assert main(["tune-all", *args, "--dry-run", "--output", str(tmp_path / "batches")]) == 0 + assert main(["generate", *args, "linear"]) == 2 + assert main(["tune-all", *args, "--agent", "fake/model", "--skip-preflight"]) == 2 + assert main(["generate", "--model", "pi05", "--model", "streamvln"]) == 2 + task = render_model(capture, tmp_path / "tasks")[0] + assert main(["run", str(task.root), "--agent", "fake/model", "--output", str(tmp_path / "runs")]) == 2 + assert not (tmp_path / "runs").exists() + + +def test_replay_rejects_input_mutation_and_wrong_output_aliases(tmp_path, registry_capture): + capture = _capture(tmp_path, registry_capture) + task = render_model(capture, tmp_path / "tasks")[0] + replay = read_json(task.root / "replay.json") + case = next(iter(replay["cases"].values())) + adapter = ReplayAdapter.__new__(ReplayAdapter) + adapter.case, adapter.replay, adapter.cfg = case, replay, task.settings + inputs = materialize(case["inputs"], task.root, 0) + local = clone_inputs(inputs) + assert adapter._check_state(local, local, inputs) == (0.0, 0.0) + local[0].add_(1) + with pytest.raises(ContractError, match="bytes"): + adapter._check_state(local, local, inputs) + with pytest.raises(ContractError, match="aliases"): + adapter._check_outputs(inputs, [inputs[0]]) + + +def test_replay_requires_every_workload(tmp_path, registry_capture): + capture = _capture(tmp_path, registry_capture) + task = render_model(capture, tmp_path / "tasks")[0] + replay = read_json(task.root / "replay.json") + replay["cases"] = {} + atomic_json(task.root / "replay.json", replay) + with pytest.raises(ContractError, match="every workload"): + TaskPackage.load(task.root) + + +def test_registry_float64_baseline_and_definition_keep_exact_values(tmp_path, registry_capture): + value = torch.tensor([[1.0 + 2**-30, 1.0 - 2**-30]], dtype=torch.float64) + weight = torch.eye(2, dtype=torch.float64) + capture = registry_capture(tmp_path, [("linear", (value, weight), {})]) + (task,) = render_model(capture, tmp_path / "tasks") + assert task.definition["inputs"]["x0"]["dtype"] == "float64" + assert task.definition["outputs"]["y0"]["dtype"] == "float64" + spec = importlib.util.spec_from_file_location( + "registry_float64_baseline", task.root / "baseline.py", submodule_search_locations=[str(task.root)] + ) + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + try: + spec.loader.exec_module(module) + result = module.run(value, weight) + assert torch.equal(result, value) + assert not torch.equal(result, value.float().double()) + finally: + for name in list(sys.modules): + if name == spec.name or name.startswith(spec.name + "."): + del sys.modules[name] + replay = read_json(task.root / "replay.json") + del replay["definition_schema"] + atomic_json(task.root / "replay.json", replay) + with pytest.raises(ContractError, match="explicit replay definition schema"): + TaskPackage.load(task.root) + + +def test_float64_replay_schema_retains_upstream_validation(tmp_path, registry_capture): + pytest.importorskip("flashinfer_bench") + from benchmarks.kernel_tuning.replay_schema import ReplayDefinition + from flashinfer_bench.data import Definition + from pydantic import ValidationError + + task = render_model(_capture(tmp_path, registry_capture), tmp_path / "tasks")[0] + definition = ReplayDefinition.model_validate(task.definition) + assert isinstance(definition, Definition) + assert definition.torch_output_dtypes == [torch.float64] + with pytest.raises(ValidationError): + Definition.model_validate(task.definition) + wrong = {**task.definition, "outputs": {"y0": {"shape": ["unknown_axis"], "dtype": "float64"}}} + with pytest.raises(ValidationError): + ReplayDefinition.model_validate(wrong) + + +@pytest.mark.gpu +@pytest.mark.parametrize("timing", ["eager", "cuda_graph"]) +def test_registry_baseline_passes_real_cuda_replay(tmp_path, registry_capture, timing): + """Exercise registry calibration, real builders, layout checks and changed-input replay.""" + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + pytest.importorskip("flashinfer_bench") + pytest.importorskip("triton") + from benchmarks.kernel_tuning.calibration import calibrate + from scripts.kernel_tuning.contracts import validate_measurement + + value = torch.ones(2, 16, device="cuda:0", dtype=torch.bfloat16) + capture = registry_capture(tmp_path, [("gelu_tanh", (value,), {})], device_type="cuda") + (task,) = render_model( + calibrate(capture, tmp_path / "calibrated"), + tmp_path / "tasks", + overrides={"timing.mode": timing, "timing.warmup": 1, "timing.iterations": 2, "timing.trials": 1}, + ) + adapter = ReplayAdapter(task) + report = adapter.evaluate({role: task.baseline() for role in ("baseline", "incumbent", "candidate")}) + validate_measurement(task, report) diff --git a/tests/kernel_tuning/test_numerics.py b/tests/kernel_tuning/test_numerics.py new file mode 100644 index 0000000..5ede410 --- /dev/null +++ b/tests/kernel_tuning/test_numerics.py @@ -0,0 +1,180 @@ +"""Byte-level correctness tests use CPU Torch; no kernel tuning is executed.""" + +from __future__ import annotations + +import importlib.util +import json +from pathlib import Path + +import pytest + +torch = pytest.importorskip("torch") + +from scripts.kernel_tuning.artifacts import REPOSITORY # noqa: E402 +from scripts.kernel_tuning.contracts import ContractError, Precision # noqa: E402 + +spec = importlib.util.spec_from_file_location( + "kernel_test_adapter", REPOSITORY / "benchmarks/kernel_tuning/flashinfer_adapter.py" +) +adapter = importlib.util.module_from_spec(spec) +spec.loader.exec_module(adapter) +compare = adapter.compare_outputs + + +def test_bit_exact_distinguishes_signed_zero_and_nan_payloads() -> None: + positive = torch.tensor([0.0]) + negative = torch.tensor([-0.0]) + assert torch.equal(positive, negative) + with pytest.raises(ContractError, match="bytes"): + compare([positive], [negative], Precision()) + first_nan = torch.tensor([0x7FC00001], dtype=torch.int32).view(torch.float32) + second_nan = torch.tensor([0x7FC00002], dtype=torch.int32).view(torch.float32) + assert compare([first_nan], [first_nan.clone()], Precision()) == (0, 0) + with pytest.raises(ContractError, match="bytes"): + compare([first_nan], [second_nan], Precision()) + + +@pytest.mark.parametrize( + "actual,expected", + [ + ([torch.ones(2)], [torch.ones(3)]), + ([torch.ones(2, dtype=torch.float16)], [torch.ones(2)]), + ([torch.ones(2)], [torch.ones(2), torch.ones(2)]), + ], +) +def test_shape_dtype_and_output_count_are_strict(actual, expected) -> None: + with pytest.raises(ContractError): + compare(actual, expected, Precision()) + + +def test_tolerance_uses_fixed_combined_atol_rtol_threshold() -> None: + precision = Precision(mode="tolerance", atol=0.1, rtol=0.1, reason="test contract") + assert compare([torch.tensor([1.19])], [torch.tensor([1.0])], precision)[0] < 0.2 + with pytest.raises(ContractError, match="threshold"): + compare([torch.tensor([1.21])], [torch.tensor([1.0])], precision) + with pytest.raises(ContractError, match="finite"): + compare([torch.tensor([float("nan")])], [torch.tensor([0.0])], precision) + with pytest.raises(ContractError, match="exactly"): + compare([torch.tensor([1], dtype=torch.int32)], [torch.tensor([2], dtype=torch.int32)], precision) + + +def test_bf16_intermediate_rounding_is_observable() -> None: + generator = torch.Generator().manual_seed(0) + residual, update, gate = [torch.randn(4096, generator=generator).to(torch.bfloat16) for _ in range(3)] + reference = residual + (update * gate).to(torch.bfloat16) + contracted = (residual.float() + update.float() * gate.float()).to(torch.bfloat16) + assert not torch.equal(reference, contracted) + with pytest.raises(ContractError, match="bytes"): + compare([contracted], [reference], Precision()) + + +@pytest.mark.gpu +@pytest.mark.parametrize("language", ["triton", "cuda"]) +def test_default_adapter_builds_multi_file_solution(example_task: Path, language: str) -> None: + """Optional target-host smoke test for both builders, with tiny declared cases.""" + import shutil + + import yaml + from scripts.kernel_tuning.artifacts import atomic_json + from scripts.kernel_tuning.contracts import TaskPackage, validate_measurement + + if not torch.cuda.is_available(): + pytest.skip("requires an NVIDIA GPU") + pytest.importorskip("flashinfer_bench") + if language == "triton": + pytest.importorskip("triton") + elif shutil.which("nvcc") is None: + pytest.skip("requires the CUDA toolkit") + task = example_task + definition = { + "name": "vector_add_test", + "op_type": "elementwise", + "axes": {"N": {"type": "var"}}, + "inputs": {"x": {"shape": ["N"], "dtype": "float32"}, "y": {"shape": ["N"], "dtype": "float32"}}, + "outputs": {"out": {"shape": ["N"], "dtype": "float32"}}, + "reference": "def run(x, y):\n return x + y\n", + } + atomic_json(task / "definition.json", definition) + (task / "baseline.py").write_text(definition["reference"]) + (task / "workloads.jsonl").write_text( + json.dumps( + { + "definition": "vector_add_test", + "workload": {"uuid": "small-tail", "axes": {"N": 257}, "inputs": {}}, + "solution": None, + "evaluation": None, + } + ) + + "\n", + encoding="utf-8", + ) + config = yaml.safe_load((task / "tuning.yaml").read_text()) + config["language"] = language + config["timing"].update(mode="cuda_graph", warmup=2, iterations=2, trials=1) + (task / "tuning.yaml").write_text(yaml.safe_dump(config)) + if language == "triton": + sources = [ + {"path": "helper.py", "content": "BLOCK = 256\n"}, + { + "path": "kernel.py", + "content": """ +import torch +import triton +import triton.language as tl +from .helper import BLOCK + +@triton.jit +def add(X, Y, Z, N: tl.constexpr, B: tl.constexpr): + i = tl.program_id(0) * B + tl.arange(0, B) + tl.store(Z + i, tl.load(X + i, i < N, 0) + tl.load(Y + i, i < N, 0), i < N) + +def run(x, y): + out = torch.empty_like(x) + add[(triton.cdiv(x.numel(), BLOCK),)](x, y, out, x.numel(), BLOCK) + return out +""", + }, + ] + entry = "kernel.py::run" + else: + sources = [ + {"path": "helper.cuh", "content": "constexpr int BLOCK = 256;\n"}, + { + "path": "kernel.cu", + "content": """ +#include +#include +#include "helper.cuh" +__global__ void add(const float* x, const float* y, float* out, int n) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) out[i] = x[i] + y[i]; +} +torch::Tensor run(torch::Tensor x, torch::Tensor y) { + auto out = torch::empty_like(x); + add<<<(x.numel()+BLOCK-1)/BLOCK, BLOCK, 0, at::cuda::getCurrentCUDAStream()>>>( + x.data_ptr(), y.data_ptr(), out.data_ptr(), x.numel()); + return out; +} +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("run", &run); } +""", + }, + ] + entry = "kernel.cu::run" + package = TaskPackage.load(task) + solution = { + "name": "add_test", + "definition": definition["name"], + "author": "test", + "spec": { + "language": language, + "target_hardware": ["cuda"], + "entry_point": entry, + "destination_passing_style": False, + **({"binding": "torch"} if language == "cuda" else {}), + }, + "sources": sources, + } + measured = adapter.FlashInferAdapter(package).evaluate( + {"baseline": package.baseline(), "incumbent": package.baseline(), "candidate": solution} + ) + validate_measurement(package, measured) diff --git a/tests/kernel_tuning/test_regions.py b/tests/kernel_tuning/test_regions.py new file mode 100644 index 0000000..25006fc --- /dev/null +++ b/tests/kernel_tuning/test_regions.py @@ -0,0 +1,167 @@ +"""CPU tests: capture-ops records the output contract of the call the model made.""" + +from __future__ import annotations + +import importlib +import sys + +import pytest + +torch = pytest.importorskip("torch") +pytest.importorskip("embodiinfer") + +from scripts.kernel_tuning.discovery import tasks # noqa: E402 +from scripts.kernel_tuning.regions import OpCapture # noqa: E402 + +from embodiinfer.backend.torch import ops as torch_ops # noqa: E402 + + +def _rope_inputs(): + # BHSD views of BSHD projections, as the policy passes them to rotate_qk. + query = torch.randn(1, 3, 2, 4).transpose(1, 2) + key = torch.randn(1, 3, 1, 4).transpose(1, 2) + cos, sin = torch.randn(1, 3, 4), torch.randn(1, 3, 4) + return query, key, cos, sin + + +class Production: + """Stand-in for a backend that allocates fresh contiguous outputs, like the Triton RoPE.""" + + def __init__(self, contiguous: bool = True) -> None: + self.calls, self.contiguous = 0, contiguous + + def __call__(self, *args): + self.calls += 1 + result = torch_ops.rotate_qk(*args) + return tuple(value.contiguous() for value in result) if self.contiguous else result + + +def _capture(tmp_path): + return OpCapture("pi05", tmp_path / "capture", ops={"rotate_qk"}) + + +def test_output_contract_follows_the_production_call(tmp_path): + capture, production = _capture(tmp_path), Production() + args = _rope_inputs() + reference = torch_ops.rotate_qk(*args) + assert not reference[0].is_contiguous(), "the reference must keep the input view layout here" + + result = capture.record("rotate_qk", torch_ops.rotate_qk, production, args, {}) + + assert production.calls == 1 + assert all(value.is_contiguous() for value in result) + assert not capture.errors + (op,) = capture.operators.values() + (case,) = op["workloads"] + assert [spec["stride"] for spec in case["outputs"]] == [list(value.stride()) for value in result] + assert [spec["stride"] for spec in case["reference_outputs"]] == [list(v.stride()) for v in reference] + assert case["output_aliases"] == [[], []] and case["output_groups"] == [[0], [1]] + + +def test_matching_layouts_keep_the_previous_case_format(tmp_path): + capture = _capture(tmp_path) + capture.record("rotate_qk", torch_ops.rotate_qk, Production(contiguous=False), _rope_inputs(), {}) + (op,) = capture.operators.values() + assert "reference_outputs" not in op["workloads"][0] + + +def test_disagreeing_structure_is_an_error_and_the_model_result_is_kept(tmp_path): + capture = _capture(tmp_path) + calls = [] + + def production(*args): + calls.append(1) + return torch_ops.rotate_qk(*args)[0] # one output instead of two + + result = capture.record("rotate_qk", torch_ops.rotate_qk, production, _rope_inputs(), {}) + assert len(calls) == 1 and isinstance(result, torch.Tensor) + assert not capture.operators + assert any("another structure" in error for error in capture.errors) + + +def test_generated_baseline_returns_the_production_layout(tmp_path): + capture = _capture(tmp_path) + args = _rope_inputs() + capture.record("rotate_qk", torch_ops.rotate_qk, Production(), args, {}) + (op,) = capture.operators.values() + code, files = tasks._sources(op) + assert "_LAYOUTS" in code + package = tmp_path / "pkg" / "relayout_baseline" + for path, text in {**files, "baseline.py": code}.items(): + target = package / path + target.parent.mkdir(parents=True, exist_ok=True) + target.write_text(text, encoding="utf-8") + sys.path.insert(0, str(tmp_path / "pkg")) + try: + run = importlib.import_module("relayout_baseline.baseline").run + finally: + sys.path.remove(str(tmp_path / "pkg")) + + expected = torch_ops.rotate_qk(*args) + outputs = run(*args) + for value, reference, spec in zip(outputs, expected, op["workloads"][0]["outputs"]): + assert list(value.stride()) == spec["stride"] + assert torch.equal(value, reference) + # Unrecorded input layouts are returned as the reference computes them. + other = tuple(value.contiguous() for value in args) + for value, reference in zip(run(*other), torch_ops.rotate_qk(*other)): + assert value.stride() == reference.stride() + + +def test_baseline_code_is_unchanged_without_layout_differences(tmp_path): + capture = _capture(tmp_path) + capture.record("rotate_qk", torch_ops.rotate_qk, Production(contiguous=False), _rope_inputs(), {}) + (op,) = capture.operators.values() + code, _ = tasks._sources(op) + assert "_LAYOUTS" not in code and "_relayout" not in code + + +def test_capture_ops_cli_keeps_registry_selection_window_and_restores_hooks(tmp_path, monkeypatch): + from scripts.kernel_tuning.__main__ import main + from scripts.kernel_tuning.discovery.contracts import ModelCapture + + from embodiinfer.engine import EngineCore + from embodiinfer.layers.ops import KernelSelection + + def execute(_): + value = torch.ones(2, 4) + selection = KernelSelection("torch") + selection.resolve("gelu_tanh")(value) + selection.resolve("linear")(value, value) + return torch.sin(value) # Unregistered ATen calls must not become tasks. + + monkeypatch.setattr(EngineCore, "execute", execute) + original_resolve = KernelSelection.resolve + script = tmp_path / "driver.py" + script.write_text( + "from embodiinfer.engine import EngineCore\n" + "from embodiinfer.policies.factory import make_policy\n" + "policy = make_policy('mock_flow_vla', preset='tiny')\n" + "for _ in range(6):\n EngineCore.execute(None)\n", + encoding="utf-8", + ) + output = tmp_path / "capture" + assert ( + main( + [ + "capture-ops", + "--model", + "mock_flow_vla", + "--ops", + "gelu_tanh", + "--warmup", + "2", + "--steps", + "3", + "--output", + str(output), + "--", + str(script), + ] + ) + == 0 + ) + (op,) = ModelCapture.load(output).operators + assert op["name"] == "registry.gelu_tanh" + assert sum(case["count"] for case in op["workloads"]) == 3 + assert EngineCore.execute is execute and KernelSelection.resolve is original_resolve diff --git a/tests/kernel_tuning/test_registration.py b/tests/kernel_tuning/test_registration.py new file mode 100644 index 0000000..38c1510 --- /dev/null +++ b/tests/kernel_tuning/test_registration.py @@ -0,0 +1,888 @@ +"""CPU registration contracts, Git isolation and real worker failure evidence.""" + +from __future__ import annotations + +import argparse +import asyncio +import copy +import json +import shutil +import subprocess +import sys +import time +from pathlib import Path +from types import SimpleNamespace + +import pytest +from scripts.kernel_tuning.artifacts import REPOSITORY, atomic_json +from scripts.kernel_tuning.contracts import ContractError, read_json +from scripts.kernel_tuning.registration import checks +from scripts.kernel_tuning.registration.cli import add_parser, dispatch +from scripts.kernel_tuning.registration.cli import test_python as parse_python +from scripts.kernel_tuning.registration.finalize import finalize +from scripts.kernel_tuning.registration.sites import locate, validate_review +from scripts.kernel_tuning.registration.spec import RegistrationSpec, load_spec +from scripts.kernel_tuning.registration.workspace import ( + REFERENCE, + REGISTRY, + ROUTED_EXAMPLE, + RegistrationRun, + check_scope, + check_unregistered, + command, + file_hash, + git, + tracked_hashes, +) + +SOURCE = "embodiinfer/policies/pi05/modeling_pi05.py" +NEW_TEST = "tests/test_registry_matmul_t.py" + + +def specification(**overrides) -> RegistrationSpec: + """Use an existing unfused registry gap without loading a model/checkpoint.""" + value = { + "model": "pi05", + "op": "matmul_t", + "sites": [{"source": SOURCE, "function": "_fused_qkv", "expression": "F.linear(x, weight)"}], + "budget": {"max_attempts": 5, "max_seconds": 120, "agent_timeout_seconds": 10}, + } + value.update(overrides) + return RegistrationSpec.parse(value) + + +def test_new_targets_selected_repository(clean_repository, tmp_path, monkeypatch, capsys) -> None: + """An uncommitted controller can create worktrees in a separate clean checkout.""" + + async def preflight(*args, **kwargs): + pass + + monkeypatch.setattr("scripts.kernel_tuning.registration.workspace.check_unregistered", preflight) + parser = argparse.ArgumentParser() + add_parser(parser.add_subparsers()) + output = tmp_path / "runs" + args = parser.parse_args( + [ + "register", + "new", + "--repository", + str(clean_repository), + "--output", + str(output), + "--model", + "pi05", + "--op", + "matmul_t", + "--site", + f"{SOURCE}::_fused_qkv::F.linear(x, weight)", + ] + ) + assert dispatch(args) == 0 + run = RegistrationRun(next(path.parent for path in output.glob("*/manifest.json"))) + assert Path(git(run.worktree, "rev-parse", "--git-common-dir").strip()) == clean_repository / ".git" + assert "Run:" in capsys.readouterr().out + + +def test_baseline_hashes_track_links_without_following_targets(tmp_path, monkeypatch) -> None: + """POSIX documentation links are valid baseline entries, but still immutable.""" + root = tmp_path / "repo" + root.mkdir() + target = tmp_path / "asset" + target.write_text("first", encoding="utf-8") + link = root / "asset-link" + try: + link.symlink_to(target) + except OSError: + pytest.skip("Creating symlinks requires platform permissions") + monkeypatch.setattr("scripts.kernel_tuning.registration.workspace.git", lambda *args: "asset-link\0") + first = tracked_hashes(root) + target.write_text("second", encoding="utf-8") + assert tracked_hashes(root) == first + link.unlink() + link.symlink_to(tmp_path / "missing") + assert tracked_hashes(root) != first + with pytest.raises(ContractError): + checks.source_path(root, "asset-link") + + +@pytest.fixture(scope="module") +def clean_repository(tmp_path_factory) -> Path: + """Reuse existing commit objects in a disposable clone; create no commits.""" + destination = tmp_path_factory.mktemp("registration-repository") / "repo" + subprocess.run( + ["git", "clone", "--shared", "--quiet", str(REPOSITORY), str(destination)], + check=True, + capture_output=True, + ) + return destination + + +@pytest.fixture +def registration_run(clean_repository, tmp_path, monkeypatch) -> RegistrationRun: + """Exercise real worktrees; registry import is covered separately without Torch.""" + + async def preflight(*args, **kwargs): + pass + + monkeypatch.setattr("scripts.kernel_tuning.registration.workspace.check_unregistered", preflight) + run = asyncio.run( + RegistrationRun.create(clean_repository, tmp_path / "runs", specification(), sys.executable) + ) + run.prepare("codex/fake:low") + return run + + +def write_candidate(run: RegistrationRun) -> None: + """Construct a valid structural patch, without executing a model.""" + source = (run.worktree / SOURCE).read_text(encoding="utf-8") + source = source.replace( + "def _fused_qkv(attn: Any, x: torch.Tensor)", + "def _fused_qkv(attn: Any, x: torch.Tensor, kernels=None)", + ) + start, end = source.index("def _fused_qkv("), source.index("def _prepare_fused_projection_weights(") + source = ( + source[:start] + + source[start:end].replace("F.linear(x, weight)", '_op(kernels, "matmul_t", False)(x, weight)') + + source[end:] + ) + source = source.replace("_fused_qkv(attn, h)", "_fused_qkv(attn, h, kernels=self._kernels)") + (run.worktree / SOURCE).write_text(source, encoding="utf-8") + with (run.worktree / REFERENCE).open("a", encoding="utf-8") as stream: + stream.write( + '\n\ndef matmul_t(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:\n """Preserve the original bias-free F.linear expression."""\n return F.linear(x, weight)\n' + ) + with (run.worktree / REGISTRY).open("a", encoding="utf-8") as stream: + stream.write('\nregister_lazy_op("matmul_t", "torch", "embodiinfer.backend.torch.ops", "matmul_t")\n') + (run.worktree / NEW_TEST).write_text("def test_new_reference():\n assert True\n", encoding="utf-8") + (run.worktree / "PLAN.md").write_text("Register the unchanged F.linear expression.", encoding="utf-8") + atomic_json( + run.worktree / "REGISTRATION.json", + { + "op": "matmul_t", + "reference": f"{REFERENCE}::matmul_t", + "sites": [ + { + "id": "S1", + "file": SOURCE, + "function": "_fused_qkv", + "call": '_op(kernels, "matmul_t", False)(x, weight)', + } + ], + "preserved": [], + "tests": [f"{NEW_TEST}::test_new_reference"], + }, + ) + + +@pytest.mark.parametrize( + "overrides", + [ + {"op": "linear.v2"}, + {"op": "for"}, + {"op": "pi05_linear"}, + {"op": "_linear"}, + {"budget": {"max_attempts": 0}}, + {"budget": {"max_seconds": True}}, + {"budget": {"unknown": 2}}, + {"tests": ["--help"]}, + {"tests": ["../test.py"]}, + {"sites": []}, + {"unexpected": True}, + ], +) +def test_spec_rejects_invalid_contracts(overrides) -> None: + with pytest.raises(ContractError): + specification(**overrides) + + +def test_yaml_and_cli_share_validation(tmp_path) -> None: + path = tmp_path / "spec.yaml" + path.write_text( + f"model: pi05\nop: matmul_t\nsites:\n - source: {SOURCE}\n function: _fused_qkv\n expression: F.linear(x, weight)\n", + encoding="utf-8", + ) + yaml_spec = load_spec(path, model=None, op=None, sites=[]) + cli_spec = load_spec( + None, model="pi05", op="matmul_t", sites=[f"{SOURCE}::_fused_qkv::F.linear(x, weight)"] + ) + assert yaml_spec == cli_spec + with pytest.raises(ContractError): + load_spec(path, model="pi05", op=None, sites=[]) + with pytest.raises(ContractError): + parse_python(["test_python=python -c bad"]) + assert parse_python([f"test_python={sys.executable}"]) == sys.executable + + +def test_locator_qualification_guards_and_call_chain(tmp_path) -> None: + path = tmp_path / SOURCE + path.parent.mkdir(parents=True) + path.write_text( + """def _op(kernels, op, triton): pass +def leaf(x, weight): + if quantized: + return module(x) + if enabled: + return F.linear(x, weight) + return x +def middle(x, weight): + return leaf(x, weight) +class Policy: + def run(self, x, weight): + self._kernels = kernels + return middle(x, weight) +""", + encoding="utf-8", + ) + spec = specification(sites=[{"source": SOURCE, "function": "leaf", "expression": "F.linear( x,weight )"}]) + found = locate(tmp_path, spec) + site = found["sites"][0] + assert site["line"] == 6 and site["column"] == 15 + assert {g["condition"] for g in site["guards"]} == {"quantized", "enabled"} + assert {c["function"] for c in site["callers"]} == {"middle", "Policy.run"} + assert site["kernels"].startswith("add kernels") + modified = copy.deepcopy(found) + modified["preserve"].append("Preserve tensor parallel module calls") + validate_review(modified, found) + modified["sites"][0]["guards"] = [] + with pytest.raises(ContractError): + validate_review(modified, found) + with path.open("a", encoding="utf-8") as stream: + stream.write("\ndef leaf(x, weight):\n return F.linear(x, weight)\n") + with pytest.raises(ContractError, match="definitions"): + locate(tmp_path, spec) + + +def test_duplicate_expression_is_not_guessed(tmp_path) -> None: + path = tmp_path / SOURCE + path.parent.mkdir(parents=True) + path.write_text( + "def _op(*args): pass\ndef _fused_qkv(x, weight):\n a = F.linear(x, weight)\n return F.linear(x, weight)\n", + encoding="utf-8", + ) + with pytest.raises(ContractError, match="2 sites"): + locate(tmp_path, specification()) + + +def test_new_requires_clean_checkout(clean_repository, tmp_path) -> None: + unrelated = clean_repository / "untracked.txt" + unrelated.write_text("keep", encoding="utf-8") + try: + with pytest.raises(ContractError, match="clean"): + asyncio.run(RegistrationRun.create(clean_repository, tmp_path, specification(), sys.executable)) + assert unrelated.read_text() == "keep" + finally: + unrelated.unlink() + + +def test_scope_catches_untracked_staged_and_protected_files(registration_run) -> None: + run = registration_run + write_candidate(run) + assert NEW_TEST in check_scope(run)["files"] + assert "new file mode" in run.patch() and "test_new_reference" in run.patch() + bad = run.worktree / "scripts/unexpected.py" + bad.write_text("pass\n", encoding="utf-8") + with pytest.raises(ContractError, match="out-of-scope"): + check_scope(run) + git(run.worktree, "add", "--", "scripts/unexpected.py") + with pytest.raises(ContractError, match="out-of-scope"): + check_scope(run) + git(run.worktree, "restore", "--staged", "--", "scripts/unexpected.py") + bad.unlink() + registry = run.worktree / REGISTRY + registry.write_text( + registry.read_text(encoding="utf-8").replace("self.mode = ", 'self.mode = "changed" or ', 1), + encoding="utf-8", + ) + with pytest.raises(ContractError, match="KernelSelection"): + check_scope(run) + + +def test_existing_test_body_cannot_be_disabled_by_insertions(registration_run) -> None: + run = registration_run + tests = run.worktree / "tests/test_kernel_ops.py" + text = tests.read_text(encoding="utf-8") + import ast + + function = next(n for n in ast.parse(text).body if isinstance(n, ast.FunctionDef)) + lines = text.splitlines(keepends=True) + lines.insert(function.body[0].lineno - 1, " return\n") + tests.write_text("".join(lines), encoding="utf-8") + with pytest.raises(ContractError, match="statements changed"): + check_scope(run) + + +@pytest.mark.parametrize( + "mutation, message", + [ + ("missing_reference", "definitions"), + ("loop", "top-level"), + ("old_expression", "original expression"), + ("triton", "Torch preference"), + ("guard", "preserve branch"), + ("backend_import", "standard library"), + ("global", "KernelSelection"), + ("old_test", "already existed"), + ], +) +def test_structure_rejects_misregistration(registration_run, mutation, message) -> None: + run = registration_run + write_candidate(run) + assert checks.check_structure(run)["sites"] == ["S1"] + if mutation == "missing_reference": + shutil.copyfile(run.base / REFERENCE, run.worktree / REFERENCE) + elif mutation == "loop": + p = run.worktree / REGISTRY + text = p.read_text(encoding="utf-8") + line = 'register_lazy_op("matmul_t", "torch", "embodiinfer.backend.torch.ops", "matmul_t")' + p.write_text(text.replace(line, "for _ in [1]:\n " + line), encoding="utf-8") + elif mutation in {"old_expression", "guard", "global"}: + p = run.worktree / SOURCE + text = p.read_text(encoding="utf-8") + if mutation == "old_expression": + text = text.replace( + 'fused = _op(kernels, "matmul_t", False)(x, weight)', + 'unused = F.linear(x, weight)\n fused = _op(kernels, "matmul_t", False)(x, weight)', + ) + elif mutation == "guard": + text = text.replace("if torch.is_grad_enabled() or attn.training:", "if attn.training:") + else: + text = text.replace("x: torch.Tensor, kernels=None)", "x: torch.Tensor)") + p.write_text(text, encoding="utf-8") + elif mutation == "backend_import": + with (run.worktree / REFERENCE).open("a", encoding="utf-8") as stream: + stream.write("\nfrom embodiinfer.engine import EngineCore\n") + else: + p = run.worktree / "REGISTRATION.json" + document = read_json(p) + if mutation == "triton": + document["sites"][0]["call"] = '_op(kernels, "matmul_t", True)(x, weight)' + else: + import ast + + name = next( + n.name + for n in ast.parse((run.base / "tests/test_kernel_ops.py").read_text(encoding="utf-8")).body + if isinstance(n, ast.FunctionDef) and n.name.startswith("test_") + ) + document["tests"] = [f"tests/test_kernel_ops.py::{name}"] + atomic_json(p, document) + with pytest.raises(ContractError, match=message): + checks.check_structure(run) + + +def test_integrity_and_baseline_overlay_recovery(registration_run) -> None: + run = registration_run + target = run.base / NEW_TEST + data = "def test_overlay(): pass\n" + target.write_text(data, encoding="utf-8") + run.manifest["base_overlay"] = {NEW_TEST: file_hash(target)} + run.save() + recovered = RegistrationRun(run.root) + recovered.recover_base() + recovered.verify() + assert not target.exists() + target.write_text("unexpected", encoding="utf-8") + with pytest.raises(ContractError, match="RUN/base"): + recovered.verify() + + +def test_context_example_comes_from_base(registration_run) -> None: + """The routed example is a BASE snapshot, never a commit that history may lose.""" + context = registration_run.root / "context" + assert (context / ROUTED_EXAMPLE).read_bytes() == (registration_run.base / ROUTED_EXAMPLE).read_bytes() + assert not (context / "migration.diff").exists() + + +def test_spec_tamper_is_fatal(registration_run) -> None: + run = registration_run + document = read_json(run.root / "spec.json") + document["spec"]["op"] = "changed" + atomic_json(run.root / "spec.json", document) + with pytest.raises(ContractError, match="specification"): + run.verify() + assert read_json(run.root / "manifest.json")["status"] == "failed" + with pytest.raises(ContractError, match="specification"): + RegistrationRun(run.root) + + +def test_finalize_only_exports_verified_bytes(registration_run) -> None: + run = registration_run + write_candidate(run) + with pytest.raises(ContractError, match="fully checked"): + finalize(run) + run.manifest.update(status="passed", verified_fingerprint=run.fingerprint()) + head = git(run.worktree, "rev-parse", "HEAD") + result = finalize(run) + assert result["commit"] is None and not result["pushed"] + assert git(run.worktree, "rev-parse", "HEAD") == head + assert "NOT been run" in (run.root / "SUMMARY.md").read_text(encoding="utf-8") + with pytest.raises(ContractError, match="requires --commit"): + finalize(run, push=True) + with (run.worktree / NEW_TEST).open("a", encoding="utf-8") as stream: + stream.write("\n# changed after checks\n") + with pytest.raises(ContractError, match="unchanged"): + finalize(run) + + +def minimal_package(root: Path) -> None: + """Make a CPU-only stand-in to test subprocess/import behavior, not model math.""" + for directory in ( + "embodiinfer", + "embodiinfer/layers", + "embodiinfer/backend", + "embodiinfer/backend/torch", + "tests", + ): + (root / directory).mkdir(parents=True, exist_ok=True) + (root / directory / "__init__.py").write_text("", encoding="utf-8") + (root / REFERENCE).write_text("", encoding="utf-8") + (root / "embodiinfer/layers/__init__.py").write_text( + "def available_ops():\n return {}\n", encoding="utf-8" + ) + + +def test_registry_preflight_uses_selected_checkout(tmp_path, monkeypatch) -> None: + minimal_package(tmp_path) + monkeypatch.setenv("PYTHONPATH", str(REPOSITORY)) + asyncio.run(check_unregistered(tmp_path, "linear", sys.executable, tmp_path / "probe.log")) + (tmp_path / "embodiinfer/layers/__init__.py").write_text( + "def available_ops():\n return {'linear': ['torch']}\n", encoding="utf-8" + ) + with pytest.raises(ContractError, match="already registered"): + asyncio.run(check_unregistered(tmp_path, "linear", sys.executable, tmp_path / "probe2.log")) + + +def test_real_registry_resolution_in_isolated_interpreter(registration_run) -> None: + pytest.importorskip("torch") + pytest.importorskip("numpy") + pytest.importorskip("PIL") + run = registration_run + asyncio.run(check_unregistered(run.base, "matmul_t", sys.executable, run.root / "real-preflight.log")) + write_candidate(run) + result = asyncio.run(checks.registry_probe(run, run.root / "real-registry.log", timeout=30)) + assert result["returncode"] == 0 + + +@pytest.mark.parametrize( + "body, effective", + [ + ("assert False", True), + ("assert True", False), + ("import pytest; pytest.skip('not tested')", False), + ("import unavailable_dependency", False), + ], +) +def test_c5_requires_each_test_to_execute_and_fail(tmp_path, monkeypatch, body, effective) -> None: + base, worktree = tmp_path / "base", tmp_path / "worktree" + minimal_package(base) + minimal_package(worktree) + (worktree / NEW_TEST).write_text(f"def test_gate():\n {body}\n", encoding="utf-8") + run = SimpleNamespace( + base=base, + worktree=worktree, + manifest={"test_python": sys.executable}, + verify=lambda: None, + save=lambda: None, + ) + + def restore(): + (base / NEW_TEST).unlink(missing_ok=True) + run.manifest["base_overlay"] = {} + + run.recover_base = restore + monkeypatch.setattr(checks, "registration", lambda _: {"tests": [f"{NEW_TEST}::test_gate"]}) + if effective: + result = asyncio.run(checks.baseline_tests(run, tmp_path / "checks", time.monotonic() + 30)) + assert result["commands"][0]["returncode"] == 1 + else: + with pytest.raises(ContractError, match="C5"): + asyncio.run(checks.baseline_tests(run, tmp_path / "checks", time.monotonic() + 30)) + assert not (base / NEW_TEST).exists() and not run.manifest["base_overlay"] + + +def test_c1_failure_stops_later_checks(registration_run, monkeypatch) -> None: + run = registration_run + (run.worktree / "unexpected.txt").write_text("bad", encoding="utf-8") + monkeypatch.setattr(checks, "check_structure", lambda _: pytest.fail("C2 must not run")) + results, feedback = asyncio.run(checks.run_checks(run, run.root / "attempts/0001", seconds=30)) + assert [r["check"] for r in results] == ["C1"] and feedback["check"] == "C1" + + +def test_worker_timeout_cleans_process(tmp_path) -> None: + with pytest.raises(TimeoutError): + asyncio.run( + command( + [sys.executable, "-c", "import time; time.sleep(30)"], + root=tmp_path, + output=tmp_path / "timeout.log", + timeout=0.2, + ) + ) + + +def test_status_never_loads_agent(registration_run, capsys, monkeypatch) -> None: + monkeypatch.setattr( + "scripts.kernel_tuning.registration.cli.execute", + lambda *a, **kw: pytest.fail("status started an agent"), + ) + assert dispatch(SimpleNamespace(registration_command="status", run=registration_run.root)) == 0 + assert json.loads(capsys.readouterr().out)["status"] == "ready" + + +@pytest.fixture +def humanize(monkeypatch): + """Use the official fake agents, adapting only POSIX journal IO on Windows.""" + pytest.importorskip("hmz") + import os + import tempfile + + from hmz.sdk import fakes + + if os.name == "nt": + from hmz.coganchor import atomic + + def writes(at, said, *, mode=None): + fd, name = tempfile.mkstemp(dir=at.parent) + try: + with os.fdopen(fd, "wb") as handle: + if isinstance(said, str): + handle.write(said.encode("utf-8")) + elif isinstance(said, bytes): + handle.write(said) + else: + handle.writelines(said) + handle.flush() + os.fsync(handle.fileno()) + os.replace(name, at) + finally: + Path(name).unlink(missing_ok=True) + + monkeypatch.setattr(atomic, "writes", writes) + return fakes + + +def fake_flow(run, humanize, coder, *, resume=False): + """Drive the real resumable flow and its framework state/session machinery.""" + from scripts.kernel_tuning.registration.flow import operator_registration + + return asyncio.run( + humanize.run_fake( + operator_registration, + "register test only", + agents={"coder": coder}, + local=humanize.FakeEnvDriver(workdir=run.root.as_posix()), + params={"run_dir": str(run.root)}, + journal=run.root / "fake-hmz.jsonl", + resume=resume, + ) + ) + + +def scripted_coder(run, humanize, *, interrupt=None, sessions=None): + """Modify a real worktree only during implementation; collect session handles.""" + + def reply(instruction, **kwargs): + if sessions is not None: + sessions.append(id(kwargs["session"])) + assert Path(str(kwargs["session"].placement.workdir)).resolve() == run.worktree + if "Write PLAN.md before" in instruction: + (run.worktree / "PLAN.md").write_text("Preserve F.linear arithmetic.", encoding="utf-8") + else: + if interrupt is not None: + raise interrupt + if not (run.worktree / "REGISTRATION.json").exists(): + write_candidate(run) + return "done" + + return humanize.FakeAgentDriver(reply=reply) + + +def test_flow_feedback_repair_same_session_then_success(registration_run, humanize, monkeypatch) -> None: + run = registration_run + sessions, calls = [], [] + + async def check(candidate, archive, *, seconds): + assert seconds > 0 + check_scope(candidate) + checks.check_structure(candidate) + calls.append(archive) + failure = ( + {"check": "C3", "summary": "synthetic formatting error", "details": {}} + if len(calls) == 1 + else None + ) + if failure: + atomic_json(archive / "feedback.json", failure) + return [{"check": "C3", "passed": failure is None, "details": {}}], failure + + monkeypatch.setattr("scripts.kernel_tuning.registration.flow.run_checks", check) + coder = scripted_coder(run, humanize, sessions=sessions) + result = fake_flow(run, humanize, coder) + assert result["status"] == "passed" + assert len(coder.prompts) == 3 and len(set(sessions)) == 1 + assert [a["status"] for a in result["attempts"]] == ["failed", "passed"] + assert "synthetic formatting error" in (run.root / "attempts/0002/repair-prompt.txt").read_text( + encoding="utf-8" + ) + assert (run.root / "patch.diff").is_file() and (run.root / "SUMMARY.md").is_file() + assert (run.root / "attempts/0001/patch.diff").is_file() + assert coder.live == 0 + + +def test_flow_replans_after_two_equal_failures_and_exhausts(registration_run, humanize, monkeypatch) -> None: + run = registration_run + sessions = [] + + async def check(*args, **kwargs): + return [{"check": "C2", "passed": False}], {"check": "C2", "summary": "synthetic", "details": {}} + + monkeypatch.setattr("scripts.kernel_tuning.registration.flow.run_checks", check) + result = fake_flow(run, humanize, scripted_coder(run, humanize, sessions=sessions)) + assert result["status"] == "exhausted" and len(result["attempts"]) == 5 + assert sessions[0] == sessions[1] == sessions[2] + assert sessions[3] != sessions[2] + assert (run.root / "attempts/0003/plan-prompt.txt").is_file() + assert (run.root / "attempts/0004/repair-prompt.txt").is_file() + assert (run.root / "attempts/0005/plan-prompt.txt").is_file() + + +def test_resume_keeps_completed_plan_and_budget(registration_run, humanize, monkeypatch) -> None: + run = registration_run + + async def check(*args, **kwargs): + return [{"check": "C1", "passed": True, "details": {}}], None + + monkeypatch.setattr("scripts.kernel_tuning.registration.flow.run_checks", check) + with pytest.raises(asyncio.CancelledError): + fake_flow(run, humanize, scripted_coder(run, humanize, interrupt=asyncio.CancelledError())) + interrupted = RegistrationRun(run.root) + assert interrupted.manifest["status"] == "interrupted" + assert interrupted.manifest["attempts"][0]["stage"] == "implement" + elapsed = interrupted.manifest["elapsed_seconds"] + coder = scripted_coder(run, humanize) + result = fake_flow(run, humanize, coder, resume=True) + assert result["status"] == "passed" and len(result["attempts"]) == 1 + assert len(coder.prompts) == 1 and "Implement PLAN.md" in coder.prompts[0] + assert result["elapsed_seconds"] > elapsed + + +def test_checks_resume_does_not_repeat_passed_gates(registration_run, monkeypatch) -> None: + run = registration_run + write_candidate(run) + archive = run.root / "attempts/0001" + atomic_json( + archive / "checks.json", + {"fingerprint": run.fingerprint(), "checks": [{"check": "C1", "passed": True}]}, + ) + monkeypatch.setattr(checks, "check_scope", lambda _: pytest.fail("C1 reran")) + monkeypatch.setattr( + checks, "check_structure", lambda _: (_ for _ in ()).throw(ContractError("C2 intentionally fails")) + ) + results, failure = asyncio.run(checks.run_checks(run, archive, seconds=30)) + assert [r["check"] for r in results] == ["C1", "C2"] and failure["check"] == "C2" + + +def test_structure_rejects_missing_kernel_forwarding(registration_run) -> None: + run = registration_run + write_candidate(run) + source = run.worktree / SOURCE + source.write_text( + source.read_text(encoding="utf-8").replace( + "_fused_qkv(attn, h, kernels=self._kernels)", "_fused_qkv(attn, h)" + ), + encoding="utf-8", + ) + with pytest.raises(ContractError, match="not forwarded"): + checks.check_structure(run) + + +def test_patch_excludes_staged_artifacts_and_preserves_no_newline(registration_run) -> None: + run = registration_run + write_candidate(run) + (run.worktree / NEW_TEST).write_text("def test_no_newline(): pass", encoding="utf-8") + git(run.worktree, "add", "--", "PLAN.md", "REGISTRATION.json") + patch = run.patch() + assert "b/PLAN.md" not in patch and "b/REGISTRATION.json" not in patch + assert "\\ No newline at end of file" in patch + patch_file = run.root / "verify.diff" + patch_file.write_text(patch, encoding="utf-8") + git(run.base, "apply", "--check", str(patch_file)) + + +def test_all_gates_run_real_cpu_workers(tmp_path) -> None: + """Exercise C1–C5 together on a tiny model-free package, with no mocked workers.""" + from scripts.kernel_tuning.contracts import digest + + base, worktree = tmp_path / "base", tmp_path / "worktree" + minimal_package(base) + (base / SOURCE).parent.mkdir(parents=True, exist_ok=True) + (base / SOURCE).write_text( + """def _op(kernels, op, use_triton): + return kernels.resolve(op) + + +def raw(x, weight): + return x * weight + + +def compute(x, weight, kernels=None, quantized=False): + if quantized: + return 99 + return raw(x, weight) +""", + encoding="utf-8", + ) + (base / REFERENCE).write_text( + "from __future__ import annotations\n\nimport typing\n\nif typing.TYPE_CHECKING:\n import torch\n", + encoding="utf-8", + ) + (base / REGISTRY).write_text( + """from importlib import import_module + +_ops = {} + + +def register_lazy_op(op, impl, module, function): + _ops[op] = (module, function) + + +def available_ops(): + return {op: ["torch"] for op in _ops} + + +def get_op(op): + module, function = _ops[op] + return getattr(import_module(module), function) + + +class KernelSelection: + def __init__(self, mode): + self.mode = mode + + def prepare(self): + for op in _ops: + self.resolve(op) + + def resolve(self, op): + return get_op(op) +""", + encoding="utf-8", + ) + (base / "embodiinfer/layers/__init__.py").write_text( + 'from .ops import KernelSelection, available_ops, get_op\n\n__all__ = ["KernelSelection", "available_ops", "get_op"]\n', + encoding="utf-8", + ) + (base / "tests/test_kernel_ops.py").write_text( + "def test_existing():\n assert 1 + 1 == 2\n", encoding="utf-8" + ) + subprocess.run([sys.executable, "-m", "ruff", "format", str(base)], check=True, capture_output=True) + shutil.copytree(base, worktree) + spec = specification(sites=[{"source": SOURCE, "function": "compute", "expression": "raw(x, weight)"}]) + atomic_json(tmp_path / "sites.json", locate(base, spec)) + source = worktree / SOURCE + source.write_text( + source.read_text(encoding="utf-8").replace( + "return raw(x, weight)", 'return _op(kernels, "matmul_t", False)(x, weight)' + ), + encoding="utf-8", + ) + with (worktree / REFERENCE).open("a", encoding="utf-8") as stream: + stream.write( + '\n\ndef matmul_t(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:\n """Synthetic multiplication reference used only to test the controller."""\n return x * weight\n' + ) + with (worktree / REGISTRY).open("a", encoding="utf-8") as stream: + stream.write('\nregister_lazy_op("matmul_t", "torch", "embodiinfer.backend.torch.ops", "matmul_t")\n') + (worktree / NEW_TEST).write_text( + """from embodiinfer.backend.torch import ops + + +def test_reference(): + assert hasattr(ops, "matmul_t") + assert ops.matmul_t(2, 3) == 6 + + +def test_route(): + from embodiinfer.policies.pi05.modeling_pi05 import compute + + calls = [] + + class Recording: + def resolve(self, op): + calls.append(op) + return lambda x, w: x * w + + assert compute(2, 3, Recording()) == 6 + assert calls == ["matmul_t"] +""", + encoding="utf-8", + ) + atomic_json( + worktree / "REGISTRATION.json", + { + "op": "matmul_t", + "reference": f"{REFERENCE}::matmul_t", + "sites": [ + { + "id": "S1", + "file": SOURCE, + "function": "compute", + "call": '_op(kernels, "matmul_t", False)(x, weight)', + } + ], + "preserved": [{"site": "S1", "branch": "quantized"}], + "tests": [f"{NEW_TEST}::test_reference", f"{NEW_TEST}::test_route"], + }, + ) + subprocess.run([sys.executable, "-m", "ruff", "format", str(worktree)], check=True, capture_output=True) + changed = [SOURCE, REFERENCE, REGISTRY, NEW_TEST, "REGISTRATION.json"] + run = SimpleNamespace( + root=tmp_path, + base=base, + worktree=worktree, + spec=spec, + base_commit="synthetic", + manifest={"test_python": sys.executable}, + verify=lambda: None, + save=lambda: None, + changes=lambda: changed, + fingerprint=lambda: digest({n: file_hash(worktree / n) for n in changed}), + ) + + def restore(): + (base / NEW_TEST).unlink(missing_ok=True) + run.manifest["base_overlay"] = {} + + run.recover_base = restore + results, feedback = asyncio.run(checks.run_checks(run, tmp_path / "checks", seconds=60)) + assert feedback is None, feedback + assert [r["check"] for r in results] == ["C1", "C2", "C3", "C4", "C5"] + assert all(r["passed"] for r in results) + assert not (base / NEW_TEST).exists() + + +@pytest.mark.parametrize("unfenced", [False, True]) +def test_flow_sessions_run_under_the_recorded_fence( + registration_run, humanize, monkeypatch, unfenced: bool +) -> None: + from scripts.kernel_tuning.permissions import FENCED, UNFENCED + + run = registration_run + run.manifest["unfenced"] = unfenced + run.save() + + async def check(*args, **kwargs): + return [{"check": "C1", "passed": True, "details": {}}], None + + monkeypatch.setattr("scripts.kernel_tuning.registration.flow.run_checks", check) + coder = scripted_coder(run, humanize) + assert fake_flow(run, humanize, coder)["status"] == "passed" + expected = UNFENCED if unfenced else FENCED + assert coder.sessions and all(session.permission == expected for session in coder.sessions) + + +def test_unfenced_choice_is_frozen_with_the_agent(registration_run) -> None: + run = registration_run + assert run.manifest["unfenced"] is False + with pytest.raises(ContractError, match="frozen"): + run.prepare("codex/fake:low", unfenced=True) + run.prepare("codex/fake:low") diff --git a/tests/kernel_tuning/test_verify.py b/tests/kernel_tuning/test_verify.py new file mode 100644 index 0000000..dc03d37 --- /dev/null +++ b/tests/kernel_tuning/test_verify.py @@ -0,0 +1,60 @@ +"""CPU tests for the model-level numerical gate.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import pytest +from scripts.kernel_tuning.contracts import ContractError +from scripts.kernel_tuning.verify import flatten_outputs, gate + + +@dataclass +class _Action: + kind: str + values: dict + + +@dataclass +class _Result: + action_space: str + actions: tuple + timing: dict + + +def _out(values, token=3): + return {".actions": ("continuous", list(values)), ".token": ("discrete", [token])} + + +def test_flatten_outputs_splits_continuous_and_discrete_leaves(): + result = _Result("joint", (_Action("chunk", {"q": [[0.5, 1.0], [2.0, 3.0]], "id": 7}),), {"ms": 1.0}) + leaves = flatten_outputs(result) + assert leaves[".actions[0].values.q[0]"] == ("continuous", [0.5, 1.0]) + assert leaves[".actions[0].values.id"] == ("discrete", [7]) + assert leaves[".action_space"] == ("discrete", ["joint"]) + assert not any("timing" in name for name in leaves) + + +def test_gate_compares_mean_and_worst_errors_over_all_inputs(): + reference = [_out([1.0, 2.0]), _out([1.0, 2.0])] + baseline = [_out([1.001, 2.0]), _out([1.003, 2.0])] # mean max-abs 2e-3, worst 3e-3 + # The first input exceeds twice its own baseline error, but the means stay within 25%. + report = gate(reference, baseline, [_out([1.0024, 2.0]), _out([1.0, 2.0])]) + assert report["passed"] and report["failed"] == [] + report = gate(reference, baseline, [_out([1.003, 2.0]), _out([1.003, 2.0])]) # mean 3e-3 > 2.5e-3 + assert not report["passed"] and report["failed"] == [".actions"] + report = gate(reference, baseline, [_out([1.0, 2.0]), _out([1.0065, 2.0])]) # worst 6.5e-3 > 6e-3 + worst = report["outputs"][".actions"]["worst_max_abs"] + assert not report["passed"] and worst["candidate"] > worst["limit"] + exact = [_out([1.0, 2.0])] + assert gate(exact, exact, [_out([1.00001, 2.0])])["passed"] # within the 1e-4 / 1e-5 floors + assert not gate(exact, exact, [_out([1.0002, 2.0])])["passed"] + + +def test_gate_requires_identical_discrete_outputs_and_comparable_reference(): + reference, baseline = [_out([1.0])], [_out([1.0])] + assert not gate(reference, baseline, [_out([1.0], token=4)])["passed"] + with pytest.raises(ContractError, match="did not share inputs/noise"): + gate(reference, [_out([5.0])], [_out([5.0])]) + with pytest.raises(ContractError, match="one output per input"): + gate(reference, baseline, []) diff --git a/tests/test_kernel_ops.py b/tests/test_kernel_ops.py new file mode 100644 index 0000000..35e63b9 --- /dev/null +++ b/tests/test_kernel_ops.py @@ -0,0 +1,183 @@ +"""CPU tests for the operator registry and kernel-profile selection.""" + +import hashlib +import json +from types import SimpleNamespace + +import pytest +import torch + +import embodiinfer.backend.triton.tuned as tuned +from embodiinfer.backend.torch import ops as torch_ops +from embodiinfer.layers import KernelSelection, available_ops, get_op +from embodiinfer.policies.pi05 import modeling_pi05 + +_TUNED_MODULE = """ +CALLS = [] + + +def gated_residual(residual, update, gate): + CALLS.append(residual.shape) + return residual + update * gate + + +def applies(residual, update, gate): + return residual.shape[-1] == 4 +""" + +_ENVIRONMENT = {"device_name": "Test GPU", "capability": [9, 9], "torch": "t", "triton": "x"} + + +def _write_profile(root, *, source=_TUNED_MODULE, applies="applies", digest=None): + root.mkdir() + (root / "gated_residual.py").write_text(source, encoding="utf-8") + entry = { + "module": "gated_residual.py", + "function": "gated_residual", + "sha256": digest or hashlib.sha256(source.encode()).hexdigest(), + } + if applies: + entry["applies"] = applies + manifest = {"schema": 1, "match": _ENVIRONMENT, "ops": {"gated_residual": entry}} + (root / "profile.json").write_text(json.dumps(manifest), encoding="utf-8") + return tuned.load_profile(root) + + +def test_core_operators_have_torch_and_triton_implementations(): + ops = available_ops() + for op in ("gated_residual", "ada_rms_norm", "gated_gelu", "rotate_qk", "split_kv_attention"): + assert {"torch", "triton"} <= set(ops[op]) + assert get_op("gated_gelu", "torch") is torch_ops.gated_gelu + + +def test_unknown_operator_or_implementation_is_explicit(): + with pytest.raises(KeyError, match="unknown operator"): + get_op("not_an_op") + with pytest.raises(KeyError, match="no 'cuda' implementation"): + get_op("gated_residual", "cuda") + + +def test_none_keeps_preference_and_wraps_only_torch(): + selection = KernelSelection("none") + wrapped = [] + + def wrapper(fn): + wrapped.append(fn) + return fn + + assert selection.resolve("gated_residual", "torch", wrapper) is torch_ops.gated_residual + assert selection.resolve("gated_residual", "triton", wrapper) is get_op("gated_residual", "triton") + assert wrapped == [torch_ops.gated_residual] + assert selection.resolve("gated_residual", "torch", wrapper) is torch_ops.gated_residual + assert len(wrapped) == 1 # cached + + +def test_torch_mode_forces_reference_implementations(): + selection = KernelSelection("torch") + assert selection.resolve("rotate_qk", "triton") is torch_ops.rotate_qk + + +def test_profile_kernel_replaces_only_applicable_calls(tmp_path): + profile = _write_profile(tmp_path / "fake") + selection = KernelSelection(profile) + op = selection.resolve("gated_residual", "torch") + module = profile.kernel("gated_residual") + calls = module.fn.__globals__["CALLS"] + x = torch.ones(1, 2, 4) + torch.testing.assert_close(op(x, x, x), x + x * x) + assert calls == [x.shape] + y = torch.ones(1, 2, 3) + torch.testing.assert_close(op(y, y, y), y + y * y) + assert calls == [x.shape] # the 3-wide call fell back + assert selection.resolve("gated_gelu", "torch") is torch_ops.gated_gelu + + +def test_profile_without_applies_always_uses_tuned_kernel(tmp_path): + profile = _write_profile(tmp_path / "fake", applies=None) + op = KernelSelection(profile).resolve("gated_residual", "triton") + assert op is profile.kernel("gated_residual").fn + + +def test_modified_tuned_kernel_is_rejected(tmp_path): + profile = _write_profile(tmp_path / "fake", digest="0" * 64) + with pytest.raises(ValueError, match="sha256"): + KernelSelection(profile).resolve("gated_residual") + + +def test_profile_selection_by_name_and_auto(tmp_path, monkeypatch): + profile = _write_profile(tmp_path / "fake") + monkeypatch.setattr(tuned, "available_profiles", lambda: {"fake": profile}) + monkeypatch.setattr(tuned, "current_environment", lambda: dict(_ENVIRONMENT)) + assert KernelSelection("fake").profile is profile + assert KernelSelection("auto").profile is profile + with pytest.raises(ValueError, match="unknown kernel_profile"): + KernelSelection("missing") + monkeypatch.setattr(tuned, "current_environment", lambda: {**_ENVIRONMENT, "triton": "other"}) + assert KernelSelection("auto").profile is None + with pytest.raises(RuntimeError, match="not this environment"): + KernelSelection("fake") + + +def test_profile_manifest_requires_match_keys(tmp_path): + root = tmp_path / "broken" + root.mkdir() + (root / "profile.json").write_text(json.dumps({"schema": 1, "match": {}, "ops": {}}), encoding="utf-8") + with pytest.raises(ValueError, match="match is missing"): + tuned.load_profile(root) + + +def test_committed_profiles_are_valid(): + for profile in tuned.available_profiles().values(): + for op in profile.ops: + assert op in available_ops() + + +def test_pi05_helpers_use_the_registry_reference(tmp_path): + x = torch.randn(2, 3, 4, dtype=torch.bfloat16) + modulation = torch.randn(2, 12, dtype=torch.bfloat16) + norm = SimpleNamespace(eps=1e-6, dense=object()) + out, gate = modeling_pi05._rmsnorm(norm, x, None, False, modulation) + expected = torch_ops.ada_rms_norm(x, modulation, 1e-6) + assert torch.equal(out, expected[0]) and torch.equal(gate, expected[1]) + + profile = _write_profile(tmp_path / "fake") + calls = profile.kernel("gated_residual").fn.__globals__["CALLS"] + gate = torch.randn(2, 1, 4) + result = modeling_pi05._gated_residual( + x.float(), x.float(), gate, False, kernels=KernelSelection(profile) + ) + torch.testing.assert_close(result, x.float() + x.float() * gate) + assert calls == [x.shape] + + +def test_prepare_resolves_everything_before_graph_tracing(monkeypatch): + import embodiinfer.layers.ops as ops_module + + selection = KernelSelection("none") + selection.prepare((None, torch_ops_wrapper := (lambda fn: fn))) + monkeypatch.setattr(ops_module, "get_op", lambda *args: pytest.fail("resolve imported after prepare")) + for op in ("gated_residual", "ada_rms_norm", "gated_gelu", "rotate_qk", "split_kv_attention"): + for prefer in ("torch", "triton"): + selection.resolve(op, prefer) + selection.resolve(op, prefer, torch_ops_wrapper) + + +def test_split_kv_attention_reference_matches_masked_gqa_attention(): + torch.manual_seed(0) + q = torch.randn(2, 4, 3, 8, dtype=torch.float64) + pk, pv = torch.randn(2, 2, 5, 8, dtype=torch.float64), torch.randn(2, 2, 5, 8, dtype=torch.float64) + k, v = torch.randn(2, 2, 3, 8, dtype=torch.float64), torch.randn(2, 2, 3, 8, dtype=torch.float64) + mask = torch.zeros(2, 1, 1, 8, dtype=torch.float64) + mask[0, ..., 1] = torch.finfo(torch.float64).min + mask[1, ..., 6] = -1.0 + actual = torch_ops.split_kv_attention(q, pk, pv, k, v, mask, 0.3) + keys = torch.cat((pk, k), dim=2).repeat_interleave(2, dim=1) + values = torch.cat((pv, v), dim=2).repeat_interleave(2, dim=1) + scores = (q @ keys.transpose(-1, -2)) * 0.3 + scores = scores.masked_fill(mask != 0, float("-inf")) + torch.testing.assert_close(actual, torch.softmax(scores, dim=-1) @ values) + torch.testing.assert_close( + torch_ops.split_kv_attention(q, pk[:, :, :0], pv[:, :, :0], k, v, None, 0.3), + torch.softmax((q @ k.repeat_interleave(2, 1).transpose(-1, -2)) * 0.3, -1) + @ v.repeat_interleave(2, 1), + ) diff --git a/tests/test_pi05_cpu_routes.py b/tests/test_pi05_cpu_routes.py index 67d2e8e..289a23c 100644 --- a/tests/test_pi05_cpu_routes.py +++ b/tests/test_pi05_cpu_routes.py @@ -5,6 +5,7 @@ import torch import embodiinfer.policies.pi05.modeling_pi05 as modeling_pi05 +from embodiinfer.layers import KernelSelection from embodiinfer.policies.pi05.modeling_pi05 import Pi05Policy, Pi05Prefix, _adarms_projection, _mlp from embodiinfer.policies.pi05.processor_pi05 import Pi05Batch @@ -35,6 +36,7 @@ def _policy_stub(*, attention: str = "eager"): compile_backend="none", denoise_attention="sdpa", _attn=_AttentionBackend(), + _kernels=KernelSelection("none"), ) return policy diff --git a/tests/test_registry_gelu_tanh.py b/tests/test_registry_gelu_tanh.py new file mode 100644 index 0000000..b060b80 --- /dev/null +++ b/tests/test_registry_gelu_tanh.py @@ -0,0 +1,535 @@ +"""CPU bit-exact parity and guarded SigLIP dispatch for tanh GELU.""" + +import functools +import math +import sys +from pathlib import Path +from types import ModuleType, SimpleNamespace + +import pytest +import torch +import torch.nn.functional as F +from torch import nn + +from embodiinfer.backend.triton.tuned import KernelProfile, TunedKernel +from embodiinfer.engine.parallel.tensor_parallel import ( + ColumnParallelLinear, + RowParallelLinear, + TensorParallelContext, +) +from embodiinfer.layers import KernelSelection, available_ops, get_op +from embodiinfer.models.linear import QuantizedLinear +from embodiinfer.policies.pi05 import modeling_pi05 as pi05 + +_DTYPES = (torch.float32, torch.bfloat16) + + +@pytest.fixture(autouse=True) +def _fixed_cpu_numerics(): + state = torch.get_rng_state() + precision = torch.get_float32_matmul_precision() + torch.random.default_generator.manual_seed(20261014) + torch.set_float32_matmul_precision("highest") + try: + yield + finally: + torch.set_rng_state(state) + torch.set_float32_matmul_precision(precision) + + +@pytest.fixture +def transformers_tanh(monkeypatch): + # The recorded CPU environment lacks Transformers. Exercise canonical class + # identity with a scoped double of its gelu_pytorch_tanh implementation. + module = ModuleType("transformers.activations") + + class PytorchGELUTanh(nn.Module): + # Mirrors transformers' GELUTanh, including its Python-formula variant. + def __init__(self, use_gelu_tanh_python=False): + super().__init__() + if use_gelu_tanh_python: + self.act = self._python + else: + self.act = functools.partial(F.gelu, approximate="tanh") + + @staticmethod + def _python(x): + return 0.5 * x * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi) * (x + 0.044715 * torch.pow(x, 3.0)))) + + def forward(self, x): + return self.act(x) + + PytorchGELUTanh.__module__ = module.__name__ + module.PytorchGELUTanh = PytorchGELUTanh + monkeypatch.setitem(sys.modules, module.__name__, module) + return PytorchGELUTanh + + +def _equal(actual, expected): + assert actual.shape == expected.shape + assert actual.dtype == expected.dtype + assert actual.device == expected.device + assert torch.equal(actual, expected) + + +def _profile(applies=None): + calls, linears = [], [] + + def gelu(x): + output = F.gelu(x, approximate="tanh") + calls.append(SimpleNamespace(x=x, metadata=(x.shape, x.stride(), x.dtype), output=output)) + return output + + def linear(x, weight, bias=None): + output = F.linear(x, weight, bias) + linears.append(SimpleNamespace(x=x, weight=weight, bias=bias, output=output)) + return output + + profile = KernelProfile( + name="gelu_tanh_cpu_test", + root=Path(__file__).parent, + match={}, + ops={"gelu_tanh": {}, "linear": {}}, + _loaded={"gelu_tanh": TunedKernel(gelu, applies), "linear": TunedKernel(linear, None)}, + ) + return KernelSelection(profile), calls, linears + + +def _dense(k, n, dtype): + return nn.Linear(k, n, dtype=dtype).eval() + + +class _Attention: + def __init__(self): + self.calls = [] + + def attend(self, q, k, v, *, scaling): + self.calls.append((q, k, v, scaling)) + return F.scaled_dot_product_attention(q, k, v, scale=scaling) + + +def _vision(dtype, selection, activation=None): + attention = SimpleNamespace( + q_proj=_dense(8, 8, dtype), + k_proj=_dense(8, 8, dtype), + v_proj=_dense(8, 8, dtype), + out_proj=_dense(8, 8, dtype), + num_heads=2, + head_dim=4, + scale=0.5, + ) + layer = SimpleNamespace( + self_attn=attention, + layer_norm1=nn.LayerNorm(8, dtype=dtype), + layer_norm2=nn.LayerNorm(8, dtype=dtype), + mlp=SimpleNamespace( + fc1=_dense(8, 12, dtype), + fc2=_dense(12, 8, dtype), + activation_fn=nn.GELU(approximate="tanh") if activation is None else activation, + ), + ) + vision = SimpleNamespace( + embeddings=SimpleNamespace( + patch_embedding=nn.Conv2d(3, 8, 2, stride=2, dtype=dtype), + position_embedding=nn.Embedding(4, 8, dtype=dtype), + position_ids=torch.arange(4).unsqueeze(0), + ), + encoder=SimpleNamespace(layers=[layer]), + post_layernorm=nn.LayerNorm(8, dtype=dtype), + ) + pg = SimpleNamespace( + vision_tower=SimpleNamespace(vision_model=vision), + multi_modal_projector=SimpleNamespace(linear=_dense(8, 7, dtype)), + config=SimpleNamespace(text_config=SimpleNamespace(hidden_size=7)), + ) + policy = SimpleNamespace( + _m=SimpleNamespace(paligemma_with_expert=SimpleNamespace(paligemma=SimpleNamespace(model=pg))), + _kernels=selection, + _get_optimizations=lambda: None, + _vision_attn=_Attention(), + ) + policy._embed_image = pi05.Pi05Policy._embed_image.__get__(policy, type(policy)) + return policy, layer + + +def _original_image(policy, image): + # Frozen baseline arithmetic, with the existing Torch linear reference + # expanded to its single F.linear call; activation still calls its module. + pg = policy._m.paligemma_with_expert.paligemma.model + vision = pg.vision_tower.vision_model + emb = vision.embeddings + h = emb.patch_embedding(image.to(emb.patch_embedding.weight.dtype)) + h = h.flatten(2).transpose(1, 2) + h = h + emb.position_embedding(emb.position_ids) + for layer in vision.encoder.layers: + x = layer.layer_norm1(h) + a = layer.self_attn + b, s, _ = x.shape + qkv = [] + for proj in (a.q_proj, a.k_proj, a.v_proj): + if type(proj) is nn.Linear: # noqa: SIM108 — retain the baseline branch structure. + projected = F.linear(x, proj.weight, proj.bias) + else: + projected = proj(x) + qkv.append(projected.view(b, s, a.num_heads, a.head_dim).transpose(1, 2)) + q, k, v = qkv + y = policy._vision_attn.attend(q, k, v, scaling=a.scale) + if type(a.out_proj) is nn.Linear: + h = h + F.linear( + y.transpose(1, 2).reshape(b, s, -1).contiguous(), a.out_proj.weight, a.out_proj.bias + ) + else: + h = h + a.out_proj(y.transpose(1, 2).reshape(b, s, -1).contiguous()) + x = layer.layer_norm2(h) + if type(layer.mlp.fc1) is nn.Linear: + activated = F.linear(x, layer.mlp.fc1.weight, layer.mlp.fc1.bias) + else: + activated = layer.mlp.fc1(x) + activated = layer.mlp.activation_fn(activated) + if type(layer.mlp.fc2) is nn.Linear: + h = h + F.linear(activated, layer.mlp.fc2.weight, layer.mlp.fc2.bias) + else: + h = h + layer.mlp.fc2(activated) + h = vision.post_layernorm(h) + h = pg.multi_modal_projector.linear(h) + scale = pg.config.text_config.hidden_size**0.5 + return ((h / scale) * scale).to(image.dtype) + + +def _positive_control(selection, calls, linears, dtype): + policy, _ = _vision(dtype, selection) + image = torch.randn(2, 3, 4, 4, dtype=dtype) + with torch.no_grad(): + _equal(policy._embed_image(image), _original_image(policy, image)) + assert len(calls) == 1, "SigLIP tanh GELU must reach the supplied profile" + calls.clear() + linears.clear() + + +def test_gelu_tanh_reference_fp32_bf16_matches_original_expression(): + from embodiinfer.backend.torch import ops + + assert callable(getattr(ops, "gelu_tanh", None)), "Torch gelu_tanh reference is missing" + from embodiinfer.backend.torch.ops import gelu_tanh + + for dtype in _DTYPES: + inputs = [torch.randn(*shape, dtype=dtype) for shape in ((17,), (2, 3, 12))] + inputs.append(torch.randn(2, 12, 3, dtype=dtype).transpose(1, 2)) + assert not inputs[-1].is_contiguous() + for x in inputs: + before = x.clone() + actual = gelu_tanh(x) + _equal(actual, F.gelu(x, approximate="tanh")) + _equal(actual, nn.GELU(approximate="tanh")(x)) + _equal(x, before) + + leaf = x.detach().clone().requires_grad_() + reference_leaf = x.detach().clone().requires_grad_() + upstream = torch.randn_like(x) + actual_grad = torch.autograd.grad(gelu_tanh(leaf), leaf, upstream)[0] + expected_grad = torch.autograd.grad( + F.gelu(reference_leaf, approximate="tanh"), reference_leaf, upstream + )[0] + _equal(actual_grad, expected_grad) + + +def test_gelu_tanh_registry_torch_only_and_profile_dispatch(): + assert "gelu_tanh" in available_ops(), "gelu_tanh registry entry is missing" + from embodiinfer.backend.torch.ops import gelu_tanh + + assert available_ops()["gelu_tanh"] == ["torch"] + assert get_op("gelu_tanh") is gelu_tanh + for mode in ("none", "torch"): + selection = KernelSelection(mode) + selection.prepare() + assert selection.resolve("gelu_tanh", "torch") is gelu_tanh + assert KernelSelection("torch").resolve("gelu_tanh", "triton") is gelu_tanh + selection, calls, _ = _profile(lambda x: x.shape[-1] == 12) + op = selection.resolve("gelu_tanh", "torch") + assert selection.resolve("gelu_tanh", "torch") is op + assert selection.resolve("gated_gelu") is get_op("gated_gelu") + for dtype in _DTYPES: + x = torch.randn(2, 12, 3, dtype=dtype).transpose(1, 2) + _equal(op(x), F.gelu(x, approximate="tanh")) + assert calls[-1].x is x + assert calls[-1].metadata == (x.shape, x.stride(), dtype) + count = len(calls) + rejected = torch.randn(2, 3, 7, dtype=dtype) + _equal(op(rejected), F.gelu(rejected, approximate="tanh")) + assert len(calls) == count + assert len(calls) == 2 + + +def test_pi05_gelu_tanh_profile_preserves_original_vision_forward(transformers_tanh): + for factory in (lambda: nn.GELU(approximate="tanh"), transformers_tanh): + for dtype, image_dtype in ( + (torch.float32, torch.float32), + (torch.bfloat16, torch.bfloat16), + (torch.bfloat16, torch.float32), + ): + selection, calls, linears = _profile() + policy, layer = _vision(dtype, selection, factory()) + image = torch.randn(2, 3, 4, 4, dtype=image_dtype) + with torch.no_grad(): + expected = _original_image(policy, image) + actual = policy._embed_image(image) + _equal(actual, expected) + assert actual.shape == (2, 4, 7) and actual.dtype == image_dtype + assert len(calls) == 1, "SigLIP tanh GELU must reach the supplied profile" + assert len(linears) == 6 + assert calls[0].x is linears[4].output + assert calls[0].metadata == (torch.Size((2, 4, 12)), (48, 12, 1), dtype) + assert linears[5].x is calls[0].output + assert linears[4].weight is layer.mlp.fc1.weight + assert linears[5].weight is layer.mlp.fc2.weight + assert linears[4].bias is layer.mlp.fc1.bias + assert linears[5].bias is layer.mlp.fc2.bias + assert linears[3].x.is_contiguous() + for original, routed in zip(*policy._vision_attn.calls, strict=True): + if isinstance(original, torch.Tensor): + _equal(routed, original) + else: + assert routed == original + for mode in ("none", "torch"): + policy._kernels = KernelSelection(mode) + _equal(policy._embed_image(image), expected) + assert len(calls) == 1 and len(linears) == 6 + + # Preserve differentiability through the complete vision path too. + policy._kernels = selection + leaf = image.detach().clone().requires_grad_() + reference_leaf = image.detach().clone().requires_grad_() + actual = policy._embed_image(leaf) + expected = _original_image(policy, reference_leaf) + _equal(actual, expected) + upstream = torch.randn_like(actual) + _equal( + torch.autograd.grad(actual, leaf, upstream)[0], + torch.autograd.grad(expected, reference_leaf, upstream)[0], + ) + assert len(calls) == 2 + + +def test_pi05_gelu_tanh_other_activations_bypass_profile(transformers_tanh): + class AlteredGELU(nn.GELU): + def forward(self, x): + return super().forward(x) + 0.5 + + class AlteredTransformersGELU(transformers_tanh): + def forward(self, x): + return super().forward(x) - 0.25 + + class RecordingCallable: + def __init__(self): + self.calls = [] + + def __call__(self, x): + self.calls.append(x) + return F.gelu(x, approximate="tanh") + 0.125 + + for dtype in _DTYPES: + selection, calls, linears = _profile() + _positive_control(selection, calls, linears, dtype) + policy, layer = _vision(dtype, selection) + image = torch.randn(2, 3, 4, 4, dtype=dtype) + recorded = RecordingCallable() + overridden = nn.GELU(approximate="tanh") + overridden.forward = recorded + hooked = nn.GELU(approximate="tanh") + hook_calls = [] + + def hook(module, args, output, hook_calls=hook_calls): + hook_calls.append(args[0]) + return output + 0.25 + + hooked.register_forward_hook(hook) + impostor = type( + "PytorchGELUTanh", + (nn.Module,), + {"__module__": "transformers.activations", "forward": lambda self, x: F.silu(x)}, + )() + custom = RecordingCallable() + for activation in ( + nn.GELU(approximate="none"), + nn.ReLU(), + custom, + AlteredGELU(approximate="tanh"), + AlteredTransformersGELU(), + transformers_tanh(use_gelu_tanh_python=True), + overridden, + hooked, + impostor, + ): + layer.mlp.activation_fn = activation + with torch.no_grad(): + expected = _original_image(policy, image) + _equal(policy._embed_image(image), expected) + assert calls == [], "a different activation must keep its own callable" + assert len(custom.calls) == len(recorded.calls) == len(hook_calls) == 2 + + +def test_pi05_gelu_tanh_linear_module_fallbacks_are_preserved(monkeypatch): + class AlteredLinear(nn.Linear): + def forward(self, x): + return super().forward(x) + 0.5 + + class QuantizedProjection(QuantizedLinear): + def __init__(self, k, n, dtype): + super().__init__() + self.in_features, self.out_features = k, n + self.weight = nn.Parameter(torch.ones(n, k, dtype=torch.int8), False) + self.bias = nn.Parameter(torch.randn(n, dtype=dtype)) + self.compute_dtype = dtype + + @classmethod + def from_linear(cls, linear, config): + raise AssertionError("test double is constructed directly") + + @classmethod + def pack(cls, linears): + raise AssertionError("vision projections must not be packed") + + def forward(self, x): + return x[..., :1].expand(*x.shape[:-1], self.out_features) + 0.25 + self.bias + + reductions = [] + group = object() + + def reduce(output, *, group): + reductions.append((output.clone(), group)) + output.add_(0.25) + + monkeypatch.setattr(torch.distributed, "all_reduce", reduce) + + def column(k, n, dtype): + return ColumnParallelLinear.from_linear( + _dense(k, n * 2, dtype), TensorParallelContext(2, 1), name="gelu_test" + ) + + def row(k, n, dtype): + return RowParallelLinear.from_linear( + _dense(k * 2, n, dtype), TensorParallelContext(2, 1, group), name="gelu_test" + ) + + for dtype in _DTYPES: + selection, calls, linears = _profile() + _positive_control(selection, calls, linears, dtype) + for factory in ( + lambda k, n, dtype: AlteredLinear(k, n, dtype=dtype), + QuantizedProjection, + column, + row, + ): + policy, layer = _vision(dtype, selection) + layer.mlp.fc1 = factory(8, 12, dtype).eval() + layer.mlp.fc2 = factory(12, 8, dtype).eval() + module_calls = [] + + def track(module, args, output, module_calls=module_calls): + module_calls.append((module, args[0], output)) + + layer.mlp.fc1.register_forward_hook(track) + layer.mlp.fc2.register_forward_hook(track) + image = torch.randn(2, 3, 4, 4, dtype=dtype) + with torch.no_grad(): + reductions.clear() + expected = _original_image(policy, image) + reference_reductions = list(reductions) + reductions.clear() + module_calls.clear() + actual = policy._embed_image(image) + _equal(actual, expected) + assert len(calls) == 1 + assert len(linears) == 4 # Q/K/V and attention output only. + assert len(module_calls) == 2 + assert module_calls[0][0] is layer.mlp.fc1 + assert module_calls[1][0] is layer.mlp.fc2 + assert calls[0].x is module_calls[0][2] + assert module_calls[1][1] is calls[0].output + assert all( + call.weight is not projection.weight + for call in linears + for projection in (layer.mlp.fc1, layer.mlp.fc2) + ) + assert len(reductions) == len(reference_reductions) == (2 if factory is row else 0) + for (actual_partial, actual_group), (expected_partial, expected_group) in zip( + reductions, reference_reductions, strict=True + ): + _equal(actual_partial, expected_partial) + assert actual_group is expected_group is group + calls.clear() + linears.clear() + + +def test_pi05_gelu_tanh_prefix_keeps_policy_local_selection(): + for dtype in _DTYPES: + selection, calls, _ = _profile() + other_selection, other_calls, _ = _profile() + policy, _ = _vision(dtype, selection) + other = SimpleNamespace(**policy.__dict__) + other._kernels = other_selection + other._embed_image = pi05.Pi05Policy._embed_image.__get__(other, type(other)) + policy._prefix_tower = other._prefix_tower = SimpleNamespace( + embed_tokens=nn.Embedding(9, 7, dtype=dtype) + ) + batch = SimpleNamespace( + images=[torch.randn(2, 3, 4, 4, dtype=dtype) for _ in range(2)], + img_masks=[torch.tensor([True, False]), torch.tensor([False, True])], + tokens=torch.tensor([[1, 2, 3], [3, 2, 1]]), + masks=torch.tensor([[True, True, False], [True, True, True]]), + ) + with torch.no_grad(): + images = [_original_image(policy, image) for image in batch.images] + lang = policy._prefix_tower.embed_tokens(batch.tokens) + expected = torch.cat([*images, lang * lang.shape[-1] ** 0.5], dim=1) + pads = torch.cat([*(mask[:, None].expand(2, 4) for mask in batch.img_masks), batch.masks], dim=1) + actual = pi05.Pi05Policy._embed_prefix(policy, batch) + for value, reference in zip(actual, (expected, pads, torch.zeros_like(pads)), strict=True): + _equal(value, reference) + assert len(calls) == 2, "each camera must route GELU through its policy's selection" + assert other_calls == [] + other_actual = pi05.Pi05Policy._embed_prefix(other, batch) + for value, reference in zip(other_actual, actual, strict=True): + _equal(value, reference) + assert len(calls) == len(other_calls) == 2 + assert calls[0].x is not other_calls[0].x + + +def test_pi05_gelu_tanh_guard_and_fallback_use_one_activation_lookup(): + class ChangingActivationMLP: + def __init__(self, mlp, activation): + self.fc1, self.fc2 = mlp.fc1, mlp.fc2 + self.activation = activation + self.lookups = 0 + + @property + def activation_fn(self): + self.lookups += 1 + return self.activation if self.lookups == 1 else nn.ReLU() + + for dtype in _DTYPES: + for approximate in ("tanh", "none"): + selection, calls, _ = _profile() + policy, layer = _vision(dtype, selection) + layer.mlp = ChangingActivationMLP(layer.mlp, nn.GELU(approximate=approximate)) + image = torch.randn(2, 3, 4, 4, dtype=dtype) + with torch.no_grad(): + expected = _original_image(policy, image) + assert layer.mlp.lookups == 1 + layer.mlp.lookups = 0 + actual = policy._embed_image(image) + _equal(actual, expected) + assert layer.mlp.lookups == 1, "the guard and fallback must use the same activation lookup" + assert len(calls) == (1 if approximate == "tanh" else 0) + + +def test_pi05_gelu_tanh_real_transformers_activations(): + """The installed Transformers' SigLIP activation routes; its Python-formula variant does not.""" + activations = pytest.importorskip("transformers.activations") + native, python = activations.ACT2FN["gelu_pytorch_tanh"], activations.ACT2FN["gelu_python_tanh"] + # A dense deterministic range: the Python formula rounds differently on part of it. + x = torch.linspace(-8, 8, 100001) + assert torch.equal(native(x), F.gelu(x, approximate="tanh")) + assert not torch.equal(python(x), F.gelu(x, approximate="tanh")) + assert pi05._is_gelu_tanh(native) + assert not pi05._is_gelu_tanh(python) diff --git a/tests/test_registry_linear.py b/tests/test_registry_linear.py new file mode 100644 index 0000000..4a8e682 --- /dev/null +++ b/tests/test_registry_linear.py @@ -0,0 +1,673 @@ +"""CPU bit-exact parity and policy routing for the linear operator registration.""" + +import inspect +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch +import torch.nn.functional as F +from torch import nn + +from embodiinfer.backend.triton.tuned import KernelProfile, TunedKernel +from embodiinfer.engine.parallel.tensor_parallel import ( + ColumnParallelLinear, + RowParallelLinear, + TensorParallelContext, +) +from embodiinfer.layers import KernelSelection, available_ops, get_op +from embodiinfer.models.linear import QuantizedLinear +from embodiinfer.policies.pi05 import modeling_pi05 as pi05 + +_DTYPES = (torch.float32, torch.bfloat16) + + +@pytest.fixture(autouse=True) +def _fixed_cpu_numerics(): + state = torch.get_rng_state() + precision = torch.get_float32_matmul_precision() + torch.random.default_generator.manual_seed(20261014) + torch.set_float32_matmul_precision("highest") + try: + yield + finally: + torch.set_rng_state(state) + torch.set_float32_matmul_precision(precision) + + +def _equal(actual, expected): + assert actual.dtype == expected.dtype + assert actual.shape == expected.shape + assert torch.equal(actual, expected) + + +def _profile(applies=None): + assert "linear" in available_ops(), "linear must be registered before profile routing" + calls = [] + + def record(x, weight, bias=None): + output = F.linear(x, weight, bias) + args = (x, weight, bias) + calls.append( + SimpleNamespace( + args=args, + metadata=[None if t is None else (t.shape, t.stride(), t.dtype) for t in args], + output=output, + ) + ) + return output + + profile = KernelProfile( + name="linear_cpu_test", + root=Path(__file__).parent, + match={}, + ops={"linear": {}}, + _loaded={"linear": TunedKernel(record, applies)}, + ) + return KernelSelection(profile), calls + + +def _call_is(call, x, weight, bias): + for actual, expected in zip(call.args, (x, weight, bias), strict=True): + assert actual is expected + assert call.metadata == [None if t is None else (t.shape, t.stride(), t.dtype) for t in (x, weight, bias)] + + +def _pair(x, first, second, selection): + assert "kernels" in inspect.signature(pi05._fused_linear_pair).parameters + return pi05._fused_linear_pair(x, first, second, kernels=selection) + + +def _dense(in_features, out_features, dtype, bias=True): + return nn.Linear(in_features, out_features, bias=bias, dtype=dtype).eval() + + +def _mlp(dtype, bias=True): + return SimpleNamespace( + gate_proj=_dense(7, 11, dtype, False), + up_proj=_dense(7, 11, dtype, False), + down_proj=_dense(11, 7, dtype, bias), + ) + + +def _original_mlp(mlp, x, fuse=True): + if fuse: + weight = torch.cat((mlp.gate_proj.weight, mlp.up_proj.weight), dim=0).contiguous() + gate, up = F.linear(x, weight).split((11, 11), dim=-1) + else: + gate, up = mlp.gate_proj(x), mlp.up_proj(x) + return mlp.down_proj(F.gelu(gate, approximate="tanh") * up) + + +def _positive_control(selection, calls, dtype): + first, second = _dense(7, 5, dtype, False), _dense(7, 9, dtype, False) + x = torch.randn(2, 3, 7, dtype=dtype) + with torch.no_grad(): + result = _pair(x, first, second, selection) + assert len(calls) == 1 + _call_is(calls[0], x, first._vvla_fused_pair_weight, None) + for actual, expected in zip(result, calls[0].output.split((5, 9), dim=-1), strict=True): + _equal(actual, expected) + calls.clear() + + +class _Attention: + def __init__(self): + self.calls = [] + + def attend(self, q, k, v, *, scaling, attn_mask=None): + self.calls.append((q, k, v, scaling)) + return F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, scale=scaling) + + +class _Activation: + def __init__(self): + self.calls = [] + + def __call__(self, x): + result = F.gelu(x, approximate="none") + self.calls.append((x, result)) + return result + + +def _vision(dtype, selection): + attention = SimpleNamespace( + q_proj=_dense(8, 8, dtype), + k_proj=_dense(8, 8, dtype), + v_proj=_dense(8, 8, dtype), + out_proj=_dense(8, 8, dtype), + num_heads=2, + head_dim=4, + scale=0.5, + ) + layer = SimpleNamespace( + self_attn=attention, + layer_norm1=nn.LayerNorm(8, dtype=dtype), + layer_norm2=nn.LayerNorm(8, dtype=dtype), + mlp=SimpleNamespace(fc1=_dense(8, 12, dtype), fc2=_dense(12, 8, dtype), activation_fn=_Activation()), + ) + vision = SimpleNamespace( + embeddings=SimpleNamespace( + patch_embedding=nn.Conv2d(3, 8, 2, stride=2, dtype=dtype), + position_embedding=nn.Embedding(4, 8, dtype=dtype), + position_ids=torch.arange(4).unsqueeze(0), + ), + encoder=SimpleNamespace(layers=[layer]), + post_layernorm=nn.LayerNorm(8, dtype=dtype), + ) + pg = SimpleNamespace( + vision_tower=SimpleNamespace(vision_model=vision), + multi_modal_projector=SimpleNamespace(linear=_dense(8, 7, dtype)), + config=SimpleNamespace(text_config=SimpleNamespace(hidden_size=7)), + ) + policy = SimpleNamespace( + _m=SimpleNamespace(paligemma_with_expert=SimpleNamespace(paligemma=SimpleNamespace(model=pg))), + _kernels=selection, + _get_optimizations=lambda: None, + _vision_attn=_Attention(), + ) + slots = [(attention, name) for name in ("q_proj", "k_proj", "v_proj", "out_proj")] + slots += [(layer.mlp, "fc1"), (layer.mlp, "fc2")] + return policy, layer, slots + + +def _original_image(policy, image): + # Frozen S3–S6 expression sequence, including entry/exit casts and layouts. + pg = policy._m.paligemma_with_expert.paligemma.model + vision = pg.vision_tower.vision_model + emb = vision.embeddings + h = emb.patch_embedding(image.to(emb.patch_embedding.weight.dtype)) + h = h.flatten(2).transpose(1, 2) + h = h + emb.position_embedding(emb.position_ids) + for layer in vision.encoder.layers: + x = layer.layer_norm1(h) + a = layer.self_attn + b, s, _ = x.shape + q, k, v = [ + proj(x).view(b, s, a.num_heads, a.head_dim).transpose(1, 2) + for proj in (a.q_proj, a.k_proj, a.v_proj) + ] + y = policy._vision_attn.attend(q, k, v, scaling=a.scale) + h = h + a.out_proj(y.transpose(1, 2).reshape(b, s, -1).contiguous()) + x = layer.layer_norm2(h) + h = h + layer.mlp.fc2(layer.mlp.activation_fn(layer.mlp.fc1(x))) + h = vision.post_layernorm(h) + h = pg.multi_modal_projector.linear(h) + scale = pg.config.text_config.hidden_size**0.5 + return ((h / scale) * scale).to(image.dtype) + + +def _track(module): + module.route_inputs = [] + module.register_forward_hook(lambda owner, args, output: owner.route_inputs.append(args[0])) + return module + + +def _check_module_bypass(make, selection, calls, dtype): + x = torch.randn(2, 3, 7, dtype=dtype) + mlp = _mlp(dtype) + mlp.down_proj = make(11, 7, dtype, True) + expected = _original_mlp(mlp, x, fuse=False) + mlp.down_proj.route_inputs.clear() + actual = pi05._mlp(mlp, x, False, fuse_projections=False, kernels=selection) + _equal(actual, expected) + assert len(mlp.down_proj.route_inputs) == 1 + assert calls == [] + + policy, _, slots = _vision(dtype, selection) + for owner, name in slots: + original = getattr(owner, name) + setattr(owner, name, make(original.in_features, original.out_features, dtype, True)) + image = torch.randn(2, 3, 4, 4, dtype=dtype) + expected = _original_image(policy, image) + for owner, name in slots: + getattr(owner, name).route_inputs.clear() + _equal(pi05.Pi05Policy._embed_image(policy, image), expected) + assert all(len(getattr(owner, name).route_inputs) == 1 for owner, name in slots) + assert calls == [] + + +def test_linear_reference_fp32_bf16_matches_original_expression(): + from embodiinfer.backend.torch import ops + + assert callable(getattr(ops, "linear", None)), "Torch linear reference is missing" + from embodiinfer.backend.torch.ops import linear + + for dtype in _DTYPES: + inputs = [torch.randn(*shape, dtype=dtype) for shape in ((7,), (4, 7), (2, 3, 7))] + inputs.append(torch.randn(2, 7, 3, dtype=dtype).transpose(1, 2)) + assert not inputs[-1].is_contiguous() + for has_bias in (False, True): + module = _dense(7, 5, dtype, has_bias) + for x in inputs: + args = (x, module.weight, module.bias) + snapshots = [None if t is None else t.detach().clone() for t in args] + actual = linear(*args) + _equal(actual, F.linear(*args)) + _equal(actual, module(x)) + assert actual.dtype == dtype + for tensor, before in zip(args, snapshots, strict=True): + if tensor is not None: + _equal(tensor, before) + + x = inputs[2].detach().clone().requires_grad_() + weight = module.weight.detach().clone().requires_grad_() + operands = [x, weight] + if has_bias: + operands.append(module.bias.detach().clone().requires_grad_()) + reference_operands = [t.detach().clone().requires_grad_() for t in operands] + actual, expected = linear(*operands), F.linear(*reference_operands) + upstream = torch.randn_like(actual) + gradients = torch.autograd.grad(actual, operands, upstream) + reference_gradients = torch.autograd.grad(expected, reference_operands, upstream) + for gradient, reference in zip(gradients, reference_gradients, strict=True): + _equal(gradient, reference) + + for x, weight, bias in ( + (torch.randn(2, 8, dtype=dtype), torch.randn(5, 7, dtype=dtype), None), + (torch.randn(2, 7, dtype=dtype), torch.randn(5, 7, dtype=dtype), torch.randn(4, dtype=dtype)), + ): + with pytest.raises(RuntimeError) as original_error: + F.linear(x, weight, bias) + with pytest.raises(RuntimeError) as routed_error: + linear(x, weight, bias) + assert str(routed_error.value) == str(original_error.value) + + +def test_linear_registry_torch_only_and_profile_dispatch(): + assert "linear" in available_ops(), "linear registry entry is missing" + from embodiinfer.backend.torch.ops import linear + + assert available_ops()["linear"] == ["torch"] + assert get_op("linear") is linear + for mode in ("none", "torch"): + assert KernelSelection(mode).resolve("linear", "torch") is linear + selection, calls = _profile(lambda x, weight, bias=None: x.shape[-1] == 7) + op = selection.resolve("linear", "torch") + assert selection.resolve("linear", "torch") is op + assert selection.resolve("gated_gelu", "torch") is get_op("gated_gelu") + for dtype in _DTYPES: + for has_bias in (False, True): + module = _dense(7, 5, dtype, has_bias) + x = torch.randn(2, 3, 7, dtype=dtype) + _equal(op(x, module.weight, module.bias), module(x)) + _call_is(calls[-1], x, module.weight, module.bias) + count = len(calls) + other = _dense(6, 5, dtype, has_bias) + rejected = torch.randn(2, 3, 6, dtype=dtype) + _equal(op(rejected, other.weight, other.bias), other(rejected)) + assert len(calls) == count + assert len(calls) == 4 + + +def test_pi05_linear_pair_profile_preserves_packed_gemm(): + selection, calls = _profile() + for dtype in _DTYPES: + first, second = _dense(7, 5, dtype, False), _dense(7, 9, dtype, False) + x = torch.randn(2, 3, 7, dtype=dtype) + with torch.no_grad(): + weight = torch.cat((first.weight, second.weight), dim=0).contiguous() + expected = F.linear(x, weight).split((5, 9), dim=-1) + actual = _pair(x, first, second, selection) + assert len(calls) == 1 + packed = first._vvla_fused_pair_weight + assert packed.shape == (14, 7) and packed.is_contiguous() + _call_is(calls[0], x, packed, None) + for result, reference in zip(actual, expected, strict=True): + _equal(result, reference) + for mode in ("none", "torch"): + plain = _pair(x, first, second, KernelSelection(mode)) + for result, reference in zip(plain, expected, strict=True): + _equal(result, reference) + repeated = _pair(x, first, second, selection) + assert len(calls) == 2 + _call_is(calls[1], x, packed, None) + for result, reference in zip(repeated, expected, strict=True): + _equal(result, reference) + calls.clear() + + +def test_pi05_mlp_linear_profile_reaches_pair_and_down(monkeypatch): + selection, calls = _profile() + other_selection, other_calls = _profile() + monkeypatch.setattr(pi05, "_rmsnorm", lambda norm, x, *args: (x, None)) + for dtype in _DTYPES: + for bias in (False, True): + mlp = _mlp(dtype, bias) + x = torch.randn(2, 3, 7, dtype=dtype) + with torch.no_grad(): + expected = _original_mlp(mlp, x) + _equal(pi05._mlp(mlp, x, False, kernels=selection), expected) + assert len(calls) == 2 + assert [call.args[1].shape for call in calls] == [(22, 7), (7, 11)] + _call_is(calls[0], x, mlp.gate_proj._vvla_fused_pair_weight, None) + gate, up = calls[0].output.split((11, 11), dim=-1) + _equal(calls[1].args[0], F.gelu(gate, approximate="tanh") * up) + _call_is(calls[1], calls[1].args[0], mlp.down_proj.weight, mlp.down_proj.bias) + for mode in ("none", "torch"): + _equal(pi05._mlp(mlp, x, False, kernels=KernelSelection(mode)), expected) + _equal(pi05._mlp(mlp, x, False, kernels=other_selection), expected) + assert len(calls) == len(other_calls) == 2 + layer = SimpleNamespace( + self_attn=SimpleNamespace(q_proj=mlp.gate_proj), + input_layernorm=object(), + post_attention_layernorm=object(), + mlp=mlp, + ) + tower = SimpleNamespace(layers=[layer], norm=object(), rotary_emb=lambda *args: (None, None)) + policy = SimpleNamespace( + native_embeddings=False, + attention="sdpa", + compile_backend="none", + _kernels=selection, + _attn_sublayer=lambda attn, hidden, *args: torch.zeros_like(hidden), + ) + actual, collected = pi05.Pi05Policy._tower_forward( + policy, tower, x, torch.zeros(2, 3, dtype=torch.long), None, None + ) + _equal(actual, (x + torch.zeros_like(x)) + expected) + assert collected is None + assert len(calls) == 4 and len(other_calls) == 2 + calls.clear() + other_calls.clear() + + +def test_pi05_vision_linear_profile_preserves_original_forward(): + selection, calls = _profile() + for dtype, image_dtype in ( + (torch.float32, torch.float32), + (torch.bfloat16, torch.bfloat16), + (torch.bfloat16, torch.float32), + ): + policy, layer, slots = _vision(dtype, selection) + image = torch.randn(2, 3, 4, 4, dtype=image_dtype) + with torch.no_grad(): + expected = _original_image(policy, image) + actual = pi05.Pi05Policy._embed_image(policy, image) + _equal(actual, expected) + assert actual.dtype == image_dtype + assert len(calls) == 6 + assert [call.args[1].shape for call in calls] == [(8, 8)] * 4 + [(12, 8), (8, 12)] + for call, (owner, name) in zip(calls, slots, strict=True): + module = getattr(owner, name) + _call_is(call, call.args[0], module.weight, module.bias) + assert all(t.dtype == dtype for t in call.args) + assert call.args[0].shape[:2] == (2, 4) + assert calls[0].args[0] is calls[1].args[0] is calls[2].args[0] + assert calls[3].args[0].is_contiguous() + activation = layer.mlp.activation_fn + assert len(activation.calls) == 2 + assert activation.calls[1][0] is calls[4].output + assert calls[5].args[0] is activation.calls[1][1] + for original, routed in zip(*policy._vision_attn.calls, strict=True): + if isinstance(original, torch.Tensor): + _equal(routed, original) + else: + assert routed == original + for mode in ("none", "torch"): + policy._kernels = KernelSelection(mode) + _equal(pi05.Pi05Policy._embed_image(policy, image), expected) + assert len(calls) == 6 + calls.clear() + + +def test_pi05_linear_pair_gradient_and_training_bypass_profile(): + selection, calls = _profile() + for dtype in _DTYPES: + _positive_control(selection, calls, dtype) + for grad_enabled, first_training, second_training in ( + (True, False, False), + (False, True, False), + (False, False, True), + ): + first = _track(_dense(7, 5, dtype, False).train(first_training)) + second = _track(_dense(7, 9, dtype, False).train(second_training)) + x = torch.randn(2, 3, 7, dtype=dtype, requires_grad=grad_enabled) + with torch.set_grad_enabled(grad_enabled): + expected = first(x), second(x) + first.route_inputs.clear() + second.route_inputs.clear() + actual = _pair(x, first, second, selection) + assert len(first.route_inputs) == len(second.route_inputs) == 1 + for result, reference in zip(actual, expected, strict=True): + _equal(result, reference) + if grad_enabled: + upstream = tuple(torch.randn_like(t) for t in actual) + operands = (x, first.weight, second.weight) + gradients = torch.autograd.grad(actual, operands, upstream) + reference_gradients = torch.autograd.grad(expected, operands, upstream) + for gradient, reference in zip(gradients, reference_gradients, strict=True): + _equal(gradient, reference) + assert not hasattr(first, "_vvla_fused_pair_weight") + assert calls == [] + + +def test_pi05_linear_quantized_branches_bypass_profile(monkeypatch): + selection, calls = _profile() + + class QuantizedProjection(QuantizedLinear): + def __init__(self, in_features, out_features, dtype, bias): + super().__init__() + self.in_features, self.out_features = in_features, out_features + self.weight = nn.Parameter(torch.ones(out_features, in_features, dtype=torch.int8), False) + self.compute_dtype = dtype + self.bias = nn.Parameter(torch.randn(out_features, dtype=dtype)) if bias else None + self.route_inputs = [] + self.eval() + + @classmethod + def from_linear(cls, linear, config): + raise AssertionError("test double is constructed directly") + + @classmethod + def pack(cls, linears): + raise AssertionError("test must use the recording packed projection") + + def forward(self, x): + self.route_inputs.append(x) + result = x[..., :1].expand(*x.shape[:-1], self.out_features) + 0.25 + return result if self.bias is None else result + self.bias + + for dtype in _DTYPES: + _positive_control(selection, calls, dtype) + with torch.no_grad(): + _check_module_bypass(QuantizedProjection, selection, calls, dtype) + first, second = QuantizedProjection(7, 5, dtype, False), QuantizedProjection(7, 9, dtype, False) + x = torch.randn(2, 3, 7, dtype=dtype) + expected = first(x), second(x) + packing, packed_calls = [], [] + + def packed(value, first=first, second=second, packed_calls=packed_calls): + packed_calls.append(value) + return torch.cat((first(value), second(value)), dim=-1) + + def pack(projections, owner, name, packing=packing, packed=packed): + packing.append((projections, owner, name)) + return packed + + monkeypatch.setattr(pi05, "_packed_quantized_projection", pack) + actual = _pair(x, first, second, selection) + assert packing == [((first, second), first, "_vvla_fused_pair_quantized")] + assert len(packed_calls) == 1 and packed_calls[0] is x + for result, reference in zip(actual, expected, strict=True): + _equal(result, reference) + assert calls == [] + + +def test_pi05_linear_subclasses_and_parallel_modules_bypass_profile(monkeypatch): + selection, calls = _profile() + + class AlteredLinear(nn.Linear): + def forward(self, x): + return super().forward(x) + 0.5 + + class CallableProjection: + def __init__(self, in_features, out_features, dtype, bias): + self.weight = torch.randn(out_features, in_features, dtype=dtype) + self.bias = torch.randn(out_features, dtype=dtype) if bias else None + self.training = False + self.route_inputs = [] + + def __call__(self, x): + self.route_inputs.append(x) + return F.linear(x, self.weight, self.bias) - 0.5 + + reductions = [] + group = object() + + def reduce(output, *, group): + reductions.append((output.clone(), group)) + output.add_(0.25) + + monkeypatch.setattr(torch.distributed, "all_reduce", reduce) + + def column(k, n, dtype, bias): + dense = _dense(k, n * 2, dtype, bias) + return _track(ColumnParallelLinear.from_linear(dense, TensorParallelContext(2, 1), name="test")) + + def row(k, n, dtype, bias): + dense = _dense(k * 2, n, dtype, bias) + context = TensorParallelContext(2, 1, group) + return _track(RowParallelLinear.from_linear(dense, context, name="test")) + + for dtype in _DTYPES: + _positive_control(selection, calls, dtype) + factories = ( + lambda k, n, dtype, bias: _track(AlteredLinear(k, n, bias=bias, dtype=dtype).eval()), + column, + row, + CallableProjection, + ) + with torch.no_grad(): + for make in factories: + _check_module_bypass(make, selection, calls, dtype) + projection = row(7, 5, dtype, True) + x = torch.randn(2, 3, 7, dtype=dtype) + mlp = SimpleNamespace(gate_proj=nn.Identity(), up_proj=nn.Identity(), down_proj=projection) + activated = F.gelu(x, approximate="tanh") * x + expected_partial = F.linear(activated, projection.weight, None) + count = len(reductions) + actual = pi05._mlp(mlp, x, False, fuse_projections=False, kernels=selection) + assert len(reductions) == count + 1 and reductions[-1][1] is group + _equal(reductions[-1][0], expected_partial) + _equal(actual, (expected_partial + 0.25) + projection.bias) + with torch.enable_grad(), pytest.raises(RuntimeError, match="inference only"): + projection.train() + pi05._mlp(mlp, x, False, fuse_projections=False, kernels=selection) + assert calls == [] + assert reductions and all(recorded_group is group for _, recorded_group in reductions) + + +def test_pi05_linear_pair_routes_any_dense_projection_like_base(): + """BASE packs every non-quantized inference pair into one GEMM, whatever its module type.""" + selection, calls = _profile() + + class AlteredLinear(nn.Linear): + def forward(self, x): + return super().forward(x) + 0.5 + + def column(k, n, dtype): + dense = _dense(k, n * 2, dtype, False) + return ColumnParallelLinear.from_linear(dense, TensorParallelContext(2, 1), name="test").eval() + + for dtype in _DTYPES: + factories = (lambda k, n, dtype: AlteredLinear(k, n, bias=False, dtype=dtype).eval(), column) + for make in factories: + for swapped in (False, True): + first, second = make(7, 5, dtype), _dense(7, 9, dtype, False) + if swapped: + first, second = second, first + x = torch.randn(2, 3, 7, dtype=dtype) + weight = torch.cat((first.weight, second.weight), dim=0).contiguous() + expected = F.linear(x, weight).split((first.out_features, second.out_features), dim=-1) + with torch.no_grad(): + actual = _pair(x, first, second, selection) + assert len(calls) == 1 + _call_is(calls[0], x, first._vvla_fused_pair_weight, None) + for result, reference in zip(actual, expected, strict=True): + _equal(result, reference) + calls.clear() + + +def test_pi05_eager_projection_branches_bypass_linear_profile(monkeypatch): + selection, calls = _profile() + + def unexpected(*args, **kwargs): + raise AssertionError("eager branch used a fused or native embedding helper") + + monkeypatch.setattr(pi05, "_fused_qkv", unexpected) + monkeypatch.setattr(pi05, "_embed_prefix_batched", unexpected) + for dtype in _DTYPES: + _positive_control(selection, calls, dtype) + policy = SimpleNamespace( + attention="eager", compile_backend="none", _kernels=selection, _attn=_Attention() + ) + attn = SimpleNamespace( + q_proj=_track(_dense(8, 8, dtype)), + k_proj=_track(_dense(8, 8, dtype)), + v_proj=_track(_dense(8, 8, dtype)), + o_proj=_track(_dense(8, 8, dtype)), + head_dim=4, + scaling=0.5, + ) + h = torch.randn(2, 3, 8, dtype=dtype) + cos, sin = torch.ones(2, 3, 4, dtype=dtype), torch.zeros(2, 3, 4, dtype=dtype) + with torch.no_grad(): + q, k, v = [ + proj(h).view(2, 3, 2, 4).transpose(1, 2) for proj in (attn.q_proj, attn.k_proj, attn.v_proj) + ] + out = F.scaled_dot_product_attention(q, k, v, scale=0.5) + expected = attn.o_proj(out.transpose(1, 2).reshape(2, 3, -1)) + for projection in (attn.q_proj, attn.k_proj, attn.v_proj, attn.o_proj): + projection.route_inputs.clear() + actual = pi05.Pi05Policy._attn_sublayer(policy, attn, h, cos, sin, None, None, []) + _equal(actual, expected) + assert all( + len(proj.route_inputs) == 1 for proj in (attn.q_proj, attn.k_proj, attn.v_proj, attn.o_proj) + ) + assert calls == [] + + mlp = _mlp(dtype) + _track(mlp.gate_proj) + _track(mlp.up_proj) + x = torch.randn(2, 3, 7, dtype=dtype) + expected = _original_mlp(mlp, x, fuse=False) + mlp.gate_proj.route_inputs.clear() + mlp.up_proj.route_inputs.clear() + with monkeypatch.context() as patch: + patch.setattr(pi05, "_fused_linear_pair", unexpected) + _equal(pi05._mlp(mlp, x, False, fuse_projections=False, kernels=selection), expected) + assert len(mlp.gate_proj.route_inputs) == len(mlp.up_proj.route_inputs) == 1 + assert len(calls) == 1 + _call_is(calls[0], calls[0].args[0], mlp.down_proj.weight, mlp.down_proj.bias) + calls.clear() + + prefix_calls = [] + hidden = torch.randn(2, 3, 8, dtype=dtype) + pads = torch.ones(2, 3, dtype=torch.bool) + + def embed_prefix(*args, prefix_calls=prefix_calls, hidden=hidden, pads=pads): + prefix_calls.append(args) + return hidden, pads, torch.zeros_like(pads) + + policy.native_embeddings = False + policy.prefix_attention = "sdpa" + policy._m = SimpleNamespace(embed_prefix=embed_prefix) + policy._embed_prefix = unexpected + policy._prefix_tower = object() + policy._make_att_2d_masks = lambda *args: torch.zeros(2, 3, 3, dtype=torch.bool) + policy._attention_mask_4d = lambda *args: None + policy._tower_forward = lambda *args, dtype=dtype, **kwargs: ( + args[1], + [(torch.zeros(2, 1, 3, 8, dtype=dtype), torch.zeros(2, 1, 3, 8, dtype=dtype))], + ) + batch = SimpleNamespace( + images=[torch.zeros(2, 3, 4, 4, dtype=dtype)], + img_masks=[torch.ones(2, dtype=torch.bool)], + tokens=torch.zeros(2, 3, dtype=torch.long), + masks=pads, + ) + prefix = pi05.Pi05Policy._encode_prefix_impl(policy, batch) + assert prefix.batch_size == 2 and len(prefix_calls) == 1 + assert prefix_calls[0][0] is batch.images + assert calls == []