From 8c1244994f7c0812c78500acc2afafc8990d15b8 Mon Sep 17 00:00:00 2001 From: owpeter Date: Sat, 3 Oct 2026 10:39:18 +0800 Subject: [PATCH 01/44] feat: add kernel design agent tuning workflow Signed-off-by: owpeter --- CONTRIBUTING.md | 64 ++ benchmarks/kernel_tuning/evaluate.py | 110 +++ .../kernel_tuning/flashinfer_adapter.py | 500 ++++++++++++++ benchmarks/kernel_tuning/requirements.txt | 4 + .../tasks/gated_residual/README.md | 47 ++ .../tasks/gated_residual/baseline.py | 9 + .../tasks/gated_residual/definition.json | 14 + .../tasks/gated_residual/tuning.yaml | 29 + .../tasks/gated_residual/workloads.jsonl | 3 + docs/proposals/0008-kernel-tuning.md | 335 +++++++++ pyproject.toml | 6 + scripts/kernel_tuning/__init__.py | 1 + scripts/kernel_tuning/__main__.py | 140 ++++ scripts/kernel_tuning/artifacts.py | 323 +++++++++ scripts/kernel_tuning/contracts.py | 498 ++++++++++++++ scripts/kernel_tuning/flow.py | 222 ++++++ scripts/kernel_tuning/prompts.py | 76 +++ scripts/kernel_tuning/pyproject.toml | 15 + scripts/kernel_tuning/runner.py | 136 ++++ scripts/kernel_tuning/uv.lock | 641 ++++++++++++++++++ tests/kernel_tuning/conftest.py | 87 +++ tests/kernel_tuning/test_cli.py | 32 + tests/kernel_tuning/test_contracts.py | 197 ++++++ tests/kernel_tuning/test_flow.py | 202 ++++++ tests/kernel_tuning/test_numerics.py | 181 +++++ 25 files changed, 3872 insertions(+) create mode 100644 benchmarks/kernel_tuning/evaluate.py create mode 100644 benchmarks/kernel_tuning/flashinfer_adapter.py create mode 100644 benchmarks/kernel_tuning/requirements.txt create mode 100644 benchmarks/kernel_tuning/tasks/gated_residual/README.md create mode 100644 benchmarks/kernel_tuning/tasks/gated_residual/baseline.py create mode 100644 benchmarks/kernel_tuning/tasks/gated_residual/definition.json create mode 100644 benchmarks/kernel_tuning/tasks/gated_residual/tuning.yaml create mode 100644 benchmarks/kernel_tuning/tasks/gated_residual/workloads.jsonl create mode 100644 docs/proposals/0008-kernel-tuning.md create mode 100644 scripts/kernel_tuning/__init__.py create mode 100644 scripts/kernel_tuning/__main__.py create mode 100644 scripts/kernel_tuning/artifacts.py create mode 100644 scripts/kernel_tuning/contracts.py create mode 100644 scripts/kernel_tuning/flow.py create mode 100644 scripts/kernel_tuning/prompts.py create mode 100644 scripts/kernel_tuning/pyproject.toml create mode 100644 scripts/kernel_tuning/runner.py create mode 100644 scripts/kernel_tuning/uv.lock create mode 100644 tests/kernel_tuning/conftest.py create mode 100644 tests/kernel_tuning/test_cli.py create mode 100644 tests/kernel_tuning/test_contracts.py create mode 100644 tests/kernel_tuning/test_flow.py create mode 100644 tests/kernel_tuning/test_numerics.py diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index b2606ee..c2b039a 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -52,6 +52,70 @@ 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 check benchmarks/kernel_tuning/tasks/gated_residual +``` + +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 [approved contract](docs/proposals/0008-kernel-tuning.md) and the +[gated residual task](benchmarks/kernel_tuning/tasks/gated_residual/README.md). +The [Chinese workflow guide](docs/proposals/0008-kernel-tuning.md#11-自动化调优使用说明) +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/`. +Saved operators require a separate review/integration before runtime use. +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. + +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 +``` + +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/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..fbe6076 --- /dev/null +++ b/benchmarks/kernel_tuning/flashinfer_adapter.py @@ -0,0 +1,500 @@ +"""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 Definition, 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 = Definition.model_validate(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"] + ) + + @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": "FlashInfer cold L2; input cloning excluded", + }, + } + + @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 do_bench, 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)) + + def reset() -> None: + import torch + + for dst, src in zip(args, inputs): + if isinstance(dst, torch.Tensor): + dst.copy_(src) + + result = do_bench( + lambda _: graph.replay(), + warmup=self.cfg.timing.warmup, + rep=self.cfg.timing.iterations, + setup=reset, + device=self.device, + ) + self._check_flags() + return float(result) + + 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 = registry.build_reference(self.definition) + 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", + } + ) + 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 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/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/kernel_tuning/tasks/gated_residual/README.md b/benchmarks/kernel_tuning/tasks/gated_residual/README.md new file mode 100644 index 0000000..e721b02 --- /dev/null +++ b/benchmarks/kernel_tuning/tasks/gated_residual/README.md @@ -0,0 +1,47 @@ +# Gated residual BF16 task + +This self-contained example covers the multiply-then-add semantics used by +`embodiinfer/backend/triton/norm.py::gated_residual`. Inputs are contiguous BF16 +`residual, update: [B,T,H]` and `gate: [B,1,H]` or `[B,T,H]`. Output has the same +shape and dtype as residual. Round the product to BF16 before adding; contracting +the expression into an FMA can violate the bit-exact contract. Inputs are immutable. + +The mathematical reference is embedded in `definition.json`. `baseline.py` is a +separate eager Torch implementation, written for this task under the repository's +Apache-2.0 license. It is an executable template baseline, **not a claim to beat +EmbodiInfer's existing Triton kernel**. Replace it with the best-known standalone +implementation before using this task to claim a production improvement. Include +all local helpers inside the task package and record their provenance here. + +The three synthetic workloads exercise broadcast gates, per-token gates, and a +non-power-of-two hidden dimension. They do not represent measured model traffic. +Use workload shapes and weights from the intended deployment for real tuning. + +On the target machine, prepare its compatible Torch/CUDA/Triton environment and +set `evaluator_python` in `tuning.yaml` to that environment's absolute Python path. +The default evaluator also needs `benchmarks/kernel_tuning/requirements.txt`. +Select `language: cuda` for a CUDA C++ search; the agent then supplies a CUDA +Solution with a supported Torch or TVM-FFI binding. The CUDA toolkit/compiler must +already be available. `target_hardware` is a Solution label, not a capability check; +the evaluator probes the actual device (including Thor) at run time. + +From the repository root, in the optional Humanize2 tooling environment: + +```bash +python -m scripts.kernel_tuning check benchmarks/kernel_tuning/tasks/gated_residual +python -m scripts.kernel_tuning check benchmarks/kernel_tuning/tasks/gated_residual --environment +python -m scripts.kernel_tuning run benchmarks/kernel_tuning/tasks/gated_residual --agent 'HARNESS/MODEL:EFFORT' +python -m scripts.kernel_tuning status results/kernel_tuning/RUN_ID +python -m scripts.kernel_tuning resume results/kernel_tuning/RUN_ID +python -m scripts.kernel_tuning export results/kernel_tuning/RUN_ID results/exported_gated_residual +``` + +Replace the agent placeholder with a configured Humanize2 agent specification. +`check` performs static validation. `--environment` additionally imports the +evaluator and probes its hardware/dependencies. `run` and `resume` execute real +coding-agent turns and measurements. Each candidate's `evaluation/command.json` +records the exact reproducible worker command and timeout. Fixed conditions, +paired measurements, and acceptance limits are in `tuning.yaml`. + +For the optional custom `benchmark.py` interface and artifact contract, see +[`0008-kernel-tuning.md`](../../../../docs/proposals/0008-kernel-tuning.md). diff --git a/benchmarks/kernel_tuning/tasks/gated_residual/baseline.py b/benchmarks/kernel_tuning/tasks/gated_residual/baseline.py new file mode 100644 index 0000000..b2c2977 --- /dev/null +++ b/benchmarks/kernel_tuning/tasks/gated_residual/baseline.py @@ -0,0 +1,9 @@ +"""Lossless eager baseline with an explicit intermediate dtype rounding.""" + +import torch + + +def run(residual: torch.Tensor, update: torch.Tensor, gate: torch.Tensor) -> torch.Tensor: + """Match the separate multiply/add semantics used by the gated residual path.""" + product = (update * gate).to(update.dtype) + return (residual + product).to(residual.dtype) diff --git a/benchmarks/kernel_tuning/tasks/gated_residual/definition.json b/benchmarks/kernel_tuning/tasks/gated_residual/definition.json new file mode 100644 index 0000000..572c2f6 --- /dev/null +++ b/benchmarks/kernel_tuning/tasks/gated_residual/definition.json @@ -0,0 +1,14 @@ +{ + "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": "import torch\n\ndef run(residual, update, gate):\n product = (update * gate).to(update.dtype)\n return (residual + product).to(residual.dtype)\n" +} diff --git a/benchmarks/kernel_tuning/tasks/gated_residual/tuning.yaml b/benchmarks/kernel_tuning/tasks/gated_residual/tuning.yaml new file mode 100644 index 0000000..bb6d9d7 --- /dev/null +++ b/benchmarks/kernel_tuning/tasks/gated_residual/tuning.yaml @@ -0,0 +1,29 @@ +schema_version: 1 +language: triton +# Set --agent explicitly when starting a real run, or fill in a valid hmz spec. +agent: "" +# Change to the absolute Python executable in the prepared Thor/runtime environment. +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 diff --git a/benchmarks/kernel_tuning/tasks/gated_residual/workloads.jsonl b/benchmarks/kernel_tuning/tasks/gated_residual/workloads.jsonl new file mode 100644 index 0000000..72193b7 --- /dev/null +++ b/benchmarks/kernel_tuning/tasks/gated_residual/workloads.jsonl @@ -0,0 +1,3 @@ +{"definition":"gated_residual_bf16","workload":{"uuid":"broadcast-small","axes":{"B":1,"T":16,"H":1024,"G":1},"inputs":{"residual":{"type":"random"},"update":{"type":"random"},"gate":{"type":"random"}}},"solution":null,"evaluation":null} +{"definition":"gated_residual_bf16","workload":{"uuid":"broadcast-wide","axes":{"B":2,"T":32,"H":2048,"G":1},"inputs":{"residual":{"type":"random"},"update":{"type":"random"},"gate":{"type":"random"}}},"solution":null,"evaluation":null} +{"definition":"gated_residual_bf16","workload":{"uuid":"token-gate-tail","axes":{"B":1,"T":17,"H":1000,"G":17},"inputs":{"residual":{"type":"random"},"update":{"type":"random"},"gate":{"type":"random"}}},"solution":null,"evaluation":null} diff --git a/docs/proposals/0008-kernel-tuning.md b/docs/proposals/0008-kernel-tuning.md new file mode 100644 index 0000000..53088d5 --- /dev/null +++ b/docs/proposals/0008-kernel-tuning.md @@ -0,0 +1,335 @@ +# 0008 — Reproducible kernel tuning with Humanize2 + +- Status: Accepted +- Author: EmbodiInfer maintainers +- Date: 2026-10-02 + +Implementation: repository CLI, Humanize2 Flow, evaluator, source/evidence archive, +and isolated tool lock are present. CPU/fake-agent coverage and NVIDIA RTX 5090 +Triton/CUDA C++ builder tests are verified. Thor compilation and performance +still require validation on that target machine. + +For the Chinese workflow guide, required task inputs, and stopping conditions, +see [Section 11](#11-自动化调优使用说明). + +## 1. Summary + +Add repository developer tooling for a serial, agent-driven kernel optimization +loop. Humanize2 (`hmz`) owns coding-agent execution, budgets, traces, and resume; +EmbodiInfer supplies the task contract, trusted evaluation, promotion rules, and +saved operator artifacts. Formal measurement belongs under `benchmarks/` and +orchestration under `scripts/`. This is model-neutral tooling, not an engine API. + +## 2. Motivation and current gap + +`embodiinfer/backend/triton/norm.py:302` and +`embodiinfer/backend/triton/split_kv_attention.py:179` contain concrete kernels and +hand-chosen launch configurations. `embodiinfer/layers/attention.py:78` provides +runtime registration, but the repository has no general operator task format, +automated search workflow, or durable archive of verified candidates. A task can +describe one of these kernels or an independent operator without importing model +or engine internals into the tuning framework. + +## 3. Goals and non-goals + +- Accept NVlabs KDA task packages and Triton or CUDA C++ multi-file solutions. +- Run on one target machine, with an explicitly selected coding agent and model. +- Freeze precision, workload, timing, and acceptance conditions before search. +- Save every attempt and export the best verified source with its evidence. +- Support bounded execution, interruptions, and explicit resume. + +Runtime registration, deployment, SSH, model-level performance claims, and +automatic installation of GPU dependencies are outside this change. Hardware +neutrality describes orchestration: a backend must support the selected hardware; +CUDA and the default FlashInfer adapter require a compatible NVIDIA environment. + +## 4. Design + +Use the [NVlabs wishlist task template](https://github.com/NVlabs/kda/tree/2186cd90dfb3b11f1ec250b4e3f321d72d7e7ef0/example) +and [minimal flow](https://github.com/NVlabs/kda/blob/ef6ce617693ef0782b3ecb9f37e39bbf10226a90/docs/agent-flow.md). +Tasks contain `README.md`, `definition.json`, `workloads.jsonl`, `baseline.py`, +and optionally `benchmark.py`. The definition embeds the mathematical reference +`run`; `baseline.py::run` supplies the initial performance baseline. The JSON +schemas remain FlashInfer Trace schemas. Project settings live in `tuning.yaml`. + +Implement a local Humanize2 Flow against +[`hmz` revision b70d4427](https://github.com/humanfia/humanize/tree/b70d4427f11308442aef20f9b2233dfbbc34a39e). +The Flow asks an existing coding agent to inspect the task, write a hypothesis +and executable plan, produce a candidate, then respond to evaluator feedback. +Agent turns perform source/JSON checks and return the candidate; GPU compilation, +numerical verification, and timing run only in the separate evaluator. A working +GPU in the evaluator does not imply that the agent's command environment grants +GPU or compiler-cache access. Plans describe validation criteria without trying +to execute that validation inside an agent turn. +It delegates all agent sessions and cancellation to Humanize2. Domain helpers +validate tasks, launch a bounded evaluator process, persist candidate evidence, +and atomically publish the best pointer. No independent coding-agent subprocess +implementation is introduced. + +The default evaluator reuses FlashInfer Bench 0.1.2 definitions, workload +materialization, and solution building. Additional gates enforce the configured +precision contract, workload completeness, graph replay correctness when selected, +and paired baseline/candidate measurements. A task-owned `benchmark.py` can +implement the documented structured evaluation interface for other hardware or +special input semantics. Exit status or printed PASS alone cannot promote a kernel. + +Each task fixes eager or CUDA Graph timing. Promotion requires the weighted mean +latency to fall by at least 3% and every workload to regress by at most 5%, against +the incumbent and initial baseline. Weights default to equal. Repeat paired +measurements before promotion. Stop at 20 attempts, two hours of active execution, +or five consecutive failures/non-improvements; all limits are configurable. + +Runs snapshot task files and record source digests, evaluator identity, hardware +and software fingerprints, precision settings, agent settings, attempts, raw +measurements, and decisions. Resume requires matching task/tool/environment +identity. Interrupted attempts remain in the archive and consume an attempt. +Export reconstructs source files plus the official Solution JSON and evidence. +Task and artifact hashes detect accidental modifications; they are not an OS +security sandbox for coding agents or kernels running as the same user. + +Humanize2 needs Python 3.12+. Install it in a separate tooling environment; the +evaluator interpreter points to the prepared target-machine runtime (for example +Thor). EmbodiInfer's Python 3.10 floor and model dependency groups stay intact. + +An independent lightweight controller was considered and rejected because the +approved design explicitly reuses Humanize2. Runtime auto-selection was rejected +because source review and separate integration are required before deployment. + +## 5. Model-agnosticism verdict + +This is developer automation and formal measurement, outside the inference +package. No engine, policy, or layer interface changes. Operator semantics belong +entirely to each task's definition and workloads. Backend and hardware details +belong to the evaluator and generated solution, not the orchestration loop. + +## 6. Losslessness and precision criterion + +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. + +## 7. Implementation plan + +Add `scripts/kernel_tuning/` for task validation, artifact storage, Humanize2 Flow, +and the `python -m scripts.kernel_tuning` developer entry point. Add +`benchmarks/kernel_tuning/` for measurement and a self-contained task example. +Document isolated setup in `CONTRIBUTING.md`; keep the optional tool's dependency +metadata separate from runtime requirements. Results default to the already +ignored `results/kernel_tuning/`. Nothing automatically registers saved kernels. + +## 8. Test plan + +CPU tests use fake coding agents and evaluators to cover actual Flow execution, +resume, failed attempts, budgets, task tampering, metric completeness, precision +contracts, regressions, atomic best publication, and export. Static checks cover +all new Python. GPU evaluator tests are explicitly marked and skip without their +optional dependencies and hardware. CPU validation does not launch a real coding +agent or tune kernels. Real agent searches and GPU validation run explicitly on +a prepared target machine, with their own saved task and measurement conditions. + +## 9. Benchmark plan + +After deployment to a prepared GPU machine, each task records hardware, software, +dtype, shapes, seeds, warmup, iterations, trials, timing mode, and baseline source. +Paired trials evaluate the full declared workload set. Record failed and neutral +attempts as well as improvements. Kernel-level results do not establish model +latency or action parity; those require the existing model benchmark separately. + +## 10. Risks and limitations + +Humanize2 and FlashInfer are evolving dependencies; use pinned versions and fail +explicitly on incompatible contracts. Arbitrary operators may need custom input +generation or timing through `benchmark.py`. Input coverage bounds correctness; +finite tests cannot prove an arbitrary generated kernel correct. Thermal state, +other GPU users, and clock changes can affect timing; paired remeasurement reduces +but does not remove this uncertainty. User-selected coding agents may execute +commands without interactive approval; run in a disposable workspace/account +appropriate to the task. Actual Thor compilation and performance remain target +machine validation work. The pinned Humanize2 journal uses POSIX primitives: real +searches run on the Linux target machine; local Windows checks and export are +supported. CPU tests on Windows adapt only upstream atomic journal file writing. + +### 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. + +### 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. + +## 11. 自动化调优使用说明 + +### 准备任务包 + +每个算子使用独立、自包含的任务目录。可复制 +`benchmarks/kernel_tuning/tasks/gated_residual/`，再替换以下五个必需文件： + +| 文件 | 人工提供的内容 | +|---|---| +| `README.md` | 调用场景、负载来源、布局与特化限制、基线来源及精度约定的理由。 | +| `definition.json` | 输入输出名称、顺序、shape、dtype、约束和可执行的数学参考 `run`。 | +| `workloads.jsonl` | 每行一个未评测的 Trace，包含唯一 UUID、实际 shape 和输入来源。 | +| `baseline.py` | 要超越的当前实现，入口为 `run`；参数顺序与 definition 的 inputs 一致。 | +| `tuning.yaml` | 语言、agent、评测 Python、设备、精度、计时、晋级规则和搜索预算。 | + +例如 RMSNorm 需要先确定普通、残差融合或自适应版本、`eps`、 +`weight` 或 `1 + weight`、归约精度和中间舍入顺序，再提供目标部署的 +batch/token/hidden 负载。正确性参考与性能基线承担不同职责，不能只给出算子名称。 +输入支持随机生成、标量和任务目录内的 safetensors；本地 helper 也应随任务保存。 +非连续输入、原地修改或特殊数据生成/硬件语义需要任务自己的 `benchmark.py`， +接口见上面的 Task-owned benchmark interface。 + +### 准备环境并启动 + +真实调优在 Linux/POSIX 目标机器上执行。先准备兼容该 GPU 的运行环境， +将 `evaluator_python` 设置为该环境的绝对 Python 路径，安装评测依赖； +独立安装并登录 coding-agent CLI。账号凭据保留在 CLI 的配置中，不写入任务包。 +Humanize2 工具环境和安装步骤见 `CONTRIBUTING.md` 的 Optional kernel tuning tools。 +Thor 应使用 Thor 兼容的 GPU 软件栈，不能直接套用其他显卡的依赖版本。 + +以下命令从仓库根目录执行。把任务路径替换为自己的任务，agent 占位符替换为 +已配置的实际 harness/model/effort： + +```bash +uv sync --project scripts/kernel_tuning --python 3.12 --frozen +uv run --project scripts/kernel_tuning python -m scripts.kernel_tuning check benchmarks/kernel_tuning/tasks/gated_residual +uv run --project scripts/kernel_tuning python -m scripts.kernel_tuning check benchmarks/kernel_tuning/tasks/gated_residual --environment +uv run --project scripts/kernel_tuning python -m scripts.kernel_tuning run benchmarks/kernel_tuning/tasks/gated_residual --agent 'HARNESS/MODEL:EFFORT' +``` + +`check` 只做结构检查；`check --environment` 额外探测依赖与硬件。 +`run` 和 `resume` 才会调用真实 agent 并执行编译、正确性检查与 GPU 测量。 + +### 自动循环与停止条件 + +1. Controller 冻结任务与工具源码，记录评测环境，检查初始基线。 +2. Humanize2 调用 coding agent 读取任务、当前最优实现和反馈，先写 `PLAN.md`。 +3. Agent 根据计划生成 Triton/CUDA 源码和完整的官方 `solution.json`，只做静态检查。 +4. 独立 evaluator 编译候选；使用相同种子对参考、基线、当前最优和候选进行校验， + 再对基线、当前最优和候选进行配对计时。 +5. Controller 根据固定门槛决定是否晋级，保存本轮源码和证据，把反馈交给下一轮。 + +默认精度为逐字节一致。允许浮点误差时，人工须事先明确设置 +`precision.mode: tolerance`、`atol`、`rtol` 和 `reason`，agent 不能修改这些条件。 +默认使用两个配对测量轮次、等权负载；每轮平均延迟须同时相对初始基线和当前 +最优改善至少 3%，且任何单个负载的回退不超过 5%。显式 `weights` 必须覆盖全部 UUID。 + +默认最多生成 20 个候选、累计执行两小时，或连续五次失败/未改善后停止； +每轮 agent 和评测也有独立超时。这些上限均在 `tuning.yaml` 中配置。 +失败和中断尝试保留并计入预算；耗尽预算后须建立新任务/run，不能修改冻结任务继续搜索。 + +### 查看、续跑与保存结果 + +`run` 会输出新建的 `results/kernel_tuning/RUN_ID/` 路径。该目录保存任务和工具快照、 +环境信息、所有候选源码/计划、评测请求及原始结果、晋级决定和 Humanize2 编排记录。 +只有候选通过晋级门槛后才生成 `best.json`；没有改善的 run 也会完整保存。 + +```bash +uv run --project scripts/kernel_tuning python -m scripts.kernel_tuning status results/kernel_tuning/RUN_ID +uv run --project scripts/kernel_tuning python -m scripts.kernel_tuning resume results/kernel_tuning/RUN_ID +uv run --project scripts/kernel_tuning python -m scripts.kernel_tuning export results/kernel_tuning/RUN_ID results/exported_operator +``` + +续跑要求任务、工具、编排环境和评测环境身份一致，并有剩余预算。 +`status` 与 `export` 不调用 agent；导出包包含最优算子的 Solution、完整源码、冻结任务、 +工具/锁文件、原始基线、胜出证据和 SHA256 清单。结果位于 Git 忽略的 `results/`， +需要另行保存。导出不会自动注册到推理 runtime，实际集成仍需代码审查和模型验证。 diff --git a/pyproject.toml b/pyproject.toml index b62785b..2eb6e82 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -147,6 +147,12 @@ 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.pytest.ini_options] testpaths = ["tests"] markers = [ 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..67ccbbd --- /dev/null +++ b/scripts/kernel_tuning/__main__.py @@ -0,0 +1,140 @@ +"""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, RunStore, atomic_json, run_lock, runtime_identity +from .contracts import HMZ_REVISION, ContractError, TaskPackage +from .runner import evaluate + + +def _execute(store: RunStore, *, resume: bool) -> dict: + 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") + 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 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) + 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") + 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) + args = parser.parse_args(argv) + try: + if 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") + store = RunStore.create(task, args.output, agent) + 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..52f69ca --- /dev/null +++ b/scripts/kernel_tuning/artifacts.py @@ -0,0 +1,323 @@ +"""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] + + +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) -> RunStore: + """Snapshot a task and baseline under a fresh, collision-resistant run ID.""" + 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, + "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") + atomic_json(scratch / "feedback.json", self.manifest["attempts"][-5:]) + 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/contracts.py b/scripts/kernel_tuning/contracts.py new file mode 100644 index 0000000..3122b60 --- /dev/null +++ b/scripts/kernel_tuning/contracts.py @@ -0,0 +1,498 @@ +"""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 + if not isinstance(data, dict): + raise ContractError("tuning.yaml must contain a mapping") + 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]: + 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") + 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") + 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 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 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/flow.py b/scripts/kernel_tuning/flow.py new file mode 100644 index 0000000..1865eed --- /dev/null +++ b/scripts/kernel_tuning/flow.py @@ -0,0 +1,222 @@ +"""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 ( + Agent, + 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.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: Agent + + +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() + 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 agents["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( + agents["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/prompts.py b/scripts/kernel_tuning/prompts.py new file mode 100644 index 0000000..da9e46b --- /dev/null +++ b/scripts/kernel_tuning/prompts.py @@ -0,0 +1,76 @@ +"""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.""" + 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. +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", + } + ], + } + 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. +Preserve dtype and intermediate rounding; bit_exact includes 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/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/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/tests/kernel_tuning/conftest.py b/tests/kernel_tuning/conftest.py new file mode 100644 index 0000000..91a2b4f --- /dev/null +++ b/tests/kernel_tuning/conftest.py @@ -0,0 +1,87 @@ +"""CPU fixtures; use --confcutdir here in the isolated, Torch-free tool environment.""" + +from __future__ import annotations + +import shutil +import sys +from pathlib import Path +from typing import Any + +import pytest +from scripts.kernel_tuning.artifacts import REPOSITORY, 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) +''' + + +@pytest.fixture +def task_dir(tmp_path: Path) -> Path: + """Make a portable task using a fake measurement adapter and explicit Python.""" + path = tmp_path / "task" + shutil.copytree(REPOSITORY / "benchmarks/kernel_tuning/tasks/gated_residual", path) + (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"} + ], + } diff --git a/tests/kernel_tuning/test_cli.py b/tests/kernel_tuning/test_cli.py new file mode 100644 index 0000000..71b8146 --- /dev/null +++ b/tests/kernel_tuning/test_cli.py @@ -0,0 +1,32 @@ +"""Read-only developer commands remain usable without Humanize2 or a GPU.""" + +from __future__ import annotations + +import json +import sys +from pathlib import Path + +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 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_flow.py b/tests/kernel_tuning/test_flow.py new file mode 100644 index 0000000..93dc835 --- /dev/null +++ b/tests/kernel_tuning/test_flow.py @@ -0,0 +1,202 @@ +"""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_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 diff --git a/tests/kernel_tuning/test_numerics.py b/tests/kernel_tuning/test_numerics.py new file mode 100644 index 0000000..87c88df --- /dev/null +++ b/tests/kernel_tuning/test_numerics.py @@ -0,0 +1,181 @@ +"""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(tmp_path: 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 = tmp_path / "task" + shutil.copytree(REPOSITORY / "benchmarks/kernel_tuning/tasks/gated_residual", 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) From a39b968c75c03ea858c2211147405856412a0b86 Mon Sep 17 00:00:00 2001 From: owpeter Date: Sat, 3 Oct 2026 10:42:26 +0800 Subject: [PATCH 02/44] chore: remove outdated sections from kernel tuning proposal Signed-off-by: owpeter --- docs/proposals/0008-kernel-tuning.md | 80 +--------------------------- 1 file changed, 1 insertion(+), 79 deletions(-) diff --git a/docs/proposals/0008-kernel-tuning.md b/docs/proposals/0008-kernel-tuning.md index 53088d5..efd6a32 100644 --- a/docs/proposals/0008-kernel-tuning.md +++ b/docs/proposals/0008-kernel-tuning.md @@ -254,82 +254,4 @@ task/solution/output locations when moving the bundle; retain task/solution dige 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. - -## 11. 自动化调优使用说明 - -### 准备任务包 - -每个算子使用独立、自包含的任务目录。可复制 -`benchmarks/kernel_tuning/tasks/gated_residual/`，再替换以下五个必需文件： - -| 文件 | 人工提供的内容 | -|---|---| -| `README.md` | 调用场景、负载来源、布局与特化限制、基线来源及精度约定的理由。 | -| `definition.json` | 输入输出名称、顺序、shape、dtype、约束和可执行的数学参考 `run`。 | -| `workloads.jsonl` | 每行一个未评测的 Trace，包含唯一 UUID、实际 shape 和输入来源。 | -| `baseline.py` | 要超越的当前实现，入口为 `run`；参数顺序与 definition 的 inputs 一致。 | -| `tuning.yaml` | 语言、agent、评测 Python、设备、精度、计时、晋级规则和搜索预算。 | - -例如 RMSNorm 需要先确定普通、残差融合或自适应版本、`eps`、 -`weight` 或 `1 + weight`、归约精度和中间舍入顺序，再提供目标部署的 -batch/token/hidden 负载。正确性参考与性能基线承担不同职责，不能只给出算子名称。 -输入支持随机生成、标量和任务目录内的 safetensors；本地 helper 也应随任务保存。 -非连续输入、原地修改或特殊数据生成/硬件语义需要任务自己的 `benchmark.py`， -接口见上面的 Task-owned benchmark interface。 - -### 准备环境并启动 - -真实调优在 Linux/POSIX 目标机器上执行。先准备兼容该 GPU 的运行环境， -将 `evaluator_python` 设置为该环境的绝对 Python 路径，安装评测依赖； -独立安装并登录 coding-agent CLI。账号凭据保留在 CLI 的配置中，不写入任务包。 -Humanize2 工具环境和安装步骤见 `CONTRIBUTING.md` 的 Optional kernel tuning tools。 -Thor 应使用 Thor 兼容的 GPU 软件栈，不能直接套用其他显卡的依赖版本。 - -以下命令从仓库根目录执行。把任务路径替换为自己的任务，agent 占位符替换为 -已配置的实际 harness/model/effort： - -```bash -uv sync --project scripts/kernel_tuning --python 3.12 --frozen -uv run --project scripts/kernel_tuning python -m scripts.kernel_tuning check benchmarks/kernel_tuning/tasks/gated_residual -uv run --project scripts/kernel_tuning python -m scripts.kernel_tuning check benchmarks/kernel_tuning/tasks/gated_residual --environment -uv run --project scripts/kernel_tuning python -m scripts.kernel_tuning run benchmarks/kernel_tuning/tasks/gated_residual --agent 'HARNESS/MODEL:EFFORT' -``` - -`check` 只做结构检查；`check --environment` 额外探测依赖与硬件。 -`run` 和 `resume` 才会调用真实 agent 并执行编译、正确性检查与 GPU 测量。 - -### 自动循环与停止条件 - -1. Controller 冻结任务与工具源码，记录评测环境，检查初始基线。 -2. Humanize2 调用 coding agent 读取任务、当前最优实现和反馈，先写 `PLAN.md`。 -3. Agent 根据计划生成 Triton/CUDA 源码和完整的官方 `solution.json`，只做静态检查。 -4. 独立 evaluator 编译候选；使用相同种子对参考、基线、当前最优和候选进行校验， - 再对基线、当前最优和候选进行配对计时。 -5. Controller 根据固定门槛决定是否晋级，保存本轮源码和证据，把反馈交给下一轮。 - -默认精度为逐字节一致。允许浮点误差时，人工须事先明确设置 -`precision.mode: tolerance`、`atol`、`rtol` 和 `reason`，agent 不能修改这些条件。 -默认使用两个配对测量轮次、等权负载；每轮平均延迟须同时相对初始基线和当前 -最优改善至少 3%，且任何单个负载的回退不超过 5%。显式 `weights` 必须覆盖全部 UUID。 - -默认最多生成 20 个候选、累计执行两小时，或连续五次失败/未改善后停止； -每轮 agent 和评测也有独立超时。这些上限均在 `tuning.yaml` 中配置。 -失败和中断尝试保留并计入预算；耗尽预算后须建立新任务/run，不能修改冻结任务继续搜索。 - -### 查看、续跑与保存结果 - -`run` 会输出新建的 `results/kernel_tuning/RUN_ID/` 路径。该目录保存任务和工具快照、 -环境信息、所有候选源码/计划、评测请求及原始结果、晋级决定和 Humanize2 编排记录。 -只有候选通过晋级门槛后才生成 `best.json`；没有改善的 run 也会完整保存。 - -```bash -uv run --project scripts/kernel_tuning python -m scripts.kernel_tuning status results/kernel_tuning/RUN_ID -uv run --project scripts/kernel_tuning python -m scripts.kernel_tuning resume results/kernel_tuning/RUN_ID -uv run --project scripts/kernel_tuning python -m scripts.kernel_tuning export results/kernel_tuning/RUN_ID results/exported_operator -``` - -续跑要求任务、工具、编排环境和评测环境身份一致，并有剩余预算。 -`status` 与 `export` 不调用 agent；导出包包含最优算子的 Solution、完整源码、冻结任务、 -工具/锁文件、原始基线、胜出证据和 SHA256 清单。结果位于 Git 忽略的 `results/`， -需要另行保存。导出不会自动注册到推理 runtime，实际集成仍需代码审查和模型验证。 +are not established by CPU/fake-agent tests. \ No newline at end of file From 6cdc04d593f96bf702ce25227bc2c6940ad6bb40 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=82=96=E4=B8=80=E6=B6=B5?= <23371301@buaa.edu.cn> Date: Sat, 3 Oct 2026 16:19:31 +0800 Subject: [PATCH 03/44] feat: generate tuning tasks from an operator catalog and tune in batches MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Declare EmbodiInfer's core Triton operators (pi0.5, ActiveVLN vision, StreamVLN decode) once in scripts/kernel_tuning/operators.py and render complete task packages from them. The baseline is the production kernel, vendored verbatim with its repository imports. Each contract follows the production kernel: bit-exact where it already is, otherwise a stated BF16 rounding-step tolerance. Add `generate` and `tune-all`: one batch generates every task, preflights all production baselines with the evaluator, runs each operator serially in its own run/resume subprocess, exports promoted kernels and writes a summary. Task validation now rejects variable axes that cannot be inferred from inputs. Signed-off-by: 肖一涵 <23371301@buaa.edu.cn> --- CONTRIBUTING.md | 13 + docs/proposals/0008-kernel-tuning.md | 49 +- scripts/kernel_tuning/__main__.py | 102 +++- scripts/kernel_tuning/batch.py | 213 ++++++++ scripts/kernel_tuning/contracts.py | 10 + scripts/kernel_tuning/generate.py | 450 +++++++++++++++++ scripts/kernel_tuning/operators.py | 513 ++++++++++++++++++++ tests/kernel_tuning/test_batch.py | 206 ++++++++ tests/kernel_tuning/test_generated_tasks.py | 232 +++++++++ 9 files changed, 1785 insertions(+), 3 deletions(-) create mode 100644 scripts/kernel_tuning/batch.py create mode 100644 scripts/kernel_tuning/generate.py create mode 100644 scripts/kernel_tuning/operators.py create mode 100644 tests/kernel_tuning/test_batch.py create mode 100644 tests/kernel_tuning/test_generated_tasks.py diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index c2b039a..6555774 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -68,6 +68,19 @@ uv run --project scripts/kernel_tuning python -m scripts.kernel_tuning --help uv run --project scripts/kernel_tuning python -m scripts.kernel_tuning check benchmarks/kernel_tuning/tasks/gated_residual ``` +To tune every core operator, generate tasks from the operator catalog and run +them as one batch (see the +[generated-task section](docs/proposals/0008-kernel-tuning.md#generated-catalog-tasks-and-batch-tuning)): + +```bash +uv run --project scripts/kernel_tuning python -m scripts.kernel_tuning generate --list +uv run --project scripts/kernel_tuning python -m scripts.kernel_tuning tune-all --preflight-only --set evaluator_python=/absolute/runtime/python +uv run --project scripts/kernel_tuning python -m scripts.kernel_tuning tune-all --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 +``` + +Select a subset with operator names or `--model pi05|activevln|streamvln`. + 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 diff --git a/docs/proposals/0008-kernel-tuning.md b/docs/proposals/0008-kernel-tuning.md index efd6a32..a546c24 100644 --- a/docs/proposals/0008-kernel-tuning.md +++ b/docs/proposals/0008-kernel-tuning.md @@ -254,4 +254,51 @@ task/solution/output locations when moving the bundle; retain task/solution dige 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. \ No newline at end of file +are not established by CPU/fake-agent tests. + +### Generated catalog tasks and batch tuning + +Hand-writing a task package per operator does not scale to every core kernel. +`scripts/kernel_tuning/operators.py` declares each core Triton operator once: +its eager Torch reference, the production call used as the baseline, the +model-derived workload shapes, and the numerical contract. `generate` renders a +complete task package from an entry. The baseline is the production kernel +itself: the defining repository module and the repository modules it imports +are copied byte-for-byte under `vendor/` (absolute `embodiinfer.` imports are +rewritten to relative ones; package `__init__` files are never copied), and +`baseline.py` calls it as the engine does. `generated.json` records the +repository revision and source digests. Structured integer inputs such as +segment offsets are written as safetensors data inside the task. + +The catalog covers the Triton paths of pi0.5 (`ada_rms_norm`, `gated_residual`, +`gated_gelu`, `rotate_qk`, `split_kv_attention`), ActiveVLN's Qwen2.5-VL vision +tower (`rotate_half_rope`, `segmented_attention`), and StreamVLN's Qwen2 decode +path (`rms_norm`, `add_rms_norm`, `swiglu`). Shapes are representative of those +models, not traced traffic. All generated tasks time with CUDA Graphs. + +Contracts follow the production kernel, not an aspiration. `gated_residual` +and `rotate_qk` are bit-exact against Torch. The other kernels already round +differently from eager Torch, so their contract is a tolerance with a stated +reason: `atol = 2**-10` and `rtol = k * 2**-8`, where `k` counts BF16 roundings +the production kernel and the eager reference do not share (2 when only the +final casts differ; 3 for `rms_norm` and `swiglu`; 4 for `gated_gelu` and +`add_rms_norm`). The `gpu`-marked +`test_production_baseline_meets_generated_contract` checks every production +baseline against its contract, including changed-input CUDA Graph replay, and +prints the fraction of the bound it uses. Every baseline passed on an RTX 4060 +Laptop GPU (Torch 2.6, Triton 3.2; at most 0.85 of its bound) and on an RTX +5090 (Torch 2.12, Triton 3.7; at most 0.90), and the evaluator preflight passed +for all ten tasks on both. Re-run both on each target before trusting a +contract there. + +`tune-all` generates the selected operators into a new batch directory under +`results/kernel_tuning/batches/`, measures every production baseline with the +evaluator (preflight), and then runs each operator serially in its own +`run`/`resume` subprocess. A failed operator does not stop the batch; an +interrupt does, and `tune-all --resume BATCH` continues it, resuming existing +run archives. Promoted kernels are exported to `exports/`, and +`summary.md`/`summary.json` report attempts, promotions, and the improvement +over the production baseline from the weakest paired round. Machine-specific +settings (`evaluator_python`, `device`, budgets) are passed with `--set` at +generation time, and `--hardware-notes` copies target-hardware notes into each +task as `HARDWARE.md` for the agent. diff --git a/scripts/kernel_tuning/__main__.py b/scripts/kernel_tuning/__main__.py index 67ccbbd..2a02caa 100644 --- a/scripts/kernel_tuning/__main__.py +++ b/scripts/kernel_tuning/__main__.py @@ -11,7 +11,7 @@ from pathlib import Path from .artifacts import REPOSITORY, RunStore, atomic_json, run_lock, runtime_identity -from .contracts import HMZ_REVISION, ContractError, TaskPackage +from .contracts import HMZ_REVISION, ContractError, TaskPackage, read_json from .runner import evaluate @@ -73,6 +73,84 @@ def _execute(store: RunStore, *, resume: bool) -> dict: ) +def _selection(parser: argparse.ArgumentParser) -> None: + parser.add_argument("operators", nargs="*", help="Catalog operator names (default: all)") + parser.add_argument("--model", action="append", help="Select operators used by a model, 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") + + +def _generate(args: argparse.Namespace) -> None: + from .generate import parse_overrides, render + from .operators import select + + operators = select(args.operators or None, args.model) + 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 .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: + raise ContractError("--resume continues the recorded batch; omit selection, --set, and --agent") + 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") + root = batch.create( + select(args.operators or None, args.model), + args.output, + agent=agent, + overrides=parse_overrides(args.set), + hardware_notes=args.hardware_notes, + ) + 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: + for name, error in batch.preflight(root).items(): + print(f"preflight {name}: {'ok' 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__) @@ -91,9 +169,29 @@ def main(argv: list[str] | None = None) -> int: 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 task packages for catalog operators; 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 catalog operators and exit") + 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") + 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") args = parser.parse_args(argv) try: - if args.command == "check": + if 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"], diff --git a/scripts/kernel_tuning/batch.py b/scripts/kernel_tuning/batch.py new file mode 100644 index 0000000..9f8d083 --- /dev/null +++ b/scripts/kernel_tuning/batch.py @@ -0,0 +1,213 @@ +"""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 subprocess +import sys +from collections.abc import Callable +from datetime import datetime, timezone +from pathlib import Path +from typing import Any + +from .artifacts import REPOSITORY, RunStore, atomic_json, run_lock +from .contracts import ContractError, TaskPackage, read_json, validate_measurement +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, +) -> 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, + "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") + results: dict[str, str | None] = {} + for name, item in batch["items"].items(): + if names is not None and name not in names or 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"]) + results[name] = None + except ContractError as exc: + results[name] = str(exc) + item.update(status="preflight_failed", error=str(exc)) + atomic_json(root / "batch.json", batch) + return results + + +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") + for name, item in batch["items"].items(): + if item["status"] in DONE: + continue + 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"), + ] + 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.", + "", + "| operator | status | attempts | promoted | improvement | baseline ms | best ms | detail |", + "|---|---|---|---|---|---|---|---|", + ] + 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/contracts.py b/scripts/kernel_tuning/contracts.py index 3122b60..5f03c05 100644 --- a/scripts/kernel_tuning/contracts.py +++ b/scripts/kernel_tuning/contracts.py @@ -214,8 +214,14 @@ def mapping(loader: Any, node: Any) -> dict[str, Any]: 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), @@ -308,6 +314,10 @@ def _load(cls, root: Path) -> TaskPackage: 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: + # 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" diff --git a/scripts/kernel_tuning/generate.py b/scripts/kernel_tuning/generate.py new file mode 100644 index 0000000..8eac2c2 --- /dev/null +++ b/scripts/kernel_tuning/generate.py @@ -0,0 +1,450 @@ +"""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": 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": 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 "" + ) + 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 + +Shapes are representative of the listed models, not traced traffic. + +| 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/operators.py b/scripts/kernel_tuning/operators.py new file mode 100644 index 0000000..a3c9c63 --- /dev/null +++ b/scripts/kernel_tuning/operators.py @@ -0,0 +1,513 @@ +"""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. + +Shapes come from the models named in each entry and are representative, not +traced traffic. Override them by editing the entry or adding workloads here. +""" + +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 + +# ActiveVLN vision tower (Qwen2.5-VL ViT): 16 heads, head_dim 80, 64-patch +# windows. 1024 patches is one 448x448 frame; 4096 is four frames. +# 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"), "float32"), + "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 FP32 AdaRMS projection in scale/shift/gate order. " + "The gate output is the FP32 gate 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 with FP32 cosine/sine tables.", + models=("activevln",), + 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"), "float32"), + "sin": Tensor(("T", "D"), "float32"), + }, + 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"), + ), + precision=_tolerance( + "The production kernel may contract the FP32 multiply-add into an FMA before the BF16 cast." + ), + notes="Every vision block applies this before attention. head_dim 80 is not a power of two.", + ), + Operator( + name="segmented_attention", + summary="Non-causal attention independently inside packed variable-length segments.", + models=("activevln",), + 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/tests/kernel_tuning/test_batch.py b/tests/kernel_tuning/test_batch.py new file mode 100644 index 0000000..46c642f --- /dev/null +++ b/tests/kernel_tuning/test_batch.py @@ -0,0 +1,206 @@ +"""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)"} + 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", "--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 + + +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", "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_generated_tasks.py b/tests/kernel_tuning/test_generated_tasks.py new file mode 100644 index 0000000..ba4639b --- /dev/null +++ b/tests/kernel_tuning/test_generated_tasks.py @@ -0,0 +1,232 @@ +"""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.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 + + task = render(catalog()[name], 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) From e76468503f2da418c6e08f3406cd17a4694b46e5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=82=96=E4=B8=80=E6=B6=B5?= <23371301@buaa.edu.cn> Date: Sat, 3 Oct 2026 16:30:16 +0800 Subject: [PATCH 04/44] fix: time CUDA Graph replays without host launch latency MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit FlashInfer's do_bench synchronizes before each sample, so the start event precedes the host launch and every sample includes launch latency and jitter. On an RTX 5090 this put microsecond kernels at about 10 us and let identical solutions differ by up to 53%, far above the 3% promotion threshold. Queue the cold-L2 flush and input reset ahead of the start event instead, so the events bracket only device execution. Generated tasks now default to 1000 iterations, and preflight reports each baseline latency so operators already at the device's launch floor can be skipped. Signed-off-by: 肖一涵 <23371301@buaa.edu.cn> --- .../kernel_tuning/flashinfer_adapter.py | 60 ++++++++++++++----- docs/proposals/0008-kernel-tuning.md | 14 +++++ scripts/kernel_tuning/__main__.py | 9 ++- scripts/kernel_tuning/batch.py | 10 +++- scripts/kernel_tuning/generate.py | 8 ++- tests/kernel_tuning/test_batch.py | 3 + 6 files changed, 84 insertions(+), 20 deletions(-) diff --git a/benchmarks/kernel_tuning/flashinfer_adapter.py b/benchmarks/kernel_tuning/flashinfer_adapter.py index fbe6076..504fd49 100644 --- a/benchmarks/kernel_tuning/flashinfer_adapter.py +++ b/benchmarks/kernel_tuning/flashinfer_adapter.py @@ -111,6 +111,7 @@ def __init__(self, task: TaskPackage) -> None: self._driver = self._command( ["nvidia-smi", "--query-gpu=uuid,driver_version", "--format=csv,noheader"] ) + self._flush: Any = None @staticmethod def _command(argv: list[str]) -> str: @@ -200,7 +201,10 @@ def environment(self) -> dict[str, Any]: "timing": asdict(self.cfg.timing), "precision": asdict(self.cfg.precision), "seeds": list(self.cfg.seeds), - "cache_policy": "FlashInfer cold L2; input cloning excluded", + "cache_policy": ( + "cold L2 per sample; input copy/clone excluded; graph samples are device " + "time without host launch latency" + ), }, } @@ -361,7 +365,7 @@ def _verify_graph(self, runnable: Any, reference: Any, workload: Any, seed: int) return error def _time(self, runnable: Any, inputs: list[Any]) -> float: - from flashinfer_bench.bench.timing import do_bench, time_runnable + from flashinfer_bench.bench.timing import time_runnable args, _ = self._call(runnable, self._clone(inputs)) if self.cfg.timing.mode == "eager": @@ -370,23 +374,47 @@ def _time(self, runnable: Any, inputs: list[Any]) -> float: ) else: graph, args, _ = self._capture(runnable, self._clone(inputs)) + result = self._time_graph(graph, args, inputs) + self._check_flags() + return float(result) - def reset() -> None: - import torch + 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 - for dst, src in zip(args, inputs): - if isinstance(dst, torch.Tensor): - dst.copy_(src) + 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)] - result = do_bench( - lambda _: graph.replay(), - warmup=self.cfg.timing.warmup, - rep=self.cfg.timing.iterations, - setup=reset, - device=self.device, - ) - self._check_flags() - return float(result) + 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.""" diff --git a/docs/proposals/0008-kernel-tuning.md b/docs/proposals/0008-kernel-tuning.md index a546c24..0abfd15 100644 --- a/docs/proposals/0008-kernel-tuning.md +++ b/docs/proposals/0008-kernel-tuning.md @@ -79,6 +79,20 @@ the incumbent and initial baseline. Weights default to equal. Repeat paired measurements before promotion. Stop at 20 attempts, two hours of active execution, or five consecutive failures/non-improvements; all limits are configurable. +Every timed sample starts from a cold L2 (a 256 MiB buffer is zeroed) with +freshly copied inputs. CUDA Graph samples queue the flush and the input copy +ahead of the start event and do not synchronize per sample, so the events bracket +only the replay's device execution. FlashInfer's `do_bench` synchronizes first; +its start event then precedes the host launch, adding several microseconds of +launch latency and scheduler jitter to every sample. On an RTX 5090 that +inflated microsecond kernels to about 10 µs and let three identical solutions +differ by up to 53%, far above the 3% promotion threshold. Without it, and with +1000 iterations (the generated-task default), identical solutions differed by at +most 4.5% and mostly under 2%. A single trivial kernel replays in about 4.1 µs +under this protocol, the device's floor; tasks whose baseline is already there +leave an agent nothing to win, so preflight prints every baseline latency. +Eager timing still uses FlashInfer's `time_runnable`, including launch cost. + Runs snapshot task files and record source digests, evaluator identity, hardware and software fingerprints, precision settings, agent settings, attempts, raw measurements, and decisions. Resume requires matching task/tool/environment diff --git a/scripts/kernel_tuning/__main__.py b/scripts/kernel_tuning/__main__.py index 2a02caa..5dd957a 100644 --- a/scripts/kernel_tuning/__main__.py +++ b/scripts/kernel_tuning/__main__.py @@ -135,8 +135,13 @@ def _tune_all(args: argparse.Namespace) -> int: print((root / "summary.md").read_text(encoding="utf-8")) return 0 if not args.skip_preflight: - for name, error in batch.preflight(root).items(): - print(f"preflight {name}: {'ok' if error is None else error}", flush=True) + 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 ( diff --git a/scripts/kernel_tuning/batch.py b/scripts/kernel_tuning/batch.py index 9f8d083..b5ef8e4 100644 --- a/scripts/kernel_tuning/batch.py +++ b/scripts/kernel_tuning/batch.py @@ -88,7 +88,7 @@ def preflight(root: Path, names: list[str] | None = None) -> dict[str, str | Non batch = read_json(root / "batch.json") results: dict[str, str | None] = {} for name, item in batch["items"].items(): - if names is not None and name not in names or item["status"] != "pending": + if (names is not None and name not in names) or item["status"] != "pending": continue task = TaskPackage.load(root / item["task"]) try: @@ -96,6 +96,14 @@ def preflight(root: Path, names: list[str] | None = None) -> dict[str, str | Non report = asyncio.run(evaluate(store, store.root / "baseline/evaluation")) validate_measurement(task, report["measurement"]) results[name] = 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)) diff --git a/scripts/kernel_tuning/generate.py b/scripts/kernel_tuning/generate.py index 8eac2c2..5c32f91 100644 --- a/scripts/kernel_tuning/generate.py +++ b/scripts/kernel_tuning/generate.py @@ -182,7 +182,13 @@ def tuning_settings(operator: Operator, overrides: dict[str, Any] | None = None) "target_hardware": ["cuda"], "seeds": [0, 1, 2], "precision": {"matmul_precision": "highest", "allow_tf32": False, **operator.precision}, - "timing": {"mode": operator.timing, "warmup": 10, "iterations": 100, "trials": 3, "paired_rounds": 2}, + "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, diff --git a/tests/kernel_tuning/test_batch.py b/tests/kernel_tuning/test_batch.py index 46c642f..13553c8 100644 --- a/tests/kernel_tuning/test_batch.py +++ b/tests/kernel_tuning/test_batch.py @@ -163,6 +163,9 @@ async def evaluate(store: RunStore, output: Path, **_: Any) -> dict[str, Any]: 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" From 1f40a621a5bd48ec36772ab546e55107486b9879 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=82=96=E4=B8=80=E6=B6=B5?= <23371301@buaa.edu.cn> Date: Sat, 3 Oct 2026 18:37:06 +0800 Subject: [PATCH 05/44] feat: capture kernel shapes from real model runs for task workloads MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `capture` runs an unmodified model script with every catalog kernel wrapped and records each call as task axes, scalars and fixed inputs with call counts; `generate --captured` and `tune-all --captured` turn the most frequent cases into weighted workloads. A pi0.5 LIBERO-10 capture is included. It shows a 560/576-token prefix rather than the estimated 968, a prefix-only split-KV attention call, and BF16 AdaRMS modulation, which the catalog now declares. Signed-off-by: 肖一涵 <23371301@buaa.edu.cn> --- CONTRIBUTING.md | 4 +- .../captures/pi05-libero10.jsonl | 8 + docs/proposals/0008-kernel-tuning.md | 29 ++ scripts/kernel_tuning/__main__.py | 42 ++- scripts/kernel_tuning/capture.py | 311 ++++++++++++++++++ scripts/kernel_tuning/generate.py | 7 +- scripts/kernel_tuning/operators.py | 7 +- tests/kernel_tuning/test_capture.py | 204 ++++++++++++ 8 files changed, 602 insertions(+), 10 deletions(-) create mode 100644 benchmarks/kernel_tuning/captures/pi05-libero10.jsonl create mode 100644 scripts/kernel_tuning/capture.py create mode 100644 tests/kernel_tuning/test_capture.py diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 6555774..d196c3c 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -79,7 +79,9 @@ uv run --project scripts/kernel_tuning python -m scripts.kernel_tuning tune-all uv run --project scripts/kernel_tuning python -m scripts.kernel_tuning tune-all --resume results/kernel_tuning/batches/BATCH_ID ``` -Select a subset with operator names or `--model pi05|activevln|streamvln`. +Select a subset with operator names or `--model pi05|activevln|streamvln`. Add +`--captured FILE` to tune with shapes recorded from real traffic by `capture`, for example +`benchmarks/kernel_tuning/captures/pi05-libero10.jsonl`. 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 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/docs/proposals/0008-kernel-tuning.md b/docs/proposals/0008-kernel-tuning.md index 0abfd15..f37b42e 100644 --- a/docs/proposals/0008-kernel-tuning.md +++ b/docs/proposals/0008-kernel-tuning.md @@ -316,3 +316,32 @@ over the production baseline from the weakest paired round. Machine-specific settings (`evaluator_python`, `device`, budgets) are passed with `--set` at generation time, and `--hardware-notes` copies target-hardware notes into each task as `HARDWARE.md` for the agent. + +### Captured workloads + +Catalog shapes are estimates. `capture` replaces them with production traffic: +it runs an unmodified model script in the model's runtime with every catalog +kernel wrapped, converts each call to the operator's task axes, scalars, and +fixed integer inputs, and writes one JSONL row per distinct case with its call +count. Calls outside a task's semantics (another dtype, broadcast tables, +custom attention scaling) are counted as skipped with the reason. Shapes depend +on the configuration and inputs, not on weight values; run with CUDA Graphs and +compilation disabled so every call reaches Python. `generate --captured` and +`tune-all --captured` then use the most frequent cases as workloads, weighted by +call count; operators the capture never saw keep their estimates. + +`benchmarks/kernel_tuning/captures/pi05-libero10.jsonl` comes from +`benchmarks/pi05-benchmark` on 200 LIBERO-10 frames (10 tasks), with the +production Triton paths enabled (`native_inference`, Triton prefix and denoise +attention). It corrected three estimates. The prefix holds two cameras plus a +48- or 64-token language bucket (560 or 576 tokens, not 968). `split_kv_attention` +also runs once per layer over the whole prefix with no cached keys (`P=0`, +`S=576`). The checkpoint's AdaRMS modulation is BF16, not FP32; the catalog now +declares BF16 (recalibrated: at most 0.79 of its bound on an RTX 5090). At B=1, +four of the five pi0.5 kernels measure about 4 µs on an RTX 5090, the device's +launch floor; only `split_kv_attention` (16-35 µs) leaves room to tune. + +```bash +PYTHONPATH=/path/to/EmbodiInfer /path/to/model-runtime/python -m scripts.kernel_tuning capture --output pi05.jsonl -- benchmark.py --config capture-config.yaml +python -m scripts.kernel_tuning tune-all --model pi05 --captured pi05.jsonl --agent ... +``` diff --git a/scripts/kernel_tuning/__main__.py b/scripts/kernel_tuning/__main__.py index 5dd957a..1c4befd 100644 --- a/scripts/kernel_tuning/__main__.py +++ b/scripts/kernel_tuning/__main__.py @@ -84,13 +84,19 @@ def _selection(parser: argparse.ArgumentParser) -> None: 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") + parser.add_argument( + "--captured", type=Path, help="Replace estimated workloads with shapes from a `capture` file" + ) def _generate(args: argparse.Namespace) -> None: + from .capture import apply_file from .generate import parse_overrides, render from .operators import select - operators = select(args.operators or None, args.model) + operators, notes = apply_file(select(args.operators or None, args.model), args.captured) + for note in notes: + print(note) if args.list: for operator in operators: print(f"{operator.name:22} {','.join(operator.models):10} {operator.summary}") @@ -111,19 +117,23 @@ def _generate(args: argparse.Namespace) -> None: 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: + if args.operators or args.model or args.set or args.hardware_notes or args.agent or args.captured: raise ContractError("--resume continues the recorded batch; omit selection, --set, and --agent") 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") + operators, notes = apply_file(select(args.operators or None, args.model), args.captured) + for note in notes: + print(note, flush=True) root = batch.create( - select(args.operators or None, args.model), + operators, args.output, agent=agent, overrides=parse_overrides(args.set), @@ -190,9 +200,31 @@ def main(argv: list[str] | None = None) -> int: 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") - args = parser.parse_args(argv) + capture = commands.add_parser( + "capture", help="Record production kernel shapes while a model script runs (model runtime)" + ) + capture.add_argument("operators", nargs="*", help="Catalog operators to record (default: all)") + capture.add_argument("--output", type=Path, required=True, help="JSONL file of captured cases") + capture.add_argument("script", nargs=argparse.REMAINDER, help="-- SCRIPT [ARGS] or -- -m MODULE [ARGS]") + raw = sys.argv[1:] if argv is None else argv + if "--" in raw and raw[0] == "capture": + split = raw.index("--") + raw, command = raw[:split], raw[split + 1 :] + else: + command = None + args = parser.parse_args(raw) try: - if args.command == "generate": + 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 == "generate": _generate(args) elif args.command == "tune-all": return _tune_all(args) diff --git a/scripts/kernel_tuning/capture.py b/scripts/kernel_tuning/capture.py new file mode 100644 index 0000000..98fdfa2 --- /dev/null +++ b/scripts/kernel_tuning/capture.py @@ -0,0 +1,311 @@ +"""Record the shapes of production kernel calls while a model actually runs. + +``python -m scripts.kernel_tuning capture --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 CUDA Graphs and compilation disabled: +inside a graph only the capture pass reaches Python, so counts would be lost. +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 .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) + if str(cos.dtype) != "torch.float32": + raise Skip(f"cos dtype {cos.dtype}") + 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}") + if _capturing(q): + raise Skip("offsets unreadable during CUDA Graph capture") + bounds = [int(x) for x in offsets.tolist()] + 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, +} + + +def _capturing(tensor: Any) -> bool: + if not tensor.is_cuda: + 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.""" + + def __init__(self) -> None: + self.cases: Counter[str] = Counter() + self.skipped: Counter[tuple[str, str]] = Counter() + + 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 + self.cases[json.dumps({"operator": name, **case}, sort_keys=True)] += 1 + + 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 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() + replaced = install(operators, 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 + + +def load(path: Path) -> list[dict[str, Any]]: + """Read a capture file written by :func:`run`.""" + rows = [json.loads(line) for line in path.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) -> tuple[list[Operator], list[str]]: + """Apply a capture file to every selected operator; ``None`` keeps the catalog shapes.""" + if path is None: + return operators, [] + 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/generate.py b/scripts/kernel_tuning/generate.py index 5c32f91..89ed760 100644 --- a/scripts/kernel_tuning/generate.py +++ b/scripts/kernel_tuning/generate.py @@ -304,6 +304,11 @@ def shape(spec: Any) -> str: 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 "" ) @@ -341,7 +346,7 @@ def shape(spec: Any) -> str: ## Workloads -Shapes are representative of the listed models, not traced traffic. +{origin} | uuid | axes and scalars | weight | where | |---|---|---|---| diff --git a/scripts/kernel_tuning/operators.py b/scripts/kernel_tuning/operators.py index a3c9c63..2890e55 100644 --- a/scripts/kernel_tuning/operators.py +++ b/scripts/kernel_tuning/operators.py @@ -96,7 +96,7 @@ def _tolerance(reason: str, steps: int = 2) -> dict[str, Any]: axes={"B": None, "T": None, "H": None, "H3": None, "ONE": 1}, inputs={ "x": Tensor(("B", "T", "H"), BF16), - "modulation": Tensor(("B", "H3"), "float32"), + "modulation": Tensor(("B", "H3"), BF16), "eps": Tensor(None, "float32"), }, outputs={"output": Tensor(("B", "T", "H"), BF16), "gate": Tensor(("B", "ONE", "H"), BF16)}, @@ -124,8 +124,9 @@ def run(x, modulation, eps): precision=_tolerance( "The production kernel reduces the FP32 variance in a different order than torch.mean." ), - notes="`modulation` is the FP32 AdaRMS projection in scale/shift/gate order. " - "The gate output is the FP32 gate cast to the hidden dtype, one row per batch item.", + 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", diff --git a/tests/kernel_tuning/test_capture.py b/tests/kernel_tuning/test_capture.py new file mode 100644 index 0000000..a8a7175 --- /dev/null +++ b/tests/kernel_tuning/test_capture.py @@ -0,0 +1,204 @@ +"""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", "--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", "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",) From f6e85386018f12c95dd12ea51ba4a48687376ed5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=82=96=E4=B8=80=E6=B6=B5?= <23371301@buaa.edu.cn> Date: Sat, 3 Oct 2026 19:13:38 +0800 Subject: [PATCH 06/44] feat: capture Qwen R2R and StreamVLN shapes on skeleton checkpoints MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Count kernels captured into CUDA Graphs on every replay, so production graph configurations can be captured, and add `skeleton`, which builds a random-weight checkpoint from safetensors headers: shapes do not depend on weight values. Committed captures now replace the catalog estimates by default (`--estimated` restores them). New captures for the Qwen R2R low-level and panoramic policies and StreamVLN correct the vision shapes (1,980/7,704 patches in ragged windows) and show decode-only StreamVLN calls. The rotary tables are BF16, which makes rotate_half_rope bit-exact, and the vision kernels belong to the Qwen R2R policies: ActiveVLN does not call them. Signed-off-by: 肖一涵 <23371301@buaa.edu.cn> --- CONTRIBUTING.md | 7 +- .../kernel_tuning/captures/qwen-r2r-low.jsonl | 2 + .../captures/qwen-r2r-panoramic.jsonl | 2 + .../captures/streamvln-r2r-rxr.jsonl | 3 + docs/proposals/0008-kernel-tuning.md | 45 ++++-- scripts/kernel_tuning/__main__.py | 41 +++++- scripts/kernel_tuning/capture.py | 115 +++++++++++++-- scripts/kernel_tuning/operators.py | 27 ++-- scripts/kernel_tuning/skeleton.py | 137 ++++++++++++++++++ tests/kernel_tuning/test_capture.py | 43 ++++++ tests/kernel_tuning/test_generated_tasks.py | 4 +- 11 files changed, 379 insertions(+), 47 deletions(-) create mode 100644 benchmarks/kernel_tuning/captures/qwen-r2r-low.jsonl create mode 100644 benchmarks/kernel_tuning/captures/qwen-r2r-panoramic.jsonl create mode 100644 benchmarks/kernel_tuning/captures/streamvln-r2r-rxr.jsonl create mode 100644 scripts/kernel_tuning/skeleton.py diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index d196c3c..80a5f9a 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -79,9 +79,10 @@ uv run --project scripts/kernel_tuning python -m scripts.kernel_tuning tune-all uv run --project scripts/kernel_tuning python -m scripts.kernel_tuning tune-all --resume results/kernel_tuning/batches/BATCH_ID ``` -Select a subset with operator names or `--model pi05|activevln|streamvln`. Add -`--captured FILE` to tune with shapes recorded from real traffic by `capture`, for example -`benchmarks/kernel_tuning/captures/pi05-libero10.jsonl`. +Select a subset with operator names or `--model pi05|qwen25-vln|streamvln`. Workloads come +from the captures in `benchmarks/kernel_tuning/captures/` (recorded by `capture`, optionally +on a random-weight `skeleton` checkpoint); `--captured FILE` selects another capture and +`--estimated` uses the catalog's estimates. 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 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/docs/proposals/0008-kernel-tuning.md b/docs/proposals/0008-kernel-tuning.md index f37b42e..e7aeab5 100644 --- a/docs/proposals/0008-kernel-tuning.md +++ b/docs/proposals/0008-kernel-tuning.md @@ -285,10 +285,13 @@ repository revision and source digests. Structured integer inputs such as segment offsets are written as safetensors data inside the task. The catalog covers the Triton paths of pi0.5 (`ada_rms_norm`, `gated_residual`, -`gated_gelu`, `rotate_qk`, `split_kv_attention`), ActiveVLN's Qwen2.5-VL vision -tower (`rotate_half_rope`, `segmented_attention`), and StreamVLN's Qwen2 decode -path (`rms_norm`, `add_rms_norm`, `swiglu`). Shapes are representative of those -models, not traced traffic. All generated tasks time with CUDA Graphs. +`gated_gelu`, `rotate_qk`, `split_kv_attention`), the Qwen2.5-VL vision tower of +the Qwen R2R low-level and panoramic policies (`rotate_half_rope`, +`segmented_attention`), and StreamVLN's Qwen2 decode path (`rms_norm`, +`add_rms_norm`, `swiglu`). ActiveVLN runs the Transformers vision tower and an +eager attention backend, so it currently calls none of these kernels. Workloads +come from committed captures (see below); the catalog's estimates are only the +fallback (`--estimated`). All generated tasks time with CUDA Graphs. Contracts follow the production kernel, not an aspiration. `gated_residual` and `rotate_qk` are bit-exact against Torch. The other kernels already round @@ -324,11 +327,21 @@ it runs an unmodified model script in the model's runtime with every catalog kernel wrapped, converts each call to the operator's task axes, scalars, and fixed integer inputs, and writes one JSONL row per distinct case with its call count. Calls outside a task's semantics (another dtype, broadcast tables, -custom attention scaling) are counted as skipped with the reason. Shapes depend -on the configuration and inputs, not on weight values; run with CUDA Graphs and -compilation disabled so every call reaches Python. `generate --captured` and -`tune-all --captured` then use the most frequent cases as workloads, weighted by -call count; operators the capture never saw keep their estimates. +custom attention scaling) are counted as skipped with the reason. Kernels +launched while a CUDA Graph is captured are attached to that graph and counted on +every replay, so production graph configurations can be captured; compilation +should be disabled. Captures committed under `benchmarks/kernel_tuning/captures/` +replace the catalog estimates by default, using each operator's most frequent +cases weighted by call count; `--captured FILE` selects another capture and +`--estimated` restores the estimates. + +Shapes depend on the configuration and the inputs, never on weight values. +`skeleton REPO --revision REV` therefore builds a random-weight copy of a +Hugging Face checkpoint: it downloads the small files and synthesizes each +safetensors file from its header, read with an HTTP range request. Benchmark +scripts then run unchanged on it, with synthetic frames where the policy resizes +images to a fixed size anyway. Random weights do not reproduce generated text, so +decode lengths are fixed in the capture configuration. `benchmarks/kernel_tuning/captures/pi05-libero10.jsonl` comes from `benchmarks/pi05-benchmark` on 200 LIBERO-10 frames (10 tasks), with the @@ -341,6 +354,20 @@ declares BF16 (recalibrated: at most 0.79 of its bound on an RTX 5090). At B=1, four of the five pi0.5 kernels measure about 4 µs on an RTX 5090, the device's launch floor; only `split_kv_attention` (16-35 µs) leaves room to tune. +The Qwen R2R captures (`qwen-r2r-low.jsonl`, `qwen-r2r-panoramic.jsonl`) use +skeleton checkpoints and synthetic R2R/RxR episodes (the trajectory release is +gated; episode lengths follow the benchmark READMEs) with `attention_backend: +triton` and CUDA Graphs on: only the graph path runs EmbodiInfer's Triton vision +forward. Each call encodes 4 history frames plus the current one: 1,980 patches +in 46 ragged windows (low-level) or 7,704 patches in 172 windows (panoramic), +not the estimated uniform 1,024/4,096. The rotary tables are BF16, which makes +the production RoPE bit-exact; the catalog now declares both. These shapes put +`segmented_attention` at 357 µs and 1,405 µs and `rotate_half_rope` at 18 µs and +61 µs on an RTX 5090. The StreamVLN capture (`streamvln-r2r-rxr.jsonl`, skeleton +weights, decode capped at the five tokens of its fast action path) sees only +single-row decode calls: `swiglu` never runs on prefill chunks, and all three +StreamVLN kernels sit at the launch floor. + ```bash PYTHONPATH=/path/to/EmbodiInfer /path/to/model-runtime/python -m scripts.kernel_tuning capture --output pi05.jsonl -- benchmark.py --config capture-config.yaml python -m scripts.kernel_tuning tune-all --model pi05 --captured pi05.jsonl --agent ... diff --git a/scripts/kernel_tuning/__main__.py b/scripts/kernel_tuning/__main__.py index 1c4befd..2417e7d 100644 --- a/scripts/kernel_tuning/__main__.py +++ b/scripts/kernel_tuning/__main__.py @@ -84,9 +84,13 @@ def _selection(parser: argparse.ArgumentParser) -> None: 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") - parser.add_argument( - "--captured", type=Path, help="Replace estimated workloads with shapes from a `capture` file" + shapes = parser.add_mutually_exclusive_group() + shapes.add_argument( + "--captured", + type=Path, + help="Shapes from this `capture` file or directory (default: benchmarks/kernel_tuning/captures)", ) + shapes.add_argument("--estimated", action="store_true", help="Use the catalog's estimated shapes") def _generate(args: argparse.Namespace) -> None: @@ -94,7 +98,9 @@ def _generate(args: argparse.Namespace) -> None: from .generate import parse_overrides, render from .operators import select - operators, notes = apply_file(select(args.operators or None, args.model), args.captured) + 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: @@ -122,14 +128,24 @@ def _tune_all(args: argparse.Namespace) -> int: 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.captured: + if ( + args.operators + or args.model + or args.set + or args.hardware_notes + or args.agent + or args.captured + or args.estimated + ): raise ContractError("--resume continues the recorded batch; omit selection, --set, and --agent") 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") - operators, notes = apply_file(select(args.operators or None, args.model), args.captured) + 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( @@ -206,6 +222,12 @@ def main(argv: list[str] | None = None) -> int: capture.add_argument("operators", nargs="*", help="Catalog operators to record (default: all)") capture.add_argument("--output", type=Path, required=True, help="JSONL file of captured cases") capture.add_argument("script", nargs=argparse.REMAINDER, help="-- SCRIPT [ARGS] or -- -m MODULE [ARGS]") + 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] == "capture": split = raw.index("--") @@ -224,6 +246,15 @@ def main(argv: list[str] | None = None) -> int: for row in rows: if "skipped" in row: print(f" skipped {row['count']} {row['operator']} calls: {row['skipped']}") + 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": diff --git a/scripts/kernel_tuning/capture.py b/scripts/kernel_tuning/capture.py index 98fdfa2..05115fb 100644 --- a/scripts/kernel_tuning/capture.py +++ b/scripts/kernel_tuning/capture.py @@ -6,8 +6,9 @@ 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 CUDA Graphs and compilation disabled: -inside a graph only the capture pass reaches Python, so counts would be lost. +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. """ @@ -26,6 +27,7 @@ from pathlib import Path from typing import Any +from .artifacts import REPOSITORY from .contracts import ContractError from .operators import Operator, Workload, catalog @@ -80,9 +82,7 @@ def _split_kv_attention(a: Bound) -> dict[str, Any]: def _rotate_half_rope(a: Bound) -> dict[str, Any]: q, k, cos = a["q"], a["k"], a["cos"] - _bf16(q, k) - if str(cos.dtype) != "torch.float32": - raise Skip(f"cos dtype {cos.dtype}") + _bf16(q, k, cos, a["sin"]) if cos.shape[0] != q.shape[0]: raise Skip("broadcast cos/sin rows") tokens, heads, width = q.shape @@ -100,9 +100,16 @@ def _segmented_attention(a: Bound) -> dict[str, Any]: 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): - raise Skip("offsets unreadable during CUDA Graph capture") - bounds = [int(x) for x in offsets.tolist()] + # 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: @@ -151,8 +158,12 @@ def _swiglu(a: Bound) -> dict[str, Any]: } +#: 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 tensor.is_cuda: + if not getattr(tensor, "is_cuda", False): return False import torch @@ -160,11 +171,18 @@ def _capturing(tensor: Any) -> bool: class Recorder: - """Convert production calls to task cases; counts identical cases instead of storing calls.""" + """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.""" @@ -176,7 +194,20 @@ def record(self, name: str, bound: Bound) -> None: except Exception as exc: # noqa: BLE001 -- recording is best effort by design self.skipped[name, f"unmapped call: {type(exc).__name__}: {exc}"] += 1 return - self.cases[json.dumps({"operator": name, **case}, sort_keys=True)] += 1 + 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.""" @@ -223,12 +254,48 @@ def install(operators: list[Operator], recorder: Recorder) -> list[tuple[Any, st 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() - replaced = install(operators, recorder) + _OFFSETS.clear() + replaced = install(operators, recorder) + install_graph_hooks(recorder) argv = sys.argv try: if command[0] == "-m": @@ -252,9 +319,19 @@ def run(command: list[str], output: Path, operators: list[Operator]) -> list[dic 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`.""" - rows = [json.loads(line) for line in path.read_text(encoding="utf-8").splitlines() if line.strip()] + """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 @@ -297,10 +374,16 @@ def apply(operator: Operator, rows: list[dict[str, Any]], *, top: int = 8) -> tu return dataclasses.replace(operator, workloads=tuple(workloads)), summary -def apply_file(operators: list[Operator], path: Path | None) -> tuple[list[Operator], list[str]]: - """Apply a capture file to every selected operator; ``None`` keeps the catalog shapes.""" +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: - return operators, [] + 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)}") diff --git a/scripts/kernel_tuning/operators.py b/scripts/kernel_tuning/operators.py index 2890e55..6590e00 100644 --- a/scripts/kernel_tuning/operators.py +++ b/scripts/kernel_tuning/operators.py @@ -4,8 +4,10 @@ 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. -Shapes come from the models named in each entry and are representative, not -traced traffic. Override them by editing the entry or adding workloads here. +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 @@ -80,8 +82,9 @@ def _tolerance(reason: str, steps: int = 2) -> dict[str, Any]: PI05_TOKENS = 50 PI05_PREFIX = 968 -# ActiveVLN vision tower (Qwen2.5-VL ViT): 16 heads, head_dim 80, 64-patch -# windows. 1024 patches is one 448x448 frame; 4096 is four frames. +# 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, ...] = ( @@ -287,8 +290,8 @@ def run(query, prefix_key, prefix_value, suffix_key, suffix_value, scaling): ), Operator( name="rotate_half_rope", - summary="Rotate-half RoPE for packed vision Q/K with FP32 cosine/sine tables.", - models=("activevln",), + 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", @@ -297,8 +300,8 @@ def run(query, prefix_key, prefix_value, suffix_key, suffix_value, scaling): inputs={ "q": Tensor(("T", "HQ", "D"), BF16), "k": Tensor(("T", "HK", "D"), BF16), - "cos": Tensor(("T", "D"), "float32"), - "sin": Tensor(("T", "D"), "float32"), + "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"), @@ -322,15 +325,13 @@ def run(q, k, cos, sin): 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"), ), - precision=_tolerance( - "The production kernel may contract the FP32 multiply-add into an FMA before the BF16 cast." - ), - notes="Every vision block applies this before attention. head_dim 80 is not a power of two.", + 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=("activevln",), + models=("qwen25-vln",), op_type="attention", source="embodiinfer/backend/triton/segmented_attention.py", entry="segmented_attention", 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/tests/kernel_tuning/test_capture.py b/tests/kernel_tuning/test_capture.py index a8a7175..112a3a9 100644 --- a/tests/kernel_tuning/test_capture.py +++ b/tests/kernel_tuning/test_capture.py @@ -202,3 +202,46 @@ def test_cli_capture_and_generate_from_capture( 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_generated_tasks.py b/tests/kernel_tuning/test_generated_tasks.py index ba4639b..6588467 100644 --- a/tests/kernel_tuning/test_generated_tasks.py +++ b/tests/kernel_tuning/test_generated_tasks.py @@ -11,6 +11,7 @@ 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, @@ -183,7 +184,8 @@ def test_production_baseline_meets_generated_contract(name: str, tmp_path: Path) spec.loader.exec_module(adapter) compare = adapter.compare_outputs - task = render(catalog()[name], tmp_path / "task") + (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) From f9bce478643f220cfa7d8352ea481115f1186a88 Mon Sep 17 00:00:00 2001 From: owpeter Date: Sun, 4 Oct 2026 10:50:15 +0800 Subject: [PATCH 07/44] feat: discover model operators for kernel tuning Capture ATen overloads and backend calls during model preparation, exclude quantized scopes, and require complete task coverage and baseline preflight before starting a model tuning batch. Preserve tensor layouts, aliases, mutations, default-generator RNG transitions and float64 values during replay. Handle zero-dropout SDPA and retain the legacy catalog workflow behind explicit flags. Validation: 104 local tooling tests and 35 server discovery/replay tests passed. Pi0.5 random-weight validation on RTX 5090 generated 94 tasks and passed all 94 baselines without starting search. The broader CPU suite retains one existing Transformers import-order test conflict. Signed-off-by: owpeter --- CONTRIBUTING.md | 50 +- .../kernel_tuning/flashinfer_adapter.py | 15 +- benchmarks/kernel_tuning/replay_adapter.py | 279 ++++++++ benchmarks/kernel_tuning/replay_schema.py | 42 ++ docs/proposals/0008-kernel-tuning.md | 154 ++++- scripts/kernel_tuning/__main__.py | 132 +++- scripts/kernel_tuning/batch.py | 178 ++++- scripts/kernel_tuning/contracts.py | 16 +- scripts/kernel_tuning/discovery/__init__.py | 5 + scripts/kernel_tuning/discovery/contracts.py | 180 +++++ scripts/kernel_tuning/discovery/recorder.py | 654 ++++++++++++++++++ scripts/kernel_tuning/discovery/rng.py | 114 +++ scripts/kernel_tuning/discovery/tasks.py | 300 ++++++++ scripts/kernel_tuning/discovery/tensors.py | 170 +++++ scripts/kernel_tuning/prompts.py | 12 + tests/kernel_tuning/test_batch.py | 12 +- tests/kernel_tuning/test_capture.py | 12 +- tests/kernel_tuning/test_discovery.py | 435 ++++++++++++ tests/kernel_tuning/test_replay_extensions.py | 286 ++++++++ 19 files changed, 2991 insertions(+), 55 deletions(-) create mode 100644 benchmarks/kernel_tuning/replay_adapter.py create mode 100644 benchmarks/kernel_tuning/replay_schema.py create mode 100644 scripts/kernel_tuning/discovery/__init__.py create mode 100644 scripts/kernel_tuning/discovery/contracts.py create mode 100644 scripts/kernel_tuning/discovery/recorder.py create mode 100644 scripts/kernel_tuning/discovery/rng.py create mode 100644 scripts/kernel_tuning/discovery/tasks.py create mode 100644 scripts/kernel_tuning/discovery/tensors.py create mode 100644 tests/kernel_tuning/test_discovery.py create mode 100644 tests/kernel_tuning/test_replay_extensions.py diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 80a5f9a..921ca82 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -68,21 +68,48 @@ uv run --project scripts/kernel_tuning python -m scripts.kernel_tuning --help uv run --project scripts/kernel_tuning python -m scripts.kernel_tuning check benchmarks/kernel_tuning/tasks/gated_residual ``` -To tune every core operator, generate tasks from the operator catalog and run -them as one batch (see the -[generated-task section](docs/proposals/0008-kernel-tuning.md#generated-catalog-tasks-and-batch-tuning)): +To tune one model, capture its existing preparation forward in its model runtime, +then generate and preflight every eligible observed call before starting agents +(see [model discovery](docs/proposals/0008-kernel-tuning.md#model-scoped-operator-discovery)): ```bash -uv run --project scripts/kernel_tuning python -m scripts.kernel_tuning generate --list -uv run --project scripts/kernel_tuning python -m scripts.kernel_tuning tune-all --preflight-only --set evaluator_python=/absolute/runtime/python -uv run --project scripts/kernel_tuning python -m scripts.kernel_tuning tune-all --agent 'HARNESS/MODEL:EFFORT' --set evaluator_python=/absolute/runtime/python +/absolute/runtime/python -m scripts.kernel_tuning capture --model pi05 --output results/kernel_tuning/captures/pi05 -- /absolute/preparation.py +uv run --project scripts/kernel_tuning python -m scripts.kernel_tuning generate --model pi05 --captured results/kernel_tuning/captures/pi05 --list +uv run --project scripts/kernel_tuning python -m scripts.kernel_tuning tune-all --model pi05 --captured results/kernel_tuning/captures/pi05 --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 ``` -Select a subset with operator names or `--model pi05|qwen25-vln|streamvln`. Workloads come -from the captures in `benchmarks/kernel_tuning/captures/` (recorded by `capture`, optionally -on a random-weight `skeleton` checkpoint); `--captured FILE` selects another capture and -`--estimated` uses the catalog's estimates. +`/absolute/preparation.py` is your existing preparation driver, constructing the +selected canonical policy through `make_policy` and running its normal forward +path. Disable compilation explicitly; select the same checkpoint, backend, +dtype, batching and decoding settings that will be tuned. Discovery covers only +paths exercised by that preparation. The capture records ATen overloads and +existing backend entrypoints, all observed shapes/layouts, and call counts. +Quantized calls and their internals are excluded; views/allocations and host work +are listed separately. No new fusion is introduced. A missing replay contract, +required fixture, or failing baseline blocks the entire model batch. `--dry-run` +only builds tasks; `--preflight-only` measures baselines without opening agents. +Model batches cannot use `--skip-preflight`, partial operator selection, or +estimated shapes. A standalone `run` cannot bypass their preflight barrier. + +Capture may store input data under a 64 MiB limit (`--fixture-bytes`): small +floating fixtures when space permits, and required integer/index fixtures. +Capture directories and tasks are hashed and checked again on resume. Large +floating inputs use seeded generation; calls needing more specific inputs must +pass preflight or receive a dedicated replay adapter. Start with a new output +directory when changing preparation/configuration. + +Default-generator `randn` tasks save required before/after RNG fixtures, check +both output bytes and generator advancement, and support eager timing only. +Explicit generators and other unsupported RNG calls remain coverage gaps. +SDPA with zero/default dropout is eligible; nonzero dropout still needs an adapter. +Discovered float64 tensors use the opt-in replay definition schema and retain +double precision through task generation and evaluation. + +The old catalog workflow remains available explicitly with `generate --catalog`, +`tune-all --catalog`, and `capture --catalog`. Only that workflow accepts legacy +JSONL captures, catalog model groups, operator subsets and `--estimated`; it +does not certify model coverage. Existing JSONL must be recaptured for model tuning. 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 @@ -125,7 +152,8 @@ uv run --project scripts/kernel_tuning ruff check scripts/kernel_tuning benchmar uv run --project scripts/kernel_tuning ruff format --check scripts/kernel_tuning benchmarks/kernel_tuning tests/kernel_tuning ``` -The `--confcutdir` option avoids loading the core test suite's Torch-dependent +CPU discovery/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 diff --git a/benchmarks/kernel_tuning/flashinfer_adapter.py b/benchmarks/kernel_tuning/flashinfer_adapter.py index 504fd49..17e694f 100644 --- a/benchmarks/kernel_tuning/flashinfer_adapter.py +++ b/benchmarks/kernel_tuning/flashinfer_adapter.py @@ -80,7 +80,7 @@ class FlashInferAdapter: def __init__(self, task: TaskPackage) -> None: import torch - from flashinfer_bench.data import Definition, Trace + from flashinfer_bench.data import Trace if importlib.metadata.version("flashinfer-bench") != FLASHINFER_VERSION: raise ContractError(f"This adapter requires flashinfer-bench=={FLASHINFER_VERSION}") @@ -95,7 +95,7 @@ def __init__(self, task: TaskPackage) -> None: self.device = str(torch.device(self.cfg.device)) torch.cuda.set_device(self.device) self.device = f"cuda:{torch.cuda.current_device()}" - self.definition = Definition.model_validate(task.definition) + 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"} @@ -113,6 +113,12 @@ def __init__(self, task: TaskPackage) -> None: ) 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]): @@ -425,7 +431,7 @@ def evaluate(self, solutions: dict[str, dict[str, Any]]) -> dict[str, Any]: if set(solutions) != set(ROLES): raise ContractError("Expected baseline, incumbent, and candidate solutions") registry = BuilderRegistry.get_instance() - reference = registry.build_reference(self.definition) + reference = self._build_reference(registry) runnables = { role: registry.build(self.definition, Solution.model_validate(solution)) for role, solution in solutions.items() @@ -471,6 +477,9 @@ def evaluate(self, solutions: dict[str, dict[str, Any]]) -> dict[str, Any]: rounds.append(measurements) return {"checks": checks, "rounds": rounds} + 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 diff --git a/benchmarks/kernel_tuning/replay_adapter.py b/benchmarks/kernel_tuning/replay_adapter.py new file mode 100644 index 0000000..5b10317 --- /dev/null +++ b/benchmarks/kernel_tuning/replay_adapter.py @@ -0,0 +1,279 @@ +"""Evaluate discovered calls with captured layouts, aliases, and mutable state.""" + +from __future__ import annotations + +from contextlib import nullcontext +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._rng: Any = None + if self.cfg.profile: + raise ContractError("Discovered tasks need a replay-aware NCU adapter before profile=true") + if self.replay.get("rng") and self.cfg.timing.mode != "eager": + raise ContractError("RNG replay currently requires eager timing") + + def _load_definition(self, definition: dict[str, Any]) -> Any: + from .replay_schema import ReplayDefinition + + return ReplayDefinition.model_validate(definition) + + def evaluate(self, solutions: dict[str, dict[str, Any]]) -> dict[str, Any]: + """Evaluate stateful RNG calls without leaking their state into the caller.""" + from scripts.kernel_tuning.discovery.rng import preserve + + rng = self.replay.get("rng") + with preserve(rng["devices"]) if rng else nullcontext(): + return super().evaluate(solutions) + + 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 _seed(self, seed: int) -> None: + # RngReplay resets only the declared generators. The ordinary adapter's + # global seed would also mutate unrelated CUDA devices. + if not self.replay.get("rng"): + super()._seed(seed) + + 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] + from scripts.kernel_tuning.discovery.rng import RngReplay + + self._rng = ( + RngReplay(self.replay["rng"], self.case, self.task.root, seed) if self.replay.get("rng") else None + ) + 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("Discovered 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]: + from scripts.kernel_tuning.discovery.rng import preserve + + rng = getattr(self, "_rng", None) + with preserve(rng.devices) if rng else nullcontext(): + 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 + from scripts.kernel_tuning.discovery.rng import snapshot + + rng = getattr(self, "_rng", None) + expected_inputs = self._clone(inputs) + self._seed(seed) + if rng: + rng.reset() + _, expected = self._call(reference, expected_inputs) + if rng: + expected_rng = snapshot(rng.devices) + rng.check_reference(expected_rng) + expected = [value.clone() for value in expected] + self._seed(seed) + local = self._clone(inputs) + if rng: + rng.reset() + _, outputs = self._call(runnable, local) + if rng: + rng.check(snapshot(rng.devices), expected_rng) + 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 = compare_outputs(outputs, expected, self.cfg.precision) + return tuple(max(a, b) for a, b in zip(error, current)) + + 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 + + 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 _time(self, runnable: Any, inputs: list[Any]) -> float: + from scripts.kernel_tuning.discovery.rng import preserve + + rng = getattr(self, "_rng", None) + with preserve(rng.devices) if rng else nullcontext(): + return self._time_samples(runnable, inputs) + + def _time_samples(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) + if self._rng: + self._rng.reset() + # 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/docs/proposals/0008-kernel-tuning.md b/docs/proposals/0008-kernel-tuning.md index e7aeab5..5c856da 100644 --- a/docs/proposals/0008-kernel-tuning.md +++ b/docs/proposals/0008-kernel-tuning.md @@ -275,7 +275,7 @@ are not established by CPU/fake-agent tests. Hand-writing a task package per operator does not scale to every core kernel. `scripts/kernel_tuning/operators.py` declares each core Triton operator once: its eager Torch reference, the production call used as the baseline, the -model-derived workload shapes, and the numerical contract. `generate` renders a +model-derived workload shapes, and the numerical contract. `generate --catalog` renders a complete task package from an entry. The baseline is the production kernel itself: the defining repository module and the repository modules it imports are copied byte-for-byte under `vendor/` (absolute `embodiinfer.` imports are @@ -308,7 +308,7 @@ Laptop GPU (Torch 2.6, Triton 3.2; at most 0.85 of its bound) and on an RTX for all ten tasks on both. Re-run both on each target before trusting a contract there. -`tune-all` generates the selected operators into a new batch directory under +`tune-all --catalog` generates the selected operators into a new batch directory under `results/kernel_tuning/batches/`, measures every production baseline with the evaluator (preflight), and then runs each operator serially in its own `run`/`resume` subprocess. A failed operator does not stop the batch; an @@ -320,9 +320,151 @@ settings (`evaluator_python`, `device`, budgets) are passed with `--set` at generation time, and `--hardware-notes` copies target-hardware notes into each task as `HARDWARE.md` for the agent. -### Captured workloads +### Model-scoped operator discovery + +The model workflow selects exactly one canonical policy ID before preparation. +During the existing preparation execution, a scoped ATen observer and wrappers +around existing backend entrypoints collect complete input/output metadata. +Existing compound kernels remain single operators; this change introduces no +new fusion, quantization, runtime registration, or engine API. + +The capture freezes the model type, supplied checkpoint/configuration descriptors, Torch version, +numerical settings, operator overloads, tensor shapes/strides, argument trees, +aliasing and mutation contracts, and every observed workload (no top-eight +truncation or estimated-shape fallback). Quantized calls and their internal +floating-point operations are excluded by invocation scope, including dynamically +created quantized projections and first-forward packing. Metadata-only operations +and host-side work are reported separately from device computation. Discovery +covers the executed preparation paths, not unvisited branches. + +Uncompiled preparation is required; compilation is rejected rather than silently +changing the model's attention or execution path. CUDA Graph metadata is collected +during warmup/capture and calls are counted on replay. Bounded fixtures may be saved +outside graph capture; cases requiring unavailable fixtures are coverage gaps. +Unknown external operators, unrepresentable arguments/state, and unsupported +contracts also remain explicit gaps. Every eligible workload must have a task and +all production baselines must pass preflight before any model batch starts an +agent. No partial-model or skip-preflight escape hatch is provided. + +ATen tasks replay the exact captured overload as their frozen bit-exact reference +and production baseline. Backend tasks snapshot their existing implementation. +The existing Humanize2 search, evaluator, paired promotion rules, archive and +export remain in use. Layout/state-aware task adapters extend evaluation without +weakening the default pure-operator checks. Inputs, current stream, numerical +settings, changed-input graph replay, and declared mutations are validated. + +This is repository tooling under scripts/kernel_tuning and measurement under +benchmarks/kernel_tuning. A static operator whitelist was rejected because it +cannot discover model calls; full-model graph compilation was rejected because +it changes operator boundaries and is outside the no-new-fusion scope. Existing +catalog tasks remain available explicitly as a legacy workflow. Legacy captures +cannot establish model coverage and must be recaptured for model tuning. + +CPU regression coverage must exercise uncatalogued operators, model identity, +all shapes/layouts, quantized scopes (including dynamic modules), unsupported +cases, immutable capture artifacts, and the all-baselines preflight barrier. +GPU/checkpoint validation remains explicit in the matching model environment; +CPU/fake-agent passes do not establish GPU correctness or performance. + +#### Commands and artifacts + +Run preparation in the selected model's own runtime. `PREPARATION.py` below is +an existing driver that constructs the specified policy through `make_policy` +and exercises its forward/decode path, with compilation disabled. The capture +does not launch an extra forward or change the attention backend. Canonical +policy IDs are those accepted by the factory, for example `pi05`, `streamvln`, +or `qwen2.5-vl-3b-r2r-low-level`; catalog group names are not model IDs. -Catalog shapes are estimates. `capture` replaces them with production traffic: +```bash +/absolute/runtime/python -m scripts.kernel_tuning capture --model pi05 --output results/kernel_tuning/captures/pi05 -- PREPARATION.py +uv run --project scripts/kernel_tuning python -m scripts.kernel_tuning generate --model pi05 --captured results/kernel_tuning/captures/pi05 --list +uv run --project scripts/kernel_tuning python -m scripts.kernel_tuning tune-all --model pi05 --captured results/kernel_tuning/captures/pi05 --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 +``` + +`--dry-run` builds and validates all tasks without evaluation. `--preflight-only` +evaluates all baselines without agents; if that batch is later resumed, supply +its agent at creation time. Model preflight failures prevent every agent from +starting. After all baselines pass, an individual search failure is recorded and +other prepared tasks may continue. Resume rechecks capture/task/tool identities +and reruns model preflight. The single-task `run`/`resume` paths also check the +parent model batch. There is no model-mode `--force`, `--estimated`, operator +filter, multiple-model selection, or `--skip-preflight` option. + +Capture artifacts are `manifest.json` (actual policy type, builder configuration, +revision, Torch version, numerical settings, completeness and hashes), +`operators.json` (overload/backend identity, signature, every workload and counts), +`coverage.json` (ready/excluded_quantized/metadata_only/host_only/blocked counts), +and bounded `fixtures/`. A model batch snapshots this directory under `capture/`, +creates a task per call signature under `tasks//`, and retains +`batch.json`, preflight records, run archives, exports and summaries. Each task +includes `replay.json` with storage/layout/mutation contracts and model identity; +its baseline calls the exact ATen overload or snapshotted backend implementation. +Different shapes and strides remain workloads of the same signature; overload, +scalar/static arguments, tensor rank/dtype and output structure distinguish tasks. + +The default 64 MiB fixture budget is explicit (`--fixture-bytes`). Integer/index +storage requires a captured fixture. Floating storage up to 1 MiB is captured +when the remaining budget permits; larger floating storage is generated from +fixed seeds. Seed 0 uses the unmodified saved fixture; other seeds perturb +floating values. First observed fixtures represent each structural case. This +is finite input validation, not proof for all values; domain-specific validity +or data-dependent output shapes may require an adapter. Budget exhaustion for a +required fixture is a coverage gap, never permission to guess indices. + +Replay preserves shared storage, strides and offsets. It checks output values, +layouts, input/output and output/output aliases, declared mutations and unchanged +input storage. Timing resets backing storage before every sample, outside timed +events, in both eager and graph modes. The default is eager; `timing.mode=cuda_graph` +also requires successful changed-input A/B/A replay. NCU profiling currently +requires a dedicated replay-aware adapter. Numerical overrides cannot weaken +the captured bit-exact production contract. + +Default-generator `aten.randn.default` and `aten.randn.generator` with +`generator=None` have an eager replay contract. Capture stores the opaque CPU +and target CUDA generator states immediately before and after the existing call, +within the same required-fixture budget. Seed 0 replays the actual captured +position; other validation seeds use independent default-generator states. +Validation compares output bytes and the complete post-call generator states, +rejecting extra draws or rewinds even when output values match. State resets +happen before each warmup/sample, outside timing events, and evaluation restores +the caller's Torch RNG states on success or failure. Explicit generator objects, +other random overloads, and RNG tasks using CUDA Graph timing remain gaps. + +SDPA's dispatcher RNG tag is conditional: `dropout_p=0` (including its schema +default) does not require an RNG replay contract. Discovery resolves positional, +keyword and omitted dropout arguments without changing the invocation, backend, +mask, scale or attention semantics. Nonzero dropout remains a coverage gap until +its own state contract is implemented; other nondeterministic tags still apply. + +Replay tasks declare `definition_schema=embodiinfer-replay-v1`. Their local +`ReplayDefinition` extends the pinned FlashInfer tensor schema with `float64`, +retains upstream structural validation, and uses the same builders with native +double-precision tensors. Captured fixtures, input materialization and output +checks preserve float64 without a cast or relaxed tolerance. Ordinary FlashInfer +definitions and their accepted dtype set are unchanged. + +Existing RoPE cache writes and greedy-sampling workspace mutations have explicit +contracts; graph GQA and prefill entries remain individual existing operators. +`GreedyWorkspace` has a typed tensor-state recipe; arbitrary Python objects still +require an adapter. Dtypes unsupported by the replay schema remain task +construction gaps rather than silently converted inputs. + +Known quantization scopes are `QuantizedLinear` (FP8/INT8/NVFP4), conversion and +packing helpers, their backend functions, and explicit native quantization +dispatcher operators. Normal integer indices are not a quantization signal. +Unknown dispatcher extensions and raw Triton launches are reported as gaps; +random operators outside the default-generator randn contract, opaque Python +arguments, and calls without tensor outputs also require an adapter. Custom +C++/CUDA calls that bypass both the dispatcher and +these entrypoints need explicit instrumentation before claiming coverage. +Capture is a single-process preparation facility; it does not certify unvisited +branches, other worker processes, or checkpoint-level action parity. Existing +fused kernels count as one call; no new fused operators are proposed or generated. + +### Captured catalog workloads + +Catalog shapes are estimates. `capture --catalog` replaces them with production traffic: it runs an unmodified model script in the model's runtime with every catalog kernel wrapped, converts each call to the operator's task axes, scalars, and fixed integer inputs, and writes one JSONL row per distinct case with its call @@ -369,6 +511,6 @@ single-row decode calls: `swiglu` never runs on prefill chunks, and all three StreamVLN kernels sit at the launch floor. ```bash -PYTHONPATH=/path/to/EmbodiInfer /path/to/model-runtime/python -m scripts.kernel_tuning capture --output pi05.jsonl -- benchmark.py --config capture-config.yaml -python -m scripts.kernel_tuning tune-all --model pi05 --captured pi05.jsonl --agent ... +PYTHONPATH=/path/to/EmbodiInfer /path/to/model-runtime/python -m scripts.kernel_tuning capture --catalog --output pi05.jsonl -- benchmark.py --config capture-config.yaml +python -m scripts.kernel_tuning tune-all --catalog --model pi05 --captured pi05.jsonl --agent 'HARNESS/MODEL:EFFORT' ``` diff --git a/scripts/kernel_tuning/__main__.py b/scripts/kernel_tuning/__main__.py index 2417e7d..03ff765 100644 --- a/scripts/kernel_tuning/__main__.py +++ b/scripts/kernel_tuning/__main__.py @@ -12,10 +12,15 @@ from .artifacts import REPOSITORY, RunStore, atomic_json, run_lock, runtime_identity 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" @@ -75,7 +80,10 @@ def _execute(store: RunStore, *, resume: bool) -> dict: def _selection(parser: argparse.ArgumentParser) -> None: parser.add_argument("operators", nargs="*", help="Catalog operator names (default: all)") - parser.add_argument("--model", action="append", help="Select operators used by a model, e.g. pi05") + 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", @@ -88,16 +96,52 @@ def _selection(parser: argparse.ArgumentParser) -> None: shapes.add_argument( "--captured", type=Path, - help="Shapes from this `capture` file or directory (default: benchmarks/kernel_tuning/captures)", + 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 --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 ) @@ -136,6 +180,7 @@ def _tune_all(args: argparse.Namespace) -> int: or args.agent or args.captured or args.estimated + or args.catalog ): raise ContractError("--resume continues the recorded batch; omit selection, --set, and --agent") root = args.resume.resolve(strict=True) @@ -143,18 +188,22 @@ def _tune_all(args: argparse.Namespace) -> int: 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") - 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, - agent=agent, - overrides=parse_overrides(args.set), - hardware_notes=args.hardware_notes, - ) + options = { + "agent": agent, + "overrides": parse_overrides(args.set), + "hardware_notes": args.hardware_notes, + } + 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) @@ -200,13 +249,13 @@ def main(argv: list[str] | None = None) -> int: 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 task packages for catalog operators; no execution") + generate = commands.add_parser("generate", help="Write every discovered model task; 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 catalog operators and exit") + generate.add_argument("--list", action="store_true", help="List discovered operators and coverage gaps") 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") @@ -220,7 +269,14 @@ def main(argv: list[str] | None = None) -> int: "capture", help="Record production kernel shapes while a model script runs (model runtime)" ) capture.add_argument("operators", nargs="*", help="Catalog operators to record (default: all)") - capture.add_argument("--output", type=Path, required=True, help="JSONL file of captured cases") + capture.add_argument("--catalog", action="store_true", help="Record legacy catalog-only JSONL") + capture.add_argument("--model", help="Canonical policy name constructed by the preparation script") + capture.add_argument( + "--fixture-bytes", type=int, default=64 * 1024 * 1024, help="Maximum stored input fixture bytes" + ) + capture.add_argument( + "--output", type=Path, required=True, help="New model capture directory (JSONL for --catalog)" + ) capture.add_argument("script", nargs=argparse.REMAINDER, help="-- SCRIPT [ARGS] or -- -m MODULE [ARGS]") skeleton = commands.add_parser( "skeleton", help="Random-weight copy of a Hugging Face checkpoint for shape capture (model runtime)" @@ -237,15 +293,31 @@ def main(argv: list[str] | None = None) -> int: args = parser.parse_args(raw) try: if args.command == "capture": - from .capture import run as run_capture - from .operators import select + if args.catalog: + 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']}") + if args.model: + raise ContractError("Legacy catalog capture does not bind a model; omit --model") + 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']}") + else: + from .discovery.recorder import run as run_model_capture + + if not args.model or args.operators: + raise ContractError( + "Capture requires one --model and no operator filter; use --catalog for legacy capture" + ) + result = run_model_capture( + command or args.script, args.output, args.model, fixture_bytes=args.fixture_bytes + ) + print(json.dumps(result, indent=2, ensure_ascii=False)) elif args.command == "skeleton": from .skeleton import build @@ -281,7 +353,17 @@ def main(argv: list[str] | None = None) -> int: 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) + 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: diff --git a/scripts/kernel_tuning/batch.py b/scripts/kernel_tuning/batch.py index b5ef8e4..7037df8 100644 --- a/scripts/kernel_tuning/batch.py +++ b/scripts/kernel_tuning/batch.py @@ -10,15 +10,18 @@ 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 -from .contracts import ContractError, TaskPackage, read_json, validate_measurement +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 @@ -86,16 +89,43 @@ def preflight(root: Path, names: list[str] | None = None) -> dict[str, str | Non 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 discovered 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 item["status"] != "pending": + 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"] = { @@ -106,11 +136,142 @@ def preflight(root: Path, names: list[str] | None = None) -> dict[str, str | Non } except ContractError as exc: results[name] = str(exc) - item.update(status="preflight_failed", error=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 +) -> 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, + "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 discovered 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 @@ -120,9 +281,11 @@ def _state(root: Path, item: dict[str, Any]) -> dict[str, Any] | None: 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"]: @@ -208,6 +371,13 @@ def write_summary(root: Path) -> list[dict[str, Any]]: "| 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 "-" diff --git a/scripts/kernel_tuning/contracts.py b/scripts/kernel_tuning/contracts.py index 5f03c05..11247df 100644 --- a/scripts/kernel_tuning/contracts.py +++ b/scripts/kernel_tuning/contracts.py @@ -306,7 +306,9 @@ def _load(cls, root: Path) -> TaskPackage: 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]: + 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") @@ -315,7 +317,9 @@ def _load(cls, root: Path) -> TaskPackage: ) 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: + 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")): @@ -364,6 +368,12 @@ def _load(cls, root: Path) -> TaskPackage: 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") return cls(root, definition, tuple(traces), settings, hashes) def verify(self) -> None: @@ -415,6 +425,8 @@ def validate_solution(solution: Any, task: TaskPackage) -> dict[str, Any]: 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("::") diff --git a/scripts/kernel_tuning/discovery/__init__.py b/scripts/kernel_tuning/discovery/__init__.py new file mode 100644 index 0000000..2176a78 --- /dev/null +++ b/scripts/kernel_tuning/discovery/__init__.py @@ -0,0 +1,5 @@ +"""Model-scoped task discovery; Torch is imported only in the capture runtime.""" + +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..cf5fdac --- /dev/null +++ b/scripts/kernel_tuning/discovery/contracts.py @@ -0,0 +1,180 @@ +"""Portable, immutable model captures, independent of Torch and coding agents.""" + +from __future__ import annotations + +import re +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", "excluded_quantized", "metadata_only", "host_only", "blocked"} +RNG_OVERLOADS = {"aten.randn.default", "aten.randn.generator"} +REPLAY_SCHEMA = "embodiinfer-replay-v1" + + +def validate_rng(contract: dict[str, Any] | None, case: dict[str, Any], files: dict[str, str]) -> None: + """Bind both RNG-state transitions to immutable files for the declared devices.""" + if contract is None: + if "rng" in case: + raise ContractError("Undeclared RNG state in workload") + return + if not isinstance(contract, dict): + raise ContractError("Invalid default-generator RNG contract") + devices = contract.get("devices") + if ( + set(contract) != {"kind", "devices"} + or contract["kind"] != "torch_default" + or not isinstance(devices, list) + or not 1 <= len(devices) <= 2 + or devices[0] != "cpu" + or (len(devices) == 2 and not re.fullmatch(r"cuda:\d+", str(devices[1]))) + ): + raise ContractError("Invalid default-generator RNG contract") + states = case.get("rng", {}) + if not isinstance(states, dict) or set(states) != {"before", "after"}: + raise ContractError("RNG replay requires before and after states") + for phase in states.values(): + if not isinstance(phase, dict) or set(phase) != set(devices): + raise ContractError("RNG state devices differ from the contract") + if any(relative_path(path) not in files for path in phase.values()): + raise ContractError("RNG fixture missing from the capture/task") + + +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("rng") is not None and replay.get("operator") not in RNG_OVERLOADS: + raise ContractError("Unsupported RNG replay operator") + 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"]] + validate_rng(replay.get("rng"), case, hashes) + 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 completed preparation run bound to one actual policy and all observed cases.""" + + 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") + 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") + 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 op["status"] == "ready": + rng = op["identity"].get("rng") + if rng is not None and op["name"] not in RNG_OVERLOADS: + raise ContractError("Unsupported RNG capture operator") + validate_rng(rng, case, hashes) + capture = cls(root, manifest, tuple(operators)) + if require_ready: + capture.require_ready() + return capture + + def require_ready(self) -> None: + """Reject incomplete runs and any non-quantized operator preparation 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/recorder.py b/scripts/kernel_tuning/discovery/recorder.py new file mode 100644 index 0000000..e965e15 --- /dev/null +++ b/scripts/kernel_tuning/discovery/recorder.py @@ -0,0 +1,654 @@ +"""Observe one model's existing execution, without replacing its numerical path.""" + +from __future__ import annotations + +import functools +import hashlib +import importlib +import inspect +import runpy +import sys +from collections import Counter +from collections.abc import Iterator +from contextlib import contextmanager +from dataclasses import fields, is_dataclass +from pathlib import Path +from typing import Any + +import torch +from torch.utils._python_dispatch import TorchDispatchMode + +from ..artifacts import REPOSITORY, atomic_json +from ..contracts import ContractError, digest, tree_hashes +from ..generate import _git_revision, vendor_sources +from .contracts import SCHEMA +from .rng import contract_for, requires_rng_adapter +from .rng import snapshot as rng_snapshot +from .tensors import describe_inputs, extent, flatten, storage_key, tensor_outputs, tensor_spec + +_VIEWS = { + "alias", + "detach", + "view", + "_unsafe_view", + "transpose", + "permute", + "as_strided", + "expand", + "squeeze", + "unsqueeze", + "slice", + "select", + "split", + "split_with_sizes", + "unbind", + "t", + "as_strided_", + "transpose_", + "squeeze_", + "unsqueeze_", + "detach_", + "resize_", + "resize_as_", + "set_", +} +_ALLOCATIONS = {"empty", "empty_like", "empty_strided", "new_empty", "new_empty_strided"} +_HOST = {"_local_scalar_dense", "is_nonzero"} +_QUANTIZED_ATEN = { + "_scaled_mm", + "_weight_int8pack_mm", + "_weight_int4pack_mm", + "_convert_weight_to_int4pack", + "quantize_per_tensor", + "quantize_per_channel", + "dequantize", +} +# These contracts describe existing public kernels, not an operator discovery whitelist. +_PURE_KERNELS = { + "rms_norm", + "add_rms_norm", + "ada_rms_norm", + "gated_residual", + "swiglu", + "gated_gelu", + "rotate_qk", + "rotate_half_rope", + "segmented_attention", + "split_kv_attention", + "flash_prefill_attention", + "flash_prefill_attention_graph", + "gqa_decode_graph", +} +_KERNEL_WRITES = { + "fused_rope_cache": ("key_cache", "value_cache"), + "fused_rope_cache_graph": ("key_cache", "value_cache"), + "fused_prefill_rope_cache_graph": ("key_cache", "value_cache"), + "fused_greedy": ("workspace",), + "fused_lm_head_greedy": ("workspace",), + "mark_penalized": ("workspace",), +} + + +def _portable(value: Any) -> Any: + if value is None or isinstance(value, (str, int, float, bool)): + return value + if isinstance(value, (torch.dtype, torch.device, Path)): + return str(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] + if is_dataclass(value) and not isinstance(value, type): + return {field.name: _portable(getattr(value, field.name)) for field in fields(value)} + return {"type": f"{type(value).__module__}.{type(value).__qualname__}"} + + +class _Observer(TorchDispatchMode): + def __init__(self, session: CaptureSession): + super().__init__() + self.session = session + + def __torch_dispatch__(self, func: Any, types: Any, args: tuple = (), kwargs: dict | None = None) -> Any: + session = self.session + kwargs = kwargs or {} + if not session.depth or session.suppressed: + return func(*args, **kwargs) + name = str(func) + short = name.split(".")[1] if name.startswith("aten.") else name + reason = None + with session.silence(): + rng = contract_for(name, kwargs) + if not name.startswith("aten."): + reason = "External dispatcher operator needs a reference adapter" + if rng is None and requires_rng_adapter(func, args, kwargs): + reason = "Random/nondeterministic operator needs an RNG replay adapter" + writes = [] + for index, argument in enumerate(func._schema.arguments): + if argument.alias_info is not None and argument.alias_info.is_write: + value = args[index] if index < len(args) else kwargs.get(argument.name) + writes.append(value) + quantized = name.split(".")[0] in {"quantized", "quantized_decomposed"} or short in _QUANTIZED_ATEN + return session.call( + {"kind": "aten", "name": name, "schema": str(func._schema)}, + func, + args, + kwargs, + writes=writes, + reason=reason, + rng=rng if rng is not None and rng["devices"][-1].split(":")[0] == session.device_type else None, + classification="excluded_quantized" + if quantized + else "metadata_only" + if short in _VIEWS | _ALLOCATIONS + else "host_only" + if short in _HOST + else None, + ) + + +class CaptureSession: + """Collect one bound model's calls; failures remain explicit coverage gaps. + + Use ``observe`` around existing preparation execution for custom drivers. The + CLI additionally binds a target constructed by the public policy factory. + ``device_type='cpu'`` is useful for deterministic observer unit tests only. + """ + + def __init__( + self, model: str, output: Path, *, fixture_bytes: int = 64 * 1024 * 1024, device_type: str = "cuda" + ) -> None: + if not model or fixture_bytes < 0: + raise ContractError("Select one model and a nonnegative fixture byte budget") + self.model, self.root = model, output.resolve() + if self.root.exists() and any(self.root.iterdir()): + raise ContractError("Capture destination must be empty; use a new directory") + self.root.mkdir(parents=True, exist_ok=True) + self.fixture_bytes, self.used_bytes, self.device_type = fixture_bytes, 0, device_type + self.target: dict[str, Any] = {} + self.instance: Any = None + self.depth, self.suppressed = 0, 0 + self.operators: dict[str, dict[str, Any]] = {} + self.errors: list[str] = [] + self.replacements: list[tuple[Any, str, Any, bool]] = [] + self._fixtures: dict[str, list[dict[str, Any]]] = {} + self.graphs: dict[int, Counter] = {} + self.graph_stack: list[Counter] = [] + self._numerics_bound = False + self._set_numerics() + + def _set_numerics(self) -> None: + self.precision = { + "mode": "bit_exact", + "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(), + } + + def _check_numerics(self) -> None: + previous = (self.precision, self.flags) + self._set_numerics() + if self._numerics_bound and previous != (self.precision, self.flags): + self.errors.append("Numerical settings changed during model preparation") + self._numerics_bound = True + + @contextmanager + def silence(self) -> Iterator[None]: + """Prevent instrumentation and atomic/quantized internals from becoming tasks.""" + self.suppressed += 1 + try: + yield + finally: + self.suppressed -= 1 + + def bind(self, instance: Any, config: dict[str, Any]) -> None: + """Bind to an actual object rather than trusting an unverified model label.""" + if self.instance is not None and self.instance is not instance: + raise ContractError("Prepare exactly one target policy instance per capture") + modules = instance.modules() if isinstance(instance, torch.nn.Module) else (instance,) + if config.get("compile_backend", "none") != "none" or any( + hasattr(module, "_orig_mod") for module in modules + ): + raise ContractError("Model discovery requires explicit uncompiled preparation") + self.instance = instance + self.target = { + "model": self.model, + "type": f"{type(instance).__module__}.{type(instance).__qualname__}", + "config": _portable(config), + } + + @contextmanager + def observe(self, instance: Any, **config: Any) -> Iterator[None]: + """Observe the existing preparation call without running it a second time.""" + self.bind(instance, config) + self.depth += 1 + try: + with _Observer(self): + yield + finally: + self.depth -= 1 + + def _snapshot( + self, key: str, tensors: list[torch.Tensor], specs: list[dict[str, Any]] + ) -> list[dict[str, Any]]: + if key in self._fixtures: + return self._fixtures[key] + if torch.cuda.is_available() and torch.cuda.is_current_stream_capturing(): + if any(not tensor.is_floating_point() for tensor in tensors): + raise ContractError("Valid non-floating fixtures must be observed during eager warmup") + return 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) + byte_count = size * tensor.element_size() + required = not tensor.is_floating_point() + # Large ordinary floating weights are generated; retain small real cases + # when budget permits. Integer/index fixtures are never guessed. + save = required or byte_count <= 1024 * 1024 + if save and self.used_bytes + byte_count > self.fixture_bytes: + if required: + raise ContractError("Required input fixture exceeds the capture byte budget") + save = False + if save: + raw = ( + tensor.detach() + .as_strided((size,), (1,), 0) + .contiguous() + .view(torch.uint8) + .cpu() + .numpy() + .tobytes() + ) + path = f"fixtures/{hashlib.sha256(raw).hexdigest()}.bin" + file = self.root / path + if not file.exists(): + file.parent.mkdir(exist_ok=True) + file.write_bytes(raw) + self.used_bytes += len(raw) + for i in indices: + specs[i]["fixture"] = path + self._fixtures[key] = specs + return specs + + def _rng_snapshot(self, devices: list[str]) -> dict[str, str]: + if torch.cuda.is_available() and torch.cuda.is_current_stream_capturing(): + raise ContractError("RNG fixtures require eager preparation outside CUDA Graph capture") + result = {} + for device, state in rng_snapshot(devices).items(): + raw = bytes(state.tolist()) + path = f"fixtures/{hashlib.sha256(raw).hexdigest()}.bin" + file = self.root / path + if not file.exists(): + if self.used_bytes + len(raw) > self.fixture_bytes: + raise ContractError("Required RNG fixture exceeds the capture byte budget") + file.parent.mkdir(exist_ok=True) + file.write_bytes(raw) + self.used_bytes += len(raw) + result[device] = path + return result + + def _add(self, identity: dict[str, Any], case: dict[str, Any], status: str, reason: str | None) -> None: + op_id = digest(identity)[:20] + op = self.operators.setdefault( + op_id, + { + "id": op_id, + "name": identity["name"], + "identity": identity, + "status": status, + "reason": reason, + "workloads": [], + }, + ) + if status == "blocked": + op.update(status=status, reason=reason) + case_id = digest({k: v for k, v in case.items() if k != "count"})[:20] + existing = next((w for w in op["workloads"] if w["id"] == case_id), None) + if existing is None: + existing = {**case, "id": case_id, "count": 0} + op["workloads"].append(existing) + if self.graph_stack: + self.graph_stack[-1][op_id, case_id] += 1 + # Retain capture-only cases too; replay traffic is counted separately. + existing["count"] += 1 + + def call( + self, + origin: dict[str, Any], + function: Any, + args: tuple, + kwargs: dict, + *, + writes: list[Any] | None = None, + reason: str | None = None, + rng: dict[str, Any] | None = None, + classification: str | None = None, + ) -> Any: + """Record before/after metadata while invoking the original callable exactly once.""" + with self.silence(): + self._check_numerics() + tensors: list[torch.Tensor] = [] + identity = dict(origin) + case: dict[str, Any] = {"inputs": [], "outputs": []} + status = classification or "ready" + try: + recipe = flatten((args, kwargs), tensors) + if any(t.is_quantized for t in tensors): + classification = status = "excluded_quantized" + specs = describe_inputs(tensors) + identity.update( + arguments=recipe, + input_dtypes=[s["dtype"] for s in specs], + input_ranks=[len(s["shape"]) for s in specs], + ) + case["inputs"] = specs + if not any(t.device.type == self.device_type for t in tensors): + status = classification or "host_only" + mutated: list[torch.Tensor] = [] + for value in writes or []: + flatten(value, mutated) + identity["mutates"] = [ + i for i, tensor in enumerate(tensors) if any(tensor is other for other in mutated) + ] + if reason: + raise ContractError(reason) + if status == "ready": + case["inputs"] = self._snapshot(digest([identity, specs]), tensors, specs) + if rng is not None: + identity["rng"] = rng + case["rng"] = {"before": self._rng_snapshot(rng["devices"])} + except (ContractError, RuntimeError, TypeError, ValueError) as exc: + status, reason = ( + ("excluded_quantized" if classification == "excluded_quantized" else "blocked"), + str(exc), + ) + try: + result = function(*args, **kwargs) + except Exception as exc: + self.errors.append(f"{origin['name']} failed during preparation: {type(exc).__name__}: {exc}") + raise + try: + if rng is not None and status != "blocked": + case["rng"]["after"] = self._rng_snapshot(rng["devices"]) + if ( + status == "host_only" + and classification is None + and any(t.device.type == self.device_type for t in tensor_outputs(result)) + ): + status = "ready" + if status not in {"metadata_only", "host_only", "excluded_quantized", "blocked"}: + outputs = tensor_outputs(result) + identity["returns"] = flatten(result, []) + # Mutated tensors are part of the task's observable outputs too. + outputs += [tensors[i] for i in identity["mutates"]] + if not outputs: + raise ContractError("Calls without tensor outputs require a dedicated task adapter") + case["outputs"] = [tensor_spec(t) for t in outputs] + identity["output_dtypes"] = [s["dtype"] for s in case["outputs"]] + identity["output_ranks"] = [len(s["shape"]) for s in case["outputs"]] + case["output_aliases"] = [ + [i for i, tensor in enumerate(tensors) if storage_key(tensor) == storage_key(output)] + for output in outputs + ] + case["output_groups"] = [ + [i for i, other in enumerate(outputs) if storage_key(output) == storage_key(other)] + for output in outputs + ] + else: + case["outputs"] = [tensor_spec(t) for t in tensor_outputs(result)] + except (ContractError, RuntimeError, TypeError, ValueError) as exc: + if rng is not None or status not in {"metadata_only", "host_only", "excluded_quantized"}: + status, reason = "blocked", str(exc) + self._add(identity, case, status, reason) + return result + + def _replace(self, holder: Any, name: str, replacement: Any) -> None: + self.replacements.append((holder, name, getattr(holder, name), name in vars(holder))) + setattr(holder, name, replacement) + + def _replace_function(self, original: Any, wrapper: Any) -> None: + for module in list(sys.modules.values()): + if module is None: + continue + for name, value in list(vars(module).items()): + if value is original: + self._replace(module, name, wrapper) + + def _atomic(self, function: Any, *, quantized: bool = False) -> Any: + source_digests: dict[str, str] | None = None + + @functools.wraps(function) + def wrapper(*args: Any, **kwargs: Any) -> Any: + nonlocal source_digests + if not self.depth or self.suppressed: + return function(*args, **kwargs) + origin = { + "kind": "backend", + "name": f"{function.__module__}.{function.__name__}", + "module": function.__module__, + "entry": function.__name__, + } + if quantized: + # A module instance is provenance, not a replayable tensor input. + call_args = args[1:] if args and isinstance(args[0], torch.nn.Module) else args + return self.call( + origin, + lambda *a, **kw: function(*args, **kwargs), + call_args, + kwargs, + classification="excluded_quantized", + ) + bound = inspect.signature(function).bind(*args, **kwargs) + bound.apply_defaults() + source = Path(inspect.getfile(function)) + origin["source_digest"] = hashlib.sha256(source.read_text(encoding="utf-8").encode()).hexdigest() + names = _KERNEL_WRITES.get(function.__name__, ()) + reason = ( + None + if function.__name__ in _PURE_KERNELS or names + else "Backend entry requires an explicit mutation/reference contract" + ) + try: + if source_digests is None: + source_digests = { + path: source_digest + for path, (_, source_digest) in vendor_sources(function.__module__).items() + } + origin["source_digests"] = source_digests + except ContractError as exc: + reason = str(exc) + return self.call( + origin, function, args, kwargs, writes=[bound.arguments[n] for n in names], reason=reason + ) + + return wrapper + + def install(self) -> None: + """Install process-local quantization, backend, factory, and graph observation hooks.""" + from embodiinfer.models import linear + from embodiinfer.policies import factory + + if self.model not in factory.available_policies(): + raise ContractError(f"Unknown policy {self.model!r}; choose from {factory.available_policies()}") + self._replace( + linear.QuantizedLinear, "forward", self._atomic(linear.QuantizedLinear.forward, quantized=True) + ) + for name in ("quantize_linear", "pack_quantized_linears"): + original = getattr(linear, name) + self._replace_function(original, self._atomic(original, quantized=True)) + for method in ("fp8", "int8", "nvfp4"): + module = importlib.import_module(f"embodiinfer.backend.torch.{method}") + for name, function in list(vars(module).items()): + if ( + inspect.isfunction(function) + and function.__module__ == module.__name__ + and not name.startswith("_") + ): + self._replace_function(function, self._atomic(function, quantized=True)) + for path in sorted((REPOSITORY / "embodiinfer/backend/triton").glob("*.py")): + if path.stem in {"__init__", "capability"}: + continue + module = importlib.import_module(f"embodiinfer.backend.triton.{path.stem}") + for name, function in list(vars(module).items()): + if ( + inspect.isfunction(function) + and function.__module__ == module.__name__ + and not name.startswith(("_", "supports_", "warmup_")) + and not name.endswith(("_available", "_capability")) + ): + self._replace_function(function, self._atomic(function, quantized=path.stem == "fp8")) + original_factory = factory.make_policy + + @functools.wraps(original_factory) + def construct(name: str, **kwargs: Any) -> Any: + if name != self.model: + raise ContractError(f"Preparation constructed {name!r}, expected {self.model!r}") + if kwargs.get("compile_backend", "none") != "none": + raise ContractError("Disable torch.compile explicitly before model discovery") + policy = original_factory(name, **kwargs) + self.bind(policy, kwargs) + for holder, methods in ( + (policy, ("encode_prefix", "denoise_step", "prepare_prefix", "forward")), + (policy.decoder, ("produce_chunk", "generate_tokens", "init_state")), + ): + for method in methods: + if not hasattr(holder, method): + continue + original = getattr(holder, method) + + def execute(*args: Any, _fn: Any = original, **kw: Any) -> Any: + self.depth += 1 + try: + return _fn(*args, **kw) + finally: + self.depth -= 1 + + self._replace(holder, method, execute) + return policy + + self._replace_function(original_factory, construct) + + def compile_guard(*args: Any, **kwargs: Any) -> Any: + raise ContractError("torch.compile is unsupported during model preparation; disable compilation") + + self._replace_function(torch.compile, compile_guard) + try: + from triton.runtime import JITFunction + except ImportError: + pass + else: + original_run = JITFunction.run + + def triton_run(kernel: Any, *args: Any, **kwargs: Any) -> Any: + if not self.depth or self.suppressed: + return original_run(kernel, *args, **kwargs) + origin = {"kind": "unregistered_triton", "name": f"{kernel.__module__}.{kernel.__name__}"} + return self.call( + origin, + lambda *a, **kw: original_run(kernel, *a, **kw), + args, + kwargs, + reason="Raw Triton entry needs a public callable and replay contract", + ) + + self._replace(JITFunction, "run", triton_run) + graph = torch.cuda.CUDAGraph + begin, end, replay = graph.capture_begin, graph.capture_end, graph.replay + + def capture_begin(instance: Any, *args: Any, **kwargs: Any) -> Any: + result = begin(instance, *args, **kwargs) + self.graph_stack.append(Counter()) + return result + + def capture_end(instance: Any, *args: Any, **kwargs: Any) -> Any: + result = end(instance, *args, **kwargs) + self.graphs[id(instance)] = self.graph_stack.pop() + return result + + def graph_replay(instance: Any, *args: Any, **kwargs: Any) -> Any: + if self.depth and not self.suppressed and id(instance) not in self.graphs: + self.errors.append("Model replayed a CUDA graph not observed during this preparation") + result = replay(instance, *args, **kwargs) + for (op_id, case_id), count in self.graphs.get(id(instance), {}).items(): + case = next(w for w in self.operators[op_id]["workloads"] if w["id"] == case_id) + case["count"] += count + return result + + for name, fn in ( + ("capture_begin", capture_begin), + ("capture_end", capture_end), + ("replay", graph_replay), + ): + self._replace(graph, name, fn) + + def close(self, *, complete: bool) -> None: + """Restore all hooks and publish a hashed success or failure capture.""" + for holder, name, original, owned in reversed(self.replacements): + if owned: + setattr(holder, name, original) + else: + delattr(holder, name) + self.replacements.clear() + if not self.target: + self.errors.append("No target policy was bound through make_policy or observe") + operators = list(self.operators.values()) + atomic_json(self.root / "operators.json", operators) + coverage = dict(Counter(op["status"] for op in operators)) + atomic_json(self.root / "coverage.json", {"operators": coverage, "errors": self.errors}) + atomic_json( + self.root / "manifest.json", + { + "schema": SCHEMA, + "target": self.target, + "complete": complete, + "errors": self.errors, + "repository_revision": _git_revision(), + "torch_version": str(torch.__version__), + "precision": self.precision, + "flags": self.flags, + "device_type": self.device_type, + "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, *, fixture_bytes: int = 64 * 1024 * 1024 +) -> dict[str, Any]: + """Execute an existing model preparation script with model-scoped observation.""" + if not command: + raise ContractError("Supply the preparation script after --") + session = CaptureSession(model, output, fixture_bytes=fixture_bytes) + argv, path = sys.argv, sys.path[:] + complete = False + try: + session.install() + with _Observer(session): + 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 + session.close(complete=complete) + from .contracts import ModelCapture + + capture = ModelCapture.load(output, model) + return {"model": model, "identity": capture.identity, "operators": len(capture.operators)} diff --git a/scripts/kernel_tuning/discovery/rng.py b/scripts/kernel_tuning/discovery/rng.py new file mode 100644 index 0000000..a82289d --- /dev/null +++ b/scripts/kernel_tuning/discovery/rng.py @@ -0,0 +1,114 @@ +"""Argument-aware RNG classification and exact default-generator randn replay.""" + +from __future__ import annotations + +from collections.abc import Iterator +from contextlib import contextmanager +from pathlib import Path +from typing import Any + +import torch + +from ..contracts import ContractError +from .contracts import RNG_OVERLOADS + + +def requires_rng_adapter(function: Any, args: tuple[Any, ...], kwargs: dict[str, Any]) -> bool: + """Treat seeded tags as conditional for SDPA with explicitly zero/default dropout.""" + tags = {tag for tag in getattr(function, "tags", ()) if "nondeterministic" in str(tag)} + if str(function) == "aten.scaled_dot_product_attention.default": + for index, argument in enumerate(function._schema.arguments): + if argument.name == "dropout_p": + dropout = kwargs.get( + argument.name, args[index] if index < len(args) else argument.default_value + ) + if type(dropout) in (int, float) and dropout == 0: + tags.discard(torch.Tag.nondeterministic_seeded) + break + return bool(tags) + + +def contract_for(name: str, kwargs: dict[str, Any]) -> dict[str, Any] | None: + """Recognize randn with the default CPU/CUDA generator; other RNG calls stay gaps.""" + if name not in RNG_OVERLOADS or kwargs.get("generator") is not None: + return None + device = torch.device(kwargs.get("device") or torch.get_default_device()) + if device.type not in {"cpu", "cuda"}: + return None + devices = ["cpu"] + if device.type == "cuda": + index = torch.cuda.current_device() if device.index is None else device.index + devices.append(f"cuda:{index}") + return {"kind": "torch_default", "devices": devices} + + +def snapshot(devices: list[str]) -> dict[str, torch.Tensor]: + """Read opaque generator states without consuming random numbers.""" + return { + device: (torch.get_rng_state() if device == "cpu" else torch.cuda.get_rng_state(device)).clone() + for device in devices + } + + +def restore(states: dict[str, torch.Tensor]) -> None: + """Restore every declared default generator from its opaque byte state.""" + for device, state in states.items(): + if device == "cpu": + torch.set_rng_state(state) + else: + torch.cuda.set_rng_state(state, device) + + +@contextmanager +def preserve(devices: list[str]) -> Iterator[None]: + """Restore the caller's RNG state on both successful and failed evaluation.""" + saved = snapshot(devices) + try: + yield + finally: + restore(saved) + + +class RngReplay: + """Replay the captured transition at seed 0 and independent states at other seeds.""" + + def __init__(self, contract: dict[str, Any], case: dict[str, Any], root: Path, seed: int) -> None: + self.devices = contract["devices"] + states = { + phase: { + device: torch.tensor(list((root / path).read_bytes()), dtype=torch.uint8) + for device, path in paths.items() + } + for phase, paths in case["rng"].items() + } + current = snapshot(self.devices) + if any( + state.shape != current[device].shape + for phase in states.values() + for device, state in phase.items() + ): + raise ContractError("RNG fixture state size differs from the current Torch runtime") + self.before = ( + states["before"] + if seed == 0 + else { + device: torch.Generator(device=device).manual_seed(seed).get_state() + for device in self.devices + } + ) + self.captured_after = states["after"] if seed == 0 else None + + def reset(self) -> None: + """Restore the selected workload state before a call, outside timed events.""" + restore(self.before) + + def check_reference(self, actual: dict[str, torch.Tensor]) -> None: + """Require the frozen overload to reproduce its observed seed-0 transition.""" + if self.captured_after is not None: + self.check(actual, self.captured_after) + + @staticmethod + def check(actual: dict[str, torch.Tensor], expected: dict[str, torch.Tensor]) -> None: + """Reject changed generator advancement even when output values are correct.""" + if set(actual) != set(expected) or any(not torch.equal(actual[d], expected[d]) for d in expected): + raise ContractError("RNG state after the call differs from the reference") diff --git a/scripts/kernel_tuning/discovery/tasks.py b/scripts/kernel_tuning/discovery/tasks.py new file mode 100644 index 0000000..51f26a2 --- /dev/null +++ b/scripts/kernel_tuning/discovery/tasks.py @@ -0,0 +1,300 @@ +"""Static task construction from verified model captures, with no Torch imports.""" + +from __future__ import annotations + +import hashlib +import json +import re +from pathlib import Path +from typing import Any + +from ..artifacts import REPOSITORY, atomic_json +from ..contracts import ContractError, TaskPackage +from ..generate import _dump_yaml, tuning_settings, vendor_sources +from ..operators import Operator, Workload +from .contracts import REPLAY_SCHEMA, ModelCapture + +# 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()) + "}" + if kind == "greedy_workspace": + return "GreedyWorkspace(**" + _expression({"dict": value}) + ")" + 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"] == "aten": + if not re.fullmatch(r"aten\.\w+\.\w+", identity["name"]): + raise ContractError("Invalid captured ATen overload") + imports, callable_name = "import torch\n", f"torch.ops.{identity['name']}" + elif identity["kind"] == "backend": + module, entry = identity["module"], identity["entry"] + if not re.fullmatch(r"embodiinfer\.backend\.\w+(?:\.\w+)*", module) or not entry.isidentifier(): + raise ContractError("Unsupported backend source identity") + source = REPOSITORY.joinpath(*module.split(".")).with_suffix(".py") + if ( + hashlib.sha256(source.read_text(encoding="utf-8").encode()).hexdigest() + != identity["source_digest"] + ): + raise ContractError(f"Backend source changed since capture: {identity['name']}") + vendored = vendor_sources(module) + if {path: source_digest for path, (_, source_digest) in vendored.items()} != identity[ + "source_digests" + ]: + raise ContractError(f"Backend dependency changed since capture: {identity['name']}") + files = {path: text for path, (text, _) in vendored.items()} + imports, callable_name = f"import torch\nfrom .vendor.{module} import {entry}\n", entry + if module == "embodiinfer.backend.triton.sampling": + imports += f"from .vendor.{module} import GreedyWorkspace\n" + else: + raise ContractError(f"No task builder for {identity['kind']}") + 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"]] + code = imports + f"\ndef run({', '.join(names)}):\n" + code += f" result = {callable_name}(*{_expression(args)}, **{_expression(kwargs)})\n" + code += " return " + (returns[0] if len(returns) == 1 else "(" + ", ".join(returns) + ")") + "\n" + return code, files + + +def _plan( + capture: ModelCapture, op: dict[str, Any], overrides: dict[str, Any], hardware_notes: Path | None +) -> dict[str, str | bytes]: + name = "discovered_" + 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() + for phase in case.get("rng", {}).values(): + for path in phase.values(): + sources[path] = (capture.root / path).read_bytes() + 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="Single observed operator; no new fusion.", + precision=capture.manifest["precision"], + timing="eager", + ) + settings = tuning_settings(descriptor, overrides) + if identity.get("rng") and settings["timing"]["mode"] != "eager": + raise ContractError("RNG replay currently requires eager timing") + if settings["precision"] != {**capture.manifest["precision"]}: + raise ContractError("Model task precision cannot override its captured numerical contract") + 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 "rng" in identity: + replay["rng"] = identity["rng"] + 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", + "README.md": f"# {op['name']}\n\nTarget: `{capture.manifest['target']['model']}`.\nCapture: `{capture.identity}`.\n\nOptimize only this captured operator, preserving all output layouts, aliases and declared input mutations in replay.json. Return tensor outputs in definition order, including the mutated input buffers. Do not fuse with neighboring operators, introduce quantization, or change numerical settings. Use return-value style (destination_passing_style=false). All recorded workloads are mandatory.\n", + "generated.json": json.dumps( + { + "generator": "scripts.kernel_tuning.discovery", + "capture_identity": capture.identity, + "operator_id": op["id"], + } + ), + } + ) + if hardware_notes is not None: + sources["HARDWARE.md"] = hardware_notes.read_text(encoding="utf-8") + if "rng" in identity: + sources["README.md"] += ( + "\nThis call consumes the default PyTorch generator. The evaluator restores the " + "opaque before-state before every call and checks the exact after-state as well as " + "the tensor outputs. Seed 0 replays the recorded transition; other seeds test fresh " + "generator states. Preserve both CPU and target-CUDA RNG state advancement. RNG " + "reset is outside timed events; CUDA Graph timing is not supported for this task.\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..f778829 --- /dev/null +++ b/scripts/kernel_tuning/discovery/tensors.py @@ -0,0 +1,170 @@ +"""Tensor argument recipes and storage-preserving replay in the model runtime.""" + +from __future__ import annotations + +import math +from dataclasses import fields +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 ( + type(value).__module__ == "embodiinfer.backend.triton.sampling" + and type(value).__name__ == "GreedyWorkspace" + ): + return { + "greedy_workspace": { + field.name: flatten(getattr(value, field.name), tensors) for field in fields(value) + } + } + 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()} + if kind == "greedy_workspace": + from embodiinfer.backend.triton.sampling import GreedyWorkspace + + return GreedyWorkspace(**{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 "greedy_workspace" in node: + raise ContractError("Workspace outputs require explicit tensor returns") + 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/prompts.py b/scripts/kernel_tuning/prompts.py index da9e46b..1eb01f3 100644 --- a/scripts/kernel_tuning/prompts.py +++ b/scripts/kernel_tuning/prompts.py @@ -10,9 +10,15 @@ 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 "" + ) 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. +{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, @@ -50,6 +56,11 @@ def implementation_prompt(task: TaskPackage) -> str: } ], } + replay = ( + "This discovered 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. @@ -59,6 +70,7 @@ def implementation_prompt(task: TaskPackage) -> str: 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 intermediate rounding; bit_exact includes 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. diff --git a/tests/kernel_tuning/test_batch.py b/tests/kernel_tuning/test_batch.py index 13553c8..b21fd6c 100644 --- a/tests/kernel_tuning/test_batch.py +++ b/tests/kernel_tuning/test_batch.py @@ -177,7 +177,7 @@ def test_cli_validates_agent_and_resume_arguments(tmp_path: Path, capsys) -> Non 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", "--model", "streamvln", "--dry-run", "--output", str(output)]) == 0 + 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 @@ -188,7 +188,15 @@ def test_cli_validates_agent_and_resume_arguments(tmp_path: Path, capsys) -> Non 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", "rms_norm", "--output", str(tmp_path / "tasks"), "--hardware-notes", str(notes)] + 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" diff --git a/tests/kernel_tuning/test_capture.py b/tests/kernel_tuning/test_capture.py index 112a3a9..4b7ff0d 100644 --- a/tests/kernel_tuning/test_capture.py +++ b/tests/kernel_tuning/test_capture.py @@ -192,13 +192,21 @@ def test_cli_capture_and_generate_from_capture( "scripts.kernel_tuning.operators.select", lambda names=None, models=None: [fake_kernel] ) output = tmp_path / "calls.jsonl" - assert main(["capture", "--output", str(output), "--", str(script)]) == 0 + 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", "rms_norm", "--output", str(tmp_path / "tasks"), "--captured", str(captured)] + 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",) diff --git a/tests/kernel_tuning/test_discovery.py b/tests/kernel_tuning/test_discovery.py new file mode 100644 index 0000000..0f6d72b --- /dev/null +++ b/tests/kernel_tuning/test_discovery.py @@ -0,0 +1,435 @@ +"""Model discovery, replay contracts and the all-task barrier, using CPU tensors.""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +torch = pytest.importorskip("torch") + +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, read_json # noqa: E402 +from scripts.kernel_tuning.discovery import ModelCapture # noqa: E402 +from scripts.kernel_tuning.discovery.recorder import CaptureSession # 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 + + +@pytest.fixture(autouse=True) +def cpu_capture(monkeypatch): + """Keep observer tests independent of CUDA initialization and numerical globals.""" + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + old = torch.get_float32_matmul_precision() + torch.set_float32_matmul_precision("highest") + yield + torch.set_float32_matmul_precision(old) + + +def capture_calls(tmp_path: Path, fn, *, fixture_bytes: int = 1024 * 1024) -> ModelCapture: + session = CaptureSession( + "mock_flow_vla", tmp_path / "capture", device_type="cpu", fixture_bytes=fixture_bytes + ) + try: + with session.observe(object(), checkpoint="test-revision"): + fn() + finally: + session.close(complete=True) + return ModelCapture.load(session.root, require_ready=False) + + +def generated_function(task: TaskPackage): + scope = {} + exec((task.root / "baseline.py").read_text(encoding="utf-8"), scope) + return scope["run"] + + +def test_discovers_uncatalogued_ops_all_shapes_and_call_counts(tmp_path: Path) -> None: + values = [torch.ones(n, 4) for n in range(1, 12)] + weight = torch.ones(4, 3) + + def forward(): + for value in [*values, values[0]]: + (value @ weight).relu() + + capture = capture_calls(tmp_path, forward) + ready = {op["name"]: op for op in capture.operators if op["status"] == "ready"} + assert set(ready) == {"aten.mm.default", "aten.relu.default"} + assert len(ready["aten.mm.default"]["workloads"]) == 11 + assert sum(case["count"] for case in ready["aten.mm.default"]["workloads"]) == 12 + tasks = render_model(capture, tmp_path / "tasks") + assert len(tasks) == 2 + assert all(len(task.workloads) == 11 for task in tasks) + for task in tasks: + replay = read_json(task.root / "replay.json") + assert task.settings.precision.mode == "bit_exact" + for case in replay["cases"].values(): + inputs = materialize(case["inputs"], task.root, 0) + result = generated_function(task)(*inputs) + assert list(result.shape) == case["outputs"][0]["shape"] + + +def test_noncontiguous_aliases_mutations_and_metadata_are_preserved(tmp_path: Path) -> None: + backing = torch.arange(24, dtype=torch.float32).reshape(4, 6) + left, right = backing[:, 1:3], backing[:, 3:5] + capture = capture_calls(tmp_path, lambda: (left.add_(right), left.transpose(0, 1))) + assert any(op["status"] == "metadata_only" for op in capture.operators) + (task,) = render_model(capture, tmp_path / "tasks") + replay = read_json(task.root / "replay.json") + (case,) = replay["cases"].values() + assert replay["mutates"] == [0] + inputs = materialize(case["inputs"], task.root, 0) + assert storage_key(inputs[0]) == storage_key(inputs[1]) + assert inputs[0].stride() == (6, 1) + assert inputs[0].storage_offset() == 1 + cloned = clone_inputs(inputs) + assert storage_key(cloned[0]) == storage_key(cloned[1]) + assert storage_key(cloned[0]) != storage_key(inputs[0]) + expected = inputs[0] + inputs[1] + output, mutated = generated_function(task)(*cloned) + assert output is mutated and output is cloned[0] + assert torch.equal(output, expected) + + +def test_required_integer_fixtures_are_valid_and_never_treated_as_quantized(tmp_path: Path) -> None: + weight = torch.arange(20, dtype=torch.float32).reshape(5, 4) + indices = torch.tensor([4, 1, 0], dtype=torch.int64) + capture = capture_calls(tmp_path, lambda: torch.nn.functional.embedding(indices, weight)) + assert not any(op["status"] == "excluded_quantized" for op in capture.operators) + (task,) = render_model(capture, tmp_path / "tasks") + (case,) = read_json(task.root / "replay.json")["cases"].values() + args = materialize(case["inputs"], task.root, 0) + assert torch.equal(generated_function(task)(*args), weight[indices]) + + +def test_tensor_factories_without_tensor_inputs_get_tasks(tmp_path: Path) -> None: + capture = capture_calls(tmp_path, lambda: torch.arange(9, device="cpu")) + (task,) = render_model(capture, tmp_path / "tasks") + assert task.definition["inputs"] == {} + assert task.settings.device == "cpu" + assert torch.equal(generated_function(task)(), torch.arange(9)) + + +def test_quantized_scope_includes_dynamic_unregistered_module_and_float_fallback(tmp_path: Path) -> None: + from embodiinfer.layers.linear import INT8Config + from embodiinfer.models import linear + + dense = torch.nn.Linear(4, 3, bias=False) + quantized = linear.INT8Linear.from_linear(dense, INT8Config(backend="torch")) + owner = torch.nn.Module() + object.__setattr__(owner, "projection", quantized) + assert list(owner.children()) == [] + value = torch.ones(2, 4) + original = linear.QuantizedLinear.forward + session = CaptureSession("mock_flow_vla", tmp_path / "capture", device_type="cpu") + try: + session.install() + with session.observe(owner): + output = owner.projection(value) + output.relu() + finally: + session.close(complete=True) + assert linear.QuantizedLinear.forward is original + capture = ModelCapture.load(session.root) + assert [op["name"] for op in capture.operators if op["status"] == "ready"] == ["aten.relu.default"] + assert sum(op["status"] == "excluded_quantized" for op in capture.operators) == 1 + assert not any("mm." in op["name"] for op in capture.operators) + + +@pytest.mark.parametrize("failure", ["random", "fixture", "external"]) +def test_any_coverage_gap_blocks_all_task_generation(tmp_path: Path, failure: str) -> None: + value = torch.ones(4) + index = torch.tensor([1, 0]) + + def forward(): + value.relu() + if failure == "random": + torch.rand_like(value) + elif failure == "fixture": + value[index] + else: + custom(value) + + library = torch.library.Library("discovery_test", "FRAGMENT") + if failure == "external": + library.define("opaque(Tensor x) -> Tensor") + library.impl("opaque", lambda x: x + 1, "CPU") + custom = torch.ops.discovery_test.opaque.default + capture = capture_calls(tmp_path, forward, fixture_bytes=0) + assert any(op["status"] == "blocked" for op in capture.operators) + with pytest.raises(ContractError, match="incomplete"): + render_model(capture, tmp_path / "tasks") + assert not (tmp_path / "tasks").exists() + + +def test_capture_model_and_immutable_fixtures_are_verified(tmp_path: Path) -> None: + value = torch.ones(4) + capture = capture_calls(tmp_path, value.relu) + with pytest.raises(ContractError, match="differs from requested"): + ModelCapture.load(capture.root, "pi05") + fixture = next((capture.root / "fixtures").iterdir()) + fixture.write_bytes(b"changed") + with pytest.raises(ContractError, match="content changed"): + render_model(capture, tmp_path / "tasks") + + +def test_single_actual_target_and_uncompiled_execution_are_required(tmp_path: Path) -> None: + session = CaptureSession("mock_flow_vla", tmp_path / "capture", device_type="cpu") + try: + session.install() + from embodiinfer.policies.factory import make_policy + + with pytest.raises(ContractError, match="expected"): + make_policy("pi05") + with pytest.raises(ContractError, match="compilation"): + torch.compile(lambda x: x) + session.bind(object(), {}) + with pytest.raises(ContractError, match="one target"): + session.bind(object(), {}) + finally: + session.close(complete=False) + + +def test_model_batch_never_starts_agent_until_every_baseline_passes(tmp_path: Path, monkeypatch) -> None: + from test_batch import measurement + + value = torch.ones(4) + capture = capture_calls(tmp_path, lambda: value.relu().square()) + 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 "pow" in read_json(store.task.root / "replay.json")["operator"]: + 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()) + assert any(error for error in 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 + # A subsequently changed task invalidates the whole model, including resume. + 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_model_cli_rejects_partial_selection_and_legacy_capture(tmp_path: Path, capsys) -> None: + assert main(["generate", "--output", str(tmp_path / "tasks")]) == 2 + assert "exactly one --model" in capsys.readouterr().err + value = torch.ones(4) + capture = capture_calls(tmp_path, value.relu) + args = ["--model", "mock_flow_vla", "--captured", str(capture.root)] + assert main(["generate", *args, "--list"]) == 0 + assert "aten.relu.default" in capsys.readouterr().out + assert main(["tune-all", *args, "--dry-run", "--output", str(tmp_path / "batches")]) == 0 + assert main(["generate", *args, "relu"]) == 2 + assert main(["tune-all", *args, "--agent", "fake/model", "--skip-preflight"]) == 2 + assert main(["generate", "--model", "pi05", "--model", "streamvln"]) == 2 + legacy = tmp_path / "legacy.jsonl" + legacy.write_text("{}", encoding="utf-8") + assert main(["generate", "--model", "pi05", "--captured", str(legacy)]) == 2 + (task,) = render_model(capture, tmp_path / "tasks") + assert main(["run", str(task.root), "--agent", "fake/model", "--output", str(tmp_path / "runs")]) == 2 + assert not (tmp_path / "runs").exists() + + +def test_factory_driver_is_scoped_and_does_not_execute_an_extra_forward(tmp_path: Path, monkeypatch) -> None: + import sys + from types import SimpleNamespace + + from scripts.kernel_tuning.discovery import recorder + + from embodiinfer.policies import factory + + class FakePolicy(torch.nn.Module): + def __init__(self): + super().__init__() + self.decoder = SimpleNamespace() + self.calls = 0 + + def forward(self, inputs): + self.calls += 1 + return inputs.relu() + + policy = FakePolicy() + monkeypatch.setitem(factory._REGISTRY, "mock_flow_vla", lambda **kwargs: policy) + # The public CLI always targets CUDA; this test substitutes the observer's + # device classification while exercising the real driver and factory hooks. + monkeypatch.setattr( + recorder, "CaptureSession", lambda *args, **kwargs: CaptureSession(*args, device_type="cpu", **kwargs) + ) + script = tmp_path / "prepare.py" + script.write_text( + "import torch\nfrom embodiinfer.policies.factory import make_policy\nx = torch.ones(3)\nx.sin() # outside the selected policy\nmodel = make_policy('mock_flow_vla', checkpoint='test')\nmodel(x)\n", + encoding="utf-8", + ) + argv, path = sys.argv, sys.path[:] + recorder.run([str(script)], tmp_path / "capture", "mock_flow_vla") + capture = ModelCapture.load(tmp_path / "capture") + assert [op["name"] for op in capture.operators] == ["aten.relu.default"] + assert policy.calls == 1 + assert "forward" not in vars(policy) + assert sys.argv is argv and sys.path == path + assert capture.manifest["target"]["config"] == {"checkpoint": "test"} + + +def test_replay_adapter_checks_output_alias_and_undeclared_mutation_on_cpu(tmp_path: Path) -> None: + from benchmarks.kernel_tuning.replay_adapter import ReplayAdapter + + value = torch.ones(4) + capture = capture_calls(tmp_path, lambda: value.add_(2)) + (task,) = render_model(capture, tmp_path / "tasks") + replay = read_json(task.root / "replay.json") + (case,) = replay["cases"].values() + # Exercise the contract logic without constructing the GPU/compiler adapter. + adapter = ReplayAdapter.__new__(ReplayAdapter) + adapter.case, adapter.replay, adapter.cfg = case, replay, task.settings + inputs = materialize(case["inputs"], task.root, 0) + expected = clone_inputs(inputs) + outputs = generated_function(task)(*expected) + adapter._check_outputs(expected, list(outputs)) + assert adapter._check_state(expected, expected, inputs) == (0.0, 0.0) + with pytest.raises(ContractError, match="aliases"): + adapter._check_outputs(expected, [x.clone() for x in outputs]) + adapter.replay = {**replay, "mutates": []} + with pytest.raises(ContractError, match="bytes"): + adapter._check_state(expected, expected, inputs) + + +def test_missing_workload_and_unbound_output_axes_are_rejected(tmp_path: Path) -> None: + value = torch.ones(4, 2) + capture = capture_calls(tmp_path, lambda: value.sum(1)) + (task,) = render_model(capture, tmp_path / "tasks") + 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) + + +@pytest.mark.parametrize("known", [True, False]) +def test_existing_backend_is_atomic_and_requires_a_replay_contract( + tmp_path: Path, monkeypatch, known: bool +) -> None: + from scripts.kernel_tuning.discovery import recorder + + from embodiinfer.backend.triton import norm + + if not known: + monkeypatch.setattr(recorder, "_PURE_KERNELS", recorder._PURE_KERNELS - {"gated_residual"}) + + values = [torch.ones(1, 2, 4) for _ in range(3)] + session = CaptureSession("mock_flow_vla", tmp_path / "capture", device_type="cpu") + original = norm.gated_residual + try: + session.install() + with session.observe(object()): + result = norm.gated_residual(*values) + finally: + session.close(complete=True) + assert norm.gated_residual is original + assert torch.equal(result, torch.full_like(result, 2)) + capture = ModelCapture.load(session.root, require_ready=False) + (op,) = capture.operators + assert op["identity"]["kind"] == "backend" + assert op["name"].endswith(".gated_residual") + if not known: + with pytest.raises(ContractError, match="explicit mutation/reference contract"): + capture.require_ready() + return + (task,) = render_model(capture, tmp_path / "tasks") + assert "vendor/embodiinfer/backend/triton/norm.py" in task.hashes + assert "from .vendor.embodiinfer.backend.triton.norm" in task.definition["reference"] + + +def test_changed_numerical_settings_block_capture(tmp_path: Path) -> None: + value = torch.ones(4) + + def forward(): + value.relu() + torch.set_float32_matmul_precision("high") + value.square() + + capture = capture_calls(tmp_path, forward) + with pytest.raises(ContractError, match="Numerical settings changed"): + capture.require_ready() + + +def test_greedy_workspace_has_an_explicit_tensor_state_recipe() -> None: + from scripts.kernel_tuning.discovery.tasks import _expression + from scripts.kernel_tuning.discovery.tensors import flatten, restore + + from embodiinfer.backend.triton.sampling import GreedyWorkspace + + workspace = GreedyWorkspace.allocate(16, torch.device("cpu")) + tensors = [] + recipe = flatten(workspace, tensors) + cloned = clone_inputs(tensors) + restored = restore(recipe, cloned) + assert restored.vocab_size == 16 + assert restored.token is cloned[recipe["greedy_workspace"]["token"]["tensor"]] + scope = {"GreedyWorkspace": GreedyWorkspace, **{f"x{i}": tensor for i, tensor in enumerate(cloned)}} + constructed = eval(_expression(recipe), scope) + assert constructed.penalized is restored.penalized + assert constructed.token is restored.token + + +@pytest.mark.gpu +@pytest.mark.parametrize("timing", ["eager", "cuda_graph"]) +def test_discovered_mutable_baseline_passes_cuda_replay(tmp_path: Path, timing: str) -> None: + """Target-runtime smoke test: real builders, alias checks, resets and graph A/B/A.""" + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + pytest.importorskip("flashinfer_bench") + pytest.importorskip("triton") + from benchmarks.kernel_tuning.replay_adapter import ReplayAdapter + from scripts.kernel_tuning.contracts import validate_measurement + + value = torch.arange(16, dtype=torch.float32, device="cuda:0").reshape(2, 8)[:, 1:5] + session = CaptureSession("mock_flow_vla", tmp_path / "capture") + try: + with session.observe(object()): + value.add_(1) + finally: + session.close(complete=True) + (task,) = render_model( + ModelCapture.load(session.root), + tmp_path / "tasks", + overrides={"timing.mode": timing, "timing.warmup": 2, "timing.iterations": 5, "timing.trials": 1}, + ) + adapter = ReplayAdapter(task) + validate_measurement( + task, adapter.evaluate({role: task.baseline() for role in ("baseline", "incumbent", "candidate")}) + ) diff --git a/tests/kernel_tuning/test_replay_extensions.py b/tests/kernel_tuning/test_replay_extensions.py new file mode 100644 index 0000000..46dd7cd --- /dev/null +++ b/tests/kernel_tuning/test_replay_extensions.py @@ -0,0 +1,286 @@ +"""Lossless random-state and float64 contracts for discovered operator tasks.""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +torch = pytest.importorskip("torch") + +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.recorder import CaptureSession # noqa: E402 +from scripts.kernel_tuning.discovery.tasks import render_model # noqa: E402 +from scripts.kernel_tuning.discovery.tensors import materialize # noqa: E402 + + +@pytest.fixture(autouse=True) +def isolated_numerics(monkeypatch): + """CPU tests must neither initialize CUDA nor leak numerical/RNG settings.""" + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + precision, state = torch.get_float32_matmul_precision(), torch.get_rng_state() + torch.set_float32_matmul_precision("highest") + yield + torch.set_float32_matmul_precision(precision) + torch.set_rng_state(state) + + +def record(tmp_path: Path, fn, *, fixture_bytes: int = 1024 * 1024) -> ModelCapture: + """Run one CPU operation through the production observer and immutable loader.""" + session = CaptureSession( + "mock_flow_vla", tmp_path / "capture", device_type="cpu", fixture_bytes=fixture_bytes + ) + try: + with session.observe(object()): + fn() + finally: + session.close(complete=True) + return ModelCapture.load(session.root, require_ready=False) + + +def baseline(task): + """Load the generated single-file ATen baseline without a GPU builder.""" + scope = {} + exec((task.root / "baseline.py").read_text(encoding="utf-8"), scope) + return scope["run"] + + +def test_default_randn_records_real_state_and_generates_a_task(tmp_path: Path) -> None: + torch.manual_seed(123) + torch.randn(7) # Include a non-initial generator position/cache. + before = torch.get_rng_state().clone() + observed = [] + capture = record(tmp_path, lambda: observed.append(torch.randn(3, 5, generator=None))) + after = torch.get_rng_state().clone() + capture.require_ready() + (task,) = render_model(capture, tmp_path / "tasks") + replay = read_json(task.root / "replay.json") + assert replay["rng"] == {"kind": "torch_default", "devices": ["cpu"]} + (case,) = replay["cases"].values() + for phase, state in (("before", before), ("after", after)): + assert (task.root / case["rng"][phase]["cpu"]).read_bytes() == bytes(state.tolist()) + torch.set_rng_state(before) + assert torch.equal(baseline(task)(), observed[0]) + assert torch.equal(torch.get_rng_state(), after) + + +def test_float64_task_preserves_dtype_and_values(tmp_path: Path) -> None: + value = torch.tensor([1.0 + 2**-30, 1.0 - 2**-30], dtype=torch.float64) + capture = record(tmp_path, value.sin) + (task,) = render_model(capture, tmp_path / "tasks") + assert task.definition["inputs"]["x0"]["dtype"] == "float64" + assert task.definition["outputs"]["y0"]["dtype"] == "float64" + replay = read_json(task.root / "replay.json") + assert replay["definition_schema"] == "embodiinfer-replay-v1" + (case,) = replay["cases"].values() + inputs = materialize(case["inputs"], task.root, 0) + actual = baseline(task)(*inputs) + assert actual.dtype == torch.float64 + assert torch.equal(actual, value.sin()) + assert not torch.equal(actual, value.float().sin().double()) + + +def test_rng_state_is_required_bounded_and_immutable(tmp_path: Path) -> None: + capture = record(tmp_path / "bounded", lambda: torch.randn(2), fixture_bytes=0) + with pytest.raises(ContractError, match="RNG fixture.*budget"): + capture.require_ready() + assert len(capture.operators) == 1 # Instrumentation does not capture its own state operations. + capture = record( + tmp_path / "after_bounded", lambda: torch.randn(2), fixture_bytes=torch.get_rng_state().numel() + ) + with pytest.raises(ContractError, match="RNG fixture.*budget"): + capture.require_ready() + capture = record(tmp_path / "valid", lambda: torch.randn(2)) + (task,) = render_model(capture, tmp_path / "tasks") + replay = read_json(task.root / "replay.json") + (case,) = replay["cases"].values() + (task.root / case["rng"]["after"]["cpu"]).unlink() + with pytest.raises(ContractError, match="RNG fixture missing"): + TaskPackage.load(task.root) + with pytest.raises(ContractError, match="eager timing"): + render_model(capture, tmp_path / "graphs", overrides={"timing.mode": "cuda_graph"}) + + +def test_unsupported_rng_and_explicit_generators_still_block(tmp_path: Path) -> None: + for name, fn in ( + ("explicit", lambda: torch.randn(2, generator=torch.Generator().manual_seed(7))), + ("uniform", lambda: torch.rand(2)), + ): + capture = record(tmp_path / name, fn) + with pytest.raises(ContractError, match="incomplete"): + capture.require_ready() + + +@pytest.mark.parametrize("style", ["omitted", "positional", "keyword"]) +def test_zero_dropout_sdpa_is_not_a_random_call(tmp_path: Path, style: str) -> None: + q = torch.randn(1, 2, 8, 16) + state = torch.get_rng_state().clone() + outputs = [] + + def forward(): + op = torch.ops.aten.scaled_dot_product_attention.default + for _ in range(2): + if style == "omitted": + outputs.append(op(q, q, q)) + elif style == "positional": + outputs.append(op(q, q, q, None, 0.0)) + else: + outputs.append(op(q, q, q, dropout_p=0.0)) + + with torch.inference_mode(): + capture = record(tmp_path, forward) + capture.require_ready() + assert len(capture.operators) == 1 + assert capture.operators[0]["name"] == "aten.scaled_dot_product_attention.default" + assert sum(c["count"] for c in capture.operators[0]["workloads"]) == 2 + assert torch.equal(state, torch.get_rng_state()) + assert torch.equal(outputs[0], outputs[1]) + (task,) = render_model(capture, tmp_path / "tasks") + assert "rng" not in read_json(task.root / "replay.json") + + +@pytest.mark.parametrize("style", ["positional", "keyword"]) +def test_nonzero_dropout_sdpa_remains_a_coverage_gap(tmp_path: Path, style: str) -> None: + q = torch.randn(1, 2, 8, 16) + + def forward(): + op = torch.ops.aten.scaled_dot_product_attention.default + return op(q, q, q, None, 0.25) if style == "positional" else op(q, q, q, dropout_p=0.25) + + with torch.inference_mode(): + capture = record(tmp_path, forward) + with pytest.raises(ContractError, match="RNG replay adapter"): + capture.require_ready() + + +def test_rng_outputs_and_advancement_are_checked_and_caller_state_restored( + tmp_path: Path, monkeypatch +) -> None: + from types import SimpleNamespace + + from benchmarks.kernel_tuning.replay_adapter import ReplayAdapter + + capture = record(tmp_path, lambda: torch.randn(15, generator=None)) + (task,) = render_model(capture, tmp_path / "tasks") + adapter = ReplayAdapter.__new__(ReplayAdapter) + adapter.task, adapter.cfg = task, task.settings + adapter.definition = SimpleNamespace( + outputs=task.definition["outputs"], torch_output_dtypes=[torch.float32] + ) + adapter.replay = read_json(task.root / "replay.json") + adapter.device = "cpu" + adapter._check_flags = lambda: None + monkeypatch.setattr(torch.cuda, "synchronize", lambda *args: None) + + def unexpected_global_seed(*args, **kwargs): + pytest.fail("RNG replay must not seed unrelated generators") + + monkeypatch.setattr(torch, "manual_seed", unexpected_global_seed) + + class Runnable: + metadata = SimpleNamespace(destination_passing_style=False) + + def __init__(self, fn): + self.fn = fn + + def __call__(self, *args): + return self.fn(*args) + + fn = baseline(task) + reference = Runnable(fn) + + def extra_draw(): + output = fn() + torch.randn(1) + return output + + def rewind(): + state = torch.get_rng_state() + output = fn() + torch.set_rng_state(state) + return output + + def crash(): + torch.randn(1) + raise RuntimeError("candidate failed") + + for seed in (0, 1, 2): + inputs = adapter._inputs(SimpleNamespace(uuid=task.workload_ids[0]), seed) + saved = torch.get_rng_state().clone() + assert adapter._verify_call(reference, reference, inputs, seed) == (0.0, 0.0) + assert torch.equal(torch.get_rng_state(), saved) + for bad in (extra_draw, rewind): + with pytest.raises(ContractError, match="RNG state"): + adapter._verify_call(Runnable(bad), reference, inputs, seed) + assert torch.equal(torch.get_rng_state(), saved) + with pytest.raises(RuntimeError, match="candidate failed"): + adapter._verify_call(Runnable(crash), reference, inputs, seed) + assert torch.equal(torch.get_rng_state(), saved) + + +def test_float64_extension_requires_replay_and_retains_upstream_validation(tmp_path: Path) -> None: + pytest.importorskip("flashinfer_bench") + from benchmarks.kernel_tuning.replay_schema import ReplayDefinition + from flashinfer_bench.data import Definition + from pydantic import ValidationError + + capture = record(tmp_path, lambda: torch.linspace(0, 1, 8, dtype=torch.float64)) + (task,) = render_model(capture, tmp_path / "tasks") + definition = ReplayDefinition.model_validate(task.definition) + assert isinstance(definition, Definition) + assert definition.torch_output_dtypes == [torch.float64] + assert definition.outputs["y0"].dtype == "float64" + assert definition.model_dump(mode="json")["outputs"]["y0"]["dtype"] == "float64" + with pytest.raises(ValidationError): + Definition.model_validate(task.definition) # No process-wide dtype monkeypatch. + wrong = {**task.definition, "outputs": {"y0": {"shape": ["unknown_axis"], "dtype": "float64"}}} + with pytest.raises(ValidationError): + ReplayDefinition.model_validate(wrong) + replay = read_json(task.root / "replay.json") + del replay["definition_schema"] + import json + + (task.root / "replay.json").write_text(json.dumps(replay), encoding="utf-8") + with pytest.raises(ContractError, match="explicit replay definition schema"): + TaskPackage.load(task.root) + + +@pytest.mark.gpu +@pytest.mark.parametrize("operation", ["randn", "float64", "sdpa"]) +def test_cuda_rng_and_float64_baselines_pass_real_evaluator(tmp_path: Path, operation: str) -> None: + """Exercise real Python builders, byte comparison, RNG transitions and eager resets.""" + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + pytest.importorskip("flashinfer_bench") + pytest.importorskip("triton") + from benchmarks.kernel_tuning.replay_adapter import ReplayAdapter + from scripts.kernel_tuning.contracts import validate_measurement + + state = torch.cuda.get_rng_state().clone() + query = torch.ones(1, 2, 8, 16, device="cuda:0", dtype=torch.bfloat16) + session = CaptureSession("mock_flow_vla", tmp_path / "capture") + try: + with torch.inference_mode(), session.observe(object()): + if operation == "randn": + torch.randn(1, 50, 32, device="cuda:0", dtype=torch.bfloat16, generator=None) + elif operation == "sdpa": + torch.nn.functional.scaled_dot_product_attention(query, query, query, dropout_p=0.0) + else: + times = torch.linspace(0, 1, 16, device="cuda:0", dtype=torch.float64) + (1000.0**times).reciprocal().sin().to(torch.bfloat16) + finally: + session.close(complete=True) + torch.cuda.set_rng_state(state) + tasks = render_model( + ModelCapture.load(session.root), + tmp_path / "tasks", + overrides={"timing.warmup": 1, "timing.iterations": 2, "timing.trials": 1}, + ) + for task in tasks: + adapter = ReplayAdapter(task) + before = torch.cuda.get_rng_state().clone() + report = adapter.evaluate({role: task.baseline() for role in ("baseline", "incumbent", "candidate")}) + validate_measurement(task, report) + if operation == "randn": + assert torch.equal(before, torch.cuda.get_rng_state()) From 33bbde53ef39b73f53718a3401979748a0af2d5a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=82=96=E4=B8=80=E6=B6=B5?= <23371301@buaa.edu.cn> Date: Wed, 7 Oct 2026 02:18:59 +0800 Subject: [PATCH 08/44] fix: close discovery capture gaps found on Qwen R2R and pi0.5 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Treat aten.item as host-side scalar work like _local_scalar_dense. Torch 2.12 delivers Tensor.item() to the observer as aten.item, which was recorded as a blocked non-tensor call and marked the Qwen R2R low-level capture incomplete. Snapshot inputs when a host-input call is promoted to ready because its output lands on the target device. Host-to-device copies such as integer index tensors moved to CUDA previously had no fixture and failed preflight with "Non-floating inputs require a valid captured fixture". Promoted calls that mutate their inputs are blocked. Observe decoder.integrate. The pi0.5 benchmark times the flow loop through integrate, whose native path reaches _denoise_step_impl without passing any observed method, so the whole denoising loop was missing while coverage still reported complete (58 instead of 96 pi0.5 tasks). Validation: three regression tests fail before and pass after the change; tests/kernel_tuning passes on CPU (97 passed, 19 skipped). On RTX 5090 the fixed recorder captured pi0.5 (96 ready) and Qwen R2R low-level (124 ready) with no blocked calls, and every baseline passed preflight. Signed-off-by: 肖一涵 <23371301@buaa.edu.cn> --- scripts/kernel_tuning/discovery/recorder.py | 8 ++- tests/kernel_tuning/test_discovery.py | 58 +++++++++++++++++++++ 2 files changed, 64 insertions(+), 2 deletions(-) diff --git a/scripts/kernel_tuning/discovery/recorder.py b/scripts/kernel_tuning/discovery/recorder.py index e965e15..13f2592 100644 --- a/scripts/kernel_tuning/discovery/recorder.py +++ b/scripts/kernel_tuning/discovery/recorder.py @@ -53,7 +53,7 @@ "set_", } _ALLOCATIONS = {"empty", "empty_like", "empty_strided", "new_empty", "new_empty_strided"} -_HOST = {"_local_scalar_dense", "is_nonzero"} +_HOST = {"item", "_local_scalar_dense", "is_nonzero"} _QUANTIZED_ATEN = { "_scaled_mm", "_weight_int8pack_mm", @@ -381,6 +381,10 @@ def call( and any(t.device.type == self.device_type for t in tensor_outputs(result)) ): status = "ready" + # Host inputs were not snapshotted before the call; capture them now. + if identity["mutates"]: + raise ContractError("Host-to-device calls mutating inputs need a dedicated adapter") + case["inputs"] = self._snapshot(digest([identity, specs]), tensors, specs) if status not in {"metadata_only", "host_only", "excluded_quantized", "blocked"}: outputs = tensor_outputs(result) identity["returns"] = flatten(result, []) @@ -514,7 +518,7 @@ def construct(name: str, **kwargs: Any) -> Any: self.bind(policy, kwargs) for holder, methods in ( (policy, ("encode_prefix", "denoise_step", "prepare_prefix", "forward")), - (policy.decoder, ("produce_chunk", "generate_tokens", "init_state")), + (policy.decoder, ("produce_chunk", "integrate", "generate_tokens", "init_state")), ): for method in methods: if not hasattr(holder, method): diff --git a/tests/kernel_tuning/test_discovery.py b/tests/kernel_tuning/test_discovery.py index 0f6d72b..a2e6ebb 100644 --- a/tests/kernel_tuning/test_discovery.py +++ b/tests/kernel_tuning/test_discovery.py @@ -112,6 +112,41 @@ def test_tensor_factories_without_tensor_inputs_get_tasks(tmp_path: Path) -> Non assert torch.equal(generated_function(task)(), torch.arange(9)) +def test_scalar_reads_are_host_work_not_coverage_gaps(tmp_path: Path) -> None: + from scripts.kernel_tuning.discovery.recorder import _Observer + + value = torch.arange(4, dtype=torch.float32).sum() + session = CaptureSession("mock_flow_vla", tmp_path / "capture", device_type="cpu") + try: + with session.observe(object(), checkpoint="test-revision"): + # Torch 2.12 delivers Tensor.item() to the observer as aten.item; older + # releases decompose it first, so deliver the overload directly. + result = _Observer(session).__torch_dispatch__(torch.ops.aten.item.default, (), (value,)) + assert result == 6.0 + finally: + session.close(complete=True) + capture = ModelCapture.load(session.root, require_ready=False) + statuses = {op["name"]: op["status"] for op in capture.operators} + assert statuses["aten.item.default"] == "host_only" + assert "blocked" not in statuses.values() + + +def test_host_to_device_copies_keep_required_input_fixtures(tmp_path: Path) -> None: + # "meta" stands in for the target device: host inputs, target-device output. + session = CaptureSession("mock_flow_vla", tmp_path / "capture", device_type="meta") + indices = torch.tensor([4, 1, 0], dtype=torch.int64) + try: + with session.observe(object(), checkpoint="test-revision"): + indices.to("meta") + finally: + session.close(complete=True) + capture = ModelCapture.load(session.root, require_ready=False) + (op,) = [op for op in capture.operators if op["status"] == "ready"] + (case,) = op["workloads"] + assert case["inputs"][0]["device"] == "cpu" + assert case["inputs"][0].get("fixture") + + def test_quantized_scope_includes_dynamic_unregistered_module_and_float_fallback(tmp_path: Path) -> None: from embodiinfer.layers.linear import INT8Config from embodiinfer.models import linear @@ -305,6 +340,29 @@ def forward(self, inputs): assert capture.manifest["target"]["config"] == {"checkpoint": "test"} +def test_decoder_integrate_entry_is_observed(tmp_path: Path, monkeypatch) -> None: + # Benchmarks may time the flow loop through decoder.integrate rather than produce_chunk. + from types import SimpleNamespace + + from scripts.kernel_tuning.discovery import recorder + + from embodiinfer.policies import factory + + policy = SimpleNamespace(decoder=SimpleNamespace(integrate=lambda x: x.relu())) + monkeypatch.setitem(factory._REGISTRY, "mock_flow_vla", lambda **kwargs: policy) + monkeypatch.setattr( + recorder, "CaptureSession", lambda *args, **kwargs: CaptureSession(*args, device_type="cpu", **kwargs) + ) + script = tmp_path / "prepare.py" + script.write_text( + "import torch\nfrom embodiinfer.policies.factory import make_policy\nmodel = make_policy('mock_flow_vla', checkpoint='test')\nmodel.decoder.integrate(torch.ones(3))\n", + encoding="utf-8", + ) + recorder.run([str(script)], tmp_path / "capture", "mock_flow_vla") + capture = ModelCapture.load(tmp_path / "capture") + assert [op["name"] for op in capture.operators] == ["aten.relu.default"] + + def test_replay_adapter_checks_output_alias_and_undeclared_mutation_on_cpu(tmp_path: Path) -> None: from benchmarks.kernel_tuning.replay_adapter import ReplayAdapter From dc3c101af0fbeffbe6f843301dd7f6315983078a Mon Sep 17 00:00:00 2001 From: owpeter Date: Wed, 7 Oct 2026 11:09:43 +0800 Subject: [PATCH 09/44] feat: calibrate operator numerical contracts Signed-off-by: owpeter --- CONTRIBUTING.md | 27 +- benchmarks/kernel_tuning/calibration.py | 380 ++++++++++++++++++ .../kernel_tuning/flashinfer_adapter.py | 5 + benchmarks/kernel_tuning/replay_adapter.py | 94 ++++- docs/proposals/0008-kernel-tuning.md | 81 +++- scripts/kernel_tuning/__main__.py | 10 + scripts/kernel_tuning/contracts.py | 28 ++ scripts/kernel_tuning/discovery/numerics.py | 138 +++++++ scripts/kernel_tuning/discovery/recorder.py | 26 +- scripts/kernel_tuning/discovery/tasks.py | 35 +- scripts/kernel_tuning/prompts.py | 4 +- tests/kernel_tuning/test_calibration.py | 330 +++++++++++++++ tests/kernel_tuning/test_discovery.py | 10 +- tests/kernel_tuning/test_replay_extensions.py | 7 +- 14 files changed, 1133 insertions(+), 42 deletions(-) create mode 100644 benchmarks/kernel_tuning/calibration.py create mode 100644 scripts/kernel_tuning/discovery/numerics.py create mode 100644 tests/kernel_tuning/test_calibration.py diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 921ca82..c421b6a 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -68,14 +68,15 @@ uv run --project scripts/kernel_tuning python -m scripts.kernel_tuning --help uv run --project scripts/kernel_tuning python -m scripts.kernel_tuning check benchmarks/kernel_tuning/tasks/gated_residual ``` -To tune one model, capture its existing preparation forward in its model runtime, -then generate and preflight every eligible observed call before starting agents +To tune one model, capture its existing preparation forward and calibrate numerical +contracts in its model runtime, then generate and preflight every eligible call (see [model discovery](docs/proposals/0008-kernel-tuning.md#model-scoped-operator-discovery)): ```bash -/absolute/runtime/python -m scripts.kernel_tuning capture --model pi05 --output results/kernel_tuning/captures/pi05 -- /absolute/preparation.py +/absolute/runtime/python -m scripts.kernel_tuning capture --model pi05 --fixture-bytes 8589934592 --output results/kernel_tuning/captures/pi05 -- /absolute/preparation.py uv run --project scripts/kernel_tuning python -m scripts.kernel_tuning generate --model pi05 --captured results/kernel_tuning/captures/pi05 --list -uv run --project scripts/kernel_tuning python -m scripts.kernel_tuning tune-all --model pi05 --captured results/kernel_tuning/captures/pi05 --agent 'HARNESS/MODEL:EFFORT' --set evaluator_python=/absolute/runtime/python +/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 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 ``` @@ -92,12 +93,18 @@ only builds tasks; `--preflight-only` measures baselines without opening agents. Model batches cannot use `--skip-preflight`, partial operator selection, or estimated shapes. A standalone `run` cannot bypass their preflight barrier. -Capture may store input data under a 64 MiB limit (`--fixture-bytes`): small -floating fixtures when space permits, and required integer/index fixtures. -Capture directories and tasks are hashed and checked again on resume. Large -floating inputs use seeded generation; calls needing more specific inputs must -pass preflight or receive a dedicated replay adapter. Start with a new output -directory when changing preparation/configuration. +The default fixture budget is 64 MiB; the example raises the upper limit to 8 GiB +without allocating it up front. Supported floating matmul/linear, reductions, +softmax and attention require complete real input fixtures, including weights. +Increase `--fixture-bytes` for the observed model; exhaustion fails explicitly. +Other exact tasks retain optional small floating fixtures and required integer +fixtures. Calibration compares only the production baseline and an FP64 reference +on recorded, random and boundary inputs. It stores per-workload/profile/output +bounds in a new hashed capture. Generation copies them to `numerics.json`; every +element must pass, and seeds/thresholds/execution settings cannot be overridden. +Calibration costs extra storage and target-device compute, before any agent runs. +Use a new capture/calibration/batch when changing settings. Old runs keep their +original contract; old captures with incomplete real inputs need recapture. Default-generator `randn` tasks save required before/after RNG fixtures, check both output bytes and generator advancement, and support eager timing only. diff --git a/benchmarks/kernel_tuning/calibration.py b/benchmarks/kernel_tuning/calibration.py new file mode 100644 index 0000000..abdbb78 --- /dev/null +++ b/benchmarks/kernel_tuning/calibration.py @@ -0,0 +1,380 @@ +"""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"] == "aten": + _, name, overload = identity["name"].split(".") + fn = getattr(getattr(torch.ops.aten, name), overload) + else: + 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]: + args, kwargs = restore(identity["arguments"], inputs) + if identity["kind"] == "aten": + _, name, overload = identity["name"].split(".") + schema = getattr(getattr(torch.ops.aten, name), overload)._schema + return { + arg.name: args[i] if i < len(args) else kwargs.get(arg.name, arg.default_value) + for i, arg in enumerate(schema.arguments) + } + names = { + "split_kv_attention": ( + "query", + "prefix_key", + "prefix_value", + "suffix_key", + "suffix_value", + "mask", + "scaling", + ), + "segmented_attention": ("q", "k", "v", "q_segment_offsets", "kv_segment_offsets"), + }[policy(identity)] + return {**dict(zip(names, 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 _attention( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + mask: Any = None, + scale: float | None = None, + causal: bool = False, + gqa: bool = False, +) -> torch.Tensor: + if gqa and q.shape[-3] != k.shape[-3]: + k = k.repeat_interleave(q.shape[-3] // k.shape[-3], dim=-3) + v = v.repeat_interleave(q.shape[-3] // v.shape[-3], dim=-3) + scale = q.shape[-1] ** -0.5 if scale is None else scale + rows = [] + # Bound oracle scratch memory without changing each row's reduction semantics. + for start in range(0, q.shape[-2], 128): + stop = min(start + 128, q.shape[-2]) + scores = (q[..., start:stop, :] @ k.transpose(-1, -2)) * scale + if causal: + visible = ( + torch.arange(k.shape[-2], device=q.device) + <= torch.arange(start, stop, device=q.device)[:, None] + ) + scores = scores.masked_fill(~visible, float("-inf")) + if mask is not None: + local = mask[..., start:stop, :] if mask.ndim >= 2 and mask.shape[-2] != 1 else mask + scores = ( + scores.masked_fill(~local, float("-inf")) if local.dtype == torch.bool else scores + local + ) + probs = torch.softmax(scores, dim=-1) + probs = torch.where(torch.isneginf(scores).all(dim=-1, keepdim=True), 0.0, probs) + rows.append(probs @ v) + return torch.cat(rows, dim=-2) if rows else q.new_empty((*q.shape[:-1], v.shape[-1])) + + +def high_precision(identity: dict[str, Any], inputs: list[torch.Tensor]) -> list[torch.Tensor]: + """Evaluate the declared mathematical operation in FP64 with unchanged attention semantics.""" + kind = policy(identity) + if kind is None: + 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] + args = _arguments(identity, promoted) + if kind == "aten": + _, name, overload = identity["name"].split(".") + if name == "scaled_dot_product_attention": + if args["dropout_p"] != 0: + raise ContractError("Attention calibration requires dropout_p=0") + return [ + _attention( + args["query"], + args["key"], + args["value"], + mask=args["attn_mask"], + scale=args["scale"], + causal=args["is_causal"], + gqa=args.get("enable_gqa", False), + ) + ] + if "dtype" in args: + args["dtype"] = torch.float64 + if "half_to_float" in args: + args["half_to_float"] = False + if name == "rms_norm" and args.get("eps") is None: + # The default epsilon belongs to the captured input dtype, not the oracle dtype. + args["eps"] = torch.finfo(inputs[0].dtype).eps + fn = getattr(getattr(torch.ops.aten, name), overload) + return tensor_outputs(fn(**args)) + if kind == "split_kv_attention": + mask = args.get("mask") + return [ + _attention( + args["query"], + torch.cat((args["prefix_key"], args["suffix_key"]), dim=2), + torch.cat((args["prefix_value"], args["suffix_value"]), dim=2), + mask=mask == 0 if mask is not None else None, + scale=args["scaling"], + gqa=True, + ) + ] + q, k, v = args["q"], args["k"], args["v"] + qb = args["q_segment_offsets"].tolist() + kb = args.get("kv_segment_offsets") + kb = qb if kb is None else kb.tolist() + output = torch.empty_like(q) + for qs, qe, ks, ke in zip(qb[:-1], qb[1:], kb[:-1], kb[1:]): + output[qs:qe] = _attention( + q[qs:qe].transpose(0, 1), + k[ks:ke].transpose(0, 1), + v[ks:ke].transpose(0, 1), + scale=args.get("scaling"), + gqa=True, + ).transpose(0, 1) + return [output] + + +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/flashinfer_adapter.py b/benchmarks/kernel_tuning/flashinfer_adapter.py index 17e694f..20bd85b 100644 --- a/benchmarks/kernel_tuning/flashinfer_adapter.py +++ b/benchmarks/kernel_tuning/flashinfer_adapter.py @@ -459,6 +459,7 @@ def evaluate(self, solutions: dict[str, dict[str, Any]]) -> dict[str, Any]: "max_abs_error": max_abs, "max_rel_error": max_rel, "graph_replay": self.cfg.timing.mode == "cuda_graph", + **self._check_metadata(), } ) rounds = [] @@ -477,6 +478,10 @@ def evaluate(self, solutions: dict[str, dict[str, Any]]) -> dict[str, Any]: 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) diff --git a/benchmarks/kernel_tuning/replay_adapter.py b/benchmarks/kernel_tuning/replay_adapter.py index 5b10317..fb4cb48 100644 --- a/benchmarks/kernel_tuning/replay_adapter.py +++ b/benchmarks/kernel_tuning/replay_adapter.py @@ -33,6 +33,14 @@ def __init__(self, task: TaskPackage) -> None: self._precision_flags = self._flags() self.case: dict[str, Any] = {} self._rng: Any = None + 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("Discovered tasks need a replay-aware NCU adapter before profile=true") if self.replay.get("rng") and self.cfg.timing.mode != "eager": @@ -74,6 +82,16 @@ def _inputs(self, workload: Any, seed: int) -> list[Any]: self._rng = ( RngReplay(self.replay["rng"], self.case, self.task.root, seed) if self.replay.get("rng") else None ) + 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 @@ -130,6 +148,22 @@ def _verify_call( ) -> tuple[float, float]: from scripts.kernel_tuning.discovery.rng import preserve + 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 rng = getattr(self, "_rng", None) with preserve(rng.devices) if rng else nullcontext(): return self._verify_transition(runnable, reference, inputs, seed) @@ -162,9 +196,28 @@ def _verify_transition( # 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 = compare_outputs(outputs, expected, self.cfg.precision) + 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]: @@ -206,6 +259,8 @@ def _capture(self, runnable: Any, inputs: list[Any]) -> tuple[Any, list[Any], li 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)): @@ -231,6 +286,43 @@ def _verify_graph(self, runnable: Any, reference: Any, workload: Any, seed: int) 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: from scripts.kernel_tuning.discovery.rng import preserve diff --git a/docs/proposals/0008-kernel-tuning.md b/docs/proposals/0008-kernel-tuning.md index 5c856da..61487df 100644 --- a/docs/proposals/0008-kernel-tuning.md +++ b/docs/proposals/0008-kernel-tuning.md @@ -118,6 +118,59 @@ belong to the evaluator and generated solution, not the orchestration loop. ## 6. Losslessness and precision criterion +### Calibrated numerical contracts for discovered operators + +Model discovery records execution settings separately from acceptance rules. +Pure copies, indexing, exact elementwise operations, RNG transitions, and unknown +semantics retain bit-exact validation. An explicit registry selects floating +matrix multiplication, linear, reductions, softmax, and supported attention/norm +calls for tolerance validation; dtype alone never relaxes an unknown operator. +Existing hand-written/catalog contracts and archived runs keep their contracts. + +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 +and RNG 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), discovery recorder/task generation, and the replay evaluator. +There are no engine, policy, backend, dependency-pin or model-level action-parity +changes. 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 @@ -376,9 +429,10 @@ policy IDs are those accepted by the factory, for example `pi05`, `streamvln`, or `qwen2.5-vl-3b-r2r-low-level`; catalog group names are not model IDs. ```bash -/absolute/runtime/python -m scripts.kernel_tuning capture --model pi05 --output results/kernel_tuning/captures/pi05 -- PREPARATION.py +/absolute/runtime/python -m scripts.kernel_tuning capture --model pi05 --fixture-bytes 8589934592 --output results/kernel_tuning/captures/pi05 -- PREPARATION.py uv run --project scripts/kernel_tuning python -m scripts.kernel_tuning generate --model pi05 --captured results/kernel_tuning/captures/pi05 --list -uv run --project scripts/kernel_tuning python -m scripts.kernel_tuning tune-all --model pi05 --captured results/kernel_tuning/captures/pi05 --agent 'HARNESS/MODEL:EFFORT' --set evaluator_python=/absolute/runtime/python +/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 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 ``` @@ -403,22 +457,25 @@ its baseline calls the exact ATen overload or snapshotted backend implementation Different shapes and strides remain workloads of the same signature; overload, scalar/static arguments, tensor rank/dtype and output structure distinguish tasks. -The default 64 MiB fixture budget is explicit (`--fixture-bytes`). Integer/index -storage requires a captured fixture. Floating storage up to 1 MiB is captured -when the remaining budget permits; larger floating storage is generated from -fixed seeds. Seed 0 uses the unmodified saved fixture; other seeds perturb -floating values. First observed fixtures represent each structural case. This -is finite input validation, not proof for all values; domain-specific validity -or data-dependent output shapes may require an adapter. Budget exhaustion for a -required fixture is a coverage gap, never permission to guess indices. +The default 64 MiB fixture budget is explicit (`--fixture-bytes`); the example +raises the maximum to 8 GiB. Integer/index storage and all inputs of supported +numerically calibrated operators require real fixtures, even for large weights. +For other exact tasks, floating storage up to 1 MiB is captured when budget +permits and other floating storage uses seeded generation. First observed +fixtures represent each structural case; calibration tests those real values, +fresh random values, zeros and cancellation inputs with masks/indices preserved. +This is finite validation, not proof for every input. Required fixture exhaustion +is a coverage gap. `calibrate` adds `calibration.json` to a new capture; generated +tasks carry each operator's frozen bounds as `numerics.json`. See the numerical +policy in Section 6 for the formula, environment binding and limitations. Replay preserves shared storage, strides and offsets. It checks output values, layouts, input/output and output/output aliases, declared mutations and unchanged input storage. Timing resets backing storage before every sample, outside timed events, in both eager and graph modes. The default is eager; `timing.mode=cuda_graph` also requires successful changed-input A/B/A replay. NCU profiling currently -requires a dedicated replay-aware adapter. Numerical overrides cannot weaken -the captured bit-exact production contract. +requires a dedicated replay-aware adapter. Numerical overrides cannot change +the selected exact contract or the frozen calibrated tolerance contract. Default-generator `aten.randn.default` and `aten.randn.generator` with `generator=None` have an eager replay contract. Capture stores the opaque CPU diff --git a/scripts/kernel_tuning/__main__.py b/scripts/kernel_tuning/__main__.py index 03ff765..60a8c80 100644 --- a/scripts/kernel_tuning/__main__.py +++ b/scripts/kernel_tuning/__main__.py @@ -278,6 +278,11 @@ def main(argv: list[str] | None = None) -> int: "--output", type=Path, required=True, help="New model capture directory (JSONL for --catalog)" ) capture.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)" ) @@ -318,6 +323,11 @@ def main(argv: list[str] | None = None) -> int: command or args.script, args.output, args.model, fixture_bytes=args.fixture_bytes ) print(json.dumps(result, indent=2, ensure_ascii=False)) + 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 diff --git a/scripts/kernel_tuning/contracts.py b/scripts/kernel_tuning/contracts.py index 11247df..4d41bc9 100644 --- a/scripts/kernel_tuning/contracts.py +++ b/scripts/kernel_tuning/contracts.py @@ -374,6 +374,27 @@ def _load(cls, root: Path) -> TaskPackage: 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: @@ -469,6 +490,13 @@ def _validate_measurement(task: TaskPackage, report: dict[str, Any]) -> None: 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"): diff --git a/scripts/kernel_tuning/discovery/numerics.py b/scripts/kernel_tuning/discovery/numerics.py new file mode 100644 index 0000000..fb5902e --- /dev/null +++ b/scripts/kernel_tuning/discovery/numerics.py @@ -0,0 +1,138 @@ +"""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)) +# Explicit semantic adapters, not substring matching on arbitrary operator names. +_ATEN = { + "mm", + "bmm", + "matmul", + "mv", + "dot", + "addmm", + "baddbmm", + "linear", + "sum", + "mean", + "_softmax", + "softmax", + "_log_softmax", + "log_softmax", + "native_layer_norm", + "layer_norm", + "rms_norm", + "scaled_dot_product_attention", +} +_BACKENDS = { + "embodiinfer.backend.triton.split_kv_attention.split_kv_attention": "split_kv_attention", + "embodiinfer.backend.triton.segmented_attention.segmented_attention": "segmented_attention", +} + + +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 + name = identity["name"] + if identity["kind"] == "aten" and name.split(".")[1] in _ATEN: + return "aten" + return _BACKENDS.get(name) + + +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/recorder.py b/scripts/kernel_tuning/discovery/recorder.py index 13f2592..2a3e465 100644 --- a/scripts/kernel_tuning/discovery/recorder.py +++ b/scripts/kernel_tuning/discovery/recorder.py @@ -22,6 +22,7 @@ from ..contracts import ContractError, digest, tree_hashes from ..generate import _git_revision, vendor_sources from .contracts import SCHEMA +from .numerics import policy from .rng import contract_for, requires_rng_adapter from .rng import snapshot as rng_snapshot from .tensors import describe_inputs, extent, flatten, storage_key, tensor_outputs, tensor_spec @@ -178,7 +179,6 @@ def __init__( def _set_numerics(self) -> None: self.precision = { - "mode": "bit_exact", "matmul_precision": torch.get_float32_matmul_precision(), "allow_tf32": torch.backends.cuda.matmul.allow_tf32, } @@ -233,7 +233,7 @@ def observe(self, instance: Any, **config: Any) -> Iterator[None]: self.depth -= 1 def _snapshot( - self, key: str, tensors: list[torch.Tensor], specs: list[dict[str, Any]] + self, key: str, tensors: list[torch.Tensor], specs: list[dict[str, Any]], *, real_inputs: bool = False ) -> list[dict[str, Any]]: if key in self._fixtures: return self._fixtures[key] @@ -246,13 +246,15 @@ def _snapshot( tensor = tensors[indices[0]] size = max(extent(specs[i]) for i in indices) byte_count = size * tensor.element_size() - required = not tensor.is_floating_point() - # Large ordinary floating weights are generated; retain small real cases - # when budget permits. Integer/index fixtures are never guessed. + required = real_inputs or not tensor.is_floating_point() + if required and byte_count > self.fixture_bytes: + raise ContractError( + "Required input fixture exceeds the capture byte budget; increase --fixture-bytes" + ) + # Calibrated arithmetic needs real weights too. Other floating tasks + # retain small fixtures when space permits; indices are never guessed. save = required or byte_count <= 1024 * 1024 - if save and self.used_bytes + byte_count > self.fixture_bytes: - if required: - raise ContractError("Required input fixture exceeds the capture byte budget") + if save and not required and self.used_bytes + byte_count > self.fixture_bytes: save = False if save: raw = ( @@ -267,6 +269,10 @@ def _snapshot( path = f"fixtures/{hashlib.sha256(raw).hexdigest()}.bin" file = self.root / path if not file.exists(): + if self.used_bytes + byte_count > self.fixture_bytes: + raise ContractError( + "Required real input fixture exceeds the capture byte budget; increase --fixture-bytes" + ) file.parent.mkdir(exist_ok=True) file.write_bytes(raw) self.used_bytes += len(raw) @@ -358,7 +364,9 @@ def call( if reason: raise ContractError(reason) if status == "ready": - case["inputs"] = self._snapshot(digest([identity, specs]), tensors, specs) + case["inputs"] = self._snapshot( + digest([identity, specs]), tensors, specs, real_inputs=policy(identity) is not None + ) if rng is not None: identity["rng"] = rng case["rng"] = {"before": self._rng_snapshot(rng["devices"])} diff --git a/scripts/kernel_tuning/discovery/tasks.py b/scripts/kernel_tuning/discovery/tasks.py index 51f26a2..9412785 100644 --- a/scripts/kernel_tuning/discovery/tasks.py +++ b/scripts/kernel_tuning/discovery/tasks.py @@ -9,10 +9,11 @@ from typing import Any from ..artifacts import REPOSITORY, atomic_json -from ..contracts import ContractError, TaskPackage +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 = { @@ -164,6 +165,20 @@ def _plan( for phase in case.get("rng", {}).values(): for path in phase.values(): sources[path] = (capture.root / path).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"], @@ -178,14 +193,16 @@ def _plan( reference=code, workloads=tuple(Workload(case["id"], {}, weight=case["count"]) for case in op["workloads"]), notes="Single observed operator; no new fusion.", - precision=capture.manifest["precision"], + precision=precision, timing="eager", ) settings = tuning_settings(descriptor, overrides) if identity.get("rng") and settings["timing"]["mode"] != "eager": raise ContractError("RNG replay currently requires eager timing") - if settings["precision"] != {**capture.manifest["precision"]}: - raise ContractError("Model task precision cannot override its captured numerical contract") + 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( @@ -213,6 +230,8 @@ def _plan( } if "rng" in identity: replay["rng"] = identity["rng"] + if numerical is not None: + replay["numerical_identity"] = identity sources.update( { "definition.json": json.dumps(definition, indent=2) + "\n", @@ -233,6 +252,14 @@ def _plan( ) 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 "rng" in identity: sources["README.md"] += ( "\nThis call consumes the default PyTorch generator. The evaluator restores the " diff --git a/scripts/kernel_tuning/prompts.py b/scripts/kernel_tuning/prompts.py index 1eb01f3..0ed820e 100644 --- a/scripts/kernel_tuning/prompts.py +++ b/scripts/kernel_tuning/prompts.py @@ -15,6 +15,8 @@ def plan_prompt(task: TaskPackage) -> str: 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. @@ -71,7 +73,7 @@ def implementation_prompt(task: TaskPackage) -> str: 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 intermediate rounding; bit_exact includes signed zero and NaN +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. diff --git a/tests/kernel_tuning/test_calibration.py b/tests/kernel_tuning/test_calibration.py new file mode 100644 index 0000000..1bc426c --- /dev/null +++ b/tests/kernel_tuning/test_calibration.py @@ -0,0 +1,330 @@ +"""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.recorder import CaptureSession # noqa: E402 +from scripts.kernel_tuning.discovery.tasks import render_model # noqa: E402 +from scripts.kernel_tuning.discovery.tensors import flatten # noqa: E402 + + +def record(root: Path, fn, *, budget: int = 1024 * 1024) -> ModelCapture: + session = CaptureSession("test", root / "raw", device_type="cpu", fixture_bytes=budget) + try: + with torch.inference_mode(), session.observe(object()): + fn() + finally: + session.close(complete=True) + return ModelCapture.load(session.root, require_ready=False) + + +def identity(name: str, *args, **kwargs) -> tuple[dict, list]: + tensors = [] + recipe = flatten((args, kwargs), tensors) + return { + "kind": "aten", + "name": name, + "arguments": recipe, + "input_dtypes": [str(t.dtype).removeprefix("torch.") for t in tensors], + "output_dtypes": ["float32"], + "mutates": [], + }, tensors + + +@pytest.fixture +def calibrated(tmp_path: Path) -> 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 = record(tmp_path, lambda: (a @ b, c @ b)) + return calibrate(raw, tmp_path / "calibrated") + + +def test_semantics_do_not_relax_copies_rng_mutation_integer_or_fp64() -> None: + x = torch.ones(3, 3) + for name in ( + "aten.mm.default", + "aten.linear.default", + "aten.sum.default", + "aten.scaled_dot_product_attention.default", + ): + spec, _ = identity(name, x) + assert policy(spec) == "aten" + for name in ("aten.clone.default", "aten.copy_.default", "aten.relu.default", "aten.index.Tensor"): + spec, _ = identity(name, x) + assert policy(spec) is None + spec, _ = identity("aten.mm.default", x) + for extra in ( + {"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) -> None: + x = torch.ones(3, 3) + capture = record(tmp_path, lambda: 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) -> None: + x = torch.ones(8, 8) + capture = record(tmp_path, lambda: x @ x, budget=0) + assert capture.operators[0]["status"] == "blocked" + assert "fixture" in capture.operators[0]["reason"] + 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 + + +@pytest.mark.parametrize( + "causal,gqa,masked", [(False, False, True), (True, False, False), (False, True, True)] +) +def test_fp64_attention_matches_math_reference(causal: bool, gqa: bool, masked: bool) -> None: + generator = torch.Generator().manual_seed(2) + q = torch.randn(1, 4, 7, 8, generator=generator) + k = torch.randn(1, 2 if gqa else 4, 9, 8, generator=generator) + v = torch.randn(1, 2 if gqa else 4, 9, 8, generator=generator) + mask = torch.ones(7, 9, dtype=torch.bool) if masked else None + if mask is not None: + mask[0] = False + mask[1:, -2:] = False + spec, tensors = identity( + "aten.scaled_dot_product_attention.default", + q, + k, + v, + attn_mask=mask, + dropout_p=0.0, + is_causal=causal, + scale=0.25, + enable_gqa=gqa, + ) + expected = torch.nn.functional.scaled_dot_product_attention( + q.double(), k.double(), v.double(), attn_mask=mask, is_causal=causal, scale=0.25, enable_gqa=gqa + ) + torch.testing.assert_close(high_precision(spec, tensors)[0], expected, atol=1e-12, rtol=1e-12) + + +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) -> None: + q = torch.ones(1, 1, 3, 4) + mask = torch.tensor([[True, False, True]]).expand(3, 3) + capture = record(tmp_path, lambda: torch.ops.aten.scaled_dot_product_attention.default(q, q, q, mask)) + # Current Torch observes the public SDPA overload atomically. + op = next(op for op in capture.operators if op["name"] == "aten.scaled_dot_product_attention.default") + 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]) + + +def test_rms_norm_default_epsilon_preserves_input_dtype_semantics() -> None: + value = torch.full((2, 4), 0.001, dtype=torch.bfloat16) + spec, tensors = identity("aten.rms_norm.default", value, [4]) + actual = high_precision(spec, tensors)[0] + expected = torch.nn.functional.rms_norm(value.double(), [4], eps=torch.finfo(value.dtype).eps) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +def test_calibrate_cli_writes_a_new_capture_before_generation(tmp_path: Path, capsys) -> None: + from scripts.kernel_tuning.__main__ import main + + value = torch.ones(2, 3) + weight = torch.ones(3, 4) + raw = record(tmp_path, lambda: value @ weight) + 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_discovery.py b/tests/kernel_tuning/test_discovery.py index a2e6ebb..a448ec0 100644 --- a/tests/kernel_tuning/test_discovery.py +++ b/tests/kernel_tuning/test_discovery.py @@ -47,6 +47,8 @@ def generated_function(task: TaskPackage): def test_discovers_uncatalogued_ops_all_shapes_and_call_counts(tmp_path: Path) -> None: + from benchmarks.kernel_tuning.calibration import calibrate + values = [torch.ones(n, 4) for n in range(1, 12)] weight = torch.ones(4, 3) @@ -59,12 +61,12 @@ def forward(): assert set(ready) == {"aten.mm.default", "aten.relu.default"} assert len(ready["aten.mm.default"]["workloads"]) == 11 assert sum(case["count"] for case in ready["aten.mm.default"]["workloads"]) == 12 - tasks = render_model(capture, tmp_path / "tasks") + tasks = render_model(calibrate(capture, tmp_path / "calibrated"), tmp_path / "tasks") assert len(tasks) == 2 assert all(len(task.workloads) == 11 for task in tasks) for task in tasks: replay = read_json(task.root / "replay.json") - assert task.settings.precision.mode == "bit_exact" + assert task.settings.precision.mode == ("tolerance" if "mm." in replay["operator"] else "bit_exact") for case in replay["cases"].values(): inputs = materialize(case["inputs"], task.root, 0) result = generated_function(task)(*inputs) @@ -387,9 +389,11 @@ def test_replay_adapter_checks_output_alias_and_undeclared_mutation_on_cpu(tmp_p def test_missing_workload_and_unbound_output_axes_are_rejected(tmp_path: Path) -> None: + from benchmarks.kernel_tuning.calibration import calibrate + value = torch.ones(4, 2) capture = capture_calls(tmp_path, lambda: value.sum(1)) - (task,) = render_model(capture, tmp_path / "tasks") + (task,) = render_model(calibrate(capture, tmp_path / "calibrated"), tmp_path / "tasks") replay = read_json(task.root / "replay.json") replay["cases"] = {} atomic_json(task.root / "replay.json", replay) diff --git a/tests/kernel_tuning/test_replay_extensions.py b/tests/kernel_tuning/test_replay_extensions.py index 46dd7cd..2cbcfd9 100644 --- a/tests/kernel_tuning/test_replay_extensions.py +++ b/tests/kernel_tuning/test_replay_extensions.py @@ -114,6 +114,8 @@ def test_unsupported_rng_and_explicit_generators_still_block(tmp_path: Path) -> @pytest.mark.parametrize("style", ["omitted", "positional", "keyword"]) def test_zero_dropout_sdpa_is_not_a_random_call(tmp_path: Path, style: str) -> None: + from benchmarks.kernel_tuning.calibration import calibrate + q = torch.randn(1, 2, 8, 16) state = torch.get_rng_state().clone() outputs = [] @@ -136,7 +138,7 @@ def forward(): assert sum(c["count"] for c in capture.operators[0]["workloads"]) == 2 assert torch.equal(state, torch.get_rng_state()) assert torch.equal(outputs[0], outputs[1]) - (task,) = render_model(capture, tmp_path / "tasks") + (task,) = render_model(calibrate(capture, tmp_path / "calibrated"), tmp_path / "tasks") assert "rng" not in read_json(task.root / "replay.json") @@ -254,6 +256,7 @@ def test_cuda_rng_and_float64_baselines_pass_real_evaluator(tmp_path: Path, oper pytest.skip("requires CUDA") pytest.importorskip("flashinfer_bench") pytest.importorskip("triton") + from benchmarks.kernel_tuning.calibration import calibrate from benchmarks.kernel_tuning.replay_adapter import ReplayAdapter from scripts.kernel_tuning.contracts import validate_measurement @@ -273,7 +276,7 @@ def test_cuda_rng_and_float64_baselines_pass_real_evaluator(tmp_path: Path, oper session.close(complete=True) torch.cuda.set_rng_state(state) tasks = render_model( - ModelCapture.load(session.root), + calibrate(ModelCapture.load(session.root), tmp_path / "calibrated"), tmp_path / "tasks", overrides={"timing.warmup": 1, "timing.iterations": 2, "timing.trials": 1}, ) From 98079f92a0df1fab6e96f6e5d9207e5b3dc9134b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=82=96=E4=B8=80=E6=B6=B5?= <23371301@buaa.edu.cn> Date: Wed, 7 Oct 2026 23:35:32 +0800 Subject: [PATCH 10/44] feat: discover model operators from end-to-end engine inference MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Operator discovery used to observe only a fixed list of policy and decoder method names. Recurrent decoding (decoder.decode), RL sampling, serving pre/post-processing, engine output packing and any work on executor threads ran outside that scope and were silently missing from complete captures. Discovery now observes everything below the model-agnostic inference entries (EngineCore, GenerationBackend, policy collate and serving adapters), with per-thread scope so async and batched serving workers are covered. Model weights read outside every entry, or a run without inference, are capture errors instead of a smaller operator set. - Add discovery.drive, which builds the policy and engine with the serving arguments and replays recorded requests or observations, one session per recurrent episode, plus benchmark-backed input exporters for pi05, the Qwen navigation profiles and StreamVLN. - Treat zero-dropout SDPA backend overloads like the composite call and skip their uninitialized philox outputs during replay. - Classify host-side preprocessing with unrepresentable arguments as host_only. Signed-off-by: 肖一涵 <23371301@buaa.edu.cn> --- benchmarks/kernel_tuning/discovery_inputs.py | 176 +++++++++++++++++ benchmarks/kernel_tuning/replay_adapter.py | 4 + docs/proposals/0008-kernel-tuning.md | 38 +++- scripts/kernel_tuning/discovery/contracts.py | 9 + scripts/kernel_tuning/discovery/drive.py | 117 +++++++++++ scripts/kernel_tuning/discovery/recorder.py | 182 ++++++++++++++--- scripts/kernel_tuning/discovery/rng.py | 37 +++- scripts/kernel_tuning/discovery/tasks.py | 2 + tests/kernel_tuning/test_discovery.py | 186 ++++++++++++++---- tests/kernel_tuning/test_replay_extensions.py | 33 ++++ 10 files changed, 701 insertions(+), 83 deletions(-) create mode 100644 benchmarks/kernel_tuning/discovery_inputs.py create mode 100644 scripts/kernel_tuning/discovery/drive.py diff --git a/benchmarks/kernel_tuning/discovery_inputs.py b/benchmarks/kernel_tuning/discovery_inputs.py new file mode 100644 index 0000000..aa8a95e --- /dev/null +++ b/benchmarks/kernel_tuning/discovery_inputs.py @@ -0,0 +1,176 @@ +"""Export benchmark data as end-to-end inputs for model operator discovery. + +The output is read by ``scripts.kernel_tuning.discovery.drive``. Data selection +reuses each benchmark's own loader, so discovery 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 discovery 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/replay_adapter.py b/benchmarks/kernel_tuning/replay_adapter.py index fb4cb48..3557fdf 100644 --- a/benchmarks/kernel_tuning/replay_adapter.py +++ b/benchmarks/kernel_tuning/replay_adapter.py @@ -200,6 +200,10 @@ def _verify_transition( 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]: + # Unspecified bytes (unused dropout RNG state) carry no result to validate. + undefined = set(self.replay.get("undefined_outputs", ())) + actual = [value for i, value in enumerate(actual) if i not in undefined] + expected = [value for i, value in enumerate(expected) if i not in undefined] if getattr(self, "numerical", None) is None: return compare_outputs(actual, expected, self.cfg.precision) from .calibration import compare_calibrated, high_precision diff --git a/docs/proposals/0008-kernel-tuning.md b/docs/proposals/0008-kernel-tuning.md index 61487df..f99ae03 100644 --- a/docs/proposals/0008-kernel-tuning.md +++ b/docs/proposals/0008-kernel-tuning.md @@ -381,6 +381,18 @@ around existing backend entrypoints collect complete input/output metadata. Existing compound kernels remain single operators; this change introduces no new fusion, quantization, runtime registration, or engine API. +The observed scope is end-to-end inference, not a list of policy methods. It +opens at the model-agnostic entries every deployment path uses: `EngineCore` +construction and execution, the `GenerationBackend` rollout calls, the policy's +`collate`, and its serving adapter's construction and `infer`/`infer_batch`. +Everything below an entry is recorded, including decoder-specific recurrent +decoding, RL sampling, request preprocessing and action postprocessing. Scope is +tracked per thread and the observer is installed on whichever thread enters, so +the async engine's executor and batched-serving workers are covered. Model +weights read by arithmetic outside every entry (a driver calling policy +internals directly) and a run without any entry are capture errors, so a +bypassed path cannot silently produce a smaller operator set. + The capture freezes the model type, supplied checkpoint/configuration descriptors, Torch version, numerical settings, operator overloads, tensor shapes/strides, argument trees, aliasing and mutation contracts, and every observed workload (no top-eight @@ -421,15 +433,22 @@ CPU/fake-agent passes do not establish GPU correctness or performance. #### Commands and artifacts -Run preparation in the selected model's own runtime. `PREPARATION.py` below is -an existing driver that constructs the specified policy through `make_policy` -and exercises its forward/decode path, with compilation disabled. The capture -does not launch an extra forward or change the attention backend. Canonical -policy IDs are those accepted by the factory, for example `pi05`, `streamvln`, -or `qwen2.5-vl-3b-r2r-low-level`; catalog group names are not model IDs. +Run preparation in the selected model's own runtime. The bundled driver +`scripts.kernel_tuning.discovery.drive` constructs the policy and engine with the +serving transports' arguments and replays recorded inputs: serving requests go +through the policy's adapter as HTTP requests would, engine observations through +`collate` and `EngineCore.execute`, with one session per recurrent episode. +`benchmarks/kernel_tuning/discovery_inputs.py` exports such inputs from the pi05, +Qwen navigation and StreamVLN benchmark data; any script that drives one of the +entries above may be used instead. Compilation must be disabled explicitly (for +example `--no-cuda-graph` for StreamVLN, whose graph mode compiles its prefill). +The capture does not launch an extra forward or change the attention backend. +Canonical policy IDs are those accepted by the factory, for example `pi05`, +`streamvln`, or `qwen2.5-vl-3b-r2r-low-level`; catalog group names are not model IDs. ```bash -/absolute/runtime/python -m scripts.kernel_tuning capture --model pi05 --fixture-bytes 8589934592 --output results/kernel_tuning/captures/pi05 -- PREPARATION.py +/absolute/runtime/python -m benchmarks.kernel_tuning.discovery_inputs pi05 --config PI05_BENCHMARK.yaml --output pi05-inputs.pt --serving-config pi05-serving.json +/absolute/runtime/python -m scripts.kernel_tuning capture --model pi05 --fixture-bytes 8589934592 --output results/kernel_tuning/captures/pi05 -- -m scripts.kernel_tuning.discovery.drive --policy pi05 --checkpoint CHECKPOINT --adapter-config pi05-serving.json --inputs pi05-inputs.pt uv run --project scripts/kernel_tuning python -m scripts.kernel_tuning generate --model pi05 --captured results/kernel_tuning/captures/pi05 --list /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 tune-all --model pi05 --captured results/kernel_tuning/captures/pi05-calibrated --agent 'HARNESS/MODEL:EFFORT' --set evaluator_python=/absolute/runtime/python @@ -493,6 +512,11 @@ default) does not require an RNG replay contract. Discovery resolves positional, keyword and omitted dropout arguments without changing the invocation, backend, mask, scale or attention semantics. Nonzero dropout remains a coverage gap until its own state contract is implemented; other nondeterministic tags still apply. +The same rule covers the backend overloads (flash, efficient, cuDNN, CPU) that +model code reaches under `no_grad`. Without dropout their `philox_seed` and +`philox_offset` outputs are left uninitialized, so replay records them as +`undefined_outputs` and checks every other output exactly; no other operator may +declare undefined outputs. Replay tasks declare `definition_schema=embodiinfer-replay-v1`. Their local `ReplayDefinition` extends the pinned FlashInfer tensor schema with `float64`, diff --git a/scripts/kernel_tuning/discovery/contracts.py b/scripts/kernel_tuning/discovery/contracts.py index cf5fdac..0030a36 100644 --- a/scripts/kernel_tuning/discovery/contracts.py +++ b/scripts/kernel_tuning/discovery/contracts.py @@ -69,6 +69,15 @@ def validate_replay( type(index) is not int or not 0 <= index < len(definition["inputs"]) for index in replay["mutates"] ): raise ContractError("Invalid declared input mutations") + undefined = replay.get("undefined_outputs", []) + if ( + not isinstance(undefined, list) + or (undefined and not str(replay.get("operator")).startswith("aten._scaled_dot_product_")) + or len(set(undefined)) != len(undefined) + or len(undefined) >= len(definition["outputs"]) + or any(type(index) is not int or not 0 <= index < len(definition["outputs"]) for index in undefined) + ): + raise ContractError("Only unused attention dropout state may be excluded from output checks") for trace in traces: work = trace["workload"] case = replay["cases"][work["uuid"]] diff --git a/scripts/kernel_tuning/discovery/drive.py b/scripts/kernel_tuning/discovery/drive.py new file mode 100644 index 0000000..dfef885 --- /dev/null +++ b/scripts/kernel_tuning/discovery/drive.py @@ -0,0 +1,117 @@ +"""Run recorded inputs through one model's production inference path. + +This is the preparation driver for model operator discovery:: + + python -m scripts.kernel_tuning capture --model POLICY --output CAPTURE -- \\ + -m scripts.kernel_tuning.discovery.drive --policy POLICY --inputs INPUTS.pt [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 capture records whatever this +inference executes; 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/recorder.py b/scripts/kernel_tuning/discovery/recorder.py index 2a3e465..7ae88d9 100644 --- a/scripts/kernel_tuning/discovery/recorder.py +++ b/scripts/kernel_tuning/discovery/recorder.py @@ -8,6 +8,7 @@ import inspect import runpy import sys +import threading from collections import Counter from collections.abc import Iterator from contextlib import contextmanager @@ -17,13 +18,14 @@ import torch from torch.utils._python_dispatch import TorchDispatchMode +from torch.utils._pytree import tree_leaves from ..artifacts import REPOSITORY, atomic_json from ..contracts import ContractError, digest, tree_hashes from ..generate import _git_revision, vendor_sources from .contracts import SCHEMA from .numerics import policy -from .rng import contract_for, requires_rng_adapter +from .rng import contract_for, requires_rng_adapter, undefined_outputs from .rng import snapshot as rng_snapshot from .tensors import describe_inputs, extent, flatten, storage_key, tensor_outputs, tensor_spec @@ -55,6 +57,18 @@ } _ALLOCATIONS = {"empty", "empty_like", "empty_strided", "new_empty", "new_empty_strided"} _HOST = {"item", "_local_scalar_dense", "is_nonzero"} +# Weight loading, casting and device placement are not inference arithmetic. +_COPIES = {"_to_copy", "copy_", "clone", "contiguous", "lift_fresh", "_copy_from", "_copy_from_and_resize"} +# Production inference enters through the model-agnostic engine contract. Every +# operator running below one of these calls, on any thread, belongs to the model. +_ENTRIES = ( + ("embodiinfer.engine.core", "EngineCore", ("__init__", "execute", "execute_pipelined", "_pipeline_step")), + ( + "embodiinfer.engine.rollout.generation_backend", + "GenerationBackend", + ("generate", "generate_with_logprob", "sample_group", "best_of_n"), + ), +) _QUANTIZED_ATEN = { "_scaled_mm", "_weight_int8pack_mm", @@ -112,7 +126,10 @@ def __init__(self, session: CaptureSession): def __torch_dispatch__(self, func: Any, types: Any, args: tuple = (), kwargs: dict | None = None) -> Any: session = self.session kwargs = kwargs or {} - if not session.depth or session.suppressed: + if session.suppressed: + return func(*args, **kwargs) + if not session.depth: + session.outside(func, args, kwargs) return func(*args, **kwargs) name = str(func) short = name.split(".")[1] if name.startswith("aten.") else name @@ -129,8 +146,11 @@ def __torch_dispatch__(self, func: Any, types: Any, args: tuple = (), kwargs: di value = args[index] if index < len(args) else kwargs.get(argument.name) writes.append(value) quantized = name.split(".")[0] in {"quantized", "quantized_decomposed"} or short in _QUANTIZED_ATEN + origin = {"kind": "aten", "name": name, "schema": str(func._schema)} + if undefined := undefined_outputs(func, args, kwargs): + origin["undefined_outputs"] = undefined return session.call( - {"kind": "aten", "name": name, "schema": str(func._schema)}, + origin, func, args, kwargs, @@ -167,7 +187,11 @@ def __init__( self.fixture_bytes, self.used_bytes, self.device_type = fixture_bytes, 0, device_type self.target: dict[str, Any] = {} self.instance: Any = None - self.depth, self.suppressed = 0, 0 + # Dispatch modes are thread-local, so inference scope is tracked per thread. + self._local = threading.local() + self.entries = 0 + self.bypassed: Counter = Counter() + self._weights: tuple[int, frozenset[int]] = (0, frozenset()) self.operators: dict[str, dict[str, Any]] = {} self.errors: list[str] = [] self.replacements: list[tuple[Any, str, Any, bool]] = [] @@ -196,6 +220,85 @@ def _check_numerics(self) -> None: self.errors.append("Numerical settings changed during model preparation") self._numerics_bound = True + @property + def depth(self) -> int: + return getattr(self._local, "depth", 0) + + @depth.setter + def depth(self, value: int) -> None: + self._local.depth = value + + @property + def suppressed(self) -> int: + return getattr(self._local, "suppressed", 0) + + @suppressed.setter + def suppressed(self, value: int) -> None: + self._local.suppressed = value + + @contextmanager + def inference(self) -> Iterator[None]: + """Observe every operator below one inference call on the current thread.""" + observer = None + if not getattr(self._local, "observing", False): + observer = _Observer(self) + observer.__enter__() + self._local.observing = True + if not self.depth: + self.entries += 1 + self.depth += 1 + try: + yield + finally: + self.depth -= 1 + if observer is not None: + self._local.observing = False + observer.__exit__(None, None, None) + + @contextmanager + def watch(self) -> Iterator[None]: + """Watch the driver thread outside inference to reject bypassed model arithmetic.""" + self._local.observing = True + try: + with _Observer(self): + yield + finally: + self._local.observing = False + + def _entry(self, function: Any) -> Any: + @functools.wraps(function) + def wrapper(*args: Any, **kwargs: Any) -> Any: + with self.inference(): + return function(*args, **kwargs) + + return wrapper + + def _weight_storages(self) -> frozenset[int]: + tensors = [*self.instance.parameters(), *self.instance.buffers()] + if not tensors: + return frozenset() + probe = tensors[0].untyped_storage().data_ptr() + # Moving or casting the model replaces storages; refresh only then. + if probe != self._weights[0]: + self._weights = (probe, frozenset(t.untyped_storage().data_ptr() for t in tensors) - {0}) + return self._weights[1] + + def outside(self, func: Any, args: tuple, kwargs: dict) -> None: + """Record model-weight arithmetic that bypassed every inference entry.""" + if not isinstance(self.instance, torch.nn.Module): + return + name = str(func) + short = name.split(".")[1] if name.startswith("aten.") else name + if short in _VIEWS | _ALLOCATIONS | _HOST | _COPIES: + return + with self.silence(): + weights = self._weight_storages() + if any( + isinstance(t, torch.Tensor) and t.untyped_storage().data_ptr() in weights + for t in tree_leaves((args, kwargs)) + ): + self.bypassed[name] += 1 + @contextmanager def silence(self) -> Iterator[None]: """Prevent instrumentation and atomic/quantized internals from becoming tasks.""" @@ -225,12 +328,8 @@ def bind(self, instance: Any, config: dict[str, Any]) -> None: def observe(self, instance: Any, **config: Any) -> Iterator[None]: """Observe the existing preparation call without running it a second time.""" self.bind(instance, config) - self.depth += 1 - try: - with _Observer(self): - yield - finally: - self.depth -= 1 + with self.inference(): + yield def _snapshot( self, key: str, tensors: list[torch.Tensor], specs: list[dict[str, Any]], *, real_inputs: bool = False @@ -342,6 +441,7 @@ def call( identity = dict(origin) case: dict[str, Any] = {"inputs": [], "outputs": []} status = classification or "ready" + explicit = reason try: recipe = flatten((args, kwargs), tensors) if any(t.is_quantized for t in tensors): @@ -375,11 +475,24 @@ def call( ("excluded_quantized" if classification == "excluded_quantized" else "blocked"), str(exc), ) + if explicit is None and classification in {"metadata_only", "host_only"}: + status = classification # Not a task, so no replay recipe is needed. try: result = function(*args, **kwargs) except Exception as exc: self.errors.append(f"{origin['name']} failed during preparation: {type(exc).__name__}: {exc}") raise + if ( + status == "blocked" + and explicit is None + and rng is None + and not any( + isinstance(t, torch.Tensor) and t.device.type == self.device_type + for t in tree_leaves((args, kwargs, result)) + ) + ): + # Host preprocessing with an unrepresentable argument is not a device task. + status = "host_only" try: if rng is not None and status != "blocked": case["rng"]["after"] = self._rng_snapshot(rng["devices"]) @@ -524,26 +637,29 @@ def construct(name: str, **kwargs: Any) -> Any: raise ContractError("Disable torch.compile explicitly before model discovery") policy = original_factory(name, **kwargs) self.bind(policy, kwargs) - for holder, methods in ( - (policy, ("encode_prefix", "denoise_step", "prepare_prefix", "forward")), - (policy.decoder, ("produce_chunk", "integrate", "generate_tokens", "init_state")), - ): - for method in methods: - if not hasattr(holder, method): - continue - original = getattr(holder, method) - - def execute(*args: Any, _fn: Any = original, **kw: Any) -> Any: - self.depth += 1 - try: - return _fn(*args, **kw) - finally: - self.depth -= 1 - - self._replace(holder, method, execute) + # Request preprocessing and serving adapters belong to end-to-end inference. + if callable(getattr(policy, "collate", None)): + self._replace(policy, "collate", self._entry(policy.collate)) + build = getattr(policy, "build_serving_adapter", None) + if callable(build): + + @functools.wraps(build) + def serving(*args: Any, **kw: Any) -> Any: + with self.inference(): + adapter = build(*args, **kw) + for method in ("infer", "infer_batch"): + if callable(getattr(adapter, method, None)): + self._replace(adapter, method, self._entry(getattr(adapter, method))) + return adapter + + self._replace(policy, "build_serving_adapter", serving) return policy self._replace_function(original_factory, construct) + for module, owner, methods in _ENTRIES: + holder = getattr(importlib.import_module(module), owner) + for method in methods: + self._replace(holder, method, self._entry(getattr(holder, method))) def compile_guard(*args: Any, **kwargs: Any) -> Any: raise ContractError("torch.compile is unsupported during model preparation; disable compilation") @@ -608,6 +724,16 @@ def close(self, *, complete: bool) -> None: self.replacements.clear() if not self.target: self.errors.append("No target policy was bound through make_policy or observe") + elif not self.entries: + self.errors.append( + "No end-to-end inference ran; drive the model through EngineCore, " + "GenerationBackend or its serving adapter" + ) + for name, count in sorted(self.bypassed.items()): + self.errors.append( + f"{name} read model weights outside an inference entry ({count} calls); " + "drive the model through EngineCore, GenerationBackend or its serving adapter" + ) operators = list(self.operators.values()) atomic_json(self.root / "operators.json", operators) coverage = dict(Counter(op["status"] for op in operators)) @@ -643,7 +769,7 @@ def run( complete = False try: session.install() - with _Observer(session): + with session.watch(): if command[0] == "-m": sys.argv = command[1:] runpy.run_module(command[1], run_name="__main__", alter_sys=True) diff --git a/scripts/kernel_tuning/discovery/rng.py b/scripts/kernel_tuning/discovery/rng.py index a82289d..f0488d3 100644 --- a/scripts/kernel_tuning/discovery/rng.py +++ b/scripts/kernel_tuning/discovery/rng.py @@ -13,21 +13,40 @@ from .contracts import RNG_OVERLOADS +def _zero_dropout_attention(function: Any, args: tuple[Any, ...], kwargs: dict[str, Any]) -> bool: + # The composite call decomposes to a backend overload (flash, efficient, + # cuDNN, CPU) before Python dispatch; none draws random state without dropout. + name = str(function) + if name != "aten.scaled_dot_product_attention.default" and not name.startswith( + "aten._scaled_dot_product_" + ): + return False + for index, argument in enumerate(function._schema.arguments): + if argument.name == "dropout_p": + dropout = kwargs.get(argument.name, args[index] if index < len(args) else argument.default_value) + return type(dropout) in (int, float) and dropout == 0 + return False + + def requires_rng_adapter(function: Any, args: tuple[Any, ...], kwargs: dict[str, Any]) -> bool: """Treat seeded tags as conditional for SDPA with explicitly zero/default dropout.""" tags = {tag for tag in getattr(function, "tags", ()) if "nondeterministic" in str(tag)} - if str(function) == "aten.scaled_dot_product_attention.default": - for index, argument in enumerate(function._schema.arguments): - if argument.name == "dropout_p": - dropout = kwargs.get( - argument.name, args[index] if index < len(args) else argument.default_value - ) - if type(dropout) in (int, float) and dropout == 0: - tags.discard(torch.Tag.nondeterministic_seeded) - break + if _zero_dropout_attention(function, args, kwargs): + tags.discard(torch.Tag.nondeterministic_seeded) return bool(tags) +def undefined_outputs(function: Any, args: tuple[Any, ...], kwargs: dict[str, Any]) -> list[int]: + """Outputs whose bytes are unspecified: dropout RNG state left uninitialized without dropout.""" + if not _zero_dropout_attention(function, args, kwargs): + return [] + return [ + i + for i, value in enumerate(function._schema.returns) + if value.name in {"philox_seed", "philox_offset"} + ] + + def contract_for(name: str, kwargs: dict[str, Any]) -> dict[str, Any] | None: """Recognize randn with the default CPU/CUDA generator; other RNG calls stay gaps.""" if name not in RNG_OVERLOADS or kwargs.get("generator") is not None: diff --git a/scripts/kernel_tuning/discovery/tasks.py b/scripts/kernel_tuning/discovery/tasks.py index 9412785..966c523 100644 --- a/scripts/kernel_tuning/discovery/tasks.py +++ b/scripts/kernel_tuning/discovery/tasks.py @@ -230,6 +230,8 @@ def _plan( } if "rng" in identity: replay["rng"] = identity["rng"] + if "undefined_outputs" in identity: + replay["undefined_outputs"] = identity["undefined_outputs"] if numerical is not None: replay["numerical_identity"] = identity sources.update( diff --git a/tests/kernel_tuning/test_discovery.py b/tests/kernel_tuning/test_discovery.py index a448ec0..7233ec1 100644 --- a/tests/kernel_tuning/test_discovery.py +++ b/tests/kernel_tuning/test_discovery.py @@ -2,6 +2,7 @@ from __future__ import annotations +import contextlib from pathlib import Path import pytest @@ -149,6 +150,21 @@ def test_host_to_device_copies_keep_required_input_fixtures(tmp_path: Path) -> N assert case["inputs"][0].get("fixture") +@pytest.mark.parametrize("device", ["cpu", "meta"]) +def test_unrepresentable_host_preprocessing_is_not_a_device_gap(tmp_path: Path, device: str) -> None: + # "meta" stands in for the target device; serving adapters check requests on the host. + session = CaptureSession("mock_flow_vla", tmp_path / "capture", device_type="meta") + try: + with session.observe(object(), checkpoint="test-revision"): + torch.ones(2, device=device).ne(float("inf")) # torch.isfinite's decomposition + finally: + session.close(complete=True) + capture = ModelCapture.load(session.root, require_ready=False) + (op,) = [op for op in capture.operators if op["name"] == "aten.ne.Scalar"] + # A device call with the same argument still needs a replay recipe. + assert op["status"] == ("host_only" if device == "cpu" else "blocked") + + def test_quantized_scope_includes_dynamic_unregistered_module_and_float_fallback(tmp_path: Path) -> None: from embodiinfer.layers.linear import INT8Config from embodiinfer.models import linear @@ -302,67 +318,159 @@ def test_model_cli_rejects_partial_selection_and_legacy_capture(tmp_path: Path, assert not (tmp_path / "runs").exists() -def test_factory_driver_is_scoped_and_does_not_execute_an_extra_forward(tmp_path: Path, monkeypatch) -> None: - import sys - from types import SimpleNamespace +_ENGINE_DRIVER = """ +from concurrent.futures import ThreadPoolExecutor + +import torch + +from embodiinfer.engine import EngineConfig, EngineCore, GenerationBackend +from embodiinfer.policies.factory import make_policy +from embodiinfer.types import Observation + +torch.manual_seed(0) +policy = make_policy("mock_flow_vla", preset="tiny") +cfg = policy.config +observation = 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,)), +) +x = torch.ones(3) +x.digamma() # driver arithmetic outside the model is not a model operator +core = EngineCore(policy, EngineConfig(device="cpu", use_cuda_graph=False, num_steps=2)) +MODE +""" +_ENGINE_MODES = { + # The async engine and batched serving execute the engine on a worker thread. + "thread": "with ThreadPoolExecutor(1) as pool:\n" + " pool.submit(core.execute, policy.collate([observation], ['r0'])).result()\n", + # The RL rollout path samples through decoder methods the engine never calls. + "rollout": "GenerationBackend(core).generate_with_logprob([observation], num_steps=2)\n", + # Benchmark-style drivers that call policy internals bypass the engine. + "bypass": "batch = policy.collate([observation], ['r0']).to('cpu', torch.float32)\n" + "policy.encode_prefix(batch)\n", +} + + +def _engine_capture(tmp_path: Path, monkeypatch, mode: str) -> ModelCapture: + from scripts.kernel_tuning.discovery import recorder + + # The public CLI always targets CUDA; this test substitutes the observer's + # device classification while exercising the real driver and engine hooks. + monkeypatch.setattr( + recorder, "CaptureSession", lambda *args, **kwargs: CaptureSession(*args, device_type="cpu", **kwargs) + ) + script = tmp_path / f"{mode}.py" + script.write_text(_ENGINE_DRIVER.replace("MODE", _ENGINE_MODES[mode]), encoding="utf-8") + with contextlib.suppress(ContractError): # Gaps are asserted from the published capture. + recorder.run([str(script)], tmp_path / mode, "mock_flow_vla") + return ModelCapture.load(tmp_path / mode, require_ready=False) + + +def _reference_operators(mode: str) -> set[str]: + """Every dispatcher operator of the same inference, collected without the recorder.""" + from torch.utils._python_dispatch import TorchDispatchMode + + from embodiinfer.engine import EngineConfig, EngineCore, GenerationBackend + from embodiinfer.policies.factory import make_policy + from embodiinfer.types import Observation + + names: set[str] = set() + + class Collect(TorchDispatchMode): + def __torch_dispatch__(self, func, types, args=(), kwargs=None): + names.add(str(func)) + return func(*args, **(kwargs or {})) + + torch.manual_seed(0) + policy = make_policy("mock_flow_vla", preset="tiny") + cfg = policy.config + observation = 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,)), + ) + with Collect(): + core = EngineCore(policy, EngineConfig(device="cpu", use_cuda_graph=False, num_steps=2)) + if mode == "thread": + core.execute(policy.collate([observation], ["r0"])) + else: + GenerationBackend(core).generate_with_logprob([observation], num_steps=2) + return names + + +@pytest.mark.parametrize("mode", ["thread", "rollout"]) +def test_engine_inference_is_captured_end_to_end(tmp_path: Path, monkeypatch, mode: str) -> None: + capture = _engine_capture(tmp_path, monkeypatch, mode) + capture.require_ready() + assert capture.manifest["target"]["config"] == {"preset": "tiny"} + captured = {op["name"] for op in capture.operators} + # No policy method list: the whole inference matches an independent collector. + assert captured == _reference_operators(mode) + assert "aten.digamma.default" not in captured + + +def test_drive_replays_recorded_observations_through_the_engine(tmp_path: Path, monkeypatch) -> None: + import json from scripts.kernel_tuning.discovery import recorder + from scripts.kernel_tuning.discovery.drive import save_inputs - from embodiinfer.policies import factory + from embodiinfer.policies.mock import preset_config + from embodiinfer.types import Observation - class FakePolicy(torch.nn.Module): - def __init__(self): - super().__init__() - self.decoder = SimpleNamespace() - self.calls = 0 + cfg = preset_config("tiny") - def forward(self, inputs): - self.calls += 1 - return inputs.relu() + def observation() -> 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,)), + ) - policy = FakePolicy() - monkeypatch.setitem(factory._REGISTRY, "mock_flow_vla", lambda **kwargs: policy) - # The public CLI always targets CUDA; this test substitutes the observer's - # device classification while exercising the real driver and factory hooks. + save_inputs(tmp_path / "inputs.pt", "observations", [[observation(), observation()], [observation()]]) + (tmp_path / "serving.json").write_text( + json.dumps({"policy_kwargs": {"preset": "tiny"}}), encoding="utf-8" + ) monkeypatch.setattr( recorder, "CaptureSession", lambda *args, **kwargs: CaptureSession(*args, device_type="cpu", **kwargs) ) - script = tmp_path / "prepare.py" - script.write_text( - "import torch\nfrom embodiinfer.policies.factory import make_policy\nx = torch.ones(3)\nx.sin() # outside the selected policy\nmodel = make_policy('mock_flow_vla', checkpoint='test')\nmodel(x)\n", - encoding="utf-8", - ) - argv, path = sys.argv, sys.path[:] - recorder.run([str(script)], tmp_path / "capture", "mock_flow_vla") - capture = ModelCapture.load(tmp_path / "capture") - assert [op["name"] for op in capture.operators] == ["aten.relu.default"] - assert policy.calls == 1 - assert "forward" not in vars(policy) - assert sys.argv is argv and sys.path == path - assert capture.manifest["target"]["config"] == {"checkpoint": "test"} - + command = ["-m", "scripts.kernel_tuning.discovery.drive", "--policy", "mock_flow_vla", "--device", "cpu"] + command += ["--adapter-config", str(tmp_path / "serving.json"), "--inputs", str(tmp_path / "inputs.pt")] + recorder.run(command, tmp_path / "capture", "mock_flow_vla") + capture = ModelCapture.load(tmp_path / "capture", "mock_flow_vla") + assert capture.manifest["target"]["config"] == {"preset": "tiny"} + addmm = next(op for op in capture.operators if op["name"] == "aten.addmm.default") + # Three requests, each with a prefix and the default denoise loop, were observed. + assert sum(case["count"] for case in addmm["workloads"]) > 3 + + +def test_model_arithmetic_outside_the_engine_is_a_coverage_gap(tmp_path: Path, monkeypatch) -> None: + capture = _engine_capture(tmp_path, monkeypatch, "bypass") + errors = capture.manifest["errors"] + assert any("read model weights outside an inference entry" in error for error in errors) + with pytest.raises(ContractError, match="outside an inference entry"): + capture.require_ready() -def test_decoder_integrate_entry_is_observed(tmp_path: Path, monkeypatch) -> None: - # Benchmarks may time the flow loop through decoder.integrate rather than produce_chunk. - from types import SimpleNamespace +def test_capture_without_inference_is_rejected(tmp_path: Path, monkeypatch) -> None: from scripts.kernel_tuning.discovery import recorder from embodiinfer.policies import factory - policy = SimpleNamespace(decoder=SimpleNamespace(integrate=lambda x: x.relu())) - monkeypatch.setitem(factory._REGISTRY, "mock_flow_vla", lambda **kwargs: policy) + monkeypatch.setitem(factory._REGISTRY, "mock_flow_vla", lambda **kwargs: torch.nn.Linear(2, 2)) monkeypatch.setattr( recorder, "CaptureSession", lambda *args, **kwargs: CaptureSession(*args, device_type="cpu", **kwargs) ) script = tmp_path / "prepare.py" script.write_text( - "import torch\nfrom embodiinfer.policies.factory import make_policy\nmodel = make_policy('mock_flow_vla', checkpoint='test')\nmodel.decoder.integrate(torch.ones(3))\n", + "from embodiinfer.policies.factory import make_policy\nmake_policy('mock_flow_vla')\n", encoding="utf-8", ) - recorder.run([str(script)], tmp_path / "capture", "mock_flow_vla") - capture = ModelCapture.load(tmp_path / "capture") - assert [op["name"] for op in capture.operators] == ["aten.relu.default"] + with pytest.raises(ContractError): + recorder.run([str(script)], tmp_path / "capture", "mock_flow_vla") + manifest = read_json(tmp_path / "capture" / "manifest.json") + assert any("No end-to-end inference ran" in error for error in manifest["errors"]) def test_replay_adapter_checks_output_alias_and_undeclared_mutation_on_cpu(tmp_path: Path) -> None: diff --git a/tests/kernel_tuning/test_replay_extensions.py b/tests/kernel_tuning/test_replay_extensions.py index 2cbcfd9..1b007f5 100644 --- a/tests/kernel_tuning/test_replay_extensions.py +++ b/tests/kernel_tuning/test_replay_extensions.py @@ -142,6 +142,39 @@ def forward(): assert "rng" not in read_json(task.root / "replay.json") +def test_zero_dropout_sdpa_backend_overload_is_not_a_random_call(tmp_path: Path) -> None: + # Under the engine's no_grad, the functional API reaches the observer as a backend overload. + q = torch.randn(1, 2, 8, 16) + with torch.no_grad(): + capture = record(tmp_path, lambda: torch.nn.functional.scaled_dot_product_attention(q, q, q)) + assert [op["name"] for op in capture.operators] == [ + "aten._scaled_dot_product_flash_attention_for_cpu.default" + ] + capture.require_ready() + + +def test_unused_attention_dropout_state_is_not_compared(tmp_path: Path) -> None: + from types import SimpleNamespace + + from benchmarks.kernel_tuning.replay_adapter import ReplayAdapter + from scripts.kernel_tuning.discovery.rng import undefined_outputs + + op = torch.ops.aten._scaled_dot_product_efficient_attention.default + q = torch.empty(1, 2, 8, 16) + # CUDA leaves philox seed/offset uninitialized when dropout is zero. + assert undefined_outputs(op, (q, q, q, None, False), {}) == [2, 3] + assert undefined_outputs(op, (q, q, q, None, False, 0.1), {}) == [] + assert undefined_outputs(torch.ops.aten.mm.default, (q, q), {}) == [] + + adapter = ReplayAdapter.__new__(ReplayAdapter) + adapter.cfg = SimpleNamespace(precision=SimpleNamespace(mode="bit_exact")) + adapter.replay = {"undefined_outputs": [1]} + value = torch.ones(3) + assert adapter._compare([value, torch.tensor(16)], [value, torch.tensor(0)], []) == (0.0, 0.0) + with pytest.raises(ContractError, match="bit-exact"): + adapter._compare([value + 1, torch.tensor(0)], [value, torch.tensor(0)], []) + + @pytest.mark.parametrize("style", ["positional", "keyword"]) def test_nonzero_dropout_sdpa_remains_a_coverage_gap(tmp_path: Path, style: str) -> None: q = torch.randn(1, 2, 8, 16) From 7967c96a534cd67c91006ba7da84b503128c12fa Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=82=96=E4=B8=80=E6=B6=B5?= <23371301@buaa.edu.cn> Date: Sun, 11 Oct 2026 13:32:26 +0800 Subject: [PATCH 11/44] fix: count native PI0.5 graphs and warm every native benchmark layout MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit With native_inference and cuda_graph enabled, the prefix and denoise-loop graphs are captured by the PI0.5 native runtime, not by EngineCore, so the warmup check reported "warmup captured no graph" and stopped. Count both. The native runtime also captures one graph per language bucket and dense/masked variant. Ten warmup calls did not cover every layout of the 1,600 LIBERO-10 frames, so a capture happened during measurement and the run was rejected after it finished. Warm each layout once before timing, as compare.py already does, and record the warmed layouts in the report. Signed-off-by: 肖一涵 <23371301@buaa.edu.cn> --- benchmarks/pi05-benchmark/benchmark.py | 40 +++++++++++++++++++++++--- 1 file changed, 36 insertions(+), 4 deletions(-) diff --git a/benchmarks/pi05-benchmark/benchmark.py b/benchmarks/pi05-benchmark/benchmark.py index 79d040b..5ea3e73 100644 --- a/benchmarks/pi05-benchmark/benchmark.py +++ b/benchmarks/pi05-benchmark/benchmark.py @@ -387,6 +387,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]], @@ -418,12 +424,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"]) @@ -439,7 +448,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,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) @@ -526,7 +535,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( @@ -537,6 +549,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 From 6efd46623f3bde5246e39d213d699cd13b791f81 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=82=96=E4=B8=80=E6=B6=B5?= <23371301@buaa.edu.cn> Date: Sun, 11 Oct 2026 13:32:37 +0800 Subject: [PATCH 12/44] feat: route operators through a registry with selectable kernel profiles MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit PI0.5 chose between Torch and Triton norm, residual, activation and RoPE by importing embodiinfer.backend.triton directly behind per-call flags, so a tuned kernel had no place to plug in. Add embodiinfer.layers.ops: operators are named by their mathematics (gated_residual, ada_rms_norm, gated_gelu, rotate_qk), each with a Torch reference in embodiinfer.backend.torch.ops and the existing Triton kernel. KernelSelection resolves an operator for the caller's preferred implementation under a kernel_profile: - none (default): the preferred implementation, as before - torch: every registered operator uses its Torch reference - auto / : verified kernels from embodiinfer/backend/triton/tuned/, matched on GPU, compute capability and Torch/Triton versions, with a per-kernel applies() fallback and a source digest check PI0.5 routes its four switch points through the registry and accepts kernel_profile; the benchmark passes it through. prepare() resolves every operator before torch.compile regions so graph tracing only reads a cache. Torch references stay uncompiled. The former _inference_helper wrappers for these helpers were first created while Dynamo traced the prefix encoder, which stores the raw function, so they never ran compiled outside that graph. Compiling them now would change the Torch path's numerics. Verified on Jetson Thor (torch 2.13.0+cu132, triton 3.7.1, BF16, B=1, pi05_libero_finetuned_v044, LIBERO-10): kernel_profile=none is bit-exact against the previous code for 1,600 native-path frames and for 20 frames of the Inductor + CUDA Graph path with TORCHINDUCTOR_DETERMINISTIC=1 and fresh compilation; native-path mean latency 99.35 ms vs 98.47 ms before. Signed-off-by: 肖一涵 <23371301@buaa.edu.cn> --- benchmarks/pi05-benchmark/benchmark.py | 8 +- embodiinfer/backend/torch/ops.py | 52 ++++++ embodiinfer/backend/triton/tuned/__init__.py | 137 +++++++++++++++ embodiinfer/layers/__init__.py | 6 + embodiinfer/layers/ops.py | 175 +++++++++++++++++++ embodiinfer/policies/pi05/modeling_pi05.py | 102 ++++++----- tests/test_kernel_ops.py | 162 +++++++++++++++++ tests/test_pi05_cpu_routes.py | 2 + 8 files changed, 601 insertions(+), 43 deletions(-) create mode 100644 embodiinfer/backend/torch/ops.py create mode 100644 embodiinfer/backend/triton/tuned/__init__.py create mode 100644 embodiinfer/layers/ops.py create mode 100644 tests/test_kernel_ops.py diff --git a/benchmarks/pi05-benchmark/benchmark.py b/benchmarks/pi05-benchmark/benchmark.py index 5ea3e73..255dd58 100644 --- a/benchmarks/pi05-benchmark/benchmark.py +++ b/benchmarks/pi05-benchmark/benchmark.py @@ -500,7 +500,13 @@ def build( compile_backend=config["compile_backend"], **{ 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), diff --git a/embodiinfer/backend/torch/ops.py b/embodiinfer/backend/torch/ops.py new file mode 100644 index 0000000..6a3283c --- /dev/null +++ b/embodiinfer/backend/torch/ops.py @@ -0,0 +1,52 @@ +"""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 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 _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) diff --git a/embodiinfer/backend/triton/tuned/__init__.py b/embodiinfer/backend/triton/tuned/__init__.py new file mode 100644 index 0000000..c5dbc2c --- /dev/null +++ b/embodiinfer/backend/triton/tuned/__init__.py @@ -0,0 +1,137 @@ +"""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 per tuned operator. +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.py", "function": "gated_residual", + "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 +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"] + digest = hashlib.sha256(path.read_bytes()).hexdigest() + if digest != entry["sha256"]: + raise ValueError(f"kernel profile {self.name!r}: {path.name} does not match its recorded sha256") + spec = importlib.util.spec_from_file_location(f"embodiinfer_tuned_{self.name}_{path.stem}", 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) + spec.loader.exec_module(module) + applies = entry.get("applies") + kernel = TunedKernel( + getattr(module, entry["function"]), getattr(module, applies) if applies else None + ) + self._loaded[op] = kernel + return kernel + + +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", +] diff --git a/embodiinfer/layers/__init__.py b/embodiinfer/layers/__init__.py index 709e2ac..caea156 100644 --- a/embodiinfer/layers/__init__.py +++ b/embodiinfer/layers/__init__.py @@ -19,6 +19,7 @@ register_linear, resolve_linear_backend, ) +from .ops import KernelSelection, available_ops, get_op, register_lazy_op, register_op __all__ = [ "AttentionBackend", @@ -36,4 +37,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..1345915 --- /dev/null +++ b/embodiinfer/layers/ops.py @@ -0,0 +1,175 @@ +"""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 + + +for _op, _module in ( + ("gated_residual", "norm"), + ("ada_rms_norm", "norm"), + ("gated_gelu", "activation"), + ("rotate_qk", "rotary"), +): + 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 42c992e..1078ca5 100644 --- a/embodiinfer/policies/pi05/modeling_pi05.py +++ b/embodiinfer/policies/pi05/modeling_pi05.py @@ -42,7 +42,7 @@ import torch.nn.functional as F 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 @@ -82,30 +82,34 @@ 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") + +# ---- 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: @@ -118,11 +122,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) @@ -151,18 +154,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: @@ -303,6 +302,7 @@ def _mlp( use_triton: bool = False, *, fuse_projections: bool = True, + kernels: KernelSelection | None = None, ) -> torch.Tensor: if fuse_projections: gate, up = _fused_linear_pair(x, mlp.gate_proj, mlp.up_proj) @@ -311,20 +311,10 @@ def _mlp( # 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: - from ...backend.triton.activation import gated_gelu - - activated = gated_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) return mlp.down_proj(activated) -def _gated_gelu_impl(gate: torch.Tensor, up: torch.Tensor) -> torch.Tensor: - return F.gelu(gate, approximate="tanh") * up - - @dataclass class Pi05Prefix: """pi0.5 prefix state: our own per-layer (K, V) cache + the pad mask. @@ -436,6 +426,7 @@ def __init__( prefix_cuda_graph: bool = False, denoise_attention: str = "sdpa", prefix_attention: str = "sdpa", + kernel_profile: str = "none", checkpoint_config: Mapping[str, Any] | None = None, tensor_parallel_size: int = 1, tensor_parallel_group=None, @@ -458,6 +449,10 @@ def __init__( self.prefix_cuda_graph = prefix_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 try: from lerobot.policies.pi05.modeling_pi05 import ( create_sinusoidal_pos_embedding, @@ -576,6 +571,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): @@ -585,6 +590,7 @@ def _clear_inference_caches(self) -> None: runtime = getattr(self, "_runtime", None) if runtime is not None: runtime.clear() + self._kernel_selection = None for key in tuple(_COMPILED_PREFIX_ENCODERS): if key[0] == id(self): del _COMPILED_PREFIX_ENCODERS[key] @@ -665,12 +671,15 @@ def _attn_sublayer( q = q.view(B, S, -1, hd).transpose(1, 2) # [B, n_head, S, hd] k = k.view(B, S, -1, hd).transpose(1, 2) # [B, n_kv, S, hd] v = v.view(B, S, -1, hd).transpose(1, 2) - if 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)) if native_expert and self.denoise_attention == "triton": @@ -735,6 +744,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, @@ -753,6 +763,7 @@ def _tower_forward( gate, self.compile_backend == "inductor", fused, + self._kernels, ) residual = hidden h, gate = _rmsnorm( @@ -762,6 +773,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, @@ -771,10 +783,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, @@ -783,6 +797,7 @@ def _tower_forward( self.compile_backend == "inductor", None if modulations is None else modulations[-1], fused, + self._kernels, ) return hidden, collected @@ -837,6 +852,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: @@ -1110,6 +1126,7 @@ def _build_pi05( prefix_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, quantization: str | Mapping[str, Any] | QuantizationConfig | None = None, @@ -1132,6 +1149,7 @@ def _build_pi05( prefix_cuda_graph=prefix_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, quantization=quantization, diff --git a/tests/test_kernel_ops.py b/tests/test_kernel_ops.py new file mode 100644 index 0000000..628939e --- /dev/null +++ b/tests/test_kernel_ops.py @@ -0,0 +1,162 @@ +"""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"): + 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"): + for prefer in ("torch", "triton"): + selection.resolve(op, prefer) + selection.resolve(op, prefer, torch_ops_wrapper) diff --git a/tests/test_pi05_cpu_routes.py b/tests/test_pi05_cpu_routes.py index 53c88b7..6b80421 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 From d4ba7e59ed75727822af76a61ff88c9143a3baf4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=82=96=E4=B8=80=E6=B6=B5?= <23371301@buaa.edu.cn> Date: Sun, 11 Oct 2026 13:32:43 +0800 Subject: [PATCH 13/44] docs: propose hotspot-driven Triton tuning with installable profiles MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Kernel tuning so far optimized isolated operators whose results never reach the model. Propose a model-level loop instead: profile the production configuration, select the regions (including new fusion boundaries) worth tuning, tune Triton against the Torch implementation as baseline, install verified kernels per hardware profile behind kernel_profile, and accept them only under a frozen model-level numerical gate and end-to-end latency. Signed-off-by: 肖一涵 <23371301@buaa.edu.cn> --- docs/proposals/0009-hotspot-kernel-tuning.md | 324 +++++++++++++++++++ 1 file changed, 324 insertions(+) create mode 100644 docs/proposals/0009-hotspot-kernel-tuning.md diff --git a/docs/proposals/0009-hotspot-kernel-tuning.md b/docs/proposals/0009-hotspot-kernel-tuning.md new file mode 100644 index 0000000..a76822a --- /dev/null +++ b/docs/proposals/0009-hotspot-kernel-tuning.md @@ -0,0 +1,324 @@ +# 0009 — Hotspot-driven Triton tuning with installable hardware profiles + +- Status: Draft +- Author: EmbodiInfer maintainers +- Date: 2026-10-09 + +## 1. Summary + +Close the loop from kernel search to model latency. Profile a model in its +production configuration, rank code regions by the time a better kernel could +recover, let a coding agent choose a few regions (including new fusion +boundaries) and write Triton implementations that must beat the model's Torch +implementation, then install the verified kernels in the repository under a +hardware profile (for example `rtx4090`). At runtime an operator registry +selects them behind an off-by-default switch and falls back to the existing +implementation whenever a kernel does not apply. Success is measured as model +end-to-end latency under a frozen model-level numerical gate, not as isolated +kernel speedups. The registry is a model-agnostic `layers/` contract; region +definitions and call sites stay in policies; tuned kernels are `backend/` code. +New work builds on this proposal instead of the per-operator workflow of +[0008](0008-kernel-tuning.md); its Humanize2 flow, evaluator timing protocol and +run archive are reused. + +## 2. Motivation and current gap + +- Tuned kernels never reach the model. 0008 exports verified sources, but + "saved operators require a separate review/integration before runtime use" + and there is no place to plug them in. +- There is no general dispatch point. Attention + (`embodiinfer/layers/attention.py:74`) and quantized linear + (`embodiinfer/layers/linear.py:253`) have registries; norm, activation and + RoPE are chosen by per-model branches that import `embodiinfer.backend.triton` + directly (`embodiinfer/policies/pi05/modeling_pi05.py:122`, `:161`, `:314`, + `:669`; `embodiinfer/models/streamvln/backbone.py:12`; + `embodiinfer/models/qwen25_vl/vision.py:47`). +- Isolated operators are the wrong unit. At B=1, four of five pi0.5 Triton + kernels and all StreamVLN kernels already run near the ~4 µs device floor + (0008, "Captured catalog workloads"). Remaining time sits in a few large + kernels and in chains of small memory-bound kernels between GEMMs; only + changing fusion boundaries reduces kernel count and intermediate traffic. +- Tuning everything is wasteful. Model discovery produced 94 pi0.5 tasks, most + without measurable model-level effect. +- The comparison target is wrong. 0008 compares a candidate with the same ATen + call or the existing Triton kernel. The goal is now for Triton to beat the + model's Torch implementation, which in production may be compiled by Inductor + (`_inference_helper`, `embodiinfer/policies/pi05/modeling_pi05.py:65`). + +## 3. Goals and non-goals + +Goals: + +- One resumable pipeline per model and hardware: profile → select → tune → + install → verify → benchmark. Only a new fusion region needs a code change, + made once and merged like any other change. +- A hotspot report that attributes device time to model regions and estimates + recoverable time; tuning only the selected regions. +- Regions may span several existing operators (new fusion boundaries). +- Baseline: the region's Torch implementation under the model's production + settings (dtype, Inductor when configured, CUDA Graphs). Existing hand-written + Triton kernels are optional starting points, not the baseline. +- Tuned kernels committed in the repository per hardware profile with + provenance, selectable at runtime, reusable by others, and later promotable to + the default Triton implementation through a normal review. + +Non-goals: + +- Making any tuned kernel the default (a separate change with evidence). +- Lower precision, quantization or changed model mathematics. A region may + reorder floating-point work, but not change dtype or matmul precision. +- Automatic merging of agent-written policy changes; they are reviewed as PRs. +- Thor/AGX profiles in the first phase; multi-GPU and batch > 1 tuning. +- Compatibility with the 0008 discovery/catalog task formats. Those commands + stay as they are and are not extended. + +## 4. Design + +### 4.1 Operator registry (`embodiinfer/layers/ops.py`) + +An operator is named by its mathematics (`ada_rms_norm`, `gated_residual`, +`add_rms_norm`, `rotate_qk`), never by a model. A policy-local fused region is +registered by its policy under a namespaced name (`pi05.residual_ada_rms_norm`). + +Every existing Torch/Triton switch point is migrated to the registry once, as an +ordinary code change: the registry call becomes the model's default code path, +with today's implementations registered as `torch` and `triton`. Tuning runs +never edit these call sites; they only add implementations to a profile. + +```python +register_op(op, impl, fn) # "torch" reference, "triton", ... +register_lazy_op(op, impl, module, function) # import on first use +selection = KernelSelection(kernel_profile) # one per policy +selection.prepare() # resolve everything eagerly +selection.resolve(op, prefer) -> Callable # cached lookup +``` + +- Torch references live in `embodiinfer/backend/torch/ops.py`; the existing + Triton kernels are registered unchanged. +- A policy builds its selection on first use on its device and drops it on + refit, device or dtype change. `prepare()` resolves every operator before any + `torch.compile` region, so graph tracing only reads the cache. +- Torch references run uncompiled. pi0.5's per-helper Inductor wrappers were + first created while Dynamo traced the prefix encoder, which stores the raw + function, so they never ran compiled outside that graph; keeping the + references uncompiled preserves those numerics. Compiling them is a separate, + measured change. +- `prefer` keeps the policy's current choice (`torch` or `triton`, from its + existing flags). A selected profile overrides only operators it contains. +- A tuned callable is wrapped with its `applies(*args)` predicate (dtype, + contiguity, shape divisibility, size limits). When it returns false the call + goes to the preferred implementation. StreamVLN's `supports_*` checks + (`embodiinfer/models/streamvln/backbone.py:125`) follow the same pattern. +- With no profile selected, resolution returns the current callables: behavior + is bit-exact with the code before this change. + +Runtime switch: `kernel_profile: none | auto | | torch` in the policy +configuration (default `none`). `auto` picks the profile whose device name, +compute capability and Triton/Torch versions match the current process, else +`none` with one log line. An explicit name that does not match raises. `torch` +forces every registered operator to its Torch reference and is the baseline +configuration for benchmarks. + +### 4.2 Hardware profiles (`embodiinfer/backend/triton/tuned//`) + +Each profile directory holds one module per tuned operator and `profile.json`: + +- match keys: GPU name, compute capability, Torch/Triton/CUDA versions tested; +- per operator: module/function, `applies` predicate, covered workloads and + weights, source digest, kernel-level and model-level gate results, the run id + and the digest of its evidence archive; +- the model, checkpoint revision and configuration the profile was verified on. + +Sources and this small manifest are tracked in Git. Evidence archives stay +outside the tracked tree, as with existing benchmark snapshots. Loading verifies +source digests, so a hand-edited kernel is reported instead of silently used. +Tuned modules are pure tensor functions and import nothing from `models`, +`policies` or `engine`. + +### 4.3 Hotspot analysis (`scripts.kernel_tuning profile`) + +Runs the model's inference driver in its production configuration on recorded +inputs, in two passes: + +1. Timing: production settings, CUDA events, end-to-end and prefill/decode + segments as the model benchmarks measure them. +2. Attribution: `torch.profiler` (CUPTI) with a `record_function` range around + every registry call and every `nn.Module.forward`. Ranges are normalized + across repeated layers (`layers.*`). When per-node kernels inside CUDA Graphs + are not reported, the pass reruns without graphs and records that it did. + +Per region the report gives device time, share of end-to-end time, calls, kernel +count, bytes moved and FLOPs from tensor metadata, and a lower bound +`max(bytes / measured bandwidth, FLOPs / measured peak)` from a device +microbenchmark. Recoverable time is `time - max(lower bound, floor × kernels)`. +Fusion candidates are adjacent producer→consumer chains of memory-bound regions; +their bound counts the eliminated intermediate traffic once. Outputs: +`hotspots.json` and `hotspots.md`. + +### 4.4 Selection and new fusion regions + +A Humanize2 selection flow gives the agent the hotspot report and the model +source. It returns `plan.json`: at most K regions (default 5), each either an +already registered operator or a proposed new fusion region, with reason and +estimated saving. Regions under 1% of end-to-end time are rejected. A person may +edit the plan or supply one with `--plan`. + +A registered operator goes straight to tuning. A new fusion region does not +exist in the model yet, so it needs a one-time code change: the agent writes the +region's Torch reference (a composition of the existing operations), registers +it, and replaces the covered call sites with one registry call. The change is +accepted only if the model with `kernel_profile: none` is bit-exact against the +unchanged model on the evaluation inputs (§6). It then lands in the repository +as a normal reviewed change; from that point the region is an ordinary +registered operator for every hardware profile and later runs reuse it. + +Workloads for every selected operator are captured by recording the registry +operator's real arguments and call counts during a production run. + +### 4.5 Region tuning + +A region task contains the Torch reference, the baseline (that reference under +the policy's compile settings, for example Inductor with the options used by +`_inference_helper`), an optional starting kernel, the captured workloads with +call-count weights, and frozen numerical bounds (§6). The 0008 Humanize2 flow, +CUDA Graph cold-L2 timing, paired promotion rule (≥ 3% weighted, ≤ 5% per +workload) and archive are reused. A candidate must also export `applies`, which +the evaluator checks is true on every workload. Expected model saving is +`Σ count × Δt`. + +### 4.6 Install, verify, benchmark + +`install` writes promoted sources and updates `profile.json`. It refuses an +operator that is not registered by the current source, so a fusion region's code +change must be merged first. `verify` runs the +model with `kernel_profile: ` against `kernel_profile: torch` on the +evaluation set and applies the model-level gate; an operator is kept only if the +whole profile passes. `benchmark` reports Torch, existing Triton and the profile +on the model's benchmark protocol, plus a per-operator ablation. The pipeline +command `optimize --model M --config C --profile P --agent A` runs all stages +with resume under `results/kernel_tuning/optimize//`. + +### 4.7 Alternatives considered + +- Add selection to the 0008 per-ATen discovery: it cannot express fusion and its + baseline is the same ATen call, not the Torch implementation. Rejected. +- Rely on Inductor (`max-autotune`) for fusion: that is the baseline being beaten. + It is reported as an extra benchmark row where it compiles. +- Keep tuned kernels outside the repository and download them: rejected, because + the team wants reviewable, shareable sources that can become defaults. +- Uncached per-call resolution: rejected for launch-path overhead; resolution + is a dictionary lookup after `prepare()` and is skipped entirely by CUDA Graph + replay. +- Apply agent policy patches without review: rejected. + +## 5. Model-agnosticism verdict + +The registry, profile loader and `kernel_profile` switch are model-agnostic +`layers/` code: operators are named by mathematics and resolution uses only the +current device and versions. Tuned kernels are backend code. Region definitions, +call sites and policy-namespaced fused operators are policy-layer and optional; +with `kernel_profile: none` no path changes. The profiler uses module hooks and +registry ranges only; model-specific inputs come from the existing drivers. + +## 6. Losslessness and precision criterion + +Triton kernels reorder accumulation, so bit-exact output is not generally +achievable. This proposal requests explicit approval of calibrated tolerances, +following the frozen thresholds already used for the pi0.5 native path +(`benchmarks/pi05-benchmark/README.md`). + +Exact gates: + +- `kernel_profile: none` is bit-exact with the code before this change. +- Integrating a region with its Torch reference as default is bit-exact against + the unpatched model (`max_i ||Δa_i||_inf = 0`). +- Precision conditions are identical: dtype, `float32_matmul_precision`, TF32 + flags and attention semantics. + +Region level (during tuning): with FP64 reference `R` of the Torch reference, +Torch baseline output `B` and candidate `C` in the production dtype with epsilon +`u`, each floating output must satisfy `|C - R| <= atol + u·|R|` with +`atol = max(u·RMS(R), 2·max(|B - R| - u·|R|, 0), tiny)`, the 0008 formula with +the Torch implementation as `B`. Inputs are the captured fixtures plus seeded +random, zero and alternating-sign cancellation inputs. Integer, boolean, layout, +aliasing and mutation results are exact; non-finite values fail. In words: a +candidate may be at most about twice as wrong as Torch. + +Model level (required for install): + +- Continuous outputs (action chunks): compute an FP32 Torch reference from the + same weights and inputs. Per observation, thresholds for max-abs and RMSE are + twice the error of the `kernel_profile: torch` run against that reference, + with floors `1e-4` and `1e-5`. Thresholds are frozen before any candidate runs + and every observation must pass. +- Discrete outputs (generated tokens or discrete actions): sequences must equal + the `kernel_profile: torch` run on every evaluation input. A difference fails; + there is no automatic waiver. +- Evaluation set: the model benchmark's frozen cross-task references (pi0.5: + the 10 references in `compare.py`) plus 200 benchmark frames. + +Inductor's runtime kernel-configuration choice makes a fresh compilation of +the same Torch path produce different outputs: on Jetson Thor, two pi0.5 +processes without a shared compile cache differed by up to 4.6e-2 in normalized +actions. Parity checks therefore run with `TORCHINDUCTOR_DETERMINISTIC=1`, which +made fresh compilations bit-exact, and model-level thresholds are derived from +the same compiled baseline that is timed. + +The thresholds, reference identity and evaluation-set digest are saved in the +optimize run and summarized in `profile.json`; `scripts.kernel_tuning verify` +reproduces them. + +## 7. Implementation plan + +1. Registry and profiles: `embodiinfer/layers/ops.py`, profile loader, + `kernel_profile` switch. One-time migration of the existing switch points to + the registry as the default code path: pi0.5's `gated_residual`, + `ada_rms_norm`, `gated_gelu` and `rotate_qk` (`split_kv_attention` already + uses the attention registry and gains profile lookup), then the Qwen R2R and + StreamVLN norm, activation and RoPE call sites. Install one kernel by hand and + run `verify`/`benchmark` on an RTX 4090 to prove the loop. +2. `profile` command and hotspot report. +3. Tasks for registered operators with Torch baselines, tuning, automated + `install`. +4. Agent selection and new fusion regions (one-time code changes, §4.4). +5. Further hardware profiles. + +New tooling lives in `scripts/kernel_tuning/` and `benchmarks/kernel_tuning/`. +Existing 0008 commands keep working and are not extended. No public symbol is +removed. + +## 8. Test plan + +CPU: registry resolution and preference; `applies` fallback; profile matching +and mismatch handling for `auto` and explicit names; digest tamper detection; +install round trip; hotspot aggregation and fusion-chain detection on synthetic +profiler traces; plan validation; region and model gate arithmetic, including +frozen thresholds and discrete-output equality. + +GPU (marked): pi0.5 with `kernel_profile: none` bit-exact against the previous +code; tuned and fallback paths inside CUDA Graph capture; the integration +bit-exact gate; `verify` on real weights. + +## 9. Benchmark plan + +pi0.5 on one RTX 4090, B=1, LIBERO-10 1,600 frames, BF16, Inductor and full +denoise-loop CUDA Graph, the protocol in `benchmarks/pi05-benchmark/README.md`. +Rows: `kernel_profile: torch`, current Triton, the tuned profile, and Inductor +`max-autotune` where it compiles; end-to-end mean/P50/P95 and prefill/decode +split; per-operator ablation. Expected gain is bounded by the recoverable time +in the hotspot report; it disappears where regions are already at the bandwidth +or launch floor. Batch > 1 and other hardware are reported only when measured. + +## 10. Risks and limitations + +- Agent kernels are often shape-specialized (a 2026-10-07 `split_kv_attention` + candidate asserted `P % PART_LENGTH == 0`); `applies` and fallback mitigate, + but uncovered shapes run the old path. +- Profiles are tied to Triton/Torch versions; a version upgrade disables them + under `auto` until re-verified. +- Attribution under CUDA Graphs depends on CUPTI support; the eager fallback pass + can shift shares slightly. +- Finite evaluation sets do not prove correctness for every input or task + success rate. +- Generated code adds review load; each install goes through a PR. +- Coding agents run commands without interactive approval, as in 0008. From 3e7df4d80f0b7707ee353f1ec8f41c10a5a33f51 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=82=96=E4=B8=80=E6=B6=B5?= <23371301@buaa.edu.cn> Date: Sun, 11 Oct 2026 14:15:41 +0800 Subject: [PATCH 14/44] feat: profile end-to-end inference and rank GPU hotspots MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add `python -m scripts.kernel_tuning profile --output DIR -- DRIVER`, run in the model runtime around an unmodified inference driver such as discovery.drive. Each top-level EngineCore.execute, GenerationBackend rollout or serving call is one step; after --warmup steps, --steps steps are profiled with torch.profiler. Each GPU kernel is attributed through its CUDA correlation id to the host launch and then to the ranges enclosing that launch: policy and decoder protocol stages, operator-registry calls, attention-backend calls and nn.Module calls with layer indices folded. Ranges are skipped while torch.compile traces, and kernels replayed from CUDA graphs keep only the ranges active at replay, so graphs and compiled regions report per stage. The report (hotspots.json, hotspots.md) ranks regions by GPU time per step with kernel counts, mean kernel time and the share of small kernels. Signed-off-by: 肖一涵 <23371301@buaa.edu.cn> --- scripts/kernel_tuning/__main__.py | 25 +- scripts/kernel_tuning/hotspots.py | 518 +++++++++++++++++++++++++++ tests/kernel_tuning/test_hotspots.py | 163 +++++++++ 3 files changed, 705 insertions(+), 1 deletion(-) create mode 100644 scripts/kernel_tuning/hotspots.py create mode 100644 tests/kernel_tuning/test_hotspots.py diff --git a/scripts/kernel_tuning/__main__.py b/scripts/kernel_tuning/__main__.py index 60a8c80..a190685 100644 --- a/scripts/kernel_tuning/__main__.py +++ b/scripts/kernel_tuning/__main__.py @@ -278,6 +278,15 @@ def main(argv: list[str] | None = None) -> int: "--output", type=Path, required=True, help="New model capture directory (JSONL for --catalog)" ) capture.add_argument("script", nargs=argparse.REMAINDER, help="-- SCRIPT [ARGS] or -- -m MODULE [ARGS]") + 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" ) @@ -290,7 +299,7 @@ def main(argv: list[str] | None = None) -> int: 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] == "capture": + if "--" in raw and raw[0] in ("capture", "profile"): split = raw.index("--") raw, command = raw[:split], raw[split + 1 :] else: @@ -323,6 +332,20 @@ def main(argv: list[str] | None = None) -> int: command or args.script, args.output, args.model, fixture_bytes=args.fixture_bytes ) print(json.dumps(result, indent=2, ensure_ascii=False)) + 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 diff --git a/scripts/kernel_tuning/hotspots.py b/scripts/kernel_tuning/hotspots.py new file mode 100644 index 0000000..2e25366 --- /dev/null +++ b/scripts/kernel_tuning/hotspots.py @@ -0,0 +1,518 @@ +"""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 keep the ranges active at replay time, so +they are attributed to their stage only; profile a run without CUDA graphs or +compilation for operator-level detail. Ranges are skipped inside +``torch.compile`` tracing, so compiled regions also report at stage level. +""" + +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:" +_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 + + +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 attribute( + ranges: Iterable[Range], launches: Iterable[Launch], kernels: Iterable[DeviceKernel] +) -> list[Kernel]: + """Attribute each device kernel to the ranges enclosing its host launch.""" + 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()) + result = [] + 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 = None + stages: list[str] = [] + innermost = None + for item in enclosing_by_id[kernel.correlation]: + name = item.name[len(RANGE) :] + if item.name.startswith(STEP): + step = int(name[len(STEP) - len(RANGE) :]) + elif name.startswith("stage:"): + if not stages or stages[-1] != name: + stages.append(name) + else: + innermost = fold_layers(name) + path = (*stages, innermost) if innermost else tuple(stages) + result.append(Kernel(step, path, kernel.name, kernel.duration_ns / 1000, "Graph" in launch.name)) + 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] + ], + } + ) + return { + "steps": steps, + "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.", + "", + "| 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() + + 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) + if self.modules: + self._install_module_hooks() + + 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 self.profiler is None or self.trace is not None: + 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}" + ) + ranges, launches, devices = [], [], [] + for event in self.trace.profiler.kineto_results.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 attribute(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/tests/kernel_tuning/test_hotspots.py b/tests/kernel_tuning/test_hotspots.py new file mode 100644 index 0000000..c4aba91 --- /dev/null +++ b/tests/kernel_tuning/test_hotspots.py @@ -0,0 +1,163 @@ +"""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",), False), + (0, ("stage:integrate",), True), + (0, ("stage:integrate",), True), + (1, (), False), + (None, (), 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 + assert session.kernels() == [] # CPU-only run: no device 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() From 04241990fb75ed15dd4148fe668822f9a38c4169 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=82=96=E4=B8=80=E6=B6=B5?= <23371301@buaa.edu.cn> Date: Sun, 11 Oct 2026 14:30:52 +0800 Subject: [PATCH 15/44] feat: label CUDA graph kernels from their capture in hotspot profiles MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Kernels replayed from a CUDA graph were attributed only to the stage that replayed the graph, which hid operator-level cost in the production configuration where most work runs inside graphs. Profile each graph capture outside the measured window and record the ranges of every node launch in capture order. Measured replays are tagged with their graph, and a replay's kernels take the capture labels in device start order when the counts match, prefixed by the replay-time stages. Regions without an operator, attention or module range (compiled code, unlabelled replays) are split by kernel category from the kernel name. On Jetson Thor, pi0.5 with native graphs: 95% of graph-replayed kernel time is now labelled; with the Inductor full-loop graph: 100%. Signed-off-by: 肖一涵 <23371301@buaa.edu.cn> --- scripts/kernel_tuning/hotspots.py | 212 ++++++++++++++++++++++----- tests/kernel_tuning/test_hotspots.py | 48 +++++- 2 files changed, 217 insertions(+), 43 deletions(-) diff --git a/scripts/kernel_tuning/hotspots.py b/scripts/kernel_tuning/hotspots.py index 2e25366..88e85e2 100644 --- a/scripts/kernel_tuning/hotspots.py +++ b/scripts/kernel_tuning/hotspots.py @@ -18,10 +18,12 @@ 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 keep the ranges active at replay time, so -they are attributed to their stage only; profile a run without CUDA graphs or -compilation for operator-level detail. Ranges are skipped inside -``torch.compile`` tracing, so compiled regions also report at stage level. +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 @@ -44,6 +46,7 @@ RANGE = "kt/" STEP = f"{RANGE}step:" +GRAPH = f"{RANGE}graph:" _ENTRIES = ( ("embodiinfer.engine.core", "EngineCore", ("execute", "execute_pipelined")), ( @@ -98,6 +101,7 @@ class Kernel: name: str duration_us: float graph: bool + labelled: bool = False def fold_layers(name: str) -> str: @@ -134,34 +138,111 @@ def _enclosing(ranges: Iterable[Range], launches: Iterable[Launch]) -> dict[int, 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] + 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.""" + """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()) - result = [] + 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 = None - stages: list[str] = [] - innermost = None - for item in enclosing_by_id[kernel.correlation]: - name = item.name[len(RANGE) :] - if item.name.startswith(STEP): - step = int(name[len(STEP) - len(RANGE) :]) - elif name.startswith("stage:"): - if not stages or stages[-1] != name: - stages.append(name) - else: - innermost = fold_layers(name) - path = (*stages, innermost) if innermost else tuple(stages) - result.append(Kernel(step, path, kernel.name, kernel.duration_ns / 1000, "Graph" in launch.name)) + 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 @@ -202,8 +283,12 @@ def summarize(kernels: Iterable[Kernel], steps: int, *, small_us: float = 5.0) - ], } ) + 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, @@ -218,6 +303,8 @@ def render(summary: dict[str, Any], *, limit: int = 30) -> str: "", 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 |", "| --- | ---: | ---: | ---: | ---: | ---: | ---: | --- |", @@ -266,6 +353,13 @@ def __init__(self, warmup: int, steps: int, *, modules: bool = True) -> 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.""" @@ -389,15 +483,56 @@ def resolve_backend(name: str) -> Any: 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 self.profiler is None or self.trace is not None: + if _compiling() or not self.recording: return name = self._module_names.get(id(module)) if name is None: @@ -432,22 +567,23 @@ def kernels(self) -> list[Kernel]: raise ContractError( f"Only {self.started} inference steps ran; profiling needs warmup {self.warmup} + {self.steps}" ) - ranges, launches, devices = [], [], [] - for event in self.trace.profiler.kineto_results.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 attribute(ranges, launches, devices) + 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 = [], [], [] + for event in profiler.profiler.kineto_results.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]: diff --git a/tests/kernel_tuning/test_hotspots.py b/tests/kernel_tuning/test_hotspots.py index c4aba91..8c2766b 100644 --- a/tests/kernel_tuning/test_hotspots.py +++ b/tests/kernel_tuning/test_hotspots.py @@ -58,11 +58,11 @@ def test_kernels_follow_their_launch_into_enclosing_ranges(): 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",), False), - (0, ("stage:integrate",), True), - (0, ("stage:integrate",), True), - (1, (), False), - (None, (), 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 @@ -161,3 +161,41 @@ def test_session_requires_enough_steps(fake_engine): 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) From dee06626e134025c3584125860d207e76d474fa3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=82=96=E4=B8=80=E6=B6=B5?= <23371301@buaa.edu.cn> Date: Sun, 11 Oct 2026 16:59:46 +0800 Subject: [PATCH 16/44] feat: route split-KV attention through the operator registry MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The denoise attention is the largest single hotspot of PI0.5's native path, but it was called on the Triton attention backend directly, so a tuned kernel could not replace it. Register split_kv_attention with the existing Triton kernel and a Torch reference (concatenated K/V, scaled_dot_product_attention with the padding mask as a boolean keep-mask, grouped-query aware). PI0.5 routes both native calls through the registry: the denoise call over cached and current K/V, and the native prefix call, which is the same kernel with an empty cache. Without a kernel profile both still call the same Triton function with the same arguments; _prepare_native_attention keeps its capability check. Signed-off-by: 肖一涵 <23371301@buaa.edu.cn> --- embodiinfer/backend/torch/ops.py | 27 ++++++++++++++++++++++ embodiinfer/layers/ops.py | 1 + embodiinfer/policies/pi05/modeling_pi05.py | 8 ++++--- tests/test_kernel_ops.py | 25 ++++++++++++++++++-- 4 files changed, 56 insertions(+), 5 deletions(-) diff --git a/embodiinfer/backend/torch/ops.py b/embodiinfer/backend/torch/ops.py index 6a3283c..3d9c866 100644 --- a/embodiinfer/backend/torch/ops.py +++ b/embodiinfer/backend/torch/ops.py @@ -50,3 +50,30 @@ def rotate_qk( 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/layers/ops.py b/embodiinfer/layers/ops.py index 1345915..fa37b13 100644 --- a/embodiinfer/layers/ops.py +++ b/embodiinfer/layers/ops.py @@ -162,6 +162,7 @@ def _select_profile(name: str) -> Any: ("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) diff --git a/embodiinfer/policies/pi05/modeling_pi05.py b/embodiinfer/policies/pi05/modeling_pi05.py index 1078ca5..1cb2a98 100644 --- a/embodiinfer/policies/pi05/modeling_pi05.py +++ b/embodiinfer/policies/pi05/modeling_pi05.py @@ -685,8 +685,8 @@ def _attn_sublayer( 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] @@ -696,7 +696,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'). diff --git a/tests/test_kernel_ops.py b/tests/test_kernel_ops.py index 628939e..35e63b9 100644 --- a/tests/test_kernel_ops.py +++ b/tests/test_kernel_ops.py @@ -45,7 +45,7 @@ def _write_profile(root, *, source=_TUNED_MODULE, applies="applies", digest=None def test_core_operators_have_torch_and_triton_implementations(): ops = available_ops() - for op in ("gated_residual", "ada_rms_norm", "gated_gelu", "rotate_qk"): + 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 @@ -156,7 +156,28 @@ def test_prepare_resolves_everything_before_graph_tracing(monkeypatch): 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"): + 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), + ) From b89d6fba9f5d37785ca0de002e1be70e9ab1d36a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=82=96=E4=B8=80=E6=B6=B5?= <23371301@buaa.edu.cn> Date: Sun, 11 Oct 2026 17:00:10 +0800 Subject: [PATCH 17/44] feat: generate tuning tasks for registry operators with Torch baselines MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Model discovery tunes isolated ATen calls against themselves. The new goal is a Triton kernel that beats the model's Torch implementation of an operator-registry entry, evaluated on the model's real inputs. Add `capture-ops` (scripts/kernel_tuning/regions.py): it runs an inference driver in eager mode and records every registry call of the measured steps, keeping the first real inputs of each layout as fixtures and counting the calls. The output is a discovery-format model capture whose operators have identity kind "registry", so calibrate, generate and tune-all apply as is: - the baseline and its FP64 oracle are the Torch reference in embodiinfer.backend.torch.ops, evaluated in FP64 for calibration; - output layouts are those of the Torch reference; - the implementation the model ran is vendored as the agent's starting point and digest-checked, like the baseline sources; - tasks time with CUDA graphs, as the operators run in production. On Jetson Thor (pi05, LIBERO-10, eager native path, 4 steps) capture-ops, calibrate and generate produced 7 tasks for split_kv_attention, ada_rms_norm, gated_residual, gated_gelu and rotate_qk. Signed-off-by: 肖一涵 <23371301@buaa.edu.cn> --- benchmarks/kernel_tuning/calibration.py | 8 + scripts/kernel_tuning/__main__.py | 30 +- scripts/kernel_tuning/discovery/numerics.py | 4 + scripts/kernel_tuning/discovery/tasks.py | 44 ++- scripts/kernel_tuning/regions.py | 386 ++++++++++++++++++++ 5 files changed, 468 insertions(+), 4 deletions(-) create mode 100644 scripts/kernel_tuning/regions.py diff --git a/benchmarks/kernel_tuning/calibration.py b/benchmarks/kernel_tuning/calibration.py index abdbb78..b76887e 100644 --- a/benchmarks/kernel_tuning/calibration.py +++ b/benchmarks/kernel_tuning/calibration.py @@ -99,6 +99,11 @@ def _arguments(identity: dict[str, Any], inputs: list[torch.Tensor]) -> dict[str arg.name: args[i] if i < len(args) else kwargs.get(arg.name, arg.default_value) for i, arg in enumerate(schema.arguments) } + if identity["kind"] == "registry": + import inspect + + function = getattr(importlib.import_module(identity["module"]), identity["entry"]) + return {**dict(zip(inspect.signature(function).parameters, args)), **kwargs} names = { "split_kv_attention": ( "query", @@ -198,6 +203,9 @@ def high_precision(identity: dict[str, Any], inputs: list[torch.Tensor]) -> list 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] args = _arguments(identity, promoted) + if kind == "registry": + # The Torch reference is dtype-generic: FP64 inputs give an FP64 oracle. + return production(identity, promoted) if kind == "aten": _, name, overload = identity["name"].split(".") if name == "scaled_dot_product_attention": diff --git a/scripts/kernel_tuning/__main__.py b/scripts/kernel_tuning/__main__.py index a190685..1e671de 100644 --- a/scripts/kernel_tuning/__main__.py +++ b/scripts/kernel_tuning/__main__.py @@ -278,6 +278,20 @@ def main(argv: list[str] | None = None) -> int: "--output", type=Path, required=True, help="New model capture directory (JSONL for --catalog)" ) 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]" + ) hotspots = commands.add_parser( "profile", help="Rank where end-to-end inference spends GPU time (model runtime)" ) @@ -299,7 +313,7 @@ def main(argv: list[str] | None = None) -> int: 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", "profile"): + if "--" in raw and raw[0] in ("capture", "capture-ops", "profile"): split = raw.index("--") raw, command = raw[:split], raw[split + 1 :] else: @@ -332,6 +346,20 @@ def main(argv: list[str] | None = None) -> int: command or args.script, args.output, args.model, fixture_bytes=args.fixture_bytes ) print(json.dumps(result, indent=2, ensure_ascii=False)) + 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 == "profile": from .hotspots import render from .hotspots import run as run_profile diff --git a/scripts/kernel_tuning/discovery/numerics.py b/scripts/kernel_tuning/discovery/numerics.py index fb5902e..5d76ec0 100644 --- a/scripts/kernel_tuning/discovery/numerics.py +++ b/scripts/kernel_tuning/discovery/numerics.py @@ -62,6 +62,10 @@ def policy(identity: dict[str, Any]) -> str | None: 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" name = identity["name"] if identity["kind"] == "aten" and name.split(".")[1] in _ATEN: return "aten" diff --git a/scripts/kernel_tuning/discovery/tasks.py b/scripts/kernel_tuning/discovery/tasks.py index 966c523..33c48ef 100644 --- a/scripts/kernel_tuning/discovery/tasks.py +++ b/scripts/kernel_tuning/discovery/tasks.py @@ -92,6 +92,16 @@ def _sources(op: dict[str, Any]) -> tuple[str, dict[str, str]]: imports, callable_name = f"import torch\nfrom .vendor.{module} import {entry}\n", entry if module == "embodiinfer.backend.triton.sampling": imports += f"from .vendor.{module} import GreedyWorkspace\n" + elif identity["kind"] == "registry": + 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 else: raise ContractError(f"No task builder for {identity['kind']}") args, kwargs = identity["arguments"]["tuple"] @@ -106,7 +116,8 @@ def _sources(op: dict[str, Any]) -> tuple[str, dict[str, str]]: def _plan( capture: ModelCapture, op: dict[str, Any], overrides: dict[str, Any], hardware_notes: Path | None ) -> dict[str, str | bytes]: - name = "discovered_" + op["id"] + registry = op["identity"]["kind"] == "registry" + name = (f"registry_{op['identity']['op']}_" if registry else "discovered_") + op["id"] unsupported = { spec["dtype"] for case in op["workloads"] for spec in case["inputs"] + case["outputs"] } - _DTYPES @@ -192,9 +203,12 @@ def _plan( outputs={}, reference=code, workloads=tuple(Workload(case["id"], {}, weight=case["count"]) for case in op["workloads"]), - notes="Single observed operator; no new fusion.", + notes="Registry operator; the Torch reference is the baseline." + if registry + else "Single observed operator; no new fusion.", precision=precision, - timing="eager", + # Registry operators run inside the model's CUDA graphs in production. + timing="cuda_graph" if registry else "eager", ) settings = tuning_settings(descriptor, overrides) if identity.get("rng") and settings["timing"]["mode"] != "eager": @@ -252,6 +266,30 @@ def _plan( ), } ) + if registry: + 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" + ) + 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: diff --git a/scripts/kernel_tuning/regions.py b/scripts/kernel_tuning/regions.py new file mode 100644 index 0000000..6346e42 --- /dev/null +++ b/scripts/kernel_tuning/regions.py @@ -0,0 +1,386 @@ +"""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 discovery 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. + +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 _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, derive its output contract from the Torch reference, then run it.""" + from embodiinfer.layers.ops import get_op + + from .discovery.tensors import ( + clone_inputs, + describe_inputs, + flatten, + restore, + storage_key, + 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) + outputs = tensor_outputs(result) + if not 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. + structure = digest([identity, specs]) + if structure not in self._fixtures: + self._fixtures[structure] = self._snapshot(tensors, specs) + case = { + "inputs": self._fixtures[structure], + "outputs": [tensor_spec(tensor) for tensor in outputs], + "output_aliases": [ + [i for i, tensor in enumerate(local) if storage_key(tensor) == storage_key(output)] + for output in outputs + ], + "output_groups": [ + [i for i, other in enumerate(outputs) if storage_key(output) == storage_key(other)] + for output in outputs + ], + } + 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: + message = f"registry.{op}: {exc}" + if message not in self.errors: + self.errors.append(message) + return function(*args, **kwargs) + + 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 + }, + } From 7f66d7f0771e95d50fc7e8150f9e051549db713b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=82=96=E4=B8=80=E6=B6=B5?= <23371301@buaa.edu.cn> Date: Sun, 11 Oct 2026 17:00:10 +0800 Subject: [PATCH 18/44] feat: install exported registry kernels into a hardware profile MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add `install EXPORT --profile NAME`, run in the model runtime on the target machine. It copies the winning sources of a registry task into embodiinfer/backend/triton/tuned/NAME// as a package and records the operator in profile.json, matched to the current GPU, compute capability and Torch/Triton versions, with the run, task and capture identities as evidence. The generated package exposes applies(), true only for argument layouts the evaluator verified (shapes, strides, dtypes and constants of every recorded workload); other calls fall back to the preferred implementation. The profile loader now accepts package directories and checks a digest over all of their files. Reinstalling needs --replace, and a profile verified for another environment is refused. Signed-off-by: 肖一涵 <23371301@buaa.edu.cn> --- embodiinfer/backend/triton/tuned/__init__.py | 40 ++++- scripts/kernel_tuning/__main__.py | 11 ++ scripts/kernel_tuning/install.py | 169 +++++++++++++++++++ tests/kernel_tuning/test_install.py | 131 ++++++++++++++ 4 files changed, 344 insertions(+), 7 deletions(-) create mode 100644 scripts/kernel_tuning/install.py create mode 100644 tests/kernel_tuning/test_install.py diff --git a/embodiinfer/backend/triton/tuned/__init__.py b/embodiinfer/backend/triton/tuned/__init__.py index c5dbc2c..6fd4693 100644 --- a/embodiinfer/backend/triton/tuned/__init__.py +++ b/embodiinfer/backend/triton/tuned/__init__.py @@ -1,7 +1,8 @@ """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 per tuned operator. +``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. @@ -12,8 +13,8 @@ "match": {"device_name": "NVIDIA Thor", "capability": [11, 0], "torch": "2.13.0+cu132", "triton": "3.7.1"}, "ops": { - "gated_residual": {"module": "gated_residual.py", "function": "gated_residual", - "applies": "applies", "sha256": ""} + "gated_residual": {"module": "gated_residual", "function": "run", + "applies": "applies", "sha256": ""} } } @@ -27,6 +28,7 @@ import hashlib import importlib.util import json +import sys from collections.abc import Callable, Mapping from dataclasses import dataclass, field from pathlib import Path @@ -65,14 +67,25 @@ def kernel(self, op: str) -> TunedKernel: return self._loaded[op] entry = self.ops[op] path = self.root / entry["module"] - digest = hashlib.sha256(path.read_bytes()).hexdigest() - if digest != entry["sha256"]: + if source_digest(path) != entry["sha256"]: raise ValueError(f"kernel profile {self.name!r}: {path.name} does not match its recorded sha256") - spec = importlib.util.spec_from_file_location(f"embodiinfer_tuned_{self.name}_{path.stem}", path) + 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) - spec.loader.exec_module(module) + # 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 @@ -81,6 +94,18 @@ def kernel(self, op: str) -> TunedKernel: 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 @@ -134,4 +159,5 @@ def current_environment() -> dict[str, Any]: "available_profiles", "current_environment", "load_profile", + "source_digest", ] diff --git a/scripts/kernel_tuning/__main__.py b/scripts/kernel_tuning/__main__.py index 1e671de..1c9a2c2 100644 --- a/scripts/kernel_tuning/__main__.py +++ b/scripts/kernel_tuning/__main__.py @@ -292,6 +292,12 @@ def main(argv: list[str] | None = None) -> int: 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") hotspots = commands.add_parser( "profile", help="Rank where end-to-end inference spends GPU time (model runtime)" ) @@ -360,6 +366,11 @@ def main(argv: list[str] | None = None) -> int: 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 == "profile": from .hotspots import render from .hotspots import run as run_profile diff --git a/scripts/kernel_tuning/install.py b/scripts/kernel_tuning/install.py new file mode 100644 index 0000000..ba71d76 --- /dev/null +++ b/scripts/kernel_tuning/install.py @@ -0,0 +1,169 @@ +"""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. +""" + +import torch + +from .{module} import {symbol} as run + +_VERIFIED = frozenset({cases}) + + +def _leaves(value, out): + if isinstance(value, torch.Tensor): + dtype = str(value.dtype).removeprefix("torch.") + out.append(("tensor", tuple(value.shape), tuple(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 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] + leaves.append(("tensor", tuple(spec["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) + ), + 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": "run", + "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/tests/kernel_tuning/test_install.py b/tests/kernel_tuning/test_install.py new file mode 100644 index 0000000..6adf093 --- /dev/null +++ b/tests/kernel_tuning/test_install.py @@ -0,0 +1,131 @@ +"""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}, + "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__["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 + 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") From f124c579dc27c9abdb346bb085ba1919e33507e9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=82=96=E4=B8=80=E6=B6=B5?= <23371301@buaa.edu.cn> Date: Sun, 11 Oct 2026 17:00:10 +0800 Subject: [PATCH 19/44] feat: add a model-level numerical gate for kernel profiles MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add `verify --profile NAME --output REPORT -- DRIVER_ARGS`, run in the model runtime with discovery.drive inputs. Every input runs three times with the same seed: an FP32 reference and a baseline in the requested dtype, both with kernel_profile=torch, and the candidate profile. For each input, continuous outputs must stay within twice the baseline's error against the reference (max-abs floor 1e-4, RMSE floor 1e-5), with thresholds fixed from the baseline alone; discrete outputs must equal the baseline. A baseline far from the reference means the runs did not share inputs or noise and is rejected rather than used as a threshold. The command exits nonzero when the gate fails. Signed-off-by: 肖一涵 <23371301@buaa.edu.cn> --- scripts/kernel_tuning/__main__.py | 20 ++- scripts/kernel_tuning/verify.py | 263 +++++++++++++++++++++++++++++ tests/kernel_tuning/test_verify.py | 55 ++++++ 3 files changed, 337 insertions(+), 1 deletion(-) create mode 100644 scripts/kernel_tuning/verify.py create mode 100644 tests/kernel_tuning/test_verify.py diff --git a/scripts/kernel_tuning/__main__.py b/scripts/kernel_tuning/__main__.py index 1c9a2c2..34d6a33 100644 --- a/scripts/kernel_tuning/__main__.py +++ b/scripts/kernel_tuning/__main__.py @@ -298,6 +298,14 @@ def main(argv: list[str] | None = None) -> int: 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)" ) @@ -319,7 +327,7 @@ def main(argv: list[str] | None = None) -> int: 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"): + if "--" in raw and raw[0] in ("capture", "capture-ops", "profile", "verify"): split = raw.index("--") raw, command = raw[:split], raw[split + 1 :] else: @@ -371,6 +379,16 @@ def main(argv: list[str] | None = None) -> int: 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", "checks", "failed", "worst_ratio")} + 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 diff --git a/scripts/kernel_tuning/verify.py b/scripts/kernel_tuning/verify.py new file mode 100644 index 0000000..435f3a1 --- /dev/null +++ b/scripts/kernel_tuning/verify.py @@ -0,0 +1,263 @@ +"""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``. + +For every input, continuous outputs must satisfy ``max|C - R| <= max(1e-4, 2 max|B - R|)`` +and ``rms(C - R) <= max(1e-5, 2 rms(B - R))``; thresholds are fixed from the +baseline before the candidate is compared. 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 +# 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]: + """Apply the frozen per-input thresholds; the baseline alone sets them.""" + if not (len(reference) == len(baseline) == len(candidate)) or not reference: + raise ContractError("All runs must produce one output per input") + rows, 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" + ) + max_abs_tol = max(MAX_ABS_FLOOR, 2 * base_abs) + rmse_tol = max(RMSE_FLOOR, 2 * base_rmse) + cand_abs, cand_rmse = _error(cand[name][1], ref[name][1]) + ok = cand_abs <= max_abs_tol and cand_rmse <= rmse_tol + passed &= ok + rows.append( + { + "input": index, + "output": name, + "kind": kind, + "passed": ok, + "candidate_max_abs": cand_abs, + "candidate_rmse": cand_rmse, + "baseline_max_abs": base_abs, + "baseline_rmse": base_rmse, + "max_abs_tolerance": max_abs_tol, + "rmse_tolerance": rmse_tol, + } + ) + worst = max( + (row for row in rows if row["kind"] == "continuous"), + key=lambda row: row["candidate_max_abs"] / row["max_abs_tolerance"], + default=None, + ) + return { + "passed": bool(passed), + "inputs": len(reference), + "checks": len(rows), + "failed": sum(not row["passed"] for row in rows), + "worst_ratio": None if worst is None else worst["candidate_max_abs"] / worst["max_abs_tolerance"], + "rows": rows, + } + + +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}, + } + ) + 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/test_verify.py b/tests/kernel_tuning/test_verify.py new file mode 100644 index 0000000..0d1394f --- /dev/null +++ b/tests/kernel_tuning/test_verify.py @@ -0,0 +1,55 @@ +"""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_uses_twice_the_baseline_error_with_floors(): + reference = [_out([1.0, 2.0])] + baseline = [_out([1.001, 2.0])] + assert gate(reference, baseline, [_out([1.0015, 2.0005])])["passed"] # rmse 1.12e-3 <= 1.41e-3 + report = gate(reference, baseline, [_out([1.0025, 2.0])]) + assert not report["passed"] and report["failed"] == 1 + exact = [_out([1.0, 2.0])] + assert gate(reference, exact, [_out([1.00001, 2.0])])["passed"] # within the 1e-4 / 1e-5 floors + assert not gate(reference, 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, []) From aafca0a54d7c645a75a4db206fbf5b61dc7a73de Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=82=96=E4=B8=80=E6=B6=B5?= <23371301@buaa.edu.cn> Date: Mon, 12 Oct 2026 18:52:39 +0800 Subject: [PATCH 20/44] fix: judge the model-level gate on errors over all inputs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A per-input bound of twice the Torch baseline's error is noise-dominated: on Jetson Thor, two Torch runs of the same 40 pi0.5 inputs differed per input, and the existing Triton path failed the bound on two inputs where Torch happened to land within 1e-3 of the FP32 reference, although its mean errors equalled Torch's. Judge each continuous output over the whole evaluation set instead: the candidate's mean max-abs and mean RMSE may exceed the baseline's by at most 25%, and its worst max-abs may be at most twice the baseline's worst (floors 1e-4 and 1e-5). Per-input errors stay in the report; discrete outputs must still equal the baseline exactly. Signed-off-by: 肖一涵 <23371301@buaa.edu.cn> --- docs/proposals/0009-hotspot-kernel-tuning.md | 13 ++-- scripts/kernel_tuning/__main__.py | 3 +- scripts/kernel_tuning/verify.py | 65 ++++++++++++++------ tests/kernel_tuning/test_verify.py | 21 ++++--- 4 files changed, 69 insertions(+), 33 deletions(-) diff --git a/docs/proposals/0009-hotspot-kernel-tuning.md b/docs/proposals/0009-hotspot-kernel-tuning.md index a76822a..0e83ecb 100644 --- a/docs/proposals/0009-hotspot-kernel-tuning.md +++ b/docs/proposals/0009-hotspot-kernel-tuning.md @@ -247,10 +247,15 @@ candidate may be at most about twice as wrong as Torch. Model level (required for install): - Continuous outputs (action chunks): compute an FP32 Torch reference from the - same weights and inputs. Per observation, thresholds for max-abs and RMSE are - twice the error of the `kernel_profile: torch` run against that reference, - with floors `1e-4` and `1e-5`. Thresholds are frozen before any candidate runs - and every observation must pass. + same weights and inputs, and measure each run's max-abs and RMSE error per + observation. Over the whole evaluation set, the candidate's mean max-abs and + mean RMSE may exceed those of the `kernel_profile: torch` run by at most 25%, + and its worst max-abs may be at most twice Torch's worst (floors `1e-4` and + `1e-5`). Thresholds are frozen before any candidate runs. A per-observation + bound of twice Torch's error was tried first and rejected: on Jetson Thor, + 40 pi0.5 observations, the existing Triton path failed it on 2 observations + where Torch happened to land within 1e-3 of the reference, while its mean + errors equalled Torch's and two Torch runs themselves differed per observation. - Discrete outputs (generated tokens or discrete actions): sequences must equal the `kernel_profile: torch` run on every evaluation input. A difference fails; there is no automatic waiver. diff --git a/scripts/kernel_tuning/__main__.py b/scripts/kernel_tuning/__main__.py index 34d6a33..0dba46c 100644 --- a/scripts/kernel_tuning/__main__.py +++ b/scripts/kernel_tuning/__main__.py @@ -385,7 +385,8 @@ def main(argv: list[str] | None = None) -> int: 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", "checks", "failed", "worst_ratio")} + 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 diff --git a/scripts/kernel_tuning/verify.py b/scripts/kernel_tuning/verify.py index 435f3a1..ffe876d 100644 --- a/scripts/kernel_tuning/verify.py +++ b/scripts/kernel_tuning/verify.py @@ -9,9 +9,12 @@ * baseline: the requested dtype with ``kernel_profile: torch``; * candidate: the requested dtype with ``kernel_profile: NAME``. -For every input, continuous outputs must satisfy ``max|C - R| <= max(1e-4, 2 max|B - R|)`` -and ``rms(C - R) <= max(1e-5, 2 rms(B - R))``; thresholds are fixed from the -baseline before the candidate is compared. Discrete outputs (integers, strings, +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. """ @@ -29,6 +32,9 @@ 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 @@ -97,10 +103,10 @@ def gate( baseline: list[dict[str, tuple[str, Any]]], candidate: list[dict[str, tuple[str, Any]]], ) -> dict[str, Any]: - """Apply the frozen per-input thresholds; the baseline alone sets them.""" + """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, passed = [], True + 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") @@ -120,40 +126,58 @@ def gate( 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" ) - max_abs_tol = max(MAX_ABS_FLOOR, 2 * base_abs) - rmse_tol = max(RMSE_FLOOR, 2 * base_rmse) cand_abs, cand_rmse = _error(cand[name][1], ref[name][1]) - ok = cand_abs <= max_abs_tol and cand_rmse <= rmse_tol - passed &= ok + errors.setdefault(name, []).append((base_abs, base_rmse, cand_abs, cand_rmse)) rows.append( { "input": index, "output": name, "kind": kind, - "passed": ok, "candidate_max_abs": cand_abs, "candidate_rmse": cand_rmse, "baseline_max_abs": base_abs, "baseline_rmse": base_rmse, - "max_abs_tolerance": max_abs_tol, - "rmse_tolerance": rmse_tol, } ) - worst = max( - (row for row in rows if row["kind"] == "continuous"), - key=lambda row: row["candidate_max_abs"] / row["max_abs_tolerance"], - default=None, - ) + 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), - "checks": len(rows), - "failed": sum(not row["passed"] for row in rows), - "worst_ratio": None if worst is None else worst["candidate_max_abs"] / worst["max_abs_tolerance"], + "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 @@ -256,6 +280,7 @@ def run(profile: str, output: Path, driver_args: list[str], *, samples: int | No "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) diff --git a/tests/kernel_tuning/test_verify.py b/tests/kernel_tuning/test_verify.py index 0d1394f..dc03d37 100644 --- a/tests/kernel_tuning/test_verify.py +++ b/tests/kernel_tuning/test_verify.py @@ -35,15 +35,20 @@ def test_flatten_outputs_splits_continuous_and_discrete_leaves(): assert not any("timing" in name for name in leaves) -def test_gate_uses_twice_the_baseline_error_with_floors(): - reference = [_out([1.0, 2.0])] - baseline = [_out([1.001, 2.0])] - assert gate(reference, baseline, [_out([1.0015, 2.0005])])["passed"] # rmse 1.12e-3 <= 1.41e-3 - report = gate(reference, baseline, [_out([1.0025, 2.0])]) - assert not report["passed"] and report["failed"] == 1 +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(reference, exact, [_out([1.00001, 2.0])])["passed"] # within the 1e-4 / 1e-5 floors - assert not gate(reference, exact, [_out([1.0002, 2.0])])["passed"] + 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(): From 5bac9db41f3914808d3d763ad5c6b5dbaeecaf4e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=82=96=E4=B8=80=E6=B6=B5?= <23371301@buaa.edu.cn> Date: Mon, 12 Oct 2026 18:52:45 +0800 Subject: [PATCH 21/44] fix: match installed kernels on canonical layouts, pass only tensor args MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit An installed kernel's `applies` compared strides exactly, so a size-one dimension recorded as (L*D, L*D, D, 1) in the eager capture never matched (L*D, D, D, 1) under Inductor, although such strides address no other element. Compare strides of size-one dimensions and of empty tensors as zero, in both the recorded signatures and the generated predicate. The evaluator calls a solution with the captured tensor arguments x0..xN only; constants such as the attention scale are part of the task. The generated package forwarded the policy's full argument list, so a matching call with a constant failed. Bind the tensors through the captured argument recipe and call the kernel as the evaluator did; `applies` already pins the constants to the verified values. Signed-off-by: 肖一涵 <23371301@buaa.edu.cn> --- scripts/kernel_tuning/install.py | 60 ++++++++++++++++++++++++++--- tests/kernel_tuning/test_install.py | 28 +++++++++++++- 2 files changed, 81 insertions(+), 7 deletions(-) diff --git a/scripts/kernel_tuning/install.py b/scripts/kernel_tuning/install.py index ba71d76..dcae27f 100644 --- a/scripts/kernel_tuning/install.py +++ b/scripts/kernel_tuning/install.py @@ -27,7 +27,9 @@ 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. +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 @@ -35,12 +37,21 @@ 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.") - out.append(("tensor", tuple(value.shape), tuple(value.stride()), dtype)) + 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) @@ -55,9 +66,40 @@ def _leaves(value, 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, ...]] = [] @@ -66,7 +108,8 @@ def walk(node: dict[str, Any]) -> None: kind, value = next(iter(node.items())) if kind == "tensor": spec = inputs[value] - leaves.append(("tensor", tuple(spec["shape"]), tuple(spec["stride"]), spec["dtype"])) + 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) @@ -137,7 +180,14 @@ def install( 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) + 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", @@ -150,7 +200,7 @@ def install( shutil.rmtree(staging, ignore_errors=True) manifest["ops"][op] = { "module": op, - "function": "run", + "function": "call", "applies": "applies", "sha256": source_digest(directory / op), "verified_layouts": len(cases), diff --git a/tests/kernel_tuning/test_install.py b/tests/kernel_tuning/test_install.py index 6adf093..650f16e 100644 --- a/tests/kernel_tuning/test_install.py +++ b/tests/kernel_tuning/test_install.py @@ -57,7 +57,11 @@ def export(tmp_path, monkeypatch): }, "task/replay.json": { "torch_version": str(torch.__version__), - "numerical_identity": {"op": "gated_residual", "arguments": recipe}, + "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"}}, @@ -91,13 +95,19 @@ def test_installed_kernel_applies_only_to_verified_layouts(export, tmp_path): 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__["CALLS"] + 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 @@ -129,3 +139,17 @@ def test_only_registry_tasks_install(export, tmp_path): (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) From ee9ac540173d7f228dc42abdc77a0fdcecf06e8b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=82=96=E4=B8=80=E6=B6=B5?= <23371301@buaa.edu.cn> Date: Mon, 12 Oct 2026 18:52:52 +0800 Subject: [PATCH 22/44] fix: keep the pi05 prefix K/V cache contiguous with or without compile MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The prefix pass cached K and V as views of the fused projection in eager mode, but Inductor returns them contiguous. Registry calls are captured eagerly while production compiles the prefix, so the denoise attention of the two runs saw different prefix V layouts and a layout-gated tuned kernel was never selected. Cache contiguous K/V in both modes. Under Inductor the copy is a no-op: with kernel_profile=none, 200 LIBERO-10 benchmark frames on Jetson Thor (BF16, B=1, TORCHINDUCTOR_DETERMINISTIC=1) were bit-identical to the previous code on both the native and the Torch path. Signed-off-by: 肖一涵 <23371301@buaa.edu.cn> --- docs/proposals/0009-hotspot-kernel-tuning.md | 9 ++++++++- embodiinfer/policies/pi05/modeling_pi05.py | 4 +++- 2 files changed, 11 insertions(+), 2 deletions(-) diff --git a/docs/proposals/0009-hotspot-kernel-tuning.md b/docs/proposals/0009-hotspot-kernel-tuning.md index 0e83ecb..7f048bd 100644 --- a/docs/proposals/0009-hotspot-kernel-tuning.md +++ b/docs/proposals/0009-hotspot-kernel-tuning.md @@ -190,7 +190,14 @@ the evaluator checks is true on every workload. Expected model saving is `install` writes promoted sources and updates `profile.json`. It refuses an operator that is not registered by the current source, so a fusion region's code -change must be merged first. `verify` runs the +change must be merged first. The installed kernel applies only to the argument +layouts the evaluator measured (shape, dtype and strides; strides of size-one +dimensions and empty tensors are ignored) and otherwise falls back. Capture runs +eagerly, so a policy must give registry arguments the same layout with and +without `torch.compile`: on Jetson Thor the pi0.5 prefix K/V cache was a view of +the fused projection in eager mode but contiguous under Inductor, and a tuned +split-KV kernel was never selected until the cache was made contiguous in both. +`verify` runs the model with `kernel_profile: ` against `kernel_profile: torch` on the evaluation set and applies the model-level gate; an operator is kept only if the whole profile passes. `benchmark` reports Torch, existing Triton and the profile diff --git a/embodiinfer/policies/pi05/modeling_pi05.py b/embodiinfer/policies/pi05/modeling_pi05.py index 1cb2a98..a81d12a 100644 --- a/embodiinfer/policies/pi05/modeling_pi05.py +++ b/embodiinfer/policies/pi05/modeling_pi05.py @@ -681,7 +681,9 @@ def _attn_sublayer( 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, :] From b58329c1edc5e41691c897238aae843423815330 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=82=96=E4=B8=80=E6=B6=B5?= <23371301@buaa.edu.cn> Date: Mon, 12 Oct 2026 18:53:02 +0800 Subject: [PATCH 23/44] feat: add a Jetson Thor kernel profile with tuned split-KV attention MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Install the first agent-tuned kernel through the full pipeline: hotspot report, registry capture, FP64 calibration, Humanize2 tuning against the Torch reference (codex/gpt-6-astra:high, run 20261012T084037Z-cf9243b720), export and install into `tuned/thor/`, model-level verify and benchmark. The kernel applies to the captured B=1 pi05 shapes (576-token prefix, 50 action tokens) and falls back otherwise. Task timing per call (CUDA Graph, cold L2): prefix 1,079 us Torch / ~178 us existing Triton / 139 us tuned; denoise 304 / ~120 / 38 us. verify on 40 recorded inputs passed (worst ratio 0.80 of the limits). Jetson Thor (MAXN, jetson_clocks), Torch 2.13.0+cu132, Triton 3.7.1, pi05_libero_finetuned_v044, B=1, BF16, 10 denoise steps, LIBERO-10 1,600 frames after 10 warmup calls, Inductor with prefix and full-loop CUDA Graphs: kernel_profile=none 98.8 ms mean (decode 49.0 ms), thor 87.6 ms (decode 37.2 ms). Generated profile packages are digest-pinned evidence, so Ruff skips them, as it does for the verbatim Wan VAE. The proposal records the measurement, implementation status and the follow-ups: shape-range tasks, automatic region selection and fusion regions. Signed-off-by: 肖一涵 <23371301@buaa.edu.cn> --- docs/proposals/0009-hotspot-kernel-tuning.md | 49 ++- .../backend/triton/tuned/thor/profile.json | 80 +++++ .../tuned/thor/split_kv_attention/__init__.py | 62 ++++ .../thor/split_kv_attention/candidate.py | 329 ++++++++++++++++++ pyproject.toml | 4 + 5 files changed, 522 insertions(+), 2 deletions(-) create mode 100644 embodiinfer/backend/triton/tuned/thor/profile.json create mode 100644 embodiinfer/backend/triton/tuned/thor/split_kv_attention/__init__.py create mode 100644 embodiinfer/backend/triton/tuned/thor/split_kv_attention/candidate.py diff --git a/docs/proposals/0009-hotspot-kernel-tuning.md b/docs/proposals/0009-hotspot-kernel-tuning.md index 7f048bd..2b0fbde 100644 --- a/docs/proposals/0009-hotspot-kernel-tuning.md +++ b/docs/proposals/0009-hotspot-kernel-tuning.md @@ -68,7 +68,7 @@ Non-goals: - Lower precision, quantization or changed model mathematics. A region may reorder floating-point work, but not change dtype or matmul precision. - Automatic merging of agent-written policy changes; they are reviewed as PRs. -- Thor/AGX profiles in the first phase; multi-GPU and batch > 1 tuning. +- Multi-GPU and batch > 1 tuning. - Compatibility with the 0008 discovery/catalog task formats. Those commands stay as they are and are not extended. @@ -295,6 +295,11 @@ reproduces them. 4. Agent selection and new fusion regions (one-time code changes, §4.4). 5. Further hardware profiles. +Status (2026-10-12): steps 1–3 are implemented for pi0.5; the loop was proven +on Jetson Thor instead of an RTX 4090 (§9), and the `thor` profile ships one +tuned operator. Qwen R2R and StreamVLN call sites, step 4 and an RTX 4090 +profile are not done; neither is shape-range tuning (§10). + New tooling lives in `scripts/kernel_tuning/` and `benchmarks/kernel_tuning/`. Existing 0008 commands keep working and are not extended. No public symbol is removed. @@ -321,11 +326,51 @@ split; per-operator ablation. Expected gain is bounded by the recoverable time in the hotspot report; it disappears where regions are already at the bandwidth or launch floor. Batch > 1 and other hardware are reported only when measured. +First measurement (2026-10-12): Jetson Thor (sm_110, MAXN, `jetson_clocks`), +Torch 2.13.0+cu132, Triton 3.7.1, `pi05_libero_finetuned_v044`, B=1, BF16, +10 denoise steps, LIBERO-10 1,600 frames after 10 warmup calls, Inductor, +prefix and full-loop CUDA Graphs, native prefix/denoise attention. + +| `kernel_profile` | Mean / P50 (ms) | Prefill (ms) | Decode (ms) | +|---|---|---|---| +| `torch` path (`native_inference: false`) | 179.9 / 179.9 | 63.9 | 111.4 | +| `none` (existing Triton) | 98.8 / 99.7 | 45.4 | 49.0 | +| `thor` | 87.6 / 85.4 | 45.2 | 37.2 | + +The `thor` profile contains one operator, `split_kv_attention`, tuned by +`codex/gpt-6-astra:high` (run `20261012T084037Z-cf9243b720`, 12 attempts). Its +task timing (CUDA Graph, cold L2) against the Torch reference and the existing +Triton kernel: + +| Workload (calls per step) | Torch (µs) | Existing Triton (µs) | Tuned (µs) | +|---|---|---|---| +| prefix, 576 keys (18) | 1,079 | ~178 | 139 | +| denoise, 576 + 50 keys (180) | 304 | ~120 | 38 | + +Under the benchmark configuration every dispatch decision selected the tuned +kernel (72 prefix and 720 denoise calls executed before graph capture; replays +reuse them). `verify` on 40 +recorded inputs passed (worst ratio 0.80 of the §6 limits). With +`kernel_profile: none`, 200 benchmark frames were bit-identical to the previous +code on both the native and the Torch path (`TORCHINDUCTOR_DETERMINISTIC=1`). + ## 10. Risks and limitations - Agent kernels are often shape-specialized (a 2026-10-07 `split_kv_attention` candidate asserted `P % PART_LENGTH == 0`); `applies` and fallback mitigate, - but uncovered shapes run the old path. + but uncovered shapes run the old path. The shipped kernel covers the captured + B=1 shapes with a 576-token prefix only; the 48-token text bucket, B > 1 and + other checkpoints fall back. Follow-up: tasks over declared shape ranges + (buckets, batch sizes, several layouts per shape), kernels with runtime strides + and scales, and range-based `applies`. +- Region selection is manual: the hotspot report ranks time, but a person picks + the operators. Follow-up: rank by `Σ count × (t − roofline floor)`, detect + chains of memory-bound kernels between GEMMs as fusion candidates, and let an + agent turn a candidate into a registered operator by PR, with a bit-exact + `kernel_profile: none` parity test as the proof. +- GEMMs (the largest share of pi0.5 time on Thor) are not registry operators. +- Agent runs depend on a relay that drops streams; failed attempts consume the + search budget. - Profiles are tied to Triton/Torch versions; a version upgrade disables them under `auto` until re-verified. - Attribution under CUDA Graphs depends on CUPTI support; the eager fallback pass diff --git a/embodiinfer/backend/triton/tuned/thor/profile.json b/embodiinfer/backend/triton/tuned/thor/profile.json new file mode 100644 index 0000000..dcd77bc --- /dev/null +++ b/embodiinfer/backend/triton/tuned/thor/profile.json @@ -0,0 +1,80 @@ +{ + "schema": 1, + "match": { + "device_name": "NVIDIA Thor", + "capability": [ + 11, + 0 + ], + "torch": "2.13.0+cu132", + "triton": "3.7.1" + }, + "ops": { + "split_kv_attention": { + "module": "split_kv_attention", + "function": "call", + "applies": "applies", + "sha256": "55e8a20b057457d63f7557aaa16f344d74397e4925047f1b6666bcf52faac2d3", + "verified_layouts": 2, + "evidence": { + "run_id": "20261012T084037Z-cf9243b720", + "task_digest": "0192979023cd12ccc2a72bb04807c258509aeb29259333a33e1f89077f51a64b", + "capture_identity": "3fda25f6804e3178470c1bc390062d402242029cc2a6fcbdf8ecd96904b22ca7", + "solution_digest": "8b1db93ae8388ca0f296151d4bbf25e4f5e8c31d43a8f6234f87630d8f92cc46", + "decision": { + "promoted": true, + "reasons": [], + "rounds": [ + { + "mean_latency_ms": { + "baseline": 0.3732264595013676, + "incumbent": 0.04873815755365473, + "candidate": 0.0469428035458832 + }, + "improvement": { + "baseline": 0.8742243419488559, + "incumbent": 0.03683672296793472 + }, + "workloads": { + "d5cf735e215bb7574e71": { + "baseline": 1.0804690765937168, + "incumbent": 0.1394609593451023, + "candidate": 0.1392327032983303 + }, + "03293aab11415ac3e8c7": { + "baseline": 0.30250219779213267, + "incumbent": 0.03966587737450997, + "candidate": 0.037713813570638494 + } + } + }, + { + "mean_latency_ms": { + "baseline": 0.37376189036260954, + "incumbent": 0.048764153013626736, + "candidate": 0.046742537992921745 + }, + "improvement": { + "baseline": 0.8749403318043637, + "incumbent": 0.041456990345758005 + }, + "workloads": { + "d5cf735e215bb7574e71": { + "baseline": 1.079287349820137, + "incumbent": 0.13775917767484983, + "candidate": 0.13755703370273112 + }, + "03293aab11415ac3e8c7": { + "baseline": 0.3032093444168568, + "incumbent": 0.039864650547504425, + "candidate": 0.037661088421940804 + } + } + } + ] + }, + "export_checksums": "4bbfc952ca6abcbce7edf7a2b3b0494c60a17354f97d5c353c115e719c58ce67" + } + } + } +} diff --git a/embodiinfer/backend/triton/tuned/thor/split_kv_attention/__init__.py b/embodiinfer/backend/triton/tuned/thor/split_kv_attention/__init__.py new file mode 100644 index 0000000..1f1c149 --- /dev/null +++ b/embodiinfer/backend/triton/tuned/thor/split_kv_attention/__init__.py @@ -0,0 +1,62 @@ +"""Tuned ``split_kv_attention`` kernel for the ``thor`` profile. + +Generated by ``scripts.kernel_tuning install`` from run 20261012T084037Z-cf9243b720; 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, 576, 256), (0, 0, 256, 1), 'bfloat16'), ('tensor', (1, 1, 576, 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, 626), (0, 0, 0, 1), 'float32'), ('constant', '0.0625')), (('tensor', (1, 8, 576, 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, 576, 256), (0, 0, 256, 1), 'bfloat16'), ('tensor', (1, 1, 576, 256), (0, 0, 2560, 1), 'bfloat16'), ('tensor', (1, 1, 1, 576), (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/thor/split_kv_attention/candidate.py b/embodiinfer/backend/triton/tuned/thor/split_kv_attention/candidate.py new file mode 100644 index 0000000..de4bc1b --- /dev/null +++ b/embodiinfer/backend/triton/tuned/thor/split_kv_attention/candidate.py @@ -0,0 +1,329 @@ +"""Direct split-KV attention for the two recorded pi05 layouts.""" + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _partial_attention( + Q, PK, PV, SK, SV, Mask, Part, Lse, + qb: tl.constexpr, qh: tl.constexpr, qs: tl.constexpr, + pkb: tl.constexpr, pkh: tl.constexpr, pks: tl.constexpr, + pvb: tl.constexpr, pvh: tl.constexpr, pvs: tl.constexpr, + skb: tl.constexpr, skh: tl.constexpr, sks: tl.constexpr, + svb: tl.constexpr, svh: tl.constexpr, svs: tl.constexpr, + mb: tl.constexpr, mn: tl.constexpr, + P: tl.constexpr, S: tl.constexpr, H: tl.constexpr, + HK: tl.constexpr, D: tl.constexpr, + SPLITS: tl.constexpr, PART_LENGTH: tl.constexpr, + BM: tl.constexpr, BN: tl.constexpr, +): + bh = tl.program_id(0) + batch = bh // H + head = bh % H + kv_head = head // (H // HK) + partition = tl.program_id(2) + rows = tl.program_id(1) * BM + tl.arange(0, BM) + dims = tl.arange(0, D) + q = tl.load( + Q + batch * qb + head * qh + rows[:, None] * qs + dims[None, :], + rows[:, None] < S, + other=0, + ) + maximum = tl.full((BM,), -float("inf"), tl.float32) + denominator = tl.zeros((BM,), tl.float32) + acc = tl.zeros((BM, D), tl.float32) + + for offset in range(PART_LENGTH // BN): + columns = partition * PART_LENGTH + offset * BN + tl.arange(0, BN) + suffix = (columns >= P) & (columns < P + S) + # Clamp inactive offsets too; only the corresponding segment is loaded. + suffix_rows = tl.maximum(columns - P, 0) + ks = tl.load( + SK + batch * skb + kv_head * skh + + suffix_rows[:, None] * sks + dims[None, :], + suffix[:, None], + other=0, + ) + if P > 0: + prefix = columns < P + kp = tl.load( + PK + batch * pkb + kv_head * pkh + + columns[:, None] * pks + dims[None, :], + prefix[:, None], + other=0, + ) + k = tl.where(prefix[:, None], kp, ks) + else: + # Empty prefix views can alias the suffix; never dereference them. + k = ks + + valid = columns < P + S + mask_value = tl.load( + Mask + batch * mb + columns * mn, valid, other=1.0 + ) + valid = valid & (mask_value == 0) + scores = tl.dot(q, tl.trans(k), input_precision="ieee", out_dtype=tl.float32) + scores = scores * (0.0625 * 1.4426950408889634) + scores = tl.where( + valid[None, :] & (rows[:, None] < S), scores, -float("inf") + ) + next_max = tl.maximum(maximum, tl.max(scores, 1)) + safe_max = tl.where(next_max == -float("inf"), 0.0, next_max) + correction = tl.exp2(maximum - safe_max) + prob = tl.exp2(scores - safe_max[:, None]) + + vs = tl.load( + SV + batch * svb + kv_head * svh + + suffix_rows[:, None] * svs + dims[None, :], + suffix[:, None], + other=0, + ) + if P > 0: + vp = tl.load( + PV + batch * pvb + kv_head * pvh + + columns[:, None] * pvs + dims[None, :], + prefix[:, None], + other=0, + ) + v = tl.where(prefix[:, None], vp, vs) + else: + v = vs + + # Preserve the starting kernel's BF16 dot operands and FP32 accumulation. + acc = acc * correction[:, None] + tl.dot( + prob.to(v.dtype), v, input_precision="ieee", out_dtype=tl.float32 + ) + denominator = denominator * correction + tl.sum(prob, 1) + maximum = next_max + + safe_denominator = tl.where(denominator > 0, denominator, 1.0) + value = acc / safe_denominator[:, None] + index = (bh * SPLITS + partition) * S + rows + tl.store(Part + index[:, None] * D + dims[None, :], value, rows[:, None] < S) + if SPLITS > 1: + logsum = tl.where( + denominator > 0, + maximum + tl.log2(safe_denominator), + -float("inf"), + ) + tl.store(Lse + index, logsum, rows < S) + + +@triton.jit +def _attention_segment( + q, rows, maximum, denominator, acc, + K, V, Mask, + ks: tl.constexpr, vs: tl.constexpr, mn: tl.constexpr, + LENGTH: tl.constexpr, MASK_OFFSET: tl.constexpr, + QUERY_ROWS: tl.constexpr, D: tl.constexpr, BN: tl.constexpr, +): + dims = tl.arange(0, D) + for offset in tl.range(0, (LENGTH + BN - 1) // BN, num_stages=2): + columns = offset * BN + tl.arange(0, BN) + in_segment = columns < LENGTH + k = tl.load( + K + columns[:, None] * ks + dims[None, :], + in_segment[:, None], + other=0, + ) + mask_value = tl.load( + Mask + (MASK_OFFSET + columns) * mn, in_segment, other=1.0 + ) + valid = in_segment & (mask_value == 0) + scores = tl.dot(q, tl.trans(k), input_precision="ieee", out_dtype=tl.float32) + scores = scores * (0.0625 * 1.4426950408889634) + scores = tl.where( + valid[None, :] & (rows[:, None] < QUERY_ROWS), scores, -float("inf") + ) + next_max = tl.maximum(maximum, tl.max(scores, 1)) + safe_max = tl.where(next_max == -float("inf"), 0.0, next_max) + correction = tl.exp2(maximum - safe_max) + prob = tl.exp2(scores - safe_max[:, None]) + + v = tl.load( + V + columns[:, None] * vs + dims[None, :], + in_segment[:, None], + other=0, + ) + acc = acc * correction[:, None] + tl.dot( + prob.to(v.dtype), v, input_precision="ieee", out_dtype=tl.float32 + ) + denominator = denominator * correction + tl.sum(prob, 1) + maximum = next_max + return maximum, denominator, acc + + +@triton.jit +def _segmented_attention( + Q, PK, PV, SK, SV, Mask, Out, + qb: tl.constexpr, qh: tl.constexpr, qs: tl.constexpr, + pkb: tl.constexpr, pkh: tl.constexpr, pks: tl.constexpr, + pvb: tl.constexpr, pvh: tl.constexpr, pvs: tl.constexpr, + skb: tl.constexpr, skh: tl.constexpr, sks: tl.constexpr, + svb: tl.constexpr, svh: tl.constexpr, svs: tl.constexpr, + mb: tl.constexpr, mn: tl.constexpr, + P: tl.constexpr, S: tl.constexpr, H: tl.constexpr, + HK: tl.constexpr, D: tl.constexpr, + BM: tl.constexpr, BN: tl.constexpr, + PACK_QUERY_HEADS: tl.constexpr, +): + # The frozen prefix ends on a tile boundary; both segments share one softmax. + tl.static_assert(P % BN == 0) + bh = tl.program_id(0) + if PACK_QUERY_HEADS: + # The host restricts this path to the contiguous, single-batch MQA case. + tl.static_assert(HK == 1) + tl.static_assert(qh == S * qs) + QUERY_ROWS: tl.constexpr = H * S + batch = 0 + head = 0 + kv_head = 0 + else: + QUERY_ROWS: tl.constexpr = S + batch = bh // H + head = bh % H + kv_head = head // (H // HK) + rows = tl.program_id(1) * BM + tl.arange(0, BM) + dims = tl.arange(0, D) + q = tl.load( + Q + batch * qb + head * qh + rows[:, None] * qs + dims[None, :], + rows[:, None] < QUERY_ROWS, + other=0, + ) + maximum = tl.full((BM,), -float("inf"), tl.float32) + denominator = tl.zeros((BM,), tl.float32) + acc = tl.zeros((BM, D), tl.float32) + + maximum, denominator, acc = _attention_segment( + q, rows, maximum, denominator, acc, + PK + batch * pkb + kv_head * pkh, + PV + batch * pvb + kv_head * pvh, + Mask + batch * mb, + pks, pvs, mn, P, 0, QUERY_ROWS, D, BN, + ) + maximum, denominator, acc = _attention_segment( + q, rows, maximum, denominator, acc, + SK + batch * skb + kv_head * skh, + SV + batch * svb + kv_head * svh, + Mask + batch * mb, + sks, svs, mn, S, P, QUERY_ROWS, D, BN, + ) + + safe_denominator = tl.where(denominator > 0, denominator, 1.0) + value = acc / safe_denominator[:, None] + if PACK_QUERY_HEADS: + out_rows = rows + else: + out_rows = bh * S + rows + tl.store( + Out + out_rows[:, None] * D + dims[None, :], + value, + rows[:, None] < QUERY_ROWS, + ) + + +@triton.jit +def _merge_attention( + Part, Lse, Out, + S: tl.constexpr, D: tl.constexpr, SPLITS: tl.constexpr, +): + row = tl.program_id(0) + bh = tl.program_id(1) + partitions = tl.arange(0, SPLITS) + dims = tl.arange(0, D) + indices = (bh * SPLITS + partitions) * S + row + logsum = tl.load(Lse + indices) + maximum = tl.max(logsum, 0) + # All-empty rows have zero weights, rather than exp2(-inf - -inf). + safe_max = tl.where(maximum == -float("inf"), 0.0, maximum) + weight = tl.exp2(logsum - safe_max) + partial = tl.load(Part + indices[:, None] * D + dims[None, :]) + denominator = tl.sum(weight, 0) + safe_denominator = tl.where(denominator > 0, denominator, 1.0) + value = tl.sum(partial * weight[:, None], 0) / safe_denominator + tl.store(Out + (bh * S + row) * D + dims, value) + + +def run(x0, x1, x2, x3, x4, x5): + """Return contiguous BF16 BHSD attention without modifying any input.""" + batch, heads, size, width = x0.shape + prefix = x1.shape[2] + if (batch, heads, width) != (1, 8, 256): + raise ValueError("Unsupported query shape") + if (prefix, size) == (0, 576): + splits, block_m, block_n, part_length = 1, 32, 64, 576 + elif (prefix, size) == (576, 50): + splits, block_m, block_n, part_length = 1, 16, 64, 640 + else: + raise ValueError("Unsupported prefix/query lengths") + + pack_query_heads = ( + (prefix, size) == (576, 50) + and x1.shape[1] == 1 + and x3.shape[1] == 1 + and tuple(x0.stride()) == (102400, 12800, 256, 1) + ) + + # Shape allocation deliberately does not preserve the input query strides. + output = torch.empty( + (batch, heads, size, width), device=x0.device, dtype=x0.dtype + ) + if (prefix, size) in ((0, 576), (576, 50)): + # Pipeline both layouts; the empty prefix performs zero iterations. + # Shared K/V lets a query tile cross head boundaries in the short case. + grid = ( + (1, triton.cdiv(heads * size, block_m), 1) + if pack_query_heads + else (batch * heads, triton.cdiv(size, block_m), 1) + ) + _segmented_attention[grid]( + x0, x1, x2, x3, x4, x5, output, + *x0.stride()[:3], + *x1.stride()[:3], + *x2.stride()[:3], + *x3.stride()[:3], + *x4.stride()[:3], + x5.stride(0), x5.stride(-1), + prefix, size, heads, x1.shape[1], width, + block_m, block_n, pack_query_heads, + num_warps=4, + num_stages=2, + ) + return output + + if splits == 1: + partial = output + logsum = None + else: + partial = torch.empty( + (batch, heads, splits, size, width), + device=x0.device, + dtype=torch.float32, + ) + logsum = torch.empty( + (batch, heads, splits, size), device=x0.device, dtype=torch.float32 + ) + + # Standard Triton launches use the caller's current CUDA stream. + _partial_attention[(batch * heads, triton.cdiv(size, block_m), splits)]( + x0, x1, x2, x3, x4, x5, partial, logsum, + *x0.stride()[:3], + *x1.stride()[:3], + *x2.stride()[:3], + *x3.stride()[:3], + *x4.stride()[:3], + x5.stride(0), x5.stride(-1), + prefix, size, heads, x1.shape[1], width, + splits, part_length, block_m, block_n, + num_warps=4, + num_stages=1, + ) + if splits > 1: + _merge_attention[(size, batch * heads)]( + partial, logsum, output, size, width, splits, + num_warps=4, + num_stages=1, + ) + return output diff --git a/pyproject.toml b/pyproject.toml index 2eb6e82..9c4773b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -171,6 +171,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. From 12a2c4eb586eb1d29fdd03fd7b7eb081b46d5030 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=82=96=E4=B8=80=E6=B6=B5?= <23371301@buaa.edu.cn> Date: Mon, 12 Oct 2026 19:03:20 +0800 Subject: [PATCH 24/44] fix: keep installed kernel profiles LF so their digests match on Windows MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Installed profile packages are pinned by the SHA-256 of their bytes in profile.json. With core.autocrlf a Windows checkout rewrote them with CRLF and the profile refused to load. Check them out with LF everywhere. Signed-off-by: 肖一涵 <23371301@buaa.edu.cn> --- .gitattributes | 3 +++ 1 file changed, 3 insertions(+) create mode 100644 .gitattributes 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 From 095e8a138ffba467e78c90760dd89c8f115c7852 Mon Sep 17 00:00:00 2001 From: owpeter Date: Tue, 13 Oct 2026 03:55:14 +0000 Subject: [PATCH 25/44] docs: propose an agent loop that registers operators for kernel tuning Kernel tuning can only install registry operators, and registering one is manual work today. Design a resumable agent-in-the-loop flow that takes a caller-selected operator or fusion region, has a coding agent write its Torch reference, register it and route the call sites through the registry, and accepts the change only when controller-owned gates pass: diff scope, structure, ruff and tests, an unchanged eager ATen sequence, registry hits per call site, and kernel_profile=none outputs bit-identical to the base revision on every parity configuration. The result is a branch and evidence for human review; tuning stays in the 0009 flow. Signed-off-by: owpeter --- docs/proposals/0009-hotspot-kernel-tuning.md | 3 +- .../0010-agent-operator-registration.md | 470 ++++++++++++++++++ 2 files changed, 472 insertions(+), 1 deletion(-) create mode 100644 docs/proposals/0010-agent-operator-registration.md diff --git a/docs/proposals/0009-hotspot-kernel-tuning.md b/docs/proposals/0009-hotspot-kernel-tuning.md index 2b0fbde..409dfce 100644 --- a/docs/proposals/0009-hotspot-kernel-tuning.md +++ b/docs/proposals/0009-hotspot-kernel-tuning.md @@ -298,7 +298,8 @@ reproduces them. Status (2026-10-12): steps 1–3 are implemented for pi0.5; the loop was proven on Jetson Thor instead of an RTX 4090 (§9), and the `thor` profile ships one tuned operator. Qwen R2R and StreamVLN call sites, step 4 and an RTX 4090 -profile are not done; neither is shape-range tuning (§10). +profile are not done; neither is shape-range tuning (§10). Step 4's registration +loop is designed in [0010](0010-agent-operator-registration.md). New tooling lives in `scripts/kernel_tuning/` and `benchmarks/kernel_tuning/`. Existing 0008 commands keep working and are not extended. No public symbol is diff --git a/docs/proposals/0010-agent-operator-registration.md b/docs/proposals/0010-agent-operator-registration.md new file mode 100644 index 0000000..5f02c85 --- /dev/null +++ b/docs/proposals/0010-agent-operator-registration.md @@ -0,0 +1,470 @@ +# 0010 — 智能体驱动的算子注册流程 + +- Status: Draft +- Author: EmbodiInfer maintainers +- Date: 2026-10-13 + +## 1. 概要 + +本提案实现 [0009](0009-hotspot-kernel-tuning.md) §4.4 和实现计划第 4 步：调用方选定一个尚未进入算子注册表的算子或融合区域后，由编码智能体完成以下改动： + +- 写出它的 Torch 参考实现； +- 在 `embodiinfer/layers/ops.py` 中注册； +- 把模型调用点改为经过注册表； +- 补齐测试。 + +控制器随后依次运行静态检查、CPU 测试、命中检查和 `kernel_profile: none` 逐位一致检查。任何一项失败，都把结构化反馈交回智能体修改，直到全部通过或预算用尽。最终产物是一个可以直接提交审阅的分支和证据包。 + +整个流程位于 `scripts/kernel_tuning/registration/`，以 `python -m scripts.kernel_tuning register …` 调用，属于开发者工具，不进入运行时。它只负责把算子变成"可调优、可安装"的注册算子，不做调优，也不自动合并。注册代码经人工审阅合并后，再进入 0009 已实现的调优流程（`capture-ops` → `calibrate` → `tune-all`/`run` → `install` → `verify`）。 + +## 2. 动机与现状缺口 + +- **只有注册算子能调优并接入模型。** `capture-ops` 只记录经过 `KernelSelection.resolve` 的调用（`scripts/kernel_tuning/regions.py:152-163`）。`install` 拒绝其他任务，也拒绝当前源码里没有注册的算子（`scripts/kernel_tuning/install.py:140-144`）。 +- **注册目前完全靠人工。** pi0.5 的 5 个算子都是在 commit `449ebe2` 中手工迁移的，每个算子要做四件事： + 1. 把 Torch 实现搬到 `embodiinfer/backend/torch/ops.py`； + 2. 在 `embodiinfer/layers/ops.py:160` 注册； + 3. 把调用点改为 `_op(kernels, name, use_triton)(...)`（`embodiinfer/policies/pi05/modeling_pi05.py:88`），并把 `self._kernels` 传下去； + 4. 在 Thor 上验证 `kernel_profile: none` 逐位一致。 +- **最大的热点在注册表之外。** 0009 §10 记录，Thor 上 pi0.5 耗时最多的是 GEMM，它们不在注册表里，新流程完全碰不到。0009 设想由智能体把热点区域写成新的注册算子并提 PR（§4.4），但这一步没有实现。 +- **改调用点的风险很具体。** 仓库多处强调形状和表达式会改变舍入： + - 合并 Q/K/V 会改变归约形状，从而改变舍入（`modeling_pi05.py:664-666`）； + - `F.linear` 和 `x @ w.T` 会走不同的 ATen 分解； + - eager 与 Inductor 下参数布局可能不同（commit `ecde612`）。 + + 这些都需要固定的门禁来兜底，不能依赖智能体自觉。 + +## 3. 目标与非目标 + +目标： + +- 调用方显式选择要注册的算子：给出算子名、种类（包装或融合）和调用点。调用点可以直接写源码位置，也可以从热点报告挑选区域，再由工具解析成源码位置。 +- 一条可恢复的流水线：选择 → 定位调用点 → 记录基线 → 智能体修改 → 静态门禁 → GPU 门禁 → 反馈循环 → 打包。每一步都有明确的产物和通过条件。 +- 尽量复用现有工具：`profile`、`capture-ops`、`verify` 的输出比较逻辑、Humanize2 flow、`run_lock`/`atomic_json`。只为现有工具做不到的事新增代码。 +- 验收由控制器负责。智能体写的测试只是补充，判定依据是控制器对比基线版本的逐位一致结果。 + +非目标： + +- 调优或安装 kernel（属于 0009 的 B、C 部分）。 +- 改变数学、精度或舍入。注册改动必须在 `kernel_profile: none` 下与基线逐位一致。 +- 自动合并、自动推送或自动开 PR。默认只生成分支和 PR 描述；提交、推送、开 PR 都需要调用方显式传参。 +- 修改 `KernelSelection`、调优档案（tuned profile）、`scripts/` 或依赖。这类改动走普通 PR。 +- Qwen R2R、StreamVLN 的首次接入。它们还没有 `kernels` 传参链，第一次需要人工建立；建立之后同样可以使用本流程。 + +## 4. 设计 + +### 4.1 总览 + +``` + 调用方 控制器（工具环境） 智能体（Humanize2） 模型运行时（GPU 子进程） + │ R0 candidates ──► 读 hotspots.json，列出未覆盖热点 + │ R1 new ─────────► 校验规格、定位调用点、建 worktree、打包上下文 + │ ◄── sites.md 等待确认 + │ R2 run ─────────► 基线录制 ─────────────────────────────────────────────────► snapshot ×2、trace-ops + │ R3 ───────────────────────────────► plan / implement + │ R4 静态门禁（范围、结构、ruff、pytest） + │ R5 GPU 门禁 ─────────────────────────────────────────────► trace-ops、capture-ops、 + │ │ snapshot、audit-layouts + │ R6 失败 → feedback.json ──────────► repair（同一会话） + │ R7 全部通过 → 证据、可选覆盖率/测速、PR.md + │ finalize ───────► （显式 --commit/--push）提交到 registry/- 分支 + │ R8 人工审阅、合并 → 进入 0009 的调优流程 +``` + +| 步骤 | 工作 | 执行者 | 现有工具 | 产物 | +|---|---|---|---|---| +| R0 | 列出未被注册表覆盖的热点 | 控制器 | 读取 `profile` 的 `hotspots.json` | `candidates.md` | +| R1 | 规格校验、调用点定位、worktree、上下文 | 控制器 | `git worktree`；运行时内 `available_ops()` | `spec.json`、`sites.json`、`context/` | +| R2 | 基线输出与 ATen 序列，确定性预检 | 运行时子进程 | 复用 `verify` 的取输出逻辑（抽成 `snapshot`） | `golden/`、`trace/base.jsonl` | +| R3 | 写计划、改代码、写测试 | 智能体 | Humanize2 flow 机制（同 `flow.py`） | worktree 改动、`REGISTRATION.json` | +| R4 | 范围、结构、ruff、pytest | 控制器 | `ruff`、`pytest` | `attempts/N/static.json` | +| R5 | ATen 序列、命中、逐位一致、布局审计 | 运行时子进程 | `capture-ops`（现有）；`snapshot`、`trace-ops`、`audit-layouts`（新增） | `attempts/N/gpu.json` | +| R6 | 汇总失败交回智能体 | 控制器 + 智能体 | 同一 Humanize2 会话 | `attempts/N/feedback.json` | +| R7 | 终检、证据、可选覆盖率和测速、PR 描述 | 控制器 | `profile`（现有）、`benchmarks/pi05-benchmark/benchmark.py`（现有） | `evidence.json`、`PR.md`、`patch.diff` | +| R8 | 审阅与合并 | 人 | GitHub | 合并后的注册代码 | + +### 4.2 调用方式：选择算子 + +所有命令都在目标 GPU 机器上运行，入口是工具环境（Python ≥ 3.12 + Humanize2，与 `tune-all` 相同）。加载模型的步骤通过 `--set runtime_python=…` 指定的模型运行时解释器在子进程里执行，与 `tune-all --set evaluator_python=…` 的做法一致。 + +```bash +TOOL="uv run --project scripts/kernel_tuning python -m scripts.kernel_tuning" + +# R0（可选）：看看哪些热点不在注册表里 +$TOOL register candidates --hotspots HOT/hotspots.json --min-share 0.01 + +# R1：为选定的算子建立一次运行（规格文件和命令行参数二选一） +$TOOL register new --spec linear.yaml --output results/kernel_tuning/register +$TOOL register new --model pi05 --op linear --kind wrap \ + --site 'embodiinfer/policies/pi05/modeling_pi05.py::_fused_qkv::F.linear(x, weight)' \ + --region 'stage:denoise_step/module:expert.layers.*.mlp.down_proj' --hotspots HOT/hotspots.json \ + --parity-config parity.yaml --output results/kernel_tuning/register +# 命令结束时打印 sites.md，由调用方确认或编辑 sites.json；--yes 跳过确认。 + +# R2–R7：运行循环；--until 可以在某个阶段后停下 +$TOOL register run RUN --agent 'codex/MODEL:EFFORT' \ + --set runtime_python=/abs/runtime/python --set test_python=/abs/dev/python \ + [--until baseline|static|parity|package] +$TOOL register resume RUN +$TOOL register status RUN + +# 收尾：只有显式传参才会提交或推送 +$TOOL register finalize RUN [--commit] [--push] +``` + +规格文件（`linear.yaml`）示例： + +```yaml +model: pi05 +op: linear # 新算子名：Python 标识符、尚未注册、按数学含义命名 +kind: wrap # wrap：把一处已有表达式包成算子；fuse：把相邻几步合成一个算子 +description: > + 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)" + - region: "stage:denoise_step/module:expert.layers.*.self_attn.o_proj" # 从 hotspots.json 选 +preserve: # 必须保持原样的分支；R1 会自动检测并补充候选 + - "QuantizedLinear：继续走 layers/linear.py 的量化注册表" + - "张量并行：RowParallelLinear / ColumnParallelLinear 保持模块调用" + - "torch.is_grad_enabled() 或 training 时的原始调用" +parity: # GPU 逐位一致检查的配置；每个配置独占一个进程和一份编译缓存 + checkpoint: /ckpt/pi05_libero_finetuned_v044 + inputs: /data/pi05-inputs.pt # discovery.drive 的录制输入 + samples: 40 # 循环里使用的输入数 + final_inputs: /data/pi05-frames.pt # R7 终检使用，建议 ≥ 200 帧 + configs: + eager: {adapter_config: pi05-eager.json, dtype: bfloat16, num_steps: 10, no_cuda_graph: true} + native: {adapter_config: pi05-native.json, dtype: bfloat16, num_steps: 10, capture_full_loop: true} + inductor: {adapter_config: pi05-inductor.json, dtype: bfloat16, num_steps: 10, capture_full_loop: true} + fp32: {adapter_config: pi05-native.json, dtype: float32, num_steps: 10, no_cuda_graph: true} +tests: # 本模型相关的 CPU 测试；不写时使用模型默认表 + - tests/test_kernel_ops.py + - tests/test_pi05_components.py + - tests/test_pi05_cpu_routes.py + - tests/test_tensor_parallel.py +budget: + max_attempts: 5 + agent_timeout_seconds: 1800 + max_seconds: 14400 +``` + +`parity.configs` 中的每一项展开为 `scripts.kernel_tuning.discovery.drive` 的参数（`--adapter-config`、`--dtype`、`--num-steps`、`--capture-full-loop`、`--no-cuda-graph` 等），因此可以覆盖策略的每条执行路径。对 pi0.5 来说，`adapter_config` 里的 `policy_kwargs` 决定走哪条路径，相关开关有 `attention`、`compile_backend`、`native_inference`、`prefix_cuda_graph`、`denoise_attention`、`prefix_attention`（`modeling_pi05.py:414-430`）。调用点处在哪条路径上，就必须有一个配置走到它。例如 `attention: eager` 走的是分开的 Q/K/V GEMM（`modeling_pi05.py:664-668`），融合投影的改动不会经过它。 + +每个模型在 `scripts/kernel_tuning/registration/models/.yaml` 里提供默认值：CPU 测试表、常见的 preserve 模式（例如 `QuantizedLinear`、TP 线性层）、parity 配置模板，以及作为示例的迁移 commit（pi0.5 是 `449ebe2`）。规格文件可以覆盖其中任意一项。 + +### 4.3 运行目录与状态 + +``` +results/kernel_tuning/register// + manifest.json 状态、各阶段结果、尝试次数、耗时；atomic_json 写入，run_lock 加锁 + spec.json 冻结的规格（含基线 commit、模型、parity 配置的摘要） + sites.json / sites.md 已确认的调用点清单 + context/ 提供给智能体的只读上下文 + base/ 基线 commit 的 git worktree（detached），只供基线录制使用 + worktree/ 候选分支 registry/- 的 git worktree，智能体只能改这里 + golden/-{1,2}.pt 基线输出（两次，用于确定性预检） + trace/base.jsonl 基线 eager 推理的 ATen 调用序列 + attempts// 每次尝试的 prompt、回复、PLAN.md、patch.diff、static.json、gpu.json、feedback.json + cache/ 每个子进程独立的 TORCHINDUCTOR_CACHE_DIR / TRITON_CACHE_DIR + evidence.json PR.md patch.diff +``` + +`resume` 从 `manifest.json` 记录的阶段继续，已经通过的阶段不会重跑。规格或基线 commit 发生变化时必须新建一次运行，不允许在原运行上续跑。 + +### 4.4 各步骤详细说明 + +#### R0 候选列表（可选，不调用智能体） + +- **工作：** 读取 `profile` 的输出，列出没有被 `op:` 范围覆盖的区域，按 `us_per_step` 排序，并标明各区域的定位精度： + - `module:` 区域可以映射到源码，比如 `expert.layers.*.mlp.down_proj` → 属性 `mlp.down_proj`； + - `kernels:gemm`、`kernels:inductor` 这类按 kernel 名分类的区域，来自编译区域或无标签的图回放，无法自动映射到源码，只能由调用方手写 `site`。 +- **复用的现有工具：** `profile` 命令和它生成的 `hotspots.json`，字段见 `scripts/kernel_tuning/hotspots.py:249-296`。 + ```bash + /abs/runtime/python -m scripts.kernel_tuning profile --output HOT -- \ + -m scripts.kernel_tuning.discovery.drive --policy pi05 --checkpoint CKPT \ + --adapter-config pi05-native.json --inputs pi05-inputs.pt --dtype bfloat16 \ + --num-steps 10 --capture-full-loop + ``` +- **新增：** `registration/candidates.py`，只读 JSON。 +- **说明：** 这一步只提供参考，选哪个算子仍由调用方决定。按 roofline 下限估算可回收时间是 0009 的后续工作，本步骤不实现。 + +#### R1 建立运行（`register new`，不调用智能体，不需要 GPU） + +1. **规格校验** + - `op` 必须是 Python 标识符。现有工具链要求参考实现函数名等于算子名，并在 `embodiinfer.backend.torch.ops` 中（`scripts/kernel_tuning/regions.py:215-217`、`scripts/kernel_tuning/discovery/tasks.py:95-98`）；安装时还用算子名作包目录名（`install.py:177`）。 + - `op` 不能已经注册：在运行时子进程里执行 `from embodiinfer.layers import available_ops`，检查 `op` 不在结果里；同时检查 `backend/torch/ops.py` 没有同名函数。 + - 模型名必须能被 `embodiinfer.policies.factory` 构造。 + - 主检出目录必须干净（`git status --porcelain` 为空），然后记录 `base = git rev-parse HEAD`。 +2. **定位调用点**（新增 `registration/sites.py`，基于 `ast`） + - `source/function/expression`：在函数体内查找与 `expression` 的 `ast.unparse` 规范化形式相同的表达式，记录文件、行号、列号。 + - `region`：取 `module:` 段的最后一级属性名（如 `down_proj`），在该模型策略源码中查找形如 `<…>.down_proj(<…>)` 的调用，列为候选。这只是启发式匹配，必须由调用方确认。 + - 对每个调用点收集包裹它的 `if` 条件，例如 `isinstance(..., QuantizedLinear)`、`torch.is_grad_enabled()`、`self.attention == "eager"`、`self.tensor_parallel.enabled`，作为 preserve 候选。再结合模型默认表，标出模块可能在运行时被替换的位置，例如 TP 会把 `o_proj`/`down_proj` 换成 `RowParallelLinear`（`embodiinfer/engine/parallel/pi05.py:48,52`）。 + - 同时检查调用点所在函数能否拿到 `KernelSelection`：方法里可以用 `self._kernels`，自由函数则需要新增 `kernels` 参数并沿调用链传递。 +3. **工作区** + ```bash + git worktree add --detach RUN/base BASE + git worktree add -b registry/- RUN/worktree BASE + ``` + 主检出目录在整个过程中保持不动。 +4. **上下文包 `context/`**（只读）： + - `AGENTS.md`； + - `layers/ops.py`、`backend/torch/ops.py`、各调用点所在文件； + - `tests/test_kernel_ops.py`，作为测试写法示例； + - 示例迁移 commit 的 diff（`git show 449ebe2 -- embodiinfer/`）； + - 0009 §4.1、§6 的摘录； + - 第 4.5 节的规则全文； + - `sites.json`。 +5. **输出与确认：** 写出 `sites.md`（每个调用点附源码片段和 preserve 条件），然后退出，等调用方确认或修改 `sites.json` 后再执行 `register run`。规格里出现 `region` 时不允许 `--yes` 跳过确认。 + +#### R2 基线录制（GPU；运行时子进程，在 `base/` worktree 中执行） + +- **工作：** + 1. **确定性预检。** 每个 parity 配置用两份全新的编译缓存各跑一次，比较两次输出。只要有任何差异就停止，报告"该配置的基线不可复现，无法做逐位一致检查"。0009 §6 记录过：不加 `TORCHINDUCTOR_DETERMINISTIC=1` 时，Inductor 两次全新编译的结果可能不同。 + 2. **保存基线输出** `golden/-1.pt`。 + 3. **记录 ATen 调用序列**：用一个只走 eager 路径的配置，对 1 个输入记录完整的 ATen 调用序列，每条记录包含算子名，以及每个张量参数的形状、步长、dtype 和设备。 +- **复用现有工具：** `verify` 里按 `kernel_profile` 收集输出的逻辑。`scripts/kernel_tuning/verify.py:181-245` 的 `_run` 已经为每个输入固定随机种子（`seed + index`），能构造 `EngineCore` 或 serving adapter，并用 `flatten_outputs` 展开输出。把它抽成公共函数 `collect_outputs(...)`，`verify` 和新命令共用，`verify` 的行为不变。 +- **新增的运行时命令：** + ```bash + # 在 RUN/base 中执行 + env PYTHONPATH=RUN/base TORCHINDUCTOR_DETERMINISTIC=1 \ + TORCHINDUCTOR_CACHE_DIR=RUN/cache/base-inductor-1 TRITON_CACHE_DIR=RUN/cache/base-inductor-1-triton \ + /abs/runtime/python -m scripts.kernel_tuning snapshot --profile none --samples 40 \ + --output RUN/golden/inductor-1.pt -- \ + --policy pi05 --checkpoint CKPT --adapter-config pi05-inductor.json \ + --inputs pi05-inputs.pt --dtype bfloat16 --num-steps 10 --capture-full-loop + + env PYTHONPATH=RUN/base /abs/runtime/python -m scripts.kernel_tuning trace-ops \ + --samples 1 --output RUN/trace/base.jsonl -- + ``` + - `snapshot`：保存每个输入的 `flatten_outputs` 结果，以及 `embodiinfer.__file__`、源码树哈希和环境信息（GPU、Torch、Triton、CUDA、`torch.get_float32_matmul_precision()`）。 + - `trace-ops`：沿用 `scripts/kernel_tuning/discovery/recorder.py:121-149` 的 `TorchDispatchMode` 观察方式，但只记录调用序列，不保存夹具。 +- **导入路径校验：** 运行时环境里 EmbodiInfer 是可编辑安装，指向主检出目录。所以每个子进程都要把 `PYTHONPATH` 设为对应 worktree，让它排在 site-packages 的 `.pth` 之前；控制器再检查输出里记录的 `embodiinfer.__file__` 确实位于该 worktree 下。不满足就直接失败，否则比较的其实是同一份代码，门禁就失效了。 + +#### R3 智能体修改（Humanize2 flow `operator_registration`） + +- **机制：** 与 `scripts/kernel_tuning/flow.py` 相同：一个可恢复的 `@flow`，一个 `coder` 智能体，工作环境是 `LocalEnv` 指向的 `worktree/`，单次超时用 `Budget`，`ctx.state` 中保存进度台账（ledger）。`__main__._execute` 中启动 `Hmz(...).run(flow=..., agents={"coder": AGENT}, params={"run_dir": RUN})` 的方式原样复用。不同之处： + - 本流程追求得到一个正确的补丁，而不是在多个候选里挑最好的一个，所以第一个通过全部门禁的修改即为结果； + - 修改在同一个 worktree 上逐步进行，失败后在同一会话里修复（见 R6）。 +- **第一轮 `plan`：** 智能体阅读 `context/` 和 `sites.json`，写出 `PLAN.md`，内容包括： + 1. 算子的数学定义，以及它与原表达式逐个 ATen 调用的对应关系； + 2. 函数签名，即哪些是张量参数、哪些是常量； + 3. 每个调用点的改法，以及 `kernels` 参数如何传到该点； + 4. 每个 preserve 分支如何保持不变； + 5. 新测试清单； + 6. 可能破坏逐位一致的风险点。 +- **第二轮 `implement`：** 在 `worktree/` 中按第 4.5 节的规则修改代码和测试，并写出 `REGISTRATION.json`： + ```json + { + "op": "linear", + "reference": "embodiinfer/backend/torch/ops.py::linear", + "signature": {"tensors": ["x", "weight", "bias?"], "constants": []}, + "sites": [{"id": "S1", "file": "embodiinfer/policies/pi05/modeling_pi05.py", "function": "_fused_qkv", + "line": 255, "call": "_op(kernels, \"linear\", False)(x, weight)"}], + "preserved": [{"site": "S1", "branch": "isinstance(projection, QuantizedLinear)", "reason": "量化注册表"}], + "expected_calls_per_inference": {"S1": 216}, + "tests": ["tests/test_kernel_ops.py::test_linear_reference_matches_f_linear"] + } + ``` + `expected_calls_per_inference` 是智能体声明的每次推理调用次数，由 R5 的命中检查核对。上例按 2 个 tower × 18 层，加上专家 tower 每个去噪步的调用估算，仅作示意。 +- **允许的操作：** 运行 `ruff` 和 CPU `pytest`。 +- **禁止的操作：** 运行 GPU 代码、`git commit`、安装依赖、修改 `context/` 或运行目录。 + +#### R4 静态门禁（控制器；CPU；按顺序执行，前一项失败时后面不再执行） + +| 门禁 | 检查内容 | 方式 | +|---|---|---| +| G1 范围 | `git -C worktree diff --name-status BASE` 只涉及允许修改的文件：`embodiinfer/backend/torch/ops.py`、`embodiinfer/layers/ops.py`、`sites.json` 涉及的源文件、`tests/test_kernel_ops.py`、`tests/test_registry_.py`。已有测试只能新增，不能删除或修改断言行。`pyproject.toml`、`uv.lock`、`scripts/`、`embodiinfer/backend/triton/tuned/` 一律不能改 | 新增 `registration/workspace.py` | +| G2 结构 | 1. `backend/torch/ops.py` 中存在 `def `，有类型注解和 docstring，且只导入 `torch` 和标准库，不依赖 `models`/`policies`/`engine`；
2. `layers/ops.py` 用 `register_lazy_op("", "torch", "embodiinfer.backend.torch.ops", "")` 注册；
3. 每个调用点改为注册表调用，原表达式只出现在 preserve 分支里；
4. 在运行时子进程中：`available_ops()[""]` 包含 `torch`，`get_op("") is backend.torch.ops.`，`KernelSelection("none").prepare()` 成功，`KernelSelection("torch").resolve("")` 返回 Torch 参考实现 | `ast` 检查 + 运行时子进程 | +| G3 格式 | 改动文件通过 `ruff check` 和 `ruff format --check` | 现有 `ruff` | +| G4 测试 | 规格中的 `tests` 和 `REGISTRATION.json` 列出的新测试全部通过。另外，把新测试放到 `base/` 上运行，必须失败，用来证明这些测试确实依赖新算子 | `test_python -m pytest … -q` | + +智能体写的测试至少要覆盖三类： + +1. 参考实现与原表达式在 CPU 随机输入（FP32 和 BF16）上 `torch.equal`； +2. 用假调优档案确认调用点经过注册表，写法见 `tests/test_kernel_ops.py:135-150`； +3. preserve 分支（量化、张量并行、训练和求梯度）不经过注册表。 + +这些测试只是补充。验收依据是 R5 中由控制器执行、对比基线版本的检查，智能体无法通过改测试绕过它。 + +#### R5 GPU 门禁（运行时子进程，在 `worktree/` 中执行；按开销从低到高排列） + +| 门禁 | 检查内容 | 调用方式 | +|---|---|---| +| G5 ATen 序列 | 候选与基线的 eager ATen 调用序列完全相同，包括算子名、形状、步长、dtype 和顺序。失败时报告第一个分歧点及前后 5 条记录 | 新增 `trace-ops`，与 R2 参数相同，`PYTHONPATH=RUN/worktree` | +| G6 命中 | 新算子被记录到；`errors` 为空；每个调用点的调用次数等于 `expected_calls_per_inference`，且不为 0 | 现有 `capture-ops`，加一个新的小选项 `--record-callers`，按调用者的文件和行号计数：
`python -m scripts.kernel_tuning capture-ops --model pi05 --ops linear --warmup 1 --steps 2 --record-callers --output RUN/attempts/N/capture -- -m scripts.kernel_tuning.discovery.drive --no-cuda-graph` | +| G7 逐位一致 | 对每个 parity 配置，候选在 `kernel_profile: none` 下的输出与 `golden/-1.pt` 完全相等（见第 6 节） | 新增 `snapshot` 和 `snapshot-compare`，环境变量与 R2 相同，使用新的缓存目录 | +| G8 布局审计（只警告） | 在生产配置下，比较 Python 实际执行（未编译、未在图捕获中）的注册调用的参数布局，与 G6 在 eager 下记录的布局是否一致。比较前按 `install.py` 的规则规范化：忽略长度为 1 的维度和空张量的步长。只在编译区域内执行的调用记为"无法观测" | 新增 `audit-layouts`，复用 `regions.py` 包装 `resolve` 的方式，只记录布局 | + +G5 不是必需的，但能把"某个输出差了多少"细化成"第 N 个 ATen 调用从 `aten.linear` 变成了 `aten.matmul`"。它的开销只有一次 eager 推理，所以放在最前面执行。 + +G8 只发警告，因为布局不一致不影响 `kernel_profile: none` 的正确性，但会让之后安装的调优 kernel 在生产中一次也选不中（commit `ecde612` 的情况）。警告会进入反馈。智能体可以把参数规范成连续布局等形式，但前提是 G5 和 G7 仍然通过；否则这个警告写进 PR 描述，由审阅者决定。 + +#### R6 反馈循环 + +- 任何门禁失败时，控制器写出 `attempts/N/feedback.json`： + ```json + {"gate": "G7", "config": "inductor", "summary": "输出与基线不一致", + "details": {"input": 3, "output": ".actions", "max_abs": 0.0078125, "mismatched": 42, "total": 350}, + "hints": ["G5 已通过：eager 序列相同，差异只出现在 Inductor 配置，检查编译区域内的参数布局或 dtype"]} + ``` + 其他门禁的 `details` 分别是：G1 的违规文件列表、G2 失败的断言、G3/G4 的输出摘要（截断到固定长度）、G5 的序列差异、G6 的调用次数对比。 +- 下一次尝试在同一个 Humanize2 会话中发出 `repair` prompt，附上 `feedback.json`，智能体在原 worktree 上继续修改。如果连续两次在同一门禁失败，就新开一个会话重新规划，worktree 保留。 +- 每次尝试结束时，把 `git -C worktree diff BASE` 存档为 `attempts/N/patch.diff`。 +- 停止条件： + - 全部门禁通过，进入 R7； + - 达到 `max_attempts` 或 `max_seconds`，状态为 `exhausted`； + - R2 确定性预检失败，或运行环境发生变化（GPU、Torch、Triton 版本，或基线 worktree 被改动），状态为 `failed`。这类问题不是智能体能修复的。 + +#### R7 终检与打包 + +1. **终检：** 用 `parity.final_inputs`（建议 ≥ 200 帧）对每个配置重跑 R2 基线和 G7，全部逐位一致才继续。循环中只用 `samples` 个输入，是为了控制每轮的 GPU 时间。 +2. **覆盖率（可选，`--coverage`）：** 用现有 `profile` 命令在生产配置下分别跑基线和候选，报告时间从哪些 `module:` 区域移到了 `op:` 区域，以及注册表覆盖的时间占比。编译区域内的调用不会有 `op:` 标签（`hotspots.py` 的模块文档字符串里有说明），报告会注明这一点。 +3. **测速（可选，`--bench`）：** 用 `benchmarks/pi05-benchmark/benchmark.py`，在相同配置和 `kernel_profile: none` 下，分别跑基线和候选。结果只报告、不卡门禁。`_op` 的分派开销只影响 eager 调用，图回放不受影响；commit `449ebe2` 测得 native 路径是 99.35 ms 对 98.47 ms。 +4. **证据 `evidence.json`：** 基线 commit；候选源码树哈希；规格摘要；每个门禁的结果；parity 配置、输入摘要和数量；环境信息；尝试次数和耗时。 +5. **`PR.md`：** 按 `.github/PULL_REQUEST_TEMPLATE.md` 的结构填写： + - Summary：算子定义、改动的调用点、保持不变的分支； + - Type of change：勾选 "Kernel, layer, or numerical change"； + - Verification：实际运行的命令和结果，以及未运行的检查； + - Parity：基线、配置、输入数、`max|Δ| = 0`； + - G8 警告和覆盖率摘要。 +6. **`finalize`：** + - 不带参数时，只打印 worktree 路径、分支名和建议的提交说明。 + - `--commit`：在 worktree 中提交到 `registry/-`。提交说明使用约定前缀，例如 `feat: route through the registry operator`，并遵守 AGENTS.md 的 Git 约定。 + - `--push`：执行 `git push -u origin `。 + - 本工具不提供合并功能。 + +#### R8 人工审阅与合并 + +审阅者重点看三件事：参考实现是否真是原数学，没有偷换表达式；preserve 分支是否完整；`expected_calls_per_inference` 是否合理。 + +合并后进入 0009 的调优流程：在合并后的代码上运行 `capture-ops --ops `，然后 `calibrate`、`tune-all`/`run`、`install`、`verify`。G6 产生的抓取结果只有在源码哈希没变时才能直接复用（`scripts/kernel_tuning/discovery/tasks.py:99-102` 会做校验）。审阅中只要改过代码，就要重新抓取。 + +### 4.5 智能体必须遵守的规则 + +以下规则写进 prompt，并由 R4 和 R5 在机制上检查： + +1. **参考实现的位置：** 写在 `embodiinfer/backend/torch/ops.py`，函数名等于算子名。参数只能是张量或常量（`int`/`float`/`bool`/`str`/`None`），返回张量或张量元组，不得原地修改输入。安装生成的 `applies()` 只能判断这两类参数（`install.py:108-122`）。 +2. **逐字保留算术：** 使用相同的 ATen 调用、相同的顺序和相同的中间 dtype 转换。不允许把 `F.linear` 改写成 `x @ w.T`，也不允许增删 `.float()`、`.to()`、`.contiguous()` 或合并、拆分 GEMM，除非 `REGISTRATION.json` 中说明理由，并且 G5、G7 仍然通过。 +3. **注册方式：** 用 `register_lazy_op("", "torch", "embodiinfer.backend.torch.ops", "")`。只有 Torch 实现的算子不能放进 `layers/ops.py:160` 那个同时注册 torch 和 triton 的循环。 +4. **调用点写法：** 使用策略现有的 helper，例如 pi0.5 的 `_op(kernels, "", False)`。只有 Torch 实现时 `prefer` 必须是 `"torch"`，否则 `get_op` 会抛出 KeyError。`KernelSelection` 用参数一路传下去，不使用全局变量。 +5. **preserve 分支：** 原样保留，只用类型判断（例如 `type(module) is nn.Linear`），不按模型名分支。 +6. **不改的内容：** 不修改 `KernelSelection`、调优档案、`scripts/`、依赖，以及已有测试的断言。 +7. **公开函数：** 带类型注解和说明契约的 docstring（AGENTS.md 的要求）。 + +### 4.6 融合区域（`kind: fuse`） + +融合区域的流程与包装相同，额外有三条要求： + +- `sites` 中必须给出相邻的一组表达式，例如 `_mlp` 中 gate/up GEMM 后接 `gated_gelu`。参考实现就是这组表达式按原顺序组成的函数，结果必须与原表达式序列逐位一致（G5、G7）。 +- 名称优先按数学含义取，例如 `gate_up_gelu`。确实只属于某个模型的区域，用 `_` 形式的标识符，例如 `pi05_residual_ada_rms_norm`。0009 §4.1 写的 `pi05.residual_ada_rms_norm` 带点号，与现有工具链对算子名的要求冲突（R1 第 1 项），因此改为下划线形式，并建议同步修订 0009 的命名示例。 +- 融合可能删掉原来单独存在的注册调用。例如被融合的 `gated_gelu` 不再单独出现，`expected_calls_per_inference` 要反映这个变化。调优档案里如果已经装过被融合掉的算子，这些 kernel 不再生效，PR 描述中需要写明。 + +### 4.7 环境与执行位置 + +- **控制器：** 工具环境，Python ≥ 3.12 + Humanize2，POSIX，与 `tune-all` 相同（`scripts/kernel_tuning/__main__.py:19-31`）。 +- **模型运行时：** 通过 `runtime_python` 指定，是模型依赖组的环境，例如 `uv sync --python 3.12 --frozen --no-dev --group pi05`。所有加载模型的步骤（R1 的注册检查、R2、G2 的运行时部分、R5、R7）都以子进程运行，带超时和进程组清理，做法同 `scripts/kernel_tuning/runner.py`。 +- **CPU 测试：** 通过 `test_python` 指定，默认等于 `runtime_python`；只能在核心开发环境运行的测试可以单独指定。 +- **机器：** 整个流程在目标 GPU 机器上运行，parity 结果只对这台机器有效，证据中会记录硬件信息。智能体的沙箱限制与 0009 相同：Thor 上 AppArmor 禁止普通用户创建用户命名空间，Codex 以无沙箱方式在低权限账号下运行。 + +### 4.8 备选方案 + +- **规则化改写（codemod），不用智能体。** 单一表达式的包装可以用 AST 替换完成，但 `kernels` 传参链、preserve 分支和融合都需要理解代码。本提案由智能体修改、由门禁兜底；将来可以给 `wrap` 加一条 codemod 快速路径，产物同样经过 R4/R5。 +- **运行时子图替换（torch.fx 或 Dynamo pattern matcher），不改源码。** 这样不需要 PR，但 eager 路径不经过 fx 图，替换也不可审阅，违背 0009"注册调用就是模型默认代码路径"的设计。拒绝。 +- **只做端到端逐位一致，不做 ATen 序列。** 失败信息太粗（只知道"差了 0.0078"），每轮都要付全部配置的 GPU 成本。两者都保留，ATen 序列放在前面执行。 +- **把注册和调优串成一条命令。** 注册改动要经过人工审阅，审阅中改动参考实现会让已有的抓取作废。所以 `register` 只做到 PR 为止，调优在合并之后进行。 + +## 5. 模型无关性结论 + +- **工具本身：** 本工具属于 `scripts/` 下的开发者工具，不进入运行时。控制器、门禁和 flow 是通用的，只依赖公开的注册表接口（`available_ops`、`get_op`、`KernelSelection`）、`discovery.drive` 的驱动参数、`profile` 和 `capture-ops` 的输出格式。 +- **模型相关的知识**（调用点、preserve 模式、parity 配置、CPU 测试表、示例 commit）放在规格文件和 `registration/models/.yaml` 中，代码里不按模型名分支。 +- **智能体产生的运行时改动** 与 `449ebe2` 的分层相同： + - `backend/torch/ops.py` 里是纯张量函数； + - `layers/ops.py` 里加一行注册； + - 调用点留在策略层。 + + 引擎层不受影响。 +- **对新模型的要求：** 必须先具备 `kernels` 传参链（一个 `KernelSelection` 和类似 `_op` 的 helper）。这个一次性的接入不属于本流程。 + +## 6. 无损判据 + +判据是逐位一致，不设容差。 + +- **参考：** 基线 commit（`spec.json` 中记录的 `base`）在同一 parity 配置、同一输入、同一随机种子（`seed + index`，与 `verify` 相同）下运行，`kernel_profile: none`。 +- **比较对象：** `flatten_outputs` 展开后的每一个输出叶子，每个输入都比较。浮点输出（如动作块）要求 `max|Δ| = 0`；整数、布尔、字符串输出要求完全相等。输出结构（叶子名称和长度）不同也算失败。 +- **配置：** `parity.configs` 中的每一项，至少覆盖每个调用点所在的执行路径，并包含生产配置（pi0.5 为 Inductor + 完整去噪 CUDA Graph）。所有运行都设置 `TORCHINDUCTOR_DETERMINISTIC=1`，每次使用全新的 `TORCHINDUCTOR_CACHE_DIR`/`TRITON_CACHE_DIR`；`float32_matmul_precision` 保持默认，并记录在证据中。 +- **前置条件：** 基线自身两次全新运行的结果完全相同（R2 确定性预检）。不满足时这次运行失败，不会放宽标准。 +- **辅助判据：** eager ATen 调用序列完全相同（G5）。 +- **复现方式：** `register run` 的 G5/G7，或者手动执行 `snapshot` + `snapshot-compare`，命令见 R2/R5。 +- **为什么可以要求逐位一致：** 注册改动只是把同一串 ATen 调用移到另一个函数里执行。Dynamo 会内联这层包装，编译出的图不变；eager 下执行的 kernel 序列也不变。所以任何差异都说明数学或布局被改了，这正是门禁要拦截的。 + +## 7. 实现计划 + +新增文件（均位于 `scripts/kernel_tuning/`）： + +| 文件 | 内容 | +|---|---| +| `registration/__init__.py` | 包说明 | +| `registration/spec.py` | 规格的 dataclass、YAML 和命令行解析、校验；冻结后写入 `spec.json` | +| `registration/sites.py` | AST 调用点定位、`region` → 源码映射、preserve 条件提取 | +| `registration/candidates.py` | R0：从 `hotspots.json` 列出未覆盖区域 | +| `registration/workspace.py` | `git worktree` 的创建和清理、diff、G1 范围检查 | +| `registration/store.py` | 运行目录、`manifest.json`、阶段状态；复用 `artifacts.atomic_json` 和 `run_lock` | +| `registration/launch.py` | 以 `runtime_python`/`test_python` 启动子进程：设置 `PYTHONPATH`、缓存目录和环境变量，处理超时，校验导入路径 | +| `registration/gates.py` | G1–G4、G5–G8 的判定和 `feedback.json` 的生成 | +| `registration/prompts.py` | plan、implement、repair 三种 prompt | +| `registration/flow.py` | Humanize2 flow `operator_registration` | +| `registration/package.py` | `evidence.json`、`PR.md`、`finalize` | +| `registration/models/pi05.yaml` | pi0.5 默认值：CPU 测试、preserve 模式、parity 配置模板、示例 commit `449ebe2` | +| `snapshot.py` | 运行时命令 `snapshot`、`snapshot-compare`（复用抽出的 `collect_outputs`） | +| `optrace.py` | 运行时命令 `trace-ops`（只记录调用序列的 `TorchDispatchMode`） | + +对现有代码的改动： + +- `verify.py`：把 `_run` 抽成公共的 `collect_outputs`，`verify` 的行为和输出不变。 +- `regions.py`：`capture-ops` 增加 `--record-callers`（按调用者的文件和行号计数），以及只记录布局的审计模式 `audit-layouts`。默认行为不变。 +- `__main__.py`：增加 `register` 子命令组（`candidates`、`new`、`run`、`resume`、`status`、`finalize`）和运行时命令（`snapshot`、`snapshot-compare`、`trace-ops`、`audit-layouts`）。 +- 文档：在 `CONTRIBUTING.md` 的 "Optional kernel tuning tools" 一节补充 `register` 用法；在 0009 §7 指向本提案，并按 §4.6 修订融合算子的命名示例。 + +不改动的部分：运行时包 `embodiinfer/`（智能体产生的改动经各自的 PR 合入）、`pyproject.toml`、`uv.lock`（工具环境已经有 `pyyaml`，见 `scripts/kernel_tuning/pyproject.toml`）。 + +## 8. 测试计划 + +CPU 测试（`tests/kernel_tuning/test_registration_*.py`，不需要 GPU，也不需要智能体）： + +- **规格：** 算子名合法性、与已注册算子重名、缺少 parity 配置、`region` 未确认时拒绝 `--yes`。 +- **调用点定位：** 在固定的示例源码上定位 `source/function/expression`；`region` → 候选调用点；提取 preserve 条件，包括 `isinstance` 和 `torch.is_grad_enabled`。 +- **G1 范围检查：** 在合成的 git 仓库上，分别验证允许的文件、越界文件、删除已有测试断言、修改 `pyproject.toml` 的情况。 +- **G2 结构检查：** 参考实现缺失、导入了 `embodiinfer.policies`、错误地放进 torch+triton 循环、原表达式残留在非 preserve 路径。 +- **G4 有效性检查：** 新测试在基线上通过时判为无效。 +- **`snapshot-compare`：** 完全相同时通过；差一个 ULP、输出结构不同、离散输出不同时都失败。 +- **`trace-ops`：** 在 CPU 上的小模块中记录序列，并在分歧点生成差异报告。 +- **flow：** 用 Humanize2 自带的假智能体（同 `tests/kernel_tuning/test_flow.py`）跑完"失败 → 反馈 → 修复 → 通过"一轮，并验证 `resume` 不会重复已完成的阶段。缺少 `hmz` 时跳过。 + +GPU 测试（标记 `gpu`，并使用 pi0.5 checkpoint 标记）： + +- 端到端回放 `449ebe2`：在 `449ebe2^` 上，以 `gated_gelu` 的同名替身（例如 `gated_gelu_v2`，避免与已注册算子冲突）为目标运行一次 `register run --until parity`，要求在预算内通过 G5–G7。 +- `linear` 作为第一个真实目标，规格见第 4.2 节，产出一个待审阅的 PR。 + +## 9. 基准计划 + +注册改动本身不应改变性能，因此不设性能门禁。 + +`--bench` 在 Jetson Thor 上用 0009 §9 的协议（B=1、BF16、10 步去噪、LIBERO-10、Inductor + 完整去噪 CUDA Graph），在 `kernel_profile: none` 下分别跑基线和候选，报告平均值、P50 和 prefill/decode 拆分。`_op` 的分派只在 eager 调用中产生开销；参考 `449ebe2` 的测量，native 路径约有 1% 的差异。 + +工具本身的成本也要在第一次运行时记录：每轮尝试的 GPU 时间（主要是 G7 各配置的模型加载和编译），以及完成一个算子所需的尝试次数，作为确定 `max_attempts` 默认值的依据。 + +## 10. 风险与限制 + +- **覆盖面有限。** parity 只覆盖配置里列出的路径。没有被执行到的分支（TP、量化、B > 1、其他 checkpoint）只靠 CPU 测试和审阅来保证。证据中会列出未覆盖的分支。 +- **调用点定位是启发式的。** `region` 到源码的映射只按属性名匹配；编译区域里的 `kernels:*` 热点无法自动定位。因此必须经过人工确认（R1）。 +- **审阅仍然必要。** 智能体的补丁可能逐位一致但结构不好，例如传参链过长或重复判断。0009 的非目标之一就是不自动合并智能体写的策略改动。 +- **GPU 成本。** 每轮尝试大约是"配置数 × 一次模型加载和编译"。通过按开销排序门禁（G5 在前）、循环中只用 `samples` 个输入、终检才用全量输入来控制。 +- **确定性依赖 Inductor 设置。** 某些配置即使设置了 `TORCHINDUCTOR_DETERMINISTIC=1` 仍无法复现时，这个配置就不能用于逐位一致检查，这次运行会失败并报告原因，不会放宽标准。 +- **导入路径。** 可编辑安装可能让子进程导入主检出目录里的代码。通过 `PYTHONPATH` 和 `embodiinfer.__file__` 校验来防范；校验不通过时直接失败。 +- **G8 观测不到编译区域内部。** 只在编译区域里执行的调用无法审计布局，只能等调优后的 `install` 和命中统计来暴露问题。 +- **工具链对算子的约束保留不变。** 算子名必须是标识符，参考实现必须在 `backend/torch/ops.py`。真正的模型专属融合如果需要放在策略目录，就要先泛化 `regions.py`、`discovery/tasks.py`、`install.py` 的相关检查，那是单独的改动。 +- **中转站断流。** 与 0009 相同，中转站断流会浪费尝试预算。 From 369fcda5a84a9a72f1cfe12288298532befbade4 Mon Sep 17 00:00:00 2001 From: owpeter Date: Tue, 13 Oct 2026 04:26:12 +0000 Subject: [PATCH 26/44] docs: withdraw the operator registration proposal The registration loop is developer tooling rather than a design proposal. Remove 0010 and its pointer from 0009; the tooling guide replaces it. Signed-off-by: owpeter --- docs/proposals/0009-hotspot-kernel-tuning.md | 3 +- .../0010-agent-operator-registration.md | 470 ------------------ 2 files changed, 1 insertion(+), 472 deletions(-) delete mode 100644 docs/proposals/0010-agent-operator-registration.md diff --git a/docs/proposals/0009-hotspot-kernel-tuning.md b/docs/proposals/0009-hotspot-kernel-tuning.md index 409dfce..2b0fbde 100644 --- a/docs/proposals/0009-hotspot-kernel-tuning.md +++ b/docs/proposals/0009-hotspot-kernel-tuning.md @@ -298,8 +298,7 @@ reproduces them. Status (2026-10-12): steps 1–3 are implemented for pi0.5; the loop was proven on Jetson Thor instead of an RTX 4090 (§9), and the `thor` profile ships one tuned operator. Qwen R2R and StreamVLN call sites, step 4 and an RTX 4090 -profile are not done; neither is shape-range tuning (§10). Step 4's registration -loop is designed in [0010](0010-agent-operator-registration.md). +profile are not done; neither is shape-range tuning (§10). New tooling lives in `scripts/kernel_tuning/` and `benchmarks/kernel_tuning/`. Existing 0008 commands keep working and are not extended. No public symbol is diff --git a/docs/proposals/0010-agent-operator-registration.md b/docs/proposals/0010-agent-operator-registration.md deleted file mode 100644 index 5f02c85..0000000 --- a/docs/proposals/0010-agent-operator-registration.md +++ /dev/null @@ -1,470 +0,0 @@ -# 0010 — 智能体驱动的算子注册流程 - -- Status: Draft -- Author: EmbodiInfer maintainers -- Date: 2026-10-13 - -## 1. 概要 - -本提案实现 [0009](0009-hotspot-kernel-tuning.md) §4.4 和实现计划第 4 步：调用方选定一个尚未进入算子注册表的算子或融合区域后，由编码智能体完成以下改动： - -- 写出它的 Torch 参考实现； -- 在 `embodiinfer/layers/ops.py` 中注册； -- 把模型调用点改为经过注册表； -- 补齐测试。 - -控制器随后依次运行静态检查、CPU 测试、命中检查和 `kernel_profile: none` 逐位一致检查。任何一项失败，都把结构化反馈交回智能体修改，直到全部通过或预算用尽。最终产物是一个可以直接提交审阅的分支和证据包。 - -整个流程位于 `scripts/kernel_tuning/registration/`，以 `python -m scripts.kernel_tuning register …` 调用，属于开发者工具，不进入运行时。它只负责把算子变成"可调优、可安装"的注册算子，不做调优，也不自动合并。注册代码经人工审阅合并后，再进入 0009 已实现的调优流程（`capture-ops` → `calibrate` → `tune-all`/`run` → `install` → `verify`）。 - -## 2. 动机与现状缺口 - -- **只有注册算子能调优并接入模型。** `capture-ops` 只记录经过 `KernelSelection.resolve` 的调用（`scripts/kernel_tuning/regions.py:152-163`）。`install` 拒绝其他任务，也拒绝当前源码里没有注册的算子（`scripts/kernel_tuning/install.py:140-144`）。 -- **注册目前完全靠人工。** pi0.5 的 5 个算子都是在 commit `449ebe2` 中手工迁移的，每个算子要做四件事： - 1. 把 Torch 实现搬到 `embodiinfer/backend/torch/ops.py`； - 2. 在 `embodiinfer/layers/ops.py:160` 注册； - 3. 把调用点改为 `_op(kernels, name, use_triton)(...)`（`embodiinfer/policies/pi05/modeling_pi05.py:88`），并把 `self._kernels` 传下去； - 4. 在 Thor 上验证 `kernel_profile: none` 逐位一致。 -- **最大的热点在注册表之外。** 0009 §10 记录，Thor 上 pi0.5 耗时最多的是 GEMM，它们不在注册表里，新流程完全碰不到。0009 设想由智能体把热点区域写成新的注册算子并提 PR（§4.4），但这一步没有实现。 -- **改调用点的风险很具体。** 仓库多处强调形状和表达式会改变舍入： - - 合并 Q/K/V 会改变归约形状，从而改变舍入（`modeling_pi05.py:664-666`）； - - `F.linear` 和 `x @ w.T` 会走不同的 ATen 分解； - - eager 与 Inductor 下参数布局可能不同（commit `ecde612`）。 - - 这些都需要固定的门禁来兜底，不能依赖智能体自觉。 - -## 3. 目标与非目标 - -目标： - -- 调用方显式选择要注册的算子：给出算子名、种类（包装或融合）和调用点。调用点可以直接写源码位置，也可以从热点报告挑选区域，再由工具解析成源码位置。 -- 一条可恢复的流水线：选择 → 定位调用点 → 记录基线 → 智能体修改 → 静态门禁 → GPU 门禁 → 反馈循环 → 打包。每一步都有明确的产物和通过条件。 -- 尽量复用现有工具：`profile`、`capture-ops`、`verify` 的输出比较逻辑、Humanize2 flow、`run_lock`/`atomic_json`。只为现有工具做不到的事新增代码。 -- 验收由控制器负责。智能体写的测试只是补充，判定依据是控制器对比基线版本的逐位一致结果。 - -非目标： - -- 调优或安装 kernel（属于 0009 的 B、C 部分）。 -- 改变数学、精度或舍入。注册改动必须在 `kernel_profile: none` 下与基线逐位一致。 -- 自动合并、自动推送或自动开 PR。默认只生成分支和 PR 描述；提交、推送、开 PR 都需要调用方显式传参。 -- 修改 `KernelSelection`、调优档案（tuned profile）、`scripts/` 或依赖。这类改动走普通 PR。 -- Qwen R2R、StreamVLN 的首次接入。它们还没有 `kernels` 传参链，第一次需要人工建立；建立之后同样可以使用本流程。 - -## 4. 设计 - -### 4.1 总览 - -``` - 调用方 控制器（工具环境） 智能体（Humanize2） 模型运行时（GPU 子进程） - │ R0 candidates ──► 读 hotspots.json，列出未覆盖热点 - │ R1 new ─────────► 校验规格、定位调用点、建 worktree、打包上下文 - │ ◄── sites.md 等待确认 - │ R2 run ─────────► 基线录制 ─────────────────────────────────────────────────► snapshot ×2、trace-ops - │ R3 ───────────────────────────────► plan / implement - │ R4 静态门禁（范围、结构、ruff、pytest） - │ R5 GPU 门禁 ─────────────────────────────────────────────► trace-ops、capture-ops、 - │ │ snapshot、audit-layouts - │ R6 失败 → feedback.json ──────────► repair（同一会话） - │ R7 全部通过 → 证据、可选覆盖率/测速、PR.md - │ finalize ───────► （显式 --commit/--push）提交到 registry/- 分支 - │ R8 人工审阅、合并 → 进入 0009 的调优流程 -``` - -| 步骤 | 工作 | 执行者 | 现有工具 | 产物 | -|---|---|---|---|---| -| R0 | 列出未被注册表覆盖的热点 | 控制器 | 读取 `profile` 的 `hotspots.json` | `candidates.md` | -| R1 | 规格校验、调用点定位、worktree、上下文 | 控制器 | `git worktree`；运行时内 `available_ops()` | `spec.json`、`sites.json`、`context/` | -| R2 | 基线输出与 ATen 序列，确定性预检 | 运行时子进程 | 复用 `verify` 的取输出逻辑（抽成 `snapshot`） | `golden/`、`trace/base.jsonl` | -| R3 | 写计划、改代码、写测试 | 智能体 | Humanize2 flow 机制（同 `flow.py`） | worktree 改动、`REGISTRATION.json` | -| R4 | 范围、结构、ruff、pytest | 控制器 | `ruff`、`pytest` | `attempts/N/static.json` | -| R5 | ATen 序列、命中、逐位一致、布局审计 | 运行时子进程 | `capture-ops`（现有）；`snapshot`、`trace-ops`、`audit-layouts`（新增） | `attempts/N/gpu.json` | -| R6 | 汇总失败交回智能体 | 控制器 + 智能体 | 同一 Humanize2 会话 | `attempts/N/feedback.json` | -| R7 | 终检、证据、可选覆盖率和测速、PR 描述 | 控制器 | `profile`（现有）、`benchmarks/pi05-benchmark/benchmark.py`（现有） | `evidence.json`、`PR.md`、`patch.diff` | -| R8 | 审阅与合并 | 人 | GitHub | 合并后的注册代码 | - -### 4.2 调用方式：选择算子 - -所有命令都在目标 GPU 机器上运行，入口是工具环境（Python ≥ 3.12 + Humanize2，与 `tune-all` 相同）。加载模型的步骤通过 `--set runtime_python=…` 指定的模型运行时解释器在子进程里执行，与 `tune-all --set evaluator_python=…` 的做法一致。 - -```bash -TOOL="uv run --project scripts/kernel_tuning python -m scripts.kernel_tuning" - -# R0（可选）：看看哪些热点不在注册表里 -$TOOL register candidates --hotspots HOT/hotspots.json --min-share 0.01 - -# R1：为选定的算子建立一次运行（规格文件和命令行参数二选一） -$TOOL register new --spec linear.yaml --output results/kernel_tuning/register -$TOOL register new --model pi05 --op linear --kind wrap \ - --site 'embodiinfer/policies/pi05/modeling_pi05.py::_fused_qkv::F.linear(x, weight)' \ - --region 'stage:denoise_step/module:expert.layers.*.mlp.down_proj' --hotspots HOT/hotspots.json \ - --parity-config parity.yaml --output results/kernel_tuning/register -# 命令结束时打印 sites.md，由调用方确认或编辑 sites.json；--yes 跳过确认。 - -# R2–R7：运行循环；--until 可以在某个阶段后停下 -$TOOL register run RUN --agent 'codex/MODEL:EFFORT' \ - --set runtime_python=/abs/runtime/python --set test_python=/abs/dev/python \ - [--until baseline|static|parity|package] -$TOOL register resume RUN -$TOOL register status RUN - -# 收尾：只有显式传参才会提交或推送 -$TOOL register finalize RUN [--commit] [--push] -``` - -规格文件（`linear.yaml`）示例： - -```yaml -model: pi05 -op: linear # 新算子名：Python 标识符、尚未注册、按数学含义命名 -kind: wrap # wrap：把一处已有表达式包成算子；fuse：把相邻几步合成一个算子 -description: > - 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)" - - region: "stage:denoise_step/module:expert.layers.*.self_attn.o_proj" # 从 hotspots.json 选 -preserve: # 必须保持原样的分支；R1 会自动检测并补充候选 - - "QuantizedLinear：继续走 layers/linear.py 的量化注册表" - - "张量并行：RowParallelLinear / ColumnParallelLinear 保持模块调用" - - "torch.is_grad_enabled() 或 training 时的原始调用" -parity: # GPU 逐位一致检查的配置；每个配置独占一个进程和一份编译缓存 - checkpoint: /ckpt/pi05_libero_finetuned_v044 - inputs: /data/pi05-inputs.pt # discovery.drive 的录制输入 - samples: 40 # 循环里使用的输入数 - final_inputs: /data/pi05-frames.pt # R7 终检使用，建议 ≥ 200 帧 - configs: - eager: {adapter_config: pi05-eager.json, dtype: bfloat16, num_steps: 10, no_cuda_graph: true} - native: {adapter_config: pi05-native.json, dtype: bfloat16, num_steps: 10, capture_full_loop: true} - inductor: {adapter_config: pi05-inductor.json, dtype: bfloat16, num_steps: 10, capture_full_loop: true} - fp32: {adapter_config: pi05-native.json, dtype: float32, num_steps: 10, no_cuda_graph: true} -tests: # 本模型相关的 CPU 测试；不写时使用模型默认表 - - tests/test_kernel_ops.py - - tests/test_pi05_components.py - - tests/test_pi05_cpu_routes.py - - tests/test_tensor_parallel.py -budget: - max_attempts: 5 - agent_timeout_seconds: 1800 - max_seconds: 14400 -``` - -`parity.configs` 中的每一项展开为 `scripts.kernel_tuning.discovery.drive` 的参数（`--adapter-config`、`--dtype`、`--num-steps`、`--capture-full-loop`、`--no-cuda-graph` 等），因此可以覆盖策略的每条执行路径。对 pi0.5 来说，`adapter_config` 里的 `policy_kwargs` 决定走哪条路径，相关开关有 `attention`、`compile_backend`、`native_inference`、`prefix_cuda_graph`、`denoise_attention`、`prefix_attention`（`modeling_pi05.py:414-430`）。调用点处在哪条路径上，就必须有一个配置走到它。例如 `attention: eager` 走的是分开的 Q/K/V GEMM（`modeling_pi05.py:664-668`），融合投影的改动不会经过它。 - -每个模型在 `scripts/kernel_tuning/registration/models/.yaml` 里提供默认值：CPU 测试表、常见的 preserve 模式（例如 `QuantizedLinear`、TP 线性层）、parity 配置模板，以及作为示例的迁移 commit（pi0.5 是 `449ebe2`）。规格文件可以覆盖其中任意一项。 - -### 4.3 运行目录与状态 - -``` -results/kernel_tuning/register// - manifest.json 状态、各阶段结果、尝试次数、耗时；atomic_json 写入，run_lock 加锁 - spec.json 冻结的规格（含基线 commit、模型、parity 配置的摘要） - sites.json / sites.md 已确认的调用点清单 - context/ 提供给智能体的只读上下文 - base/ 基线 commit 的 git worktree（detached），只供基线录制使用 - worktree/ 候选分支 registry/- 的 git worktree，智能体只能改这里 - golden/-{1,2}.pt 基线输出（两次，用于确定性预检） - trace/base.jsonl 基线 eager 推理的 ATen 调用序列 - attempts// 每次尝试的 prompt、回复、PLAN.md、patch.diff、static.json、gpu.json、feedback.json - cache/ 每个子进程独立的 TORCHINDUCTOR_CACHE_DIR / TRITON_CACHE_DIR - evidence.json PR.md patch.diff -``` - -`resume` 从 `manifest.json` 记录的阶段继续，已经通过的阶段不会重跑。规格或基线 commit 发生变化时必须新建一次运行，不允许在原运行上续跑。 - -### 4.4 各步骤详细说明 - -#### R0 候选列表（可选，不调用智能体） - -- **工作：** 读取 `profile` 的输出，列出没有被 `op:` 范围覆盖的区域，按 `us_per_step` 排序，并标明各区域的定位精度： - - `module:` 区域可以映射到源码，比如 `expert.layers.*.mlp.down_proj` → 属性 `mlp.down_proj`； - - `kernels:gemm`、`kernels:inductor` 这类按 kernel 名分类的区域，来自编译区域或无标签的图回放，无法自动映射到源码，只能由调用方手写 `site`。 -- **复用的现有工具：** `profile` 命令和它生成的 `hotspots.json`，字段见 `scripts/kernel_tuning/hotspots.py:249-296`。 - ```bash - /abs/runtime/python -m scripts.kernel_tuning profile --output HOT -- \ - -m scripts.kernel_tuning.discovery.drive --policy pi05 --checkpoint CKPT \ - --adapter-config pi05-native.json --inputs pi05-inputs.pt --dtype bfloat16 \ - --num-steps 10 --capture-full-loop - ``` -- **新增：** `registration/candidates.py`，只读 JSON。 -- **说明：** 这一步只提供参考，选哪个算子仍由调用方决定。按 roofline 下限估算可回收时间是 0009 的后续工作，本步骤不实现。 - -#### R1 建立运行（`register new`，不调用智能体，不需要 GPU） - -1. **规格校验** - - `op` 必须是 Python 标识符。现有工具链要求参考实现函数名等于算子名，并在 `embodiinfer.backend.torch.ops` 中（`scripts/kernel_tuning/regions.py:215-217`、`scripts/kernel_tuning/discovery/tasks.py:95-98`）；安装时还用算子名作包目录名（`install.py:177`）。 - - `op` 不能已经注册：在运行时子进程里执行 `from embodiinfer.layers import available_ops`，检查 `op` 不在结果里；同时检查 `backend/torch/ops.py` 没有同名函数。 - - 模型名必须能被 `embodiinfer.policies.factory` 构造。 - - 主检出目录必须干净（`git status --porcelain` 为空），然后记录 `base = git rev-parse HEAD`。 -2. **定位调用点**（新增 `registration/sites.py`，基于 `ast`） - - `source/function/expression`：在函数体内查找与 `expression` 的 `ast.unparse` 规范化形式相同的表达式，记录文件、行号、列号。 - - `region`：取 `module:` 段的最后一级属性名（如 `down_proj`），在该模型策略源码中查找形如 `<…>.down_proj(<…>)` 的调用，列为候选。这只是启发式匹配，必须由调用方确认。 - - 对每个调用点收集包裹它的 `if` 条件，例如 `isinstance(..., QuantizedLinear)`、`torch.is_grad_enabled()`、`self.attention == "eager"`、`self.tensor_parallel.enabled`，作为 preserve 候选。再结合模型默认表，标出模块可能在运行时被替换的位置，例如 TP 会把 `o_proj`/`down_proj` 换成 `RowParallelLinear`（`embodiinfer/engine/parallel/pi05.py:48,52`）。 - - 同时检查调用点所在函数能否拿到 `KernelSelection`：方法里可以用 `self._kernels`，自由函数则需要新增 `kernels` 参数并沿调用链传递。 -3. **工作区** - ```bash - git worktree add --detach RUN/base BASE - git worktree add -b registry/- RUN/worktree BASE - ``` - 主检出目录在整个过程中保持不动。 -4. **上下文包 `context/`**（只读）： - - `AGENTS.md`； - - `layers/ops.py`、`backend/torch/ops.py`、各调用点所在文件； - - `tests/test_kernel_ops.py`，作为测试写法示例； - - 示例迁移 commit 的 diff（`git show 449ebe2 -- embodiinfer/`）； - - 0009 §4.1、§6 的摘录； - - 第 4.5 节的规则全文； - - `sites.json`。 -5. **输出与确认：** 写出 `sites.md`（每个调用点附源码片段和 preserve 条件），然后退出，等调用方确认或修改 `sites.json` 后再执行 `register run`。规格里出现 `region` 时不允许 `--yes` 跳过确认。 - -#### R2 基线录制（GPU；运行时子进程，在 `base/` worktree 中执行） - -- **工作：** - 1. **确定性预检。** 每个 parity 配置用两份全新的编译缓存各跑一次，比较两次输出。只要有任何差异就停止，报告"该配置的基线不可复现，无法做逐位一致检查"。0009 §6 记录过：不加 `TORCHINDUCTOR_DETERMINISTIC=1` 时，Inductor 两次全新编译的结果可能不同。 - 2. **保存基线输出** `golden/-1.pt`。 - 3. **记录 ATen 调用序列**：用一个只走 eager 路径的配置，对 1 个输入记录完整的 ATen 调用序列，每条记录包含算子名，以及每个张量参数的形状、步长、dtype 和设备。 -- **复用现有工具：** `verify` 里按 `kernel_profile` 收集输出的逻辑。`scripts/kernel_tuning/verify.py:181-245` 的 `_run` 已经为每个输入固定随机种子（`seed + index`），能构造 `EngineCore` 或 serving adapter，并用 `flatten_outputs` 展开输出。把它抽成公共函数 `collect_outputs(...)`，`verify` 和新命令共用，`verify` 的行为不变。 -- **新增的运行时命令：** - ```bash - # 在 RUN/base 中执行 - env PYTHONPATH=RUN/base TORCHINDUCTOR_DETERMINISTIC=1 \ - TORCHINDUCTOR_CACHE_DIR=RUN/cache/base-inductor-1 TRITON_CACHE_DIR=RUN/cache/base-inductor-1-triton \ - /abs/runtime/python -m scripts.kernel_tuning snapshot --profile none --samples 40 \ - --output RUN/golden/inductor-1.pt -- \ - --policy pi05 --checkpoint CKPT --adapter-config pi05-inductor.json \ - --inputs pi05-inputs.pt --dtype bfloat16 --num-steps 10 --capture-full-loop - - env PYTHONPATH=RUN/base /abs/runtime/python -m scripts.kernel_tuning trace-ops \ - --samples 1 --output RUN/trace/base.jsonl -- - ``` - - `snapshot`：保存每个输入的 `flatten_outputs` 结果，以及 `embodiinfer.__file__`、源码树哈希和环境信息（GPU、Torch、Triton、CUDA、`torch.get_float32_matmul_precision()`）。 - - `trace-ops`：沿用 `scripts/kernel_tuning/discovery/recorder.py:121-149` 的 `TorchDispatchMode` 观察方式，但只记录调用序列，不保存夹具。 -- **导入路径校验：** 运行时环境里 EmbodiInfer 是可编辑安装，指向主检出目录。所以每个子进程都要把 `PYTHONPATH` 设为对应 worktree，让它排在 site-packages 的 `.pth` 之前；控制器再检查输出里记录的 `embodiinfer.__file__` 确实位于该 worktree 下。不满足就直接失败，否则比较的其实是同一份代码，门禁就失效了。 - -#### R3 智能体修改（Humanize2 flow `operator_registration`） - -- **机制：** 与 `scripts/kernel_tuning/flow.py` 相同：一个可恢复的 `@flow`，一个 `coder` 智能体，工作环境是 `LocalEnv` 指向的 `worktree/`，单次超时用 `Budget`，`ctx.state` 中保存进度台账（ledger）。`__main__._execute` 中启动 `Hmz(...).run(flow=..., agents={"coder": AGENT}, params={"run_dir": RUN})` 的方式原样复用。不同之处： - - 本流程追求得到一个正确的补丁，而不是在多个候选里挑最好的一个，所以第一个通过全部门禁的修改即为结果； - - 修改在同一个 worktree 上逐步进行，失败后在同一会话里修复（见 R6）。 -- **第一轮 `plan`：** 智能体阅读 `context/` 和 `sites.json`，写出 `PLAN.md`，内容包括： - 1. 算子的数学定义，以及它与原表达式逐个 ATen 调用的对应关系； - 2. 函数签名，即哪些是张量参数、哪些是常量； - 3. 每个调用点的改法，以及 `kernels` 参数如何传到该点； - 4. 每个 preserve 分支如何保持不变； - 5. 新测试清单； - 6. 可能破坏逐位一致的风险点。 -- **第二轮 `implement`：** 在 `worktree/` 中按第 4.5 节的规则修改代码和测试，并写出 `REGISTRATION.json`： - ```json - { - "op": "linear", - "reference": "embodiinfer/backend/torch/ops.py::linear", - "signature": {"tensors": ["x", "weight", "bias?"], "constants": []}, - "sites": [{"id": "S1", "file": "embodiinfer/policies/pi05/modeling_pi05.py", "function": "_fused_qkv", - "line": 255, "call": "_op(kernels, \"linear\", False)(x, weight)"}], - "preserved": [{"site": "S1", "branch": "isinstance(projection, QuantizedLinear)", "reason": "量化注册表"}], - "expected_calls_per_inference": {"S1": 216}, - "tests": ["tests/test_kernel_ops.py::test_linear_reference_matches_f_linear"] - } - ``` - `expected_calls_per_inference` 是智能体声明的每次推理调用次数，由 R5 的命中检查核对。上例按 2 个 tower × 18 层，加上专家 tower 每个去噪步的调用估算，仅作示意。 -- **允许的操作：** 运行 `ruff` 和 CPU `pytest`。 -- **禁止的操作：** 运行 GPU 代码、`git commit`、安装依赖、修改 `context/` 或运行目录。 - -#### R4 静态门禁（控制器；CPU；按顺序执行，前一项失败时后面不再执行） - -| 门禁 | 检查内容 | 方式 | -|---|---|---| -| G1 范围 | `git -C worktree diff --name-status BASE` 只涉及允许修改的文件：`embodiinfer/backend/torch/ops.py`、`embodiinfer/layers/ops.py`、`sites.json` 涉及的源文件、`tests/test_kernel_ops.py`、`tests/test_registry_.py`。已有测试只能新增，不能删除或修改断言行。`pyproject.toml`、`uv.lock`、`scripts/`、`embodiinfer/backend/triton/tuned/` 一律不能改 | 新增 `registration/workspace.py` | -| G2 结构 | 1. `backend/torch/ops.py` 中存在 `def `，有类型注解和 docstring，且只导入 `torch` 和标准库，不依赖 `models`/`policies`/`engine`；
2. `layers/ops.py` 用 `register_lazy_op("", "torch", "embodiinfer.backend.torch.ops", "")` 注册；
3. 每个调用点改为注册表调用，原表达式只出现在 preserve 分支里；
4. 在运行时子进程中：`available_ops()[""]` 包含 `torch`，`get_op("") is backend.torch.ops.`，`KernelSelection("none").prepare()` 成功，`KernelSelection("torch").resolve("")` 返回 Torch 参考实现 | `ast` 检查 + 运行时子进程 | -| G3 格式 | 改动文件通过 `ruff check` 和 `ruff format --check` | 现有 `ruff` | -| G4 测试 | 规格中的 `tests` 和 `REGISTRATION.json` 列出的新测试全部通过。另外，把新测试放到 `base/` 上运行，必须失败，用来证明这些测试确实依赖新算子 | `test_python -m pytest … -q` | - -智能体写的测试至少要覆盖三类： - -1. 参考实现与原表达式在 CPU 随机输入（FP32 和 BF16）上 `torch.equal`； -2. 用假调优档案确认调用点经过注册表，写法见 `tests/test_kernel_ops.py:135-150`； -3. preserve 分支（量化、张量并行、训练和求梯度）不经过注册表。 - -这些测试只是补充。验收依据是 R5 中由控制器执行、对比基线版本的检查，智能体无法通过改测试绕过它。 - -#### R5 GPU 门禁（运行时子进程，在 `worktree/` 中执行；按开销从低到高排列） - -| 门禁 | 检查内容 | 调用方式 | -|---|---|---| -| G5 ATen 序列 | 候选与基线的 eager ATen 调用序列完全相同，包括算子名、形状、步长、dtype 和顺序。失败时报告第一个分歧点及前后 5 条记录 | 新增 `trace-ops`，与 R2 参数相同，`PYTHONPATH=RUN/worktree` | -| G6 命中 | 新算子被记录到；`errors` 为空；每个调用点的调用次数等于 `expected_calls_per_inference`，且不为 0 | 现有 `capture-ops`，加一个新的小选项 `--record-callers`，按调用者的文件和行号计数：
`python -m scripts.kernel_tuning capture-ops --model pi05 --ops linear --warmup 1 --steps 2 --record-callers --output RUN/attempts/N/capture -- -m scripts.kernel_tuning.discovery.drive --no-cuda-graph` | -| G7 逐位一致 | 对每个 parity 配置，候选在 `kernel_profile: none` 下的输出与 `golden/-1.pt` 完全相等（见第 6 节） | 新增 `snapshot` 和 `snapshot-compare`，环境变量与 R2 相同，使用新的缓存目录 | -| G8 布局审计（只警告） | 在生产配置下，比较 Python 实际执行（未编译、未在图捕获中）的注册调用的参数布局，与 G6 在 eager 下记录的布局是否一致。比较前按 `install.py` 的规则规范化：忽略长度为 1 的维度和空张量的步长。只在编译区域内执行的调用记为"无法观测" | 新增 `audit-layouts`，复用 `regions.py` 包装 `resolve` 的方式，只记录布局 | - -G5 不是必需的，但能把"某个输出差了多少"细化成"第 N 个 ATen 调用从 `aten.linear` 变成了 `aten.matmul`"。它的开销只有一次 eager 推理，所以放在最前面执行。 - -G8 只发警告，因为布局不一致不影响 `kernel_profile: none` 的正确性，但会让之后安装的调优 kernel 在生产中一次也选不中（commit `ecde612` 的情况）。警告会进入反馈。智能体可以把参数规范成连续布局等形式，但前提是 G5 和 G7 仍然通过；否则这个警告写进 PR 描述，由审阅者决定。 - -#### R6 反馈循环 - -- 任何门禁失败时，控制器写出 `attempts/N/feedback.json`： - ```json - {"gate": "G7", "config": "inductor", "summary": "输出与基线不一致", - "details": {"input": 3, "output": ".actions", "max_abs": 0.0078125, "mismatched": 42, "total": 350}, - "hints": ["G5 已通过：eager 序列相同，差异只出现在 Inductor 配置，检查编译区域内的参数布局或 dtype"]} - ``` - 其他门禁的 `details` 分别是：G1 的违规文件列表、G2 失败的断言、G3/G4 的输出摘要（截断到固定长度）、G5 的序列差异、G6 的调用次数对比。 -- 下一次尝试在同一个 Humanize2 会话中发出 `repair` prompt，附上 `feedback.json`，智能体在原 worktree 上继续修改。如果连续两次在同一门禁失败，就新开一个会话重新规划，worktree 保留。 -- 每次尝试结束时，把 `git -C worktree diff BASE` 存档为 `attempts/N/patch.diff`。 -- 停止条件： - - 全部门禁通过，进入 R7； - - 达到 `max_attempts` 或 `max_seconds`，状态为 `exhausted`； - - R2 确定性预检失败，或运行环境发生变化（GPU、Torch、Triton 版本，或基线 worktree 被改动），状态为 `failed`。这类问题不是智能体能修复的。 - -#### R7 终检与打包 - -1. **终检：** 用 `parity.final_inputs`（建议 ≥ 200 帧）对每个配置重跑 R2 基线和 G7，全部逐位一致才继续。循环中只用 `samples` 个输入，是为了控制每轮的 GPU 时间。 -2. **覆盖率（可选，`--coverage`）：** 用现有 `profile` 命令在生产配置下分别跑基线和候选，报告时间从哪些 `module:` 区域移到了 `op:` 区域，以及注册表覆盖的时间占比。编译区域内的调用不会有 `op:` 标签（`hotspots.py` 的模块文档字符串里有说明），报告会注明这一点。 -3. **测速（可选，`--bench`）：** 用 `benchmarks/pi05-benchmark/benchmark.py`，在相同配置和 `kernel_profile: none` 下，分别跑基线和候选。结果只报告、不卡门禁。`_op` 的分派开销只影响 eager 调用，图回放不受影响；commit `449ebe2` 测得 native 路径是 99.35 ms 对 98.47 ms。 -4. **证据 `evidence.json`：** 基线 commit；候选源码树哈希；规格摘要；每个门禁的结果；parity 配置、输入摘要和数量；环境信息；尝试次数和耗时。 -5. **`PR.md`：** 按 `.github/PULL_REQUEST_TEMPLATE.md` 的结构填写： - - Summary：算子定义、改动的调用点、保持不变的分支； - - Type of change：勾选 "Kernel, layer, or numerical change"； - - Verification：实际运行的命令和结果，以及未运行的检查； - - Parity：基线、配置、输入数、`max|Δ| = 0`； - - G8 警告和覆盖率摘要。 -6. **`finalize`：** - - 不带参数时，只打印 worktree 路径、分支名和建议的提交说明。 - - `--commit`：在 worktree 中提交到 `registry/-`。提交说明使用约定前缀，例如 `feat: route through the registry operator`，并遵守 AGENTS.md 的 Git 约定。 - - `--push`：执行 `git push -u origin `。 - - 本工具不提供合并功能。 - -#### R8 人工审阅与合并 - -审阅者重点看三件事：参考实现是否真是原数学，没有偷换表达式；preserve 分支是否完整；`expected_calls_per_inference` 是否合理。 - -合并后进入 0009 的调优流程：在合并后的代码上运行 `capture-ops --ops `，然后 `calibrate`、`tune-all`/`run`、`install`、`verify`。G6 产生的抓取结果只有在源码哈希没变时才能直接复用（`scripts/kernel_tuning/discovery/tasks.py:99-102` 会做校验）。审阅中只要改过代码，就要重新抓取。 - -### 4.5 智能体必须遵守的规则 - -以下规则写进 prompt，并由 R4 和 R5 在机制上检查： - -1. **参考实现的位置：** 写在 `embodiinfer/backend/torch/ops.py`，函数名等于算子名。参数只能是张量或常量（`int`/`float`/`bool`/`str`/`None`），返回张量或张量元组，不得原地修改输入。安装生成的 `applies()` 只能判断这两类参数（`install.py:108-122`）。 -2. **逐字保留算术：** 使用相同的 ATen 调用、相同的顺序和相同的中间 dtype 转换。不允许把 `F.linear` 改写成 `x @ w.T`，也不允许增删 `.float()`、`.to()`、`.contiguous()` 或合并、拆分 GEMM，除非 `REGISTRATION.json` 中说明理由，并且 G5、G7 仍然通过。 -3. **注册方式：** 用 `register_lazy_op("", "torch", "embodiinfer.backend.torch.ops", "")`。只有 Torch 实现的算子不能放进 `layers/ops.py:160` 那个同时注册 torch 和 triton 的循环。 -4. **调用点写法：** 使用策略现有的 helper，例如 pi0.5 的 `_op(kernels, "", False)`。只有 Torch 实现时 `prefer` 必须是 `"torch"`，否则 `get_op` 会抛出 KeyError。`KernelSelection` 用参数一路传下去，不使用全局变量。 -5. **preserve 分支：** 原样保留，只用类型判断（例如 `type(module) is nn.Linear`），不按模型名分支。 -6. **不改的内容：** 不修改 `KernelSelection`、调优档案、`scripts/`、依赖，以及已有测试的断言。 -7. **公开函数：** 带类型注解和说明契约的 docstring（AGENTS.md 的要求）。 - -### 4.6 融合区域（`kind: fuse`） - -融合区域的流程与包装相同，额外有三条要求： - -- `sites` 中必须给出相邻的一组表达式，例如 `_mlp` 中 gate/up GEMM 后接 `gated_gelu`。参考实现就是这组表达式按原顺序组成的函数，结果必须与原表达式序列逐位一致（G5、G7）。 -- 名称优先按数学含义取，例如 `gate_up_gelu`。确实只属于某个模型的区域，用 `_` 形式的标识符，例如 `pi05_residual_ada_rms_norm`。0009 §4.1 写的 `pi05.residual_ada_rms_norm` 带点号，与现有工具链对算子名的要求冲突（R1 第 1 项），因此改为下划线形式，并建议同步修订 0009 的命名示例。 -- 融合可能删掉原来单独存在的注册调用。例如被融合的 `gated_gelu` 不再单独出现，`expected_calls_per_inference` 要反映这个变化。调优档案里如果已经装过被融合掉的算子，这些 kernel 不再生效，PR 描述中需要写明。 - -### 4.7 环境与执行位置 - -- **控制器：** 工具环境，Python ≥ 3.12 + Humanize2，POSIX，与 `tune-all` 相同（`scripts/kernel_tuning/__main__.py:19-31`）。 -- **模型运行时：** 通过 `runtime_python` 指定，是模型依赖组的环境，例如 `uv sync --python 3.12 --frozen --no-dev --group pi05`。所有加载模型的步骤（R1 的注册检查、R2、G2 的运行时部分、R5、R7）都以子进程运行，带超时和进程组清理，做法同 `scripts/kernel_tuning/runner.py`。 -- **CPU 测试：** 通过 `test_python` 指定，默认等于 `runtime_python`；只能在核心开发环境运行的测试可以单独指定。 -- **机器：** 整个流程在目标 GPU 机器上运行，parity 结果只对这台机器有效，证据中会记录硬件信息。智能体的沙箱限制与 0009 相同：Thor 上 AppArmor 禁止普通用户创建用户命名空间，Codex 以无沙箱方式在低权限账号下运行。 - -### 4.8 备选方案 - -- **规则化改写（codemod），不用智能体。** 单一表达式的包装可以用 AST 替换完成，但 `kernels` 传参链、preserve 分支和融合都需要理解代码。本提案由智能体修改、由门禁兜底；将来可以给 `wrap` 加一条 codemod 快速路径，产物同样经过 R4/R5。 -- **运行时子图替换（torch.fx 或 Dynamo pattern matcher），不改源码。** 这样不需要 PR，但 eager 路径不经过 fx 图，替换也不可审阅，违背 0009"注册调用就是模型默认代码路径"的设计。拒绝。 -- **只做端到端逐位一致，不做 ATen 序列。** 失败信息太粗（只知道"差了 0.0078"），每轮都要付全部配置的 GPU 成本。两者都保留，ATen 序列放在前面执行。 -- **把注册和调优串成一条命令。** 注册改动要经过人工审阅，审阅中改动参考实现会让已有的抓取作废。所以 `register` 只做到 PR 为止，调优在合并之后进行。 - -## 5. 模型无关性结论 - -- **工具本身：** 本工具属于 `scripts/` 下的开发者工具，不进入运行时。控制器、门禁和 flow 是通用的，只依赖公开的注册表接口（`available_ops`、`get_op`、`KernelSelection`）、`discovery.drive` 的驱动参数、`profile` 和 `capture-ops` 的输出格式。 -- **模型相关的知识**（调用点、preserve 模式、parity 配置、CPU 测试表、示例 commit）放在规格文件和 `registration/models/.yaml` 中，代码里不按模型名分支。 -- **智能体产生的运行时改动** 与 `449ebe2` 的分层相同： - - `backend/torch/ops.py` 里是纯张量函数； - - `layers/ops.py` 里加一行注册； - - 调用点留在策略层。 - - 引擎层不受影响。 -- **对新模型的要求：** 必须先具备 `kernels` 传参链（一个 `KernelSelection` 和类似 `_op` 的 helper）。这个一次性的接入不属于本流程。 - -## 6. 无损判据 - -判据是逐位一致，不设容差。 - -- **参考：** 基线 commit（`spec.json` 中记录的 `base`）在同一 parity 配置、同一输入、同一随机种子（`seed + index`，与 `verify` 相同）下运行，`kernel_profile: none`。 -- **比较对象：** `flatten_outputs` 展开后的每一个输出叶子，每个输入都比较。浮点输出（如动作块）要求 `max|Δ| = 0`；整数、布尔、字符串输出要求完全相等。输出结构（叶子名称和长度）不同也算失败。 -- **配置：** `parity.configs` 中的每一项，至少覆盖每个调用点所在的执行路径，并包含生产配置（pi0.5 为 Inductor + 完整去噪 CUDA Graph）。所有运行都设置 `TORCHINDUCTOR_DETERMINISTIC=1`，每次使用全新的 `TORCHINDUCTOR_CACHE_DIR`/`TRITON_CACHE_DIR`；`float32_matmul_precision` 保持默认，并记录在证据中。 -- **前置条件：** 基线自身两次全新运行的结果完全相同（R2 确定性预检）。不满足时这次运行失败，不会放宽标准。 -- **辅助判据：** eager ATen 调用序列完全相同（G5）。 -- **复现方式：** `register run` 的 G5/G7，或者手动执行 `snapshot` + `snapshot-compare`，命令见 R2/R5。 -- **为什么可以要求逐位一致：** 注册改动只是把同一串 ATen 调用移到另一个函数里执行。Dynamo 会内联这层包装，编译出的图不变；eager 下执行的 kernel 序列也不变。所以任何差异都说明数学或布局被改了，这正是门禁要拦截的。 - -## 7. 实现计划 - -新增文件（均位于 `scripts/kernel_tuning/`）： - -| 文件 | 内容 | -|---|---| -| `registration/__init__.py` | 包说明 | -| `registration/spec.py` | 规格的 dataclass、YAML 和命令行解析、校验；冻结后写入 `spec.json` | -| `registration/sites.py` | AST 调用点定位、`region` → 源码映射、preserve 条件提取 | -| `registration/candidates.py` | R0：从 `hotspots.json` 列出未覆盖区域 | -| `registration/workspace.py` | `git worktree` 的创建和清理、diff、G1 范围检查 | -| `registration/store.py` | 运行目录、`manifest.json`、阶段状态；复用 `artifacts.atomic_json` 和 `run_lock` | -| `registration/launch.py` | 以 `runtime_python`/`test_python` 启动子进程：设置 `PYTHONPATH`、缓存目录和环境变量，处理超时，校验导入路径 | -| `registration/gates.py` | G1–G4、G5–G8 的判定和 `feedback.json` 的生成 | -| `registration/prompts.py` | plan、implement、repair 三种 prompt | -| `registration/flow.py` | Humanize2 flow `operator_registration` | -| `registration/package.py` | `evidence.json`、`PR.md`、`finalize` | -| `registration/models/pi05.yaml` | pi0.5 默认值：CPU 测试、preserve 模式、parity 配置模板、示例 commit `449ebe2` | -| `snapshot.py` | 运行时命令 `snapshot`、`snapshot-compare`（复用抽出的 `collect_outputs`） | -| `optrace.py` | 运行时命令 `trace-ops`（只记录调用序列的 `TorchDispatchMode`） | - -对现有代码的改动： - -- `verify.py`：把 `_run` 抽成公共的 `collect_outputs`，`verify` 的行为和输出不变。 -- `regions.py`：`capture-ops` 增加 `--record-callers`（按调用者的文件和行号计数），以及只记录布局的审计模式 `audit-layouts`。默认行为不变。 -- `__main__.py`：增加 `register` 子命令组（`candidates`、`new`、`run`、`resume`、`status`、`finalize`）和运行时命令（`snapshot`、`snapshot-compare`、`trace-ops`、`audit-layouts`）。 -- 文档：在 `CONTRIBUTING.md` 的 "Optional kernel tuning tools" 一节补充 `register` 用法；在 0009 §7 指向本提案，并按 §4.6 修订融合算子的命名示例。 - -不改动的部分：运行时包 `embodiinfer/`（智能体产生的改动经各自的 PR 合入）、`pyproject.toml`、`uv.lock`（工具环境已经有 `pyyaml`，见 `scripts/kernel_tuning/pyproject.toml`）。 - -## 8. 测试计划 - -CPU 测试（`tests/kernel_tuning/test_registration_*.py`，不需要 GPU，也不需要智能体）： - -- **规格：** 算子名合法性、与已注册算子重名、缺少 parity 配置、`region` 未确认时拒绝 `--yes`。 -- **调用点定位：** 在固定的示例源码上定位 `source/function/expression`；`region` → 候选调用点；提取 preserve 条件，包括 `isinstance` 和 `torch.is_grad_enabled`。 -- **G1 范围检查：** 在合成的 git 仓库上，分别验证允许的文件、越界文件、删除已有测试断言、修改 `pyproject.toml` 的情况。 -- **G2 结构检查：** 参考实现缺失、导入了 `embodiinfer.policies`、错误地放进 torch+triton 循环、原表达式残留在非 preserve 路径。 -- **G4 有效性检查：** 新测试在基线上通过时判为无效。 -- **`snapshot-compare`：** 完全相同时通过；差一个 ULP、输出结构不同、离散输出不同时都失败。 -- **`trace-ops`：** 在 CPU 上的小模块中记录序列，并在分歧点生成差异报告。 -- **flow：** 用 Humanize2 自带的假智能体（同 `tests/kernel_tuning/test_flow.py`）跑完"失败 → 反馈 → 修复 → 通过"一轮，并验证 `resume` 不会重复已完成的阶段。缺少 `hmz` 时跳过。 - -GPU 测试（标记 `gpu`，并使用 pi0.5 checkpoint 标记）： - -- 端到端回放 `449ebe2`：在 `449ebe2^` 上，以 `gated_gelu` 的同名替身（例如 `gated_gelu_v2`，避免与已注册算子冲突）为目标运行一次 `register run --until parity`，要求在预算内通过 G5–G7。 -- `linear` 作为第一个真实目标，规格见第 4.2 节，产出一个待审阅的 PR。 - -## 9. 基准计划 - -注册改动本身不应改变性能，因此不设性能门禁。 - -`--bench` 在 Jetson Thor 上用 0009 §9 的协议（B=1、BF16、10 步去噪、LIBERO-10、Inductor + 完整去噪 CUDA Graph），在 `kernel_profile: none` 下分别跑基线和候选，报告平均值、P50 和 prefill/decode 拆分。`_op` 的分派只在 eager 调用中产生开销；参考 `449ebe2` 的测量，native 路径约有 1% 的差异。 - -工具本身的成本也要在第一次运行时记录：每轮尝试的 GPU 时间（主要是 G7 各配置的模型加载和编译），以及完成一个算子所需的尝试次数，作为确定 `max_attempts` 默认值的依据。 - -## 10. 风险与限制 - -- **覆盖面有限。** parity 只覆盖配置里列出的路径。没有被执行到的分支（TP、量化、B > 1、其他 checkpoint）只靠 CPU 测试和审阅来保证。证据中会列出未覆盖的分支。 -- **调用点定位是启发式的。** `region` 到源码的映射只按属性名匹配；编译区域里的 `kernels:*` 热点无法自动定位。因此必须经过人工确认（R1）。 -- **审阅仍然必要。** 智能体的补丁可能逐位一致但结构不好，例如传参链过长或重复判断。0009 的非目标之一就是不自动合并智能体写的策略改动。 -- **GPU 成本。** 每轮尝试大约是"配置数 × 一次模型加载和编译"。通过按开销排序门禁（G5 在前）、循环中只用 `samples` 个输入、终检才用全量输入来控制。 -- **确定性依赖 Inductor 设置。** 某些配置即使设置了 `TORCHINDUCTOR_DETERMINISTIC=1` 仍无法复现时，这个配置就不能用于逐位一致检查，这次运行会失败并报告原因，不会放宽标准。 -- **导入路径。** 可编辑安装可能让子进程导入主检出目录里的代码。通过 `PYTHONPATH` 和 `embodiinfer.__file__` 校验来防范；校验不通过时直接失败。 -- **G8 观测不到编译区域内部。** 只在编译区域里执行的调用无法审计布局，只能等调优后的 `install` 和命中统计来暴露问题。 -- **工具链对算子的约束保留不变。** 算子名必须是标识符，参考实现必须在 `backend/torch/ops.py`。真正的模型专属融合如果需要放在策略目录，就要先泛化 `regions.py`、`discovery/tasks.py`、`install.py` 的相关检查，那是单独的改动。 -- **中转站断流。** 与 0009 相同，中转站断流会浪费尝试预算。 From d4b1a267c64ccbef9fc5a5fa44bb804e5d1cf614 Mon Sep 17 00:00:00 2001 From: owpeter Date: Tue, 13 Oct 2026 04:26:12 +0000 Subject: [PATCH 27/44] docs: describe the agent loop that registers operators for kernel tuning Specify the register tooling under scripts/kernel_tuning/registration: the caller selects an operator and its call sites, a coding agent writes the Torch reference, registers it, routes the call sites through the registry and adds tests, and CPU checks on scope, structure, ruff, tests and test effectiveness feed failures back until the change passes and a review branch is produced. GPU parity and tuning stay outside the loop. Signed-off-by: owpeter --- scripts/kernel_tuning/registration/README.md | 308 +++++++++++++++++++ 1 file changed, 308 insertions(+) create mode 100644 scripts/kernel_tuning/registration/README.md diff --git a/scripts/kernel_tuning/registration/README.md b/scripts/kernel_tuning/registration/README.md new file mode 100644 index 0000000..cd6d679 --- /dev/null +++ b/scripts/kernel_tuning/registration/README.md @@ -0,0 +1,308 @@ +# 算子注册流程（智能体闭环） + +本目录是 `python -m scripts.kernel_tuning register …` 工具链的规格说明，代码尚未实现。 + +本工具链自动完成"把一个算子接入算子注册表"这一步： + +1. 调用方选定算子； +2. 工具定位调用点； +3. 智能体修改代码、注册算子、编写测试； +4. 控制器运行检查； +5. 失败时把信息反馈给智能体重做； +6. 全部通过后产出一个待审阅的分支。 + +不在本流程内的工作： + +- 调优和安装 kernel：使用现有的 `capture-ops`、`calibrate`、`tune-all`/`run`、`install`、`verify`。 +- GPU 上的整模型逐位一致验证：按 [0009 §6](../../../docs/proposals/0009-hotspot-kernel-tuning.md) 的要求，合并前在目标 GPU 上单独做一次。 +- 合并：由人工完成。 + +本流程的检查全部在 CPU 上运行，因此不需要 GPU。 + +## 1. 流程总览 + +``` +选定算子 ──► 定位调用点 ──► 智能体修改 ──► 检查 ──通过──► 产出分支 + (步骤1) (步骤2) (步骤3) (步骤4) (步骤6) + │ ▲ │ + 调用方确认 └── 反馈 ◄───┘ 失败 (步骤5) +``` + +| 步骤 | 工作 | 执行者 | 使用的现有工具 | 产物 | +|---|---|---|---|---| +| 1 选定算子 | 校验规格，记录基线 commit，建立 worktree | 控制器 | `git worktree`；在测试环境中调用 `embodiinfer.layers.available_ops` | `spec.json`、`base/`、`worktree/` | +| 2 定位调用点 | 找到要改的表达式，以及需要保持原样的分支 | 控制器，调用方确认 | 无（`ast`） | `sites.json`、`sites.md` | +| 3 智能体修改 | 写参考实现、注册、改调用点、写测试 | 智能体 | Humanize2 flow 机制（同 `scripts/kernel_tuning/flow.py`） | worktree 中的改动、`PLAN.md`、`REGISTRATION.json` | +| 4 检查 | 范围、结构、格式、测试、测试有效性 | 控制器 | `ruff`、`pytest` | `attempts//checks.json` | +| 5 反馈重做 | 把失败信息交回智能体 | 控制器 + 智能体 | 同一个 Humanize2 会话 | `attempts//feedback.json` | +| 6 产出分支 | 汇总改动，按需提交和推送 | 控制器 | `git` | 分支 `registry/-`、`patch.diff`、`SUMMARY.md` | + +## 2. 命令 + +控制器运行在调优工具环境里（Python ≥ 3.12 + Humanize2，POSIX），与 `tune-all` 相同。需要导入 `embodiinfer` 的检查，通过 `test_python` 指定的解释器以子进程运行。 + +```bash +TOOL="uv run --project scripts/kernel_tuning python -m scripts.kernel_tuning" + +# 步骤 1、2：为选定的算子建立一次运行，打印 sites.md 后停下，等待调用方确认 +$TOOL register new --spec linear.yaml --output results/kernel_tuning/register + +# 步骤 3–5：智能体修改、检查、反馈，直到通过或预算用尽 +$TOOL register run RUN --agent 'HARNESS/MODEL:EFFORT' --set test_python=/abs/dev/python + +$TOOL register resume RUN # 中断后继续 +$TOOL register status RUN # 查看当前阶段和每次尝试的结果 + +# 步骤 6：默认只打印分支和改动；提交、推送都需要显式传参 +$TOOL register finalize RUN [--commit] [--push] +``` + +`register new` 也可以不用规格文件，直接用命令行参数： + +```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)' \ + --output results/kernel_tuning/register +``` + +## 3. 规格文件 + +规格文件决定这次运行注册哪个算子、改哪些调用点。 + +```yaml +model: pi05 # 调用点所属的策略 +op: linear # 新算子名，要求见步骤 1 +description: > # 算子的数学含义，写进智能体的上下文 + 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: # 必须保持原样的分支；步骤 2 会自动补充候选 + - "QuantizedLinear 继续走 embodiinfer/layers/linear.py 的量化注册表" + - "张量并行下的 RowParallelLinear / ColumnParallelLinear 保持模块调用" +tests: # 每次检查都要通过的已有 CPU 测试 + - tests/test_kernel_ops.py + - tests/test_pi05_components.py + - tests/test_pi05_cpu_routes.py + - tests/test_tensor_parallel.py +budget: + max_attempts: 5 # 最多尝试次数 + agent_timeout_seconds: 1800 # 每轮智能体调用的超时 + max_seconds: 7200 # 整次运行的时间上限 +``` + +## 4. 各步骤 + +### 步骤 1：选定算子（`register new`） + +控制器依次检查以下各项，任何一项不满足就报错退出： + +1. **算子名合法。** `op` 必须是 Python 标识符，并按数学含义命名，不能用模型名。原因是现有调优工具要求参考实现位于 `embodiinfer.backend.torch.ops`，且函数名与算子名相同（`scripts/kernel_tuning/regions.py:215-217`、`scripts/kernel_tuning/discovery/tasks.py:95-98`）；`install` 还会用算子名作包目录名（`scripts/kernel_tuning/install.py:177`）。 +2. **算子尚未注册。** 以 `PYTHONPATH=<仓库根目录>` 在 `test_python` 中运行下面的命令，`op` 不能出现在结果里： + ```bash + python -c "from embodiinfer.layers import available_ops; print(sorted(available_ops()))" + ``` + 同时 `embodiinfer/backend/torch/ops.py` 中不能已有同名函数。 +3. **工作目录干净。** `git status --porcelain` 为空，然后记录 `BASE=$(git rev-parse HEAD)`。 +4. **建立两个 worktree。** 主检出目录在整个过程中不被修改： + ```bash + git worktree add --detach RUN/base "$BASE" # 基线，步骤 4 的有效性检查用 + git worktree add -b registry/- RUN/worktree "$BASE" # 智能体只改这里 + ``` +5. **冻结规格。** 规格和 `BASE` 一起写入 `RUN/spec.json`。 + +### 步骤 2：定位调用点（`register new` 的后半段） + +对每个 `site`，控制器用 `ast` 完成以下工作： + +- **定位表达式。** 在 `function` 的函数体中查找与 `expression` 规范化后（`ast.unparse`）相同的表达式，记录行号和列号。找不到，或者找到多处又无法区分时，直接报错。 +- **收集 preserve 候选。** 收集包裹该表达式的 `if` 条件，例如 `isinstance(..., QuantizedLinear)`、`torch.is_grad_enabled()`、`self.attention == "eager"`、`self.tensor_parallel.enabled`，与规格中的 `preserve` 合并。 +- **检查能否拿到 `KernelSelection`。** + - 调用点在策略方法里时，可以直接用 `self._kernels`（`embodiinfer/policies/pi05/modeling_pi05.py:574`）。 + - 调用点在自由函数里时，需要新增 `kernels` 参数。控制器会列出同一文件中调用这个函数的位置，智能体要把参数一路传过去。 + +结果写入 `sites.json`，并生成给人看的 `sites.md`，内容是每个调用点的源码片段和 preserve 条件。然后命令停下。调用方确认无误、或直接修改 `sites.json` 之后，再运行 `register run`。传 `--yes` 可以跳过确认。 + +### 步骤 3：智能体修改（`register run`） + +**运行机制**与现有调优流程相同：控制器用 Humanize2 运行一个可恢复的 flow `operator_registration`，结构参照 `scripts/kernel_tuning/flow.py`；启动方式与 `scripts/kernel_tuning/__main__.py` 中的 `_execute` 一致： + +```python +Hmz(workspace=RUN / "workspace").run( + flow=Path("scripts/kernel_tuning/registration/flow.py"), + task=f"Register {op} and route its call sites through the registry", + agents={"coder": AGENT}, + params={"run_dir": str(RUN)}, + budget={"duration": remaining_seconds, "graceful": False}, + resume=resume, +) +``` + +智能体的工作环境（`LocalEnv`）指向 `RUN/worktree`。只读上下文放在 `RUN/context/`，包括： + +- `AGENTS.md`； +- `embodiinfer/layers/ops.py`、`embodiinfer/backend/torch/ops.py`，以及各调用点所在的文件； +- `tests/test_kernel_ops.py`，作为测试写法示例； +- 一次已完成迁移的 diff 作为示例：`git show 449ebe2 -- embodiinfer/`； +- `sites.json`； +- 下面的规则全文。 + +**每次尝试分两轮：** + +1. **`plan`**：写出 `PLAN.md`，包括： + - 算子的数学定义，以及它与原表达式的逐项对应关系； + - 函数签名，即哪些参数是张量、哪些是常量； + - 每个调用点的改法，以及 `kernels` 如何传到该点； + - preserve 分支如何保持不变； + - 要新增的测试。 +2. **`implement`**：在 worktree 中完成修改，并写出 `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": "isinstance(projection, QuantizedLinear)"}], + "tests": ["tests/test_kernel_ops.py::test_linear_reference_matches_f_linear"] + } + ``` + +**智能体必须遵守的规则**（写进 prompt，并由步骤 4 检查）： + +1. **参考实现。** 写在 `embodiinfer/backend/torch/ops.py`，函数名等于算子名，带类型注解和说明契约的 docstring，只导入 `torch` 和标准库。参数只能是张量或常量（`int`/`float`/`bool`/`str`/`None`），返回张量或张量元组，不得原地修改输入。 +2. **逐字保留原算术。** 使用相同的 ATen 调用、顺序和中间 dtype 转换。例如不能把 `F.linear` 改写成 `x @ w.T`，不能增删 `.float()`、`.to()`、`.contiguous()`，也不能合并或拆分 GEMM。仓库多处说明过：表达式或形状一变，舍入就会变（`embodiinfer/policies/pi05/modeling_pi05.py:664-666`）。 +3. **注册。** 单独写一行： + ```python + register_lazy_op("", "torch", "embodiinfer.backend.torch.ops", "") + ``` + 只有 Torch 实现的算子不能放进 `embodiinfer/layers/ops.py:160` 那个同时注册 torch 和 triton 的循环。 +4. **调用点。** 使用策略现有的 helper，pi0.5 是 `_op(kernels, "", False)(...)`。只有 Torch 实现时，偏好的实现必须是 `torch`，否则 `get_op` 会抛出 KeyError。`KernelSelection` 用参数传递，不使用全局变量。 +5. **preserve 分支原样保留。** 判断方式只用类型（例如 `type(module) is nn.Linear`），不按模型名分支。 +6. **测试至少覆盖三类：** + - 参考实现与原表达式在 CPU 随机输入（FP32 和 BF16）上 `torch.equal`； + - 调用点经过注册表，做法是用一个会记录调用的假调优档案，写法见 `tests/test_kernel_ops.py:135-150`； + - preserve 分支不经过注册表。 +7. **不能改的内容：** 已有测试的断言、`KernelSelection`、`embodiinfer/backend/triton/tuned/`、`scripts/`、`pyproject.toml`、`uv.lock`。 +8. **允许和禁止的操作：** 可以运行 `ruff` 和 CPU `pytest`；不能运行 GPU 代码，不能 `git commit`，不能安装依赖。 + +### 步骤 4：检查 + +控制器按顺序执行以下检查，前一项失败时后面的不再执行。所有检查都在 CPU 上运行。凡是需要导入 `embodiinfer` 的检查，都以 `PYTHONPATH=RUN/worktree` 在 `test_python` 中运行，并确认导入的 `embodiinfer.__file__` 位于 worktree 内。这是因为开发环境里 EmbodiInfer 是可编辑安装，不设置的话会导入主检出目录中的旧代码。 + +| 检查 | 内容 | 方式 | +|---|---|---| +| C1 范围 | `git -C RUN/worktree diff --name-status $BASE` 只涉及：`embodiinfer/backend/torch/ops.py`、`embodiinfer/layers/ops.py`、`sites.json` 中的源文件、`tests/test_kernel_ops.py` 或新的 `tests/test_registry_.py`。已有测试文件只允许新增内容，不允许删改原有行 | `git diff` | +| C2 结构 | ① `backend/torch/ops.py` 中有 `def `，带注解和 docstring，且不导入 `embodiinfer.models`/`policies`/`engine`；② `layers/ops.py` 中有对应的 `register_lazy_op` 行；③ 每个调用点已改为注册表调用，原表达式只出现在 preserve 分支中；④ 在 `test_python` 中验证：`available_ops()[""]` 包含 `torch`、`get_op("")` 就是参考实现、`KernelSelection("none").prepare()` 成功、`KernelSelection("torch").resolve("")` 返回参考实现 | `ast` + 子进程 | +| C3 格式 | 改动文件通过 `ruff check` 和 `ruff format --check` | `ruff` | +| C4 测试 | 规格 `tests` 中的测试和 `REGISTRATION.json` 列出的新测试全部通过 | `test_python -m pytest -q`，在 `RUN/worktree` 中运行 | +| C5 测试有效性 | 把新测试复制到 `RUN/base`，以 `PYTHONPATH=RUN/base` 运行，结果必须失败，以此证明这些测试确实依赖本次改动 | `test_python -m pytest` | + +### 步骤 5：反馈重做 + +任何检查失败时，控制器写出 `attempts//feedback.json`： + +```json +{ + "check": "C2", + "summary": "调用点 S2 仍直接调用 mlp.down_proj(activated)，且不在 preserve 分支内", + "details": {"site": "S2", "file": "embodiinfer/policies/pi05/modeling_pi05.py", "line": 315} +} +``` + +`details` 的内容随检查而定： +- C1：越界的文件或被改动的已有测试行； +- C2：失败的具体检查项； +- C3、C4：命令输出摘要，截断到固定长度； +- C5：在基线上仍然通过的测试名。 + +**下一次尝试：** +- 默认在同一个 Humanize2 会话中发出 `repair` prompt，附上 `feedback.json`，智能体在原 worktree 上继续修改； +- 连续两次在同一项检查失败时，新开一个会话，从 `plan` 重新开始，但 worktree 保留。 + +每次尝试结束时，控制器把 `git -C RUN/worktree diff $BASE` 存为 `attempts//patch.diff`。 + +**停止条件：** +- 全部检查通过：进入步骤 6； +- 达到 `max_attempts` 或 `max_seconds`：状态记为 `exhausted`，worktree 保留供人工接手； +- 规格、`BASE` 或 `RUN/base` 被改动：状态记为 `failed`。 + +### 步骤 6：产出分支（`register finalize`） + +全部检查通过后，控制器生成两个文件： + +- `RUN/patch.diff`； +- `RUN/SUMMARY.md`，内容包括： + - 算子定义； + - 改动的调用点和保持不变的分支； + - 每项检查的命令和结果； + - 尝试次数； + - 明确写出未做的验证：GPU 整模型逐位一致。 + +`finalize` 的行为： + +- 不带参数时，只打印 worktree 路径、分支名和建议的提交说明。 +- `--commit`：在 worktree 中提交到 `registry/-`。提交说明遵守 AGENTS.md 的 Git 约定，例如 `feat: route pi0.5 projections through the linear registry operator`。 +- `--push`：执行 `git push -u origin registry/-`。 + +本工具不提供合并功能。合并前，还需要在目标 GPU 上按 0009 §6 确认 `kernel_profile: none` 时整模型输出与 `BASE` 逐位一致，方法与 commit `449ebe2` 相同。 + +## 5. 运行目录 + +``` +results/kernel_tuning/register// + manifest.json 当前阶段、每次尝试的结果、耗时（atomic_json 写入，run_lock 加锁） + spec.json 冻结的规格和 BASE + sites.json sites.md + context/ 智能体的只读上下文 + base/ BASE 的 worktree（只读） + worktree/ 分支 registry/- 的 worktree + workspace/ Humanize2 的工作区和日志 + attempts// plan/implement/repair 的 prompt 和回复、PLAN.md、REGISTRATION.json、 + checks.json、feedback.json、patch.diff + patch.diff SUMMARY.md +``` + +`resume` 会从 `manifest.json` 记录的阶段继续，已经通过的阶段不再重跑。规格变了就要新建一次运行。 + +## 6. 现有工具链的复用 + +| 现有部件 | 用在 | 用法 | +|---|---|---| +| Humanize2 flow 机制（`scripts/kernel_tuning/flow.py`、`__main__.py` 中的 `_execute`） | 步骤 3、5 | 新写一个 flow，沿用 `@flow(..., resumable=True)`、`AgentCollection`、`LocalEnv`、`Budget`、`ctx.state` 中的进度台账，以及 `Hmz(...).run(...)` 的启动方式 | +| `scripts/kernel_tuning/artifacts.py` 中的 `atomic_json`、`run_lock` | 运行目录 | 直接导入 | +| `scripts/kernel_tuning/runner.py` 的子进程处理（超时、进程组清理） | 步骤 1、4 的子进程 | 照搬这套做法，但启动的是 `test_python` | +| `embodiinfer.layers`（`available_ops`、`get_op`、`KernelSelection`） | 步骤 1、C2 | 在 `test_python` 子进程中导入 | +| `ruff`、`pytest` | C3–C5 | 命令行调用 | +| `git worktree` | 步骤 1、6 | 命令行调用 | + +`profile`、`capture-ops`、`calibrate`、`tune-all`、`install`、`verify` 不在本流程内使用。它们属于调优阶段，或者需要 GPU。 + +## 7. 需要新写的代码 + +| 文件 | 内容 | +|---|---| +| `registration/__init__.py` | 包说明 | +| `registration/spec.py` | 规格的解析（YAML 或命令行）和校验，生成 `spec.json` | +| `registration/sites.py` | 用 `ast` 定位调用点、提取 preserve 条件、查找调用链 | +| `registration/workspace.py` | worktree 的创建和清理、diff、C1 检查 | +| `registration/checks.py` | C2–C5 检查，生成 `feedback.json` | +| `registration/prompts.py` | `plan`、`implement`、`repair` 三种 prompt | +| `registration/flow.py` | Humanize2 flow `operator_registration` | +| `registration/finalize.py` | 生成 `SUMMARY.md`，提交和推送 | +| `__main__.py` | 增加 `register` 子命令组：`new`、`run`、`resume`、`status`、`finalize` | +| `tests/kernel_tuning/test_registration.py` | CPU 测试，覆盖：规格校验；在示例源码上定位调用点；用合成仓库测 C1；构造缺参考实现、错进循环、原表达式残留等情况测 C2；C5 能识别在基线上也通过的测试；用 Humanize2 自带的假智能体（同 `tests/kernel_tuning/test_flow.py`）跑通"失败 → 反馈 → 修复 → 通过"并验证 `resume` | + +工具环境已经包含 `pyyaml`（见 `scripts/kernel_tuning/pyproject.toml`），不需要新增依赖。 + +## 8. 限制 + +- **CPU 检查不能证明 GPU 上逐位一致。** 例如 pi0.5 的 native 路径只在 CUDA 上启用（`embodiinfer/policies/pi05/modeling_pi05.py:565-572`），它上面的调用点在 CPU 测试中执行不到。因此合并前仍需在 GPU 上验证一次（见步骤 6）。 +- **调用点由调用方指定。** 工具只负责定位和确认，不会从热点报告中自动挑选。 +- **每次只注册一个算子，不支持把多个相邻算子融合成一个。** +- **策略需要已有 `kernels` 传参链。** 必须已经有 `KernelSelection` 和类似 `_op` 的 helper，目前只有 pi0.5 具备。其他策略要先人工接入一次。 +- **智能体写出的补丁可能检查全过但结构不好**（例如传参链过长），仍然需要人工审阅。 From adefd22f8236d05f08fc7845921d222eb0e51873 Mon Sep 17 00:00:00 2001 From: owpeter Date: Tue, 13 Oct 2026 00:24:50 -0700 Subject: [PATCH 28/44] feat: add operator registration toolchain Signed-off-by: owpeter --- CONTRIBUTING.md | 8 + docs/proposals/0009-hotspot-kernel-tuning.md | 65 +- scripts/kernel_tuning/__main__.py | 7 + scripts/kernel_tuning/artifacts.py | 11 +- scripts/kernel_tuning/prompts.py | 4 + scripts/kernel_tuning/registration/README.md | 41 +- .../kernel_tuning/registration/__init__.py | 1 + scripts/kernel_tuning/registration/checks.py | 417 +++++++++ scripts/kernel_tuning/registration/cli.py | 169 ++++ .../kernel_tuning/registration/finalize.py | 99 ++ scripts/kernel_tuning/registration/flow.py | 213 +++++ scripts/kernel_tuning/registration/prompts.py | 97 ++ scripts/kernel_tuning/registration/sites.py | 210 +++++ scripts/kernel_tuning/registration/spec.py | 141 +++ .../kernel_tuning/registration/workspace.py | 436 +++++++++ tests/kernel_tuning/test_flow.py | 27 + tests/kernel_tuning/test_registration.py | 852 ++++++++++++++++++ 17 files changed, 2789 insertions(+), 9 deletions(-) create mode 100644 scripts/kernel_tuning/registration/__init__.py create mode 100644 scripts/kernel_tuning/registration/checks.py create mode 100644 scripts/kernel_tuning/registration/cli.py create mode 100644 scripts/kernel_tuning/registration/finalize.py create mode 100644 scripts/kernel_tuning/registration/flow.py create mode 100644 scripts/kernel_tuning/registration/prompts.py create mode 100644 scripts/kernel_tuning/registration/sites.py create mode 100644 scripts/kernel_tuning/registration/spec.py create mode 100644 scripts/kernel_tuning/registration/workspace.py create mode 100644 tests/kernel_tuning/test_registration.py diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index c421b6a..5894ed2 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -146,6 +146,14 @@ explains task inputs, the full tuning loop, stopping conditions, and saved resul The controller snapshots task inputs, saves all attempts, and exports improving solutions with their correctness and timing evidence under `results/kernel_tuning/`. Saved operators require a separate review/integration before runtime use. +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. diff --git a/docs/proposals/0009-hotspot-kernel-tuning.md b/docs/proposals/0009-hotspot-kernel-tuning.md index 2b0fbde..fd0c839 100644 --- a/docs/proposals/0009-hotspot-kernel-tuning.md +++ b/docs/proposals/0009-hotspot-kernel-tuning.md @@ -186,6 +186,11 @@ workload) and archive are reused. A candidate must also export `applies`, which the evaluator checks is true on every workload. Expected model saving is `Σ count × Δt`. +Each planning turn receives the recent candidates' verified source alongside +their measured failures, including rejected candidates that never became the +incumbent. This lets a fresh agent inspect failed configurations before choosing +a new change. Optional hardware notes are part of the frozen task context. + ### 4.6 Install, verify, benchmark `install` writes promoted sources and updates `profile.json`. It refuses an @@ -218,6 +223,62 @@ with resume under `results/kernel_tuning/optimize//`. replay. - Apply agent policy patches without review: rejected. +### 4.8 Agent-assisted operator registration + +`python -m scripts.kernel_tuning register` implements the registration step +between manual operator selection and tuning. Its developer-only controller is +in `scripts/kernel_tuning/registration/`; it changes no inference API or default +runtime behavior. The caller supplies one mathematical operator and exact source +expressions. The initial supported routing convention is pi0.5's existing `_op` +helper and explicit `KernelSelection` parameters. Other policies must first gain +an equivalent integration; this tool does not create model-specific engine logic. + +The [registration contract](../../scripts/kernel_tuning/registration/README.md) +defines the YAML specification and `new`, `run`, `resume`, `status`, and `finalize` +commands. `new` requires a clean checkout and probes the public registry in a +separately selected `test_python`. `--repository` may select a clean model checkout +separate from an uncommitted controller checkout. It records BASE, creates a detached baseline +and a `registry/-` worktree, and stops after AST call-site review. +Explicit `run` confirms the reviewed sites and freezes the execution settings. +Humanize2 plans and implements in the branch worktree, keeping a session for +repair turns and opening a fresh planning session after two equal check failures. +An interrupted run retains its completed stages and elapsed budget. + +The controller applies five fail-fast CPU gates: file scope (including untracked +and staged changes); annotated Torch reference, standalone lazy registration, +preserved branches, parameter forwarding and public registry resolution; Ruff; +declared existing tests and new test nodes; and baseline effectiveness. Every new +test node must execute on BASE and fail by assertion, with all parametrizations +covered. Collection errors, missing dependencies and skips are not proof. The +baseline overlay is journaled before copying and restored after failures or +interruptions. Check evidence is bound to source bytes and the registration +manifest, so changing a patch invalidates completed gates. +Baseline integrity hashes tracked symbolic links by their link text without +following their targets; selected source and test files must still be regular files. + +This is lossless routing, not kernel optimization: the reference must preserve +ATen calls, arithmetic order, dtype conversions and shapes. Generated tests must +compare original and reference expressions with `torch.equal` on CPU FP32 and +BF16, exercise a recording profile, and verify preservation branches. The §6 +whole-model target-GPU bit-exact gate remains mandatory before merge. Static +checks and finite CPU tests cannot prove arbitrary numerical equivalence or +that agent-authored tests cover every behavior. + +Passing runs export `patch.diff` and `SUMMARY.md`. Finalization prints the branch +by default; commits and pushes each need their explicit flags. No merge is +provided. Operator selection, fusion, GPU testing and tuning remain separate. +The alternative of adding registration to the GPU tuning flow was rejected: +registration needs baseline code checkouts and routing tests, whereas tuning +measures candidate kernels against a frozen numerical/performance contract. + +Validation uses disposable clones of existing commit objects, structural +counterexamples, real CPU subprocesses, and official Humanize2 fake agents for +repair/session restart/resume. It requires no model checkpoint and makes no +performance claim. Coding agents and tests run with the user's process +permissions; directory bounds, integrity checks and hidden accelerator devices +are not an OS security sandbox. Actual Humanize2 execution requires POSIX and +Python 3.12+, while runtime imports use the separately prepared test interpreter. + ## 5. Model-agnosticism verdict The registry, profile loader and `kernel_profile` switch are model-agnostic @@ -298,7 +359,9 @@ reproduces them. Status (2026-10-12): steps 1–3 are implemented for pi0.5; the loop was proven on Jetson Thor instead of an RTX 4090 (§9), and the `thor` profile ships one tuned operator. Qwen R2R and StreamVLN call sites, step 4 and an RTX 4090 -profile are not done; neither is shape-range tuning (§10). +profile are not done; neither is shape-range tuning (§10). The manually selected, +single-operator registration portion of step 4 now has the CPU controller in +§4.8; automatic selection and new fusion remain future work. New tooling lives in `scripts/kernel_tuning/` and `benchmarks/kernel_tuning/`. Existing 0008 commands keep working and are not extended. No public symbol is diff --git a/scripts/kernel_tuning/__main__.py b/scripts/kernel_tuning/__main__.py index 0dba46c..f28de4e 100644 --- a/scripts/kernel_tuning/__main__.py +++ b/scripts/kernel_tuning/__main__.py @@ -235,6 +235,9 @@ 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( @@ -334,6 +337,10 @@ def main(argv: list[str] | None = None) -> int: command = None args = parser.parse_args(raw) try: + if args.command == "register": + from .registration.cli import dispatch + + return dispatch(args) if args.command == "capture": if args.catalog: from .capture import run as run_capture diff --git a/scripts/kernel_tuning/artifacts.py b/scripts/kernel_tuning/artifacts.py index 52f69ca..d54d23c 100644 --- a/scripts/kernel_tuning/artifacts.py +++ b/scripts/kernel_tuning/artifacts.py @@ -240,7 +240,16 @@ def begin_attempt(self, *, persist: Callable[[], None] | None = None) -> tuple[s 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") - atomic_json(scratch / "feedback.json", self.manifest["attempts"][-5:]) + 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 diff --git a/scripts/kernel_tuning/prompts.py b/scripts/kernel_tuning/prompts.py index 0ed820e..eac6024 100644 --- a/scripts/kernel_tuning/prompts.py +++ b/scripts/kernel_tuning/prompts.py @@ -20,6 +20,10 @@ def plan_prompt(task: TaskPackage) -> str: 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. diff --git a/scripts/kernel_tuning/registration/README.md b/scripts/kernel_tuning/registration/README.md index cd6d679..114d66b 100644 --- a/scripts/kernel_tuning/registration/README.md +++ b/scripts/kernel_tuning/registration/README.md @@ -1,6 +1,7 @@ # 算子注册流程（智能体闭环） -本目录是 `python -m scripts.kernel_tuning register …` 工具链的规格说明，代码尚未实现。 +本目录实现 `python -m scripts.kernel_tuning register …` 工具链。设计归属见 +[0009 §4.8](../../../docs/proposals/0009-hotspot-kernel-tuning.md#48-agent-assisted-operator-registration)。 本工具链自动完成"把一个算子接入算子注册表"这一步： @@ -45,7 +46,8 @@ TOOL="uv run --project scripts/kernel_tuning python -m scripts.kernel_tuning" # 步骤 1、2：为选定的算子建立一次运行，打印 sites.md 后停下，等待调用方确认 -$TOOL register new --spec linear.yaml --output results/kernel_tuning/register +$TOOL register new --spec linear.yaml --output results/kernel_tuning/register \ + --set test_python=/abs/dev/python # 步骤 3–5：智能体修改、检查、反馈，直到通过或预算用尽 $TOOL register run RUN --agent 'HARNESS/MODEL:EFFORT' --set test_python=/abs/dev/python @@ -57,6 +59,14 @@ $TOOL register status RUN # 查看当前阶段和每次尝试的结果 $TOOL register finalize RUN [--commit] [--push] ``` +`new` 默认用当前解释器探测注册表；工具环境没有 Torch 时，必须指定准备好的 +`test_python`。它也需要 pytest 及规格中 CPU 测试所需的模型依赖。Ruff 在工具环境中运行。 +控制器与待注册代码可以分开存放：`new --repository /abs/clean-checkout` 指定干净的模型仓库， +默认使用控制器所在仓库。这样可以验证尚未提交的工具链，同时从已提交的模型基线建立 worktree。 +`run` 命令本身表示调用方已经确认 `sites.md`；`--yes` 可显式记录确认。 +首次 `run` 可以更换 `test_python` 并重新探测，此后解释器和 agent 固定，`resume` 不接受覆盖。 +`--push` 需要同时传 `--commit`，或此前已经完成提交。 + `register new` 也可以不用规格文件，直接用命令行参数： ```bash @@ -109,6 +119,7 @@ budget: ``` 同时 `embodiinfer/backend/torch/ops.py` 中不能已有同名函数。 3. **工作目录干净。** `git status --porcelain` 为空，然后记录 `BASE=$(git rev-parse HEAD)`。 + 基线中的已有符号链接按链接目标文本记录完整性，不跟随目标；选定的源码和测试仍须是普通文件。 4. **建立两个 worktree。** 主检出目录在整个过程中不被修改： ```bash git worktree add --detach RUN/base "$BASE" # 基线，步骤 4 的有效性检查用 @@ -126,14 +137,17 @@ budget: - 调用点在策略方法里时，可以直接用 `self._kernels`（`embodiinfer/policies/pi05/modeling_pi05.py:574`）。 - 调用点在自由函数里时，需要新增 `kernels` 参数。控制器会列出同一文件中调用这个函数的位置，智能体要把参数一路传过去。 -结果写入 `sites.json`，并生成给人看的 `sites.md`，内容是每个调用点的源码片段和 preserve 条件。然后命令停下。调用方确认无误、或直接修改 `sites.json` 之后，再运行 `register run`。传 `--yes` 可以跳过确认。 +结果写入 `sites.json`，并生成给人看的 `sites.md`，内容是每个调用点的源码片段和 preserve 条件。 +然后命令停下。调用方确认无误、或直接修改 `sites.json` 的补充说明之后，再运行 `register run`。 +源文件、函数、表达式、站点 ID 和自动提取的基线 guard 不可删改；更换这些输入需要新建运行。 +传 `--yes` 可显式记录确认，`new` 始终在此停下，不会自动打开智能体。 ### 步骤 3：智能体修改（`register run`） **运行机制**与现有调优流程相同：控制器用 Humanize2 运行一个可恢复的 flow `operator_registration`，结构参照 `scripts/kernel_tuning/flow.py`；启动方式与 `scripts/kernel_tuning/__main__.py` 中的 `_execute` 一致： ```python -Hmz(workspace=RUN / "workspace").run( +Hmz(workspace=RUN).run( flow=Path("scripts/kernel_tuning/registration/flow.py"), task=f"Register {op} and route its call sites through the registry", agents={"coder": AGENT}, @@ -152,6 +166,11 @@ Hmz(workspace=RUN / "workspace").run( - `sites.json`； - 下面的规则全文。 +固定版 Humanize2 的 `LocalEnv.derive_subdir` 只接受子目录，所以 SDK 的 workspace +设为 `RUN`，再派生 `worktree` 给智能体。框架日志归档到 `RUN/workspace/`； +`humanize.json` 记录实际 epic 及归档路径。源码快照和规则文件会校验哈希。 +浅克隆中缺少示例提交 `449ebe2` 时，明确记录缺失，继续使用当前源码和测试示例。 + **每次尝试分两轮：** 1. **`plan`**：写出 `PLAN.md`，包括： @@ -200,7 +219,7 @@ Hmz(workspace=RUN / "workspace").run( | C2 结构 | ① `backend/torch/ops.py` 中有 `def `，带注解和 docstring，且不导入 `embodiinfer.models`/`policies`/`engine`；② `layers/ops.py` 中有对应的 `register_lazy_op` 行；③ 每个调用点已改为注册表调用，原表达式只出现在 preserve 分支中；④ 在 `test_python` 中验证：`available_ops()[""]` 包含 `torch`、`get_op("")` 就是参考实现、`KernelSelection("none").prepare()` 成功、`KernelSelection("torch").resolve("")` 返回参考实现 | `ast` + 子进程 | | C3 格式 | 改动文件通过 `ruff check` 和 `ruff format --check` | `ruff` | | C4 测试 | 规格 `tests` 中的测试和 `REGISTRATION.json` 列出的新测试全部通过 | `test_python -m pytest -q`，在 `RUN/worktree` 中运行 | -| C5 测试有效性 | 把新测试复制到 `RUN/base`，以 `PYTHONPATH=RUN/base` 运行，结果必须失败，以此证明这些测试确实依赖本次改动 | `test_python -m pytest` | +| C5 测试有效性 | 把新测试复制到 `RUN/base`，以 `PYTHONPATH=RUN/base` 逐个运行，每个测试的所有参数化实例必须执行并断言失败；通过、跳过、收集失败或依赖异常均不算有效证据 | `test_python` 内调用 `pytest.main`，保存 JUnit 报告 | ### 步骤 5：反馈重做 @@ -267,7 +286,10 @@ results/kernel_tuning/register// patch.diff SUMMARY.md ``` -`resume` 会从 `manifest.json` 记录的阶段继续，已经通过的阶段不再重跑。规格变了就要新建一次运行。 +`resume` 会结合 `manifest.json` 和 Humanize2 的状态台账继续，已经通过的阶段在源码和 +`REGISTRATION.json` 哈希未变时不再重跑。规格变了就要新建一次运行。 +同次执行中的 repair 沿用会话；进程中断后从记录的阶段继续，新会话会收到规格、已有计划和反馈。 +基线测试复制前先记入恢复台账，失败、取消和下次恢复时均还原这些测试文件。 ## 6. 现有工具链的复用 @@ -282,7 +304,7 @@ results/kernel_tuning/register// `profile`、`capture-ops`、`calibrate`、`tune-all`、`install`、`verify` 不在本流程内使用。它们属于调优阶段，或者需要 GPU。 -## 7. 需要新写的代码 +## 7. 代码布局 | 文件 | 内容 | |---|---| @@ -294,6 +316,7 @@ results/kernel_tuning/register// | `registration/prompts.py` | `plan`、`implement`、`repair` 三种 prompt | | `registration/flow.py` | Humanize2 flow `operator_registration` | | `registration/finalize.py` | 生成 `SUMMARY.md`，提交和推送 | +| `registration/cli.py` | 参数解析、运行环境校验、Humanize2 启动 | | `__main__.py` | 增加 `register` 子命令组：`new`、`run`、`resume`、`status`、`finalize` | | `tests/kernel_tuning/test_registration.py` | CPU 测试，覆盖：规格校验；在示例源码上定位调用点；用合成仓库测 C1；构造缺参考实现、错进循环、原表达式残留等情况测 C2；C5 能识别在基线上也通过的测试；用 Humanize2 自带的假智能体（同 `tests/kernel_tuning/test_flow.py`）跑通"失败 → 反馈 → 修复 → 通过"并验证 `resume` | @@ -306,3 +329,7 @@ results/kernel_tuning/register// - **每次只注册一个算子，不支持把多个相邻算子融合成一个。** - **策略需要已有 `kernels` 传参链。** 必须已经有 `KernelSelection` 和类似 `_op` 的 helper，目前只有 pi0.5 具备。其他策略要先人工接入一次。 - **智能体写出的补丁可能检查全过但结构不好**（例如传参链过长），仍然需要人工审阅。 +- **新测试必须单独列出函数 node ID。** 不接受整个文件、已有测试或挑选部分参数化实例。 + 新符号的导入放在函数内部，先断言符号存在，避免用缺失导入代替有效的基线断言失败。 +- **目录和哈希检查不等同于安全沙箱。** 智能体和测试继承运行用户的权限；控制器隐藏 CUDA/ROCm + 设备，禁止运行 GPU 测试，但这不是对任意不受信代码的操作系统隔离。 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..d550d19 --- /dev/null +++ b/scripts/kernel_tuning/registration/cli.py @@ -0,0 +1,169 @@ +"""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, atomic_json, run_lock, runtime_identity +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") + 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 + # 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) + 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..d8e5c00 --- /dev/null +++ b/scripts/kernel_tuning/registration/flow.py @@ -0,0 +1,213 @@ +"""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 ( + Agent, + 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.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: Agent + + +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 + 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 agents["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( + agents["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..294a4a2 --- /dev/null +++ b/scripts/kernel_tuning/registration/prompts.py @@ -0,0 +1,97 @@ +"""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, test_kernel_ops.py and +migration.diff. 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..9c85aee --- /dev/null +++ b/scripts/kernel_tuning/registration/workspace.py @@ -0,0 +1,436 @@ +"""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" +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) -> None: + """Freeze reviewed call sites and build context before the first agent turn.""" + self.recover_base() + self.verify() + if self.manifest["sites_hash"]: + if agent != self.manifest["agent"] or (python and python != self.manifest["test_python"]): + raise ContractError("Agent/test_python 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, + "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") + try: + example = git(self.base, "show", "449ebe2", "--", "embodiinfer/") + except ContractError: + example = "The optional migration example 449ebe2 is unavailable in this repository.\n" + (context / "migration.diff").write_text(example, encoding="utf-8") + self.manifest.update( + agent=agent, + 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/tests/kernel_tuning/test_flow.py b/tests/kernel_tuning/test_flow.py index 93dc835..916e812 100644 --- a/tests/kernel_tuning/test_flow.py +++ b/tests/kernel_tuning/test_flow.py @@ -109,6 +109,33 @@ def test_full_loop_promotes_then_preserves_failures_and_exports(store, candidate 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): diff --git a/tests/kernel_tuning/test_registration.py b/tests/kernel_tuning/test_registration.py new file mode 100644 index 0000000..855eaf0 --- /dev/null +++ b/tests/kernel_tuning/test_registration.py @@ -0,0 +1,852 @@ +"""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, + RegistrationRun, + check_scope, + check_unregistered, + command, + file_hash, + git, + tracked_hashes, +) + +SOURCE = "embodiinfer/policies/pi05/modeling_pi05.py" +NEW_TEST = "tests/test_registry_linear.py" + + +def specification(**overrides) -> RegistrationSpec: + """Use an existing unfused registry gap without loading a model/checkpoint.""" + value = { + "model": "pi05", + "op": "linear", + "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", + "linear", + "--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, "linear", 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 linear(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("linear", "torch", "embodiinfer.backend.torch.ops", "linear")\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": "linear", + "reference": f"{REFERENCE}::linear", + "sites": [ + { + "id": "S1", + "file": SOURCE, + "function": "_fused_qkv", + "call": '_op(kernels, "linear", 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: linear\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="linear", 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("linear", "torch", "embodiinfer.backend.torch.ops", "linear")' + 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, "linear", False)(x, weight)', + 'unused = F.linear(x, weight)\n fused = _op(kernels, "linear", 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, "linear", 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_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, "linear", 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, "linear", False)(x, weight)' + ), + encoding="utf-8", + ) + with (worktree / REFERENCE).open("a", encoding="utf-8") as stream: + stream.write( + '\n\ndef linear(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("linear", "torch", "embodiinfer.backend.torch.ops", "linear")\n') + (worktree / NEW_TEST).write_text( + """from embodiinfer.backend.torch import ops + + +def test_reference(): + assert hasattr(ops, "linear") + assert ops.linear(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 == ["linear"] +""", + encoding="utf-8", + ) + atomic_json( + worktree / "REGISTRATION.json", + { + "op": "linear", + "reference": f"{REFERENCE}::linear", + "sites": [ + { + "id": "S1", + "file": SOURCE, + "function": "compute", + "call": '_op(kernels, "linear", 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() From 412fc6963579f7f1e9abd9e2d9fb59d36bd62443 Mon Sep 17 00:00:00 2001 From: owpeter Date: Wed, 14 Oct 2026 04:08:25 +0000 Subject: [PATCH 29/44] refactor: read the routed registry example from BASE, not a fixed commit The registration agent's worked example came from `git show` of a hard-coded commit, which disappears from history under a squash or rebase merge. Copy the already-routed pi0.5 policy from BASE into the read-only context instead, and rewrite the registration README to match the current controller. Signed-off-by: owpeter --- scripts/kernel_tuning/registration/README.md | 548 +++++++++--------- scripts/kernel_tuning/registration/prompts.py | 5 +- .../kernel_tuning/registration/workspace.py | 9 +- tests/kernel_tuning/test_registration.py | 8 + 4 files changed, 300 insertions(+), 270 deletions(-) diff --git a/scripts/kernel_tuning/registration/README.md b/scripts/kernel_tuning/registration/README.md index 114d66b..5ef8b8b 100644 --- a/scripts/kernel_tuning/registration/README.md +++ b/scripts/kernel_tuning/registration/README.md @@ -1,89 +1,65 @@ -# 算子注册流程（智能体闭环） +# 算子注册工具（`register`） -本目录实现 `python -m scripts.kernel_tuning register …` 工具链。设计归属见 +`python -m scripts.kernel_tuning register …` 把一个尚未注册的数学算子接入算子注册表 +（`embodiinfer/layers/ops.py`），并把选定的调用点改为经过注册表调用。设计归属见 [0009 §4.8](../../../docs/proposals/0009-hotspot-kernel-tuning.md#48-agent-assisted-operator-registration)。 -本工具链自动完成"把一个算子接入算子注册表"这一步： +工具由一个控制器驱动编码智能体，循环执行"计划 → 实现 → CPU 检查 → 反馈修复"，直到全部检查通过或预算用尽。 +最终产物是一个待人工审阅的分支和 `patch.diff`。 -1. 调用方选定算子； -2. 工具定位调用点； -3. 智能体修改代码、注册算子、编写测试； -4. 控制器运行检查； -5. 失败时把信息反馈给智能体重做； -6. 全部通过后产出一个待审阅的分支。 +注册只改调用路径，不改数值：参考实现必须逐字保留原表达式的算术。以下工作不在本工具范围内： -不在本流程内的工作： +- 选择要注册的算子和调用点（由调用方指定）； +- 融合多个相邻算子； +- 调优、安装 kernel（`tune-all`/`run`、`install`、`verify` 等）； +- GPU 上的整模型逐位一致验证（合并前必须单独完成，见[合并前](#合并前)）； +- 合并分支（由人工完成）。 -- 调优和安装 kernel：使用现有的 `capture-ops`、`calibrate`、`tune-all`/`run`、`install`、`verify`。 -- GPU 上的整模型逐位一致验证：按 [0009 §6](../../../docs/proposals/0009-hotspot-kernel-tuning.md) 的要求，合并前在目标 GPU 上单独做一次。 -- 合并：由人工完成。 +## 1. 前提条件 -本流程的检查全部在 CPU 上运行，因此不需要 GPU。 +- **工具环境**：Python ≥ 3.12，并已安装 Humanize2 和 Ruff，与 `tune-all` 共用： + `uv sync --project scripts/kernel_tuning --python 3.12 --frozen`。 + `run` 和 `resume` 只能在 POSIX 系统上执行；`new`、`status`、`finalize` 不启动智能体。 + 启动后工具环境的包版本（`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`）满足。其他策略需要先人工接入一次。 -## 1. 流程总览 - -``` -选定算子 ──► 定位调用点 ──► 智能体修改 ──► 检查 ──通过──► 产出分支 - (步骤1) (步骤2) (步骤3) (步骤4) (步骤6) - │ ▲ │ - 调用方确认 └── 反馈 ◄───┘ 失败 (步骤5) -``` - -| 步骤 | 工作 | 执行者 | 使用的现有工具 | 产物 | -|---|---|---|---|---| -| 1 选定算子 | 校验规格，记录基线 commit，建立 worktree | 控制器 | `git worktree`；在测试环境中调用 `embodiinfer.layers.available_ops` | `spec.json`、`base/`、`worktree/` | -| 2 定位调用点 | 找到要改的表达式，以及需要保持原样的分支 | 控制器，调用方确认 | 无（`ast`） | `sites.json`、`sites.md` | -| 3 智能体修改 | 写参考实现、注册、改调用点、写测试 | 智能体 | Humanize2 flow 机制（同 `scripts/kernel_tuning/flow.py`） | worktree 中的改动、`PLAN.md`、`REGISTRATION.json` | -| 4 检查 | 范围、结构、格式、测试、测试有效性 | 控制器 | `ruff`、`pytest` | `attempts//checks.json` | -| 5 反馈重做 | 把失败信息交回智能体 | 控制器 + 智能体 | 同一个 Humanize2 会话 | `attempts//feedback.json` | -| 6 产出分支 | 汇总改动，按需提交和推送 | 控制器 | `git` | 分支 `registry/-`、`patch.diff`、`SUMMARY.md` | - -## 2. 命令 - -控制器运行在调优工具环境里（Python ≥ 3.12 + Humanize2，POSIX），与 `tune-all` 相同。需要导入 `embodiinfer` 的检查，通过 `test_python` 指定的解释器以子进程运行。 +## 2. 快速开始 ```bash TOOL="uv run --project scripts/kernel_tuning python -m scripts.kernel_tuning" -# 步骤 1、2：为选定的算子建立一次运行，打印 sites.md 后停下，等待调用方确认 -$TOOL register new --spec linear.yaml --output results/kernel_tuning/register \ - --set test_python=/abs/dev/python +# 1. 建立运行：校验规格、建立 worktree、定位调用点，打印 sites.md 后停下 +$TOOL register new --spec linear.yaml --set test_python=/abs/runtime/python +# 输出最后一行为 "Run: " -# 步骤 3–5：智能体修改、检查、反馈，直到通过或预算用尽 -$TOOL register run RUN --agent 'HARNESS/MODEL:EFFORT' --set test_python=/abs/dev/python +# 2. 审阅 RUN/sites.md（需要时补充 RUN/sites.json），然后启动智能体循环 +$TOOL register run RUN --agent 'HARNESS/MODEL:EFFORT' -$TOOL register resume RUN # 中断后继续 -$TOOL register status RUN # 查看当前阶段和每次尝试的结果 +# 中断后继续；查看状态 +$TOOL register resume RUN +$TOOL register status RUN -# 步骤 6：默认只打印分支和改动；提交、推送都需要显式传参 +# 3. 导出已验证的补丁；提交和推送都需要显式参数 $TOOL register finalize RUN [--commit] [--push] ``` -`new` 默认用当前解释器探测注册表；工具环境没有 Torch 时，必须指定准备好的 -`test_python`。它也需要 pytest 及规格中 CPU 测试所需的模型依赖。Ruff 在工具环境中运行。 -控制器与待注册代码可以分开存放：`new --repository /abs/clean-checkout` 指定干净的模型仓库， -默认使用控制器所在仓库。这样可以验证尚未提交的工具链，同时从已提交的模型基线建立 worktree。 -`run` 命令本身表示调用方已经确认 `sites.md`；`--yes` 可显式记录确认。 -首次 `run` 可以更换 `test_python` 并重新探测，此后解释器和 agent 固定，`resume` 不接受覆盖。 -`--push` 需要同时传 `--commit`，或此前已经完成提交。 - -`register new` 也可以不用规格文件，直接用命令行参数： - -```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)' \ - --output results/kernel_tuning/register -``` +所有命令都在仓库根目录执行。`run`/`resume` 的退出码在状态为 `passed` 或 `finalized` 时为 0， +否则为 1。 -## 3. 规格文件 +## 3. 规格 -规格文件决定这次运行注册哪个算子、改哪些调用点。 +规格决定本次运行注册哪个算子、修改哪些调用点、用哪些测试回归，以及预算。 ```yaml -model: pi05 # 调用点所属的策略 -op: linear # 新算子名，要求见步骤 1 -description: > # 算子的数学含义，写进智能体的上下文 +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 @@ -92,244 +68,290 @@ sites: # 要改的调用点：文件、所在函数、 - source: embodiinfer/policies/pi05/modeling_pi05.py function: _mlp expression: "mlp.down_proj(activated)" -preserve: # 必须保持原样的分支；步骤 2 会自动补充候选 - - "QuantizedLinear 继续走 embodiinfer/layers/linear.py 的量化注册表" - - "张量并行下的 RowParallelLinear / ColumnParallelLinear 保持模块调用" +preserve: # 必须保持原样的路径，作为审阅说明交给智能体 + - "QuantizedLinear 继续走量化注册表" tests: # 每次检查都要通过的已有 CPU 测试 - tests/test_kernel_ops.py - - tests/test_pi05_components.py - tests/test_pi05_cpu_routes.py - - tests/test_tensor_parallel.py budget: max_attempts: 5 # 最多尝试次数 - agent_timeout_seconds: 1800 # 每轮智能体调用的超时 - max_seconds: 7200 # 整次运行的时间上限 + agent_timeout_seconds: 1800 # 单次智能体调用的超时 + max_seconds: 7200 # 整次运行的累计时长上限，中断后继续累计 ``` -## 4. 各步骤 - -### 步骤 1：选定算子（`register new`） - -控制器依次检查以下各项，任何一项不满足就报错退出： - -1. **算子名合法。** `op` 必须是 Python 标识符，并按数学含义命名，不能用模型名。原因是现有调优工具要求参考实现位于 `embodiinfer.backend.torch.ops`，且函数名与算子名相同（`scripts/kernel_tuning/regions.py:215-217`、`scripts/kernel_tuning/discovery/tasks.py:95-98`）；`install` 还会用算子名作包目录名（`scripts/kernel_tuning/install.py:177`）。 -2. **算子尚未注册。** 以 `PYTHONPATH=<仓库根目录>` 在 `test_python` 中运行下面的命令，`op` 不能出现在结果里： - ```bash - python -c "from embodiinfer.layers import available_ops; print(sorted(available_ops()))" - ``` - 同时 `embodiinfer/backend/torch/ops.py` 中不能已有同名函数。 -3. **工作目录干净。** `git status --porcelain` 为空，然后记录 `BASE=$(git rev-parse HEAD)`。 - 基线中的已有符号链接按链接目标文本记录完整性，不跟随目标；选定的源码和测试仍须是普通文件。 -4. **建立两个 worktree。** 主检出目录在整个过程中不被修改： - ```bash - git worktree add --detach RUN/base "$BASE" # 基线，步骤 4 的有效性检查用 - git worktree add -b registry/- RUN/worktree "$BASE" # 智能体只改这里 - ``` -5. **冻结规格。** 规格和 `BASE` 一起写入 `RUN/spec.json`。 - -### 步骤 2：定位调用点（`register new` 的后半段） - -对每个 `site`，控制器用 `ast` 完成以下工作： - -- **定位表达式。** 在 `function` 的函数体中查找与 `expression` 规范化后（`ast.unparse`）相同的表达式，记录行号和列号。找不到，或者找到多处又无法区分时，直接报错。 -- **收集 preserve 候选。** 收集包裹该表达式的 `if` 条件，例如 `isinstance(..., QuantizedLinear)`、`torch.is_grad_enabled()`、`self.attention == "eager"`、`self.tensor_parallel.enabled`，与规格中的 `preserve` 合并。 -- **检查能否拿到 `KernelSelection`。** - - 调用点在策略方法里时，可以直接用 `self._kernels`（`embodiinfer/policies/pi05/modeling_pi05.py:574`）。 - - 调用点在自由函数里时，需要新增 `kernels` 参数。控制器会列出同一文件中调用这个函数的位置，智能体要把参数一路传过去。 - -结果写入 `sites.json`，并生成给人看的 `sites.md`，内容是每个调用点的源码片段和 preserve 条件。 -然后命令停下。调用方确认无误、或直接修改 `sites.json` 的补充说明之后，再运行 `register run`。 -源文件、函数、表达式、站点 ID 和自动提取的基线 guard 不可删改；更换这些输入需要新建运行。 -传 `--yes` 可显式记录确认，`new` 始终在此停下，不会自动打开智能体。 - -### 步骤 3：智能体修改（`register run`） - -**运行机制**与现有调优流程相同：控制器用 Humanize2 运行一个可恢复的 flow `operator_registration`，结构参照 `scripts/kernel_tuning/flow.py`；启动方式与 `scripts/kernel_tuning/__main__.py` 中的 `_execute` 一致： - -```python -Hmz(workspace=RUN).run( - flow=Path("scripts/kernel_tuning/registration/flow.py"), - task=f"Register {op} and route its call sites through the registry", - agents={"coder": AGENT}, - params={"run_dir": str(RUN)}, - budget={"duration": remaining_seconds, "graceful": False}, - resume=resume, -) +校验规则（`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)' ``` -智能体的工作环境（`LocalEnv`）指向 `RUN/worktree`。只读上下文放在 `RUN/context/`，包括： - -- `AGENTS.md`； -- `embodiinfer/layers/ops.py`、`embodiinfer/backend/torch/ops.py`，以及各调用点所在的文件； -- `tests/test_kernel_ops.py`，作为测试写法示例； -- 一次已完成迁移的 diff 作为示例：`git show 449ebe2 -- embodiinfer/`； -- `sites.json`； -- 下面的规则全文。 - -固定版 Humanize2 的 `LocalEnv.derive_subdir` 只接受子目录，所以 SDK 的 workspace -设为 `RUN`，再派生 `worktree` 给智能体。框架日志归档到 `RUN/workspace/`； -`humanize.json` 记录实际 epic 及归档路径。源码快照和规则文件会校验哈希。 -浅克隆中缺少示例提交 `449ebe2` 时，明确记录缺失，继续使用当前源码和测试示例。 - -**每次尝试分两轮：** - -1. **`plan`**：写出 `PLAN.md`，包括： - - 算子的数学定义，以及它与原表达式的逐项对应关系； - - 函数签名，即哪些参数是张量、哪些是常量； - - 每个调用点的改法，以及 `kernels` 如何传到该点； - - preserve 分支如何保持不变； - - 要新增的测试。 -2. **`implement`**：在 worktree 中完成修改，并写出 `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": "isinstance(projection, QuantizedLinear)"}], - "tests": ["tests/test_kernel_ops.py::test_linear_reference_matches_f_linear"] - } - ``` - -**智能体必须遵守的规则**（写进 prompt，并由步骤 4 检查）： - -1. **参考实现。** 写在 `embodiinfer/backend/torch/ops.py`，函数名等于算子名，带类型注解和说明契约的 docstring，只导入 `torch` 和标准库。参数只能是张量或常量（`int`/`float`/`bool`/`str`/`None`），返回张量或张量元组，不得原地修改输入。 -2. **逐字保留原算术。** 使用相同的 ATen 调用、顺序和中间 dtype 转换。例如不能把 `F.linear` 改写成 `x @ w.T`，不能增删 `.float()`、`.to()`、`.contiguous()`，也不能合并或拆分 GEMM。仓库多处说明过：表达式或形状一变，舍入就会变（`embodiinfer/policies/pi05/modeling_pi05.py:664-666`）。 -3. **注册。** 单独写一行： - ```python - register_lazy_op("", "torch", "embodiinfer.backend.torch.ops", "") - ``` - 只有 Torch 实现的算子不能放进 `embodiinfer/layers/ops.py:160` 那个同时注册 torch 和 triton 的循环。 -4. **调用点。** 使用策略现有的 helper，pi0.5 是 `_op(kernels, "", False)(...)`。只有 Torch 实现时，偏好的实现必须是 `torch`，否则 `get_op` 会抛出 KeyError。`KernelSelection` 用参数传递，不使用全局变量。 -5. **preserve 分支原样保留。** 判断方式只用类型（例如 `type(module) is nn.Linear`），不按模型名分支。 -6. **测试至少覆盖三类：** - - 参考实现与原表达式在 CPU 随机输入（FP32 和 BF16）上 `torch.equal`； - - 调用点经过注册表，做法是用一个会记录调用的假调优档案，写法见 `tests/test_kernel_ops.py:135-150`； - - preserve 分支不经过注册表。 -7. **不能改的内容：** 已有测试的断言、`KernelSelection`、`embodiinfer/backend/triton/tuned/`、`scripts/`、`pyproject.toml`、`uv.lock`。 -8. **允许和禁止的操作：** 可以运行 `ruff` 和 CPU `pytest`；不能运行 GPU 代码，不能 `git commit`，不能安装依赖。 +## 4. 流程 -### 步骤 4：检查 +### 4.1 `register new`：建立运行 -控制器按顺序执行以下检查，前一项失败时后面的不再执行。所有检查都在 CPU 上运行。凡是需要导入 `embodiinfer` 的检查，都以 `PYTHONPATH=RUN/worktree` 在 `test_python` 中运行，并确认导入的 `embodiinfer.__file__` 位于 worktree 内。这是因为开发环境里 EmbodiInfer 是可编辑安装，不设置的话会导入主检出目录中的旧代码。 +`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]` 执行一次即视为确认调用点，然后： + +1. 恢复并校验运行目录（见 4.6）。 +2. 用当前或新指定的 `test_python` 在 `RUN/base` 上重新做预检。 +3. 将 `sites.json` 与 BASE 上重新定位的结果比对，然后冻结 agent、`test_python` 和 `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`。 + +**每次尝试的阶段：** + +| 尝试 | 阶段 | 会话 | |---|---|---| -| C1 范围 | `git -C RUN/worktree diff --name-status $BASE` 只涉及：`embodiinfer/backend/torch/ops.py`、`embodiinfer/layers/ops.py`、`sites.json` 中的源文件、`tests/test_kernel_ops.py` 或新的 `tests/test_registry_.py`。已有测试文件只允许新增内容，不允许删改原有行 | `git diff` | -| C2 结构 | ① `backend/torch/ops.py` 中有 `def `，带注解和 docstring，且不导入 `embodiinfer.models`/`policies`/`engine`；② `layers/ops.py` 中有对应的 `register_lazy_op` 行；③ 每个调用点已改为注册表调用，原表达式只出现在 preserve 分支中；④ 在 `test_python` 中验证：`available_ops()[""]` 包含 `torch`、`get_op("")` 就是参考实现、`KernelSelection("none").prepare()` 成功、`KernelSelection("torch").resolve("")` 返回参考实现 | `ast` + 子进程 | -| C3 格式 | 改动文件通过 `ruff check` 和 `ruff format --check` | `ruff` | -| C4 测试 | 规格 `tests` 中的测试和 `REGISTRATION.json` 列出的新测试全部通过 | `test_python -m pytest -q`，在 `RUN/worktree` 中运行 | -| C5 测试有效性 | 把新测试复制到 `RUN/base`，以 `PYTHONPATH=RUN/base` 逐个运行，每个测试的所有参数化实例必须执行并断言失败；通过、跳过、收集失败或依赖异常均不算有效证据 | `test_python` 内调用 `pytest.main`，保存 JUnit 报告 | +| 第一次 | `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`，并与当时源码的指纹绑定。中断后恢复时，只有源码字节未变， +已通过的检查才会复用。检查过程中源码发生变化会被判为失败。 -### 步骤 5：反馈重做 +### 4.5 反馈、重试和停止 -任何检查失败时，控制器写出 `attempts//feedback.json`： +检查失败时，控制器写出 `attempts//feedback.json`，下一次尝试的 prompt 会附上它： ```json -{ - "check": "C2", - "summary": "调用点 S2 仍直接调用 mlp.down_proj(activated)，且不在 preserve 分支内", - "details": {"site": "S2", "file": "embodiinfer/policies/pi05/modeling_pi05.py", "line": 315} -} +{"check": "C2", "summary": "S2: original expression remains outside a preserve branch at line 315", "details": {}} ``` -`details` 的内容随检查而定： -- C1：越界的文件或被改动的已有测试行； -- C2：失败的具体检查项； -- C3、C4：命令输出摘要，截断到固定长度； -- C5：在基线上仍然通过的测试名。 +智能体调用本身出错时，`check` 为 `agent`。每次尝试结束时，当时的补丁都会保存为 `attempts//patch.diff`。 -**下一次尝试：** -- 默认在同一个 Humanize2 会话中发出 `repair` prompt，附上 `feedback.json`，智能体在原 worktree 上继续修改； -- 连续两次在同一项检查失败时，新开一个会话，从 `plan` 重新开始，但 worktree 保留。 +| 结果 | 状态 | +|---|---| +| C1–C5 全部通过 | `passed`，记录已验证的源码指纹，并写出 `RUN/patch.diff` 和 `RUN/SUMMARY.md` | +| 达到 `max_attempts` 或 `max_seconds` | `exhausted`，worktree 保留，供人工接手 | +| 冻结输入被改动（规格、BASE、`RUN/base`、分支、提交历史、已审阅的调用点、只读上下文），或控制器、运行环境出错 | `failed`，不会转为反馈交给智能体重试 | +| Ctrl-C 或取消 | `interrupted`，可用 `resume` 继续 | -每次尝试结束时，控制器把 `git -C RUN/worktree diff $BASE` 存为 `attempts//patch.diff`。 +### 4.6 `resume` 和 `status` -**停止条件：** -- 全部检查通过：进入步骤 6； -- 达到 `max_attempts` 或 `max_seconds`：状态记为 `exhausted`，worktree 保留供人工接手； -- 规格、`BASE` 或 `RUN/base` 被改动：状态记为 `failed`。 +- `status RUN` 打印 `manifest.json`，不加锁，也不启动智能体。 +- `resume RUN` 根据 `manifest.json` 和 Humanize2 的状态台账继续： + - 已完成的阶段不重跑； + - 已用时长每秒落盘一次，中断不会重置预算； + - 状态已是 `passed`、`finalized`、`failed` 或 `exhausted` 时直接返回。 -### 步骤 6：产出分支（`register finalize`） +所有修改运行目录的命令都持有 `RUN/.run.lock`。C5 往 `RUN/base` 复制测试之前，会先把这些文件记入 +`manifest.json` 的 `base_overlay`；任何失败、取消或下一次启动时，都会先用 BASE 还原这些文件。 -全部检查通过后，控制器生成两个文件： +### 4.7 `register finalize`：导出和提交 -- `RUN/patch.diff`； -- `RUN/SUMMARY.md`，内容包括： - - 算子定义； - - 改动的调用点和保持不变的分支； - - 每项检查的命令和结果； - - 尝试次数； - - 明确写出未做的验证：GPU 整模型逐位一致。 +`finalize` 只接受状态为 `passed` 或 `finalized`、且源码指纹与通过检查时一致的运行。 +它会重新写出 `RUN/patch.diff` 和 `RUN/SUMMARY.md`。补丁不包含 `PLAN.md` 和 `REGISTRATION.json`。 -`finalize` 的行为： +| 参数 | 行为 | +|---|---| +| 无 | 打印 worktree 路径、分支名、建议的提交说明和产物路径，不修改 Git | +| `--commit` | 只暂存通过 C1 的文件（暂存区里有其他文件时报错），在 worktree 中以 `feat: route calls through the registry operator` 提交。提交使用当前 Git 身份，会运行仓库的 hook。状态变为 `finalized` | +| `--push` | 需要同时传 `--commit`，或此前已经提交。推送前再次确认文件字节与已验证的一致，然后执行 `git push -u origin registry/-` | -- 不带参数时，只打印 worktree 路径、分支名和建议的提交说明。 -- `--commit`：在 worktree 中提交到 `registry/-`。提交说明遵守 AGENTS.md 的 Git 约定，例如 `feat: route pi0.5 projections through the linear registry operator`。 -- `--push`：执行 `git push -u origin registry/-`。 +本工具不提供合并功能。 -本工具不提供合并功能。合并前，还需要在目标 GPU 上按 0009 §6 确认 `kernel_profile: none` 时整模型输出与 `BASE` 逐位一致，方法与 commit `449ebe2` 相同。 +#### 合并前 -## 5. 运行目录 +CPU 检查不能证明 GPU 上逐位一致。例如 pi0.5 的 native 路径只在 CUDA 上启用，CPU 测试执行不到其中的调用点。 +合并前必须在目标 GPU 上按 [0009 §6](../../../docs/proposals/0009-hotspot-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` 的末尾也会写明这一项尚未完成。 -``` -results/kernel_tuning/register// - manifest.json 当前阶段、每次尝试的结果、耗时（atomic_json 写入，run_lock 加锁） - spec.json 冻结的规格和 BASE - sites.json sites.md - context/ 智能体的只读上下文 - base/ BASE 的 worktree（只读） - worktree/ 分支 registry/- 的 worktree - workspace/ Humanize2 的工作区和日志 - attempts// plan/implement/repair 的 prompt 和回复、PLAN.md、REGISTRATION.json、 - checks.json、feedback.json、patch.diff - patch.diff SUMMARY.md -``` +## 5. Git 的使用和清理 -`resume` 会结合 `manifest.json` 和 Humanize2 的状态台账继续，已经通过的阶段在源码和 -`REGISTRATION.json` 哈希未变时不再重跑。规格变了就要新建一次运行。 -同次执行中的 repair 沿用会话；进程中断后从记录的阶段继续，新会话会收到规格、已有计划和反馈。 -基线测试复制前先记入恢复台账，失败、取消和下次恢复时均还原这些测试文件。 +智能体执行命令时没有交互确认，控制器因此用 Git 做隔离和取证： -## 6. 现有工具链的复用 +- **隔离**：两个 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` 时才提交或推送。 -| 现有部件 | 用在 | 用法 | -|---|---|---| -| Humanize2 flow 机制（`scripts/kernel_tuning/flow.py`、`__main__.py` 中的 `_execute`） | 步骤 3、5 | 新写一个 flow，沿用 `@flow(..., resumable=True)`、`AgentCollection`、`LocalEnv`、`Budget`、`ctx.state` 中的进度台账，以及 `Hmz(...).run(...)` 的启动方式 | -| `scripts/kernel_tuning/artifacts.py` 中的 `atomic_json`、`run_lock` | 运行目录 | 直接导入 | -| `scripts/kernel_tuning/runner.py` 的子进程处理（超时、进程组清理） | 步骤 1、4 的子进程 | 照搬这套做法，但启动的是 `test_python` | -| `embodiinfer.layers`（`available_ops`、`get_op`、`KernelSelection`） | 步骤 1、C2 | 在 `test_python` 子进程中导入 | -| `ruff`、`pytest` | C3–C5 | 命令行调用 | -| `git worktree` | 步骤 1、6 | 命令行调用 | +运行结束后，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 +``` -`profile`、`capture-ops`、`calibrate`、`tune-all`、`install`、`verify` 不在本流程内使用。它们属于调优阶段，或者需要 GPU。 +状态的变化顺序为：`awaiting_confirmation` → `ready` → `running` → `passed` → `finalized`。 +中途可能变为 `exhausted`、`failed` 或 `interrupted`。 ## 7. 代码布局 | 文件 | 内容 | |---|---| -| `registration/__init__.py` | 包说明 | -| `registration/spec.py` | 规格的解析（YAML 或命令行）和校验，生成 `spec.json` | -| `registration/sites.py` | 用 `ast` 定位调用点、提取 preserve 条件、查找调用链 | -| `registration/workspace.py` | worktree 的创建和清理、diff、C1 检查 | -| `registration/checks.py` | C2–C5 检查，生成 `feedback.json` | -| `registration/prompts.py` | `plan`、`implement`、`repair` 三种 prompt | -| `registration/flow.py` | Humanize2 flow `operator_registration` | -| `registration/finalize.py` | 生成 `SUMMARY.md`，提交和推送 | -| `registration/cli.py` | 参数解析、运行环境校验、Humanize2 启动 | -| `__main__.py` | 增加 `register` 子命令组：`new`、`run`、`resume`、`status`、`finalize` | -| `tests/kernel_tuning/test_registration.py` | CPU 测试，覆盖：规格校验；在示例源码上定位调用点；用合成仓库测 C1；构造缺参考实现、错进循环、原表达式残留等情况测 C2；C5 能识别在基线上也通过的测试；用 Humanize2 自带的假智能体（同 `tests/kernel_tuning/test_flow.py`）跑通"失败 → 反馈 → 修复 → 通过"并验证 `resume` | - -工具环境已经包含 `pyyaml`（见 `scripts/kernel_tuning/pyproject.toml`），不需要新增依赖。 +| `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 上逐位一致。** 例如 pi0.5 的 native 路径只在 CUDA 上启用（`embodiinfer/policies/pi05/modeling_pi05.py:565-572`），它上面的调用点在 CPU 测试中执行不到。因此合并前仍需在 GPU 上验证一次（见步骤 6）。 -- **调用点由调用方指定。** 工具只负责定位和确认，不会从热点报告中自动挑选。 -- **每次只注册一个算子，不支持把多个相邻算子融合成一个。** -- **策略需要已有 `kernels` 传参链。** 必须已经有 `KernelSelection` 和类似 `_op` 的 helper，目前只有 pi0.5 具备。其他策略要先人工接入一次。 -- **智能体写出的补丁可能检查全过但结构不好**（例如传参链过长），仍然需要人工审阅。 -- **新测试必须单独列出函数 node ID。** 不接受整个文件、已有测试或挑选部分参数化实例。 - 新符号的导入放在函数内部，先断言符号存在，避免用缺失导入代替有效的基线断言失败。 -- **目录和哈希检查不等同于安全沙箱。** 智能体和测试继承运行用户的权限；控制器隐藏 CUDA/ROCm - 设备，禁止运行 GPU 测试，但这不是对任意不受信代码的操作系统隔离。 +- CPU 检查和有限的测试不能证明 GPU 上逐位一致，也不能保证智能体写的测试覆盖了全部行为。 +- 每次运行只注册一个算子，不支持融合。调用点只能位于 `embodiinfer/policies//` 下， + 且所在文件必须已有 `_op` 集成。`kernels` 传递链的检查只覆盖同一文件内的调用方。 +- 检查全部通过的补丁仍可能结构不佳（例如传参链过长），需要人工审阅。 +- 目录边界、哈希校验和隐藏加速设备都不构成操作系统级的沙箱。智能体和测试以当前用户的权限运行， + 必要时请在隔离的账户或容器中运行。 diff --git a/scripts/kernel_tuning/registration/prompts.py b/scripts/kernel_tuning/registration/prompts.py index 294a4a2..62cdeff 100644 --- a/scripts/kernel_tuning/registration/prompts.py +++ b/scripts/kernel_tuning/registration/prompts.py @@ -55,8 +55,9 @@ def prompt(run: RegistrationRun, stage: str, feedback: dict | None = None) -> st {json.dumps(run.spec.as_dict(), indent=2, ensure_ascii=False)} {RULES} -Read context/sites.json, AGENTS.md, the source snapshots, test_kernel_ops.py and -migration.diff. Existing PLAN.md and REGISTRATION.json describe previous work. +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 ( diff --git a/scripts/kernel_tuning/registration/workspace.py b/scripts/kernel_tuning/registration/workspace.py index 9c85aee..df6d9fb 100644 --- a/scripts/kernel_tuning/registration/workspace.py +++ b/scripts/kernel_tuning/registration/workspace.py @@ -24,6 +24,9 @@ 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"} @@ -303,6 +306,7 @@ def prepare(self, agent: str, python: str | None = None) -> None: "AGENTS.md", REFERENCE, REGISTRY, + ROUTED_EXAMPLE, "tests/test_kernel_ops.py", *(s.source for s in self.spec.sites), } @@ -314,11 +318,6 @@ def prepare(self, agent: str, python: str | None = None) -> None: from .prompts import RULES (context / "RULES.md").write_text(RULES, encoding="utf-8") - try: - example = git(self.base, "show", "449ebe2", "--", "embodiinfer/") - except ContractError: - example = "The optional migration example 449ebe2 is unavailable in this repository.\n" - (context / "migration.diff").write_text(example, encoding="utf-8") self.manifest.update( agent=agent, test_python=python or self.manifest["test_python"], diff --git a/tests/kernel_tuning/test_registration.py b/tests/kernel_tuning/test_registration.py index 855eaf0..4a610eb 100644 --- a/tests/kernel_tuning/test_registration.py +++ b/tests/kernel_tuning/test_registration.py @@ -25,6 +25,7 @@ from scripts.kernel_tuning.registration.workspace import ( REFERENCE, REGISTRY, + ROUTED_EXAMPLE, RegistrationRun, check_scope, check_unregistered, @@ -385,6 +386,13 @@ def test_integrity_and_baseline_overlay_recovery(registration_run) -> None: 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") From b288d384fa3fbab6a6e899587a803ede694cd1d7 Mon Sep 17 00:00:00 2001 From: owpeter Date: Tue, 13 Oct 2026 22:46:20 -0700 Subject: [PATCH 30/44] chore: drop the gated residual example task Tests now write the hand-written task package into a temporary directory, and the quick start generates a catalog task to check instead. Signed-off-by: owpeter --- CONTRIBUTING.md | 6 +- .../tasks/gated_residual/README.md | 47 ---------- .../tasks/gated_residual/baseline.py | 9 -- .../tasks/gated_residual/definition.json | 14 --- .../tasks/gated_residual/tuning.yaml | 29 ------- .../tasks/gated_residual/workloads.jsonl | 3 - tests/kernel_tuning/conftest.py | 87 +++++++++++++++++-- tests/kernel_tuning/test_numerics.py | 5 +- 8 files changed, 87 insertions(+), 113 deletions(-) delete mode 100644 benchmarks/kernel_tuning/tasks/gated_residual/README.md delete mode 100644 benchmarks/kernel_tuning/tasks/gated_residual/baseline.py delete mode 100644 benchmarks/kernel_tuning/tasks/gated_residual/definition.json delete mode 100644 benchmarks/kernel_tuning/tasks/gated_residual/tuning.yaml delete mode 100644 benchmarks/kernel_tuning/tasks/gated_residual/workloads.jsonl diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 5894ed2..fb4c382 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -65,7 +65,8 @@ 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 check benchmarks/kernel_tuning/tasks/gated_residual +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 ``` To tune one model, capture its existing preparation forward and calibrate numerical @@ -139,8 +140,7 @@ plan before applying it. When inheriting system packages, verify that the instal 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 [approved contract](docs/proposals/0008-kernel-tuning.md) and the -[gated residual task](benchmarks/kernel_tuning/tasks/gated_residual/README.md). +See the [approved contract](docs/proposals/0008-kernel-tuning.md). The [Chinese workflow guide](docs/proposals/0008-kernel-tuning.md#11-自动化调优使用说明) explains task inputs, the full tuning loop, stopping conditions, and saved results. The controller snapshots task inputs, saves all attempts, and exports improving diff --git a/benchmarks/kernel_tuning/tasks/gated_residual/README.md b/benchmarks/kernel_tuning/tasks/gated_residual/README.md deleted file mode 100644 index e721b02..0000000 --- a/benchmarks/kernel_tuning/tasks/gated_residual/README.md +++ /dev/null @@ -1,47 +0,0 @@ -# Gated residual BF16 task - -This self-contained example covers the multiply-then-add semantics used by -`embodiinfer/backend/triton/norm.py::gated_residual`. Inputs are contiguous BF16 -`residual, update: [B,T,H]` and `gate: [B,1,H]` or `[B,T,H]`. Output has the same -shape and dtype as residual. Round the product to BF16 before adding; contracting -the expression into an FMA can violate the bit-exact contract. Inputs are immutable. - -The mathematical reference is embedded in `definition.json`. `baseline.py` is a -separate eager Torch implementation, written for this task under the repository's -Apache-2.0 license. It is an executable template baseline, **not a claim to beat -EmbodiInfer's existing Triton kernel**. Replace it with the best-known standalone -implementation before using this task to claim a production improvement. Include -all local helpers inside the task package and record their provenance here. - -The three synthetic workloads exercise broadcast gates, per-token gates, and a -non-power-of-two hidden dimension. They do not represent measured model traffic. -Use workload shapes and weights from the intended deployment for real tuning. - -On the target machine, prepare its compatible Torch/CUDA/Triton environment and -set `evaluator_python` in `tuning.yaml` to that environment's absolute Python path. -The default evaluator also needs `benchmarks/kernel_tuning/requirements.txt`. -Select `language: cuda` for a CUDA C++ search; the agent then supplies a CUDA -Solution with a supported Torch or TVM-FFI binding. The CUDA toolkit/compiler must -already be available. `target_hardware` is a Solution label, not a capability check; -the evaluator probes the actual device (including Thor) at run time. - -From the repository root, in the optional Humanize2 tooling environment: - -```bash -python -m scripts.kernel_tuning check benchmarks/kernel_tuning/tasks/gated_residual -python -m scripts.kernel_tuning check benchmarks/kernel_tuning/tasks/gated_residual --environment -python -m scripts.kernel_tuning run benchmarks/kernel_tuning/tasks/gated_residual --agent 'HARNESS/MODEL:EFFORT' -python -m scripts.kernel_tuning status results/kernel_tuning/RUN_ID -python -m scripts.kernel_tuning resume results/kernel_tuning/RUN_ID -python -m scripts.kernel_tuning export results/kernel_tuning/RUN_ID results/exported_gated_residual -``` - -Replace the agent placeholder with a configured Humanize2 agent specification. -`check` performs static validation. `--environment` additionally imports the -evaluator and probes its hardware/dependencies. `run` and `resume` execute real -coding-agent turns and measurements. Each candidate's `evaluation/command.json` -records the exact reproducible worker command and timeout. Fixed conditions, -paired measurements, and acceptance limits are in `tuning.yaml`. - -For the optional custom `benchmark.py` interface and artifact contract, see -[`0008-kernel-tuning.md`](../../../../docs/proposals/0008-kernel-tuning.md). diff --git a/benchmarks/kernel_tuning/tasks/gated_residual/baseline.py b/benchmarks/kernel_tuning/tasks/gated_residual/baseline.py deleted file mode 100644 index b2c2977..0000000 --- a/benchmarks/kernel_tuning/tasks/gated_residual/baseline.py +++ /dev/null @@ -1,9 +0,0 @@ -"""Lossless eager baseline with an explicit intermediate dtype rounding.""" - -import torch - - -def run(residual: torch.Tensor, update: torch.Tensor, gate: torch.Tensor) -> torch.Tensor: - """Match the separate multiply/add semantics used by the gated residual path.""" - product = (update * gate).to(update.dtype) - return (residual + product).to(residual.dtype) diff --git a/benchmarks/kernel_tuning/tasks/gated_residual/definition.json b/benchmarks/kernel_tuning/tasks/gated_residual/definition.json deleted file mode 100644 index 572c2f6..0000000 --- a/benchmarks/kernel_tuning/tasks/gated_residual/definition.json +++ /dev/null @@ -1,14 +0,0 @@ -{ - "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": "import torch\n\ndef run(residual, update, gate):\n product = (update * gate).to(update.dtype)\n return (residual + product).to(residual.dtype)\n" -} diff --git a/benchmarks/kernel_tuning/tasks/gated_residual/tuning.yaml b/benchmarks/kernel_tuning/tasks/gated_residual/tuning.yaml deleted file mode 100644 index bb6d9d7..0000000 --- a/benchmarks/kernel_tuning/tasks/gated_residual/tuning.yaml +++ /dev/null @@ -1,29 +0,0 @@ -schema_version: 1 -language: triton -# Set --agent explicitly when starting a real run, or fill in a valid hmz spec. -agent: "" -# Change to the absolute Python executable in the prepared Thor/runtime environment. -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 diff --git a/benchmarks/kernel_tuning/tasks/gated_residual/workloads.jsonl b/benchmarks/kernel_tuning/tasks/gated_residual/workloads.jsonl deleted file mode 100644 index 72193b7..0000000 --- a/benchmarks/kernel_tuning/tasks/gated_residual/workloads.jsonl +++ /dev/null @@ -1,3 +0,0 @@ -{"definition":"gated_residual_bf16","workload":{"uuid":"broadcast-small","axes":{"B":1,"T":16,"H":1024,"G":1},"inputs":{"residual":{"type":"random"},"update":{"type":"random"},"gate":{"type":"random"}}},"solution":null,"evaluation":null} -{"definition":"gated_residual_bf16","workload":{"uuid":"broadcast-wide","axes":{"B":2,"T":32,"H":2048,"G":1},"inputs":{"residual":{"type":"random"},"update":{"type":"random"},"gate":{"type":"random"}}},"solution":null,"evaluation":null} -{"definition":"gated_residual_bf16","workload":{"uuid":"token-gate-tail","axes":{"B":1,"T":17,"H":1000,"G":17},"inputs":{"residual":{"type":"random"},"update":{"type":"random"},"gate":{"type":"random"}}},"solution":null,"evaluation":null} diff --git a/tests/kernel_tuning/conftest.py b/tests/kernel_tuning/conftest.py index 91a2b4f..e5c8542 100644 --- a/tests/kernel_tuning/conftest.py +++ b/tests/kernel_tuning/conftest.py @@ -2,13 +2,13 @@ from __future__ import annotations -import shutil +import json import sys from pathlib import Path from typing import Any import pytest -from scripts.kernel_tuning.artifacts import REPOSITORY, RunStore +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") @@ -46,11 +46,88 @@ def create_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 task_dir(tmp_path: Path) -> Path: - """Make a portable task using a fake measurement adapter and explicit Python.""" +def example_task(tmp_path: Path) -> Path: + """Write a hand-written BF16 gated residual task package with synthetic workloads.""" path = tmp_path / "task" - shutil.copytree(REPOSITORY / "benchmarks/kernel_tuning/tasks/gated_residual", path) + 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 diff --git a/tests/kernel_tuning/test_numerics.py b/tests/kernel_tuning/test_numerics.py index 87c88df..5ede410 100644 --- a/tests/kernel_tuning/test_numerics.py +++ b/tests/kernel_tuning/test_numerics.py @@ -70,7 +70,7 @@ def test_bf16_intermediate_rounding_is_observable() -> None: @pytest.mark.gpu @pytest.mark.parametrize("language", ["triton", "cuda"]) -def test_default_adapter_builds_multi_file_solution(tmp_path: Path, language: str) -> None: +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 @@ -85,8 +85,7 @@ def test_default_adapter_builds_multi_file_solution(tmp_path: Path, language: st pytest.importorskip("triton") elif shutil.which("nvcc") is None: pytest.skip("requires the CUDA toolkit") - task = tmp_path / "task" - shutil.copytree(REPOSITORY / "benchmarks/kernel_tuning/tasks/gated_residual", task) + task = example_task definition = { "name": "vector_add_test", "op_type": "elementwise", From 6354c5aa831a100890c4a3a636350944aab6ab3a Mon Sep 17 00:00:00 2001 From: owpeter Date: Tue, 13 Oct 2026 22:46:20 -0700 Subject: [PATCH 31/44] chore: drop the Thor kernel profile and its measurements Keep the profile loader and install tooling; remove the committed profile, the Thor benchmark tables, run identity, and the relay note from 0009. Signed-off-by: owpeter --- docs/proposals/0009-hotspot-kernel-tuning.md | 44 +-- .../backend/triton/tuned/thor/profile.json | 80 ----- .../tuned/thor/split_kv_attention/__init__.py | 62 ---- .../thor/split_kv_attention/candidate.py | 329 ------------------ 4 files changed, 9 insertions(+), 506 deletions(-) delete mode 100644 embodiinfer/backend/triton/tuned/thor/profile.json delete mode 100644 embodiinfer/backend/triton/tuned/thor/split_kv_attention/__init__.py delete mode 100644 embodiinfer/backend/triton/tuned/thor/split_kv_attention/candidate.py diff --git a/docs/proposals/0009-hotspot-kernel-tuning.md b/docs/proposals/0009-hotspot-kernel-tuning.md index fd0c839..214b42c 100644 --- a/docs/proposals/0009-hotspot-kernel-tuning.md +++ b/docs/proposals/0009-hotspot-kernel-tuning.md @@ -356,10 +356,10 @@ reproduces them. 4. Agent selection and new fusion regions (one-time code changes, §4.4). 5. Further hardware profiles. -Status (2026-10-12): steps 1–3 are implemented for pi0.5; the loop was proven -on Jetson Thor instead of an RTX 4090 (§9), and the `thor` profile ships one -tuned operator. Qwen R2R and StreamVLN call sites, step 4 and an RTX 4090 -profile are not done; neither is shape-range tuning (§10). The manually selected, +Status (2026-10-12): steps 1–3 are implemented for pi0.5; the loop was +exercised end to end on Jetson Thor instead of an RTX 4090. No hardware profile +is committed yet. Qwen R2R and StreamVLN call sites, step 4 and committed +profiles are not done; neither is shape-range tuning (§10). The manually selected, single-operator registration portion of step 4 now has the CPU controller in §4.8; automatic selection and new fusion remain future work. @@ -389,40 +389,16 @@ split; per-operator ablation. Expected gain is bounded by the recoverable time in the hotspot report; it disappears where regions are already at the bandwidth or launch floor. Batch > 1 and other hardware are reported only when measured. -First measurement (2026-10-12): Jetson Thor (sm_110, MAXN, `jetson_clocks`), -Torch 2.13.0+cu132, Triton 3.7.1, `pi05_libero_finetuned_v044`, B=1, BF16, -10 denoise steps, LIBERO-10 1,600 frames after 10 warmup calls, Inductor, -prefix and full-loop CUDA Graphs, native prefix/denoise attention. - -| `kernel_profile` | Mean / P50 (ms) | Prefill (ms) | Decode (ms) | -|---|---|---|---| -| `torch` path (`native_inference: false`) | 179.9 / 179.9 | 63.9 | 111.4 | -| `none` (existing Triton) | 98.8 / 99.7 | 45.4 | 49.0 | -| `thor` | 87.6 / 85.4 | 45.2 | 37.2 | - -The `thor` profile contains one operator, `split_kv_attention`, tuned by -`codex/gpt-6-astra:high` (run `20261012T084037Z-cf9243b720`, 12 attempts). Its -task timing (CUDA Graph, cold L2) against the Torch reference and the existing -Triton kernel: - -| Workload (calls per step) | Torch (µs) | Existing Triton (µs) | Tuned (µs) | -|---|---|---|---| -| prefix, 576 keys (18) | 1,079 | ~178 | 139 | -| denoise, 576 + 50 keys (180) | 304 | ~120 | 38 | - -Under the benchmark configuration every dispatch decision selected the tuned -kernel (72 prefix and 720 denoise calls executed before graph capture; replays -reuse them). `verify` on 40 -recorded inputs passed (worst ratio 0.80 of the §6 limits). With -`kernel_profile: none`, 200 benchmark frames were bit-identical to the previous -code on both the native and the Torch path (`TORCHINDUCTOR_DETERMINISTIC=1`). +With `kernel_profile: none`, 200 benchmark frames were bit-identical to the +previous code on both the native and the Torch path +(`TORCHINDUCTOR_DETERMINISTIC=1`). ## 10. Risks and limitations - Agent kernels are often shape-specialized (a 2026-10-07 `split_kv_attention` candidate asserted `P % PART_LENGTH == 0`); `applies` and fallback mitigate, - but uncovered shapes run the old path. The shipped kernel covers the captured - B=1 shapes with a 576-token prefix only; the 48-token text bucket, B > 1 and + but uncovered shapes run the old path. A kernel tuned on captured B=1 shapes + with a 576-token prefix covers only those; the 48-token text bucket, B > 1 and other checkpoints fall back. Follow-up: tasks over declared shape ranges (buckets, batch sizes, several layouts per shape), kernels with runtime strides and scales, and range-based `applies`. @@ -432,8 +408,6 @@ code on both the native and the Torch path (`TORCHINDUCTOR_DETERMINISTIC=1`). agent turn a candidate into a registered operator by PR, with a bit-exact `kernel_profile: none` parity test as the proof. - GEMMs (the largest share of pi0.5 time on Thor) are not registry operators. -- Agent runs depend on a relay that drops streams; failed attempts consume the - search budget. - Profiles are tied to Triton/Torch versions; a version upgrade disables them under `auto` until re-verified. - Attribution under CUDA Graphs depends on CUPTI support; the eager fallback pass diff --git a/embodiinfer/backend/triton/tuned/thor/profile.json b/embodiinfer/backend/triton/tuned/thor/profile.json deleted file mode 100644 index dcd77bc..0000000 --- a/embodiinfer/backend/triton/tuned/thor/profile.json +++ /dev/null @@ -1,80 +0,0 @@ -{ - "schema": 1, - "match": { - "device_name": "NVIDIA Thor", - "capability": [ - 11, - 0 - ], - "torch": "2.13.0+cu132", - "triton": "3.7.1" - }, - "ops": { - "split_kv_attention": { - "module": "split_kv_attention", - "function": "call", - "applies": "applies", - "sha256": "55e8a20b057457d63f7557aaa16f344d74397e4925047f1b6666bcf52faac2d3", - "verified_layouts": 2, - "evidence": { - "run_id": "20261012T084037Z-cf9243b720", - "task_digest": "0192979023cd12ccc2a72bb04807c258509aeb29259333a33e1f89077f51a64b", - "capture_identity": "3fda25f6804e3178470c1bc390062d402242029cc2a6fcbdf8ecd96904b22ca7", - "solution_digest": "8b1db93ae8388ca0f296151d4bbf25e4f5e8c31d43a8f6234f87630d8f92cc46", - "decision": { - "promoted": true, - "reasons": [], - "rounds": [ - { - "mean_latency_ms": { - "baseline": 0.3732264595013676, - "incumbent": 0.04873815755365473, - "candidate": 0.0469428035458832 - }, - "improvement": { - "baseline": 0.8742243419488559, - "incumbent": 0.03683672296793472 - }, - "workloads": { - "d5cf735e215bb7574e71": { - "baseline": 1.0804690765937168, - "incumbent": 0.1394609593451023, - "candidate": 0.1392327032983303 - }, - "03293aab11415ac3e8c7": { - "baseline": 0.30250219779213267, - "incumbent": 0.03966587737450997, - "candidate": 0.037713813570638494 - } - } - }, - { - "mean_latency_ms": { - "baseline": 0.37376189036260954, - "incumbent": 0.048764153013626736, - "candidate": 0.046742537992921745 - }, - "improvement": { - "baseline": 0.8749403318043637, - "incumbent": 0.041456990345758005 - }, - "workloads": { - "d5cf735e215bb7574e71": { - "baseline": 1.079287349820137, - "incumbent": 0.13775917767484983, - "candidate": 0.13755703370273112 - }, - "03293aab11415ac3e8c7": { - "baseline": 0.3032093444168568, - "incumbent": 0.039864650547504425, - "candidate": 0.037661088421940804 - } - } - } - ] - }, - "export_checksums": "4bbfc952ca6abcbce7edf7a2b3b0494c60a17354f97d5c353c115e719c58ce67" - } - } - } -} diff --git a/embodiinfer/backend/triton/tuned/thor/split_kv_attention/__init__.py b/embodiinfer/backend/triton/tuned/thor/split_kv_attention/__init__.py deleted file mode 100644 index 1f1c149..0000000 --- a/embodiinfer/backend/triton/tuned/thor/split_kv_attention/__init__.py +++ /dev/null @@ -1,62 +0,0 @@ -"""Tuned ``split_kv_attention`` kernel for the ``thor`` profile. - -Generated by ``scripts.kernel_tuning install`` from run 20261012T084037Z-cf9243b720; 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, 576, 256), (0, 0, 256, 1), 'bfloat16'), ('tensor', (1, 1, 576, 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, 626), (0, 0, 0, 1), 'float32'), ('constant', '0.0625')), (('tensor', (1, 8, 576, 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, 576, 256), (0, 0, 256, 1), 'bfloat16'), ('tensor', (1, 1, 576, 256), (0, 0, 2560, 1), 'bfloat16'), ('tensor', (1, 1, 1, 576), (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/thor/split_kv_attention/candidate.py b/embodiinfer/backend/triton/tuned/thor/split_kv_attention/candidate.py deleted file mode 100644 index de4bc1b..0000000 --- a/embodiinfer/backend/triton/tuned/thor/split_kv_attention/candidate.py +++ /dev/null @@ -1,329 +0,0 @@ -"""Direct split-KV attention for the two recorded pi05 layouts.""" - -import torch -import triton -import triton.language as tl - - -@triton.jit -def _partial_attention( - Q, PK, PV, SK, SV, Mask, Part, Lse, - qb: tl.constexpr, qh: tl.constexpr, qs: tl.constexpr, - pkb: tl.constexpr, pkh: tl.constexpr, pks: tl.constexpr, - pvb: tl.constexpr, pvh: tl.constexpr, pvs: tl.constexpr, - skb: tl.constexpr, skh: tl.constexpr, sks: tl.constexpr, - svb: tl.constexpr, svh: tl.constexpr, svs: tl.constexpr, - mb: tl.constexpr, mn: tl.constexpr, - P: tl.constexpr, S: tl.constexpr, H: tl.constexpr, - HK: tl.constexpr, D: tl.constexpr, - SPLITS: tl.constexpr, PART_LENGTH: tl.constexpr, - BM: tl.constexpr, BN: tl.constexpr, -): - bh = tl.program_id(0) - batch = bh // H - head = bh % H - kv_head = head // (H // HK) - partition = tl.program_id(2) - rows = tl.program_id(1) * BM + tl.arange(0, BM) - dims = tl.arange(0, D) - q = tl.load( - Q + batch * qb + head * qh + rows[:, None] * qs + dims[None, :], - rows[:, None] < S, - other=0, - ) - maximum = tl.full((BM,), -float("inf"), tl.float32) - denominator = tl.zeros((BM,), tl.float32) - acc = tl.zeros((BM, D), tl.float32) - - for offset in range(PART_LENGTH // BN): - columns = partition * PART_LENGTH + offset * BN + tl.arange(0, BN) - suffix = (columns >= P) & (columns < P + S) - # Clamp inactive offsets too; only the corresponding segment is loaded. - suffix_rows = tl.maximum(columns - P, 0) - ks = tl.load( - SK + batch * skb + kv_head * skh - + suffix_rows[:, None] * sks + dims[None, :], - suffix[:, None], - other=0, - ) - if P > 0: - prefix = columns < P - kp = tl.load( - PK + batch * pkb + kv_head * pkh - + columns[:, None] * pks + dims[None, :], - prefix[:, None], - other=0, - ) - k = tl.where(prefix[:, None], kp, ks) - else: - # Empty prefix views can alias the suffix; never dereference them. - k = ks - - valid = columns < P + S - mask_value = tl.load( - Mask + batch * mb + columns * mn, valid, other=1.0 - ) - valid = valid & (mask_value == 0) - scores = tl.dot(q, tl.trans(k), input_precision="ieee", out_dtype=tl.float32) - scores = scores * (0.0625 * 1.4426950408889634) - scores = tl.where( - valid[None, :] & (rows[:, None] < S), scores, -float("inf") - ) - next_max = tl.maximum(maximum, tl.max(scores, 1)) - safe_max = tl.where(next_max == -float("inf"), 0.0, next_max) - correction = tl.exp2(maximum - safe_max) - prob = tl.exp2(scores - safe_max[:, None]) - - vs = tl.load( - SV + batch * svb + kv_head * svh - + suffix_rows[:, None] * svs + dims[None, :], - suffix[:, None], - other=0, - ) - if P > 0: - vp = tl.load( - PV + batch * pvb + kv_head * pvh - + columns[:, None] * pvs + dims[None, :], - prefix[:, None], - other=0, - ) - v = tl.where(prefix[:, None], vp, vs) - else: - v = vs - - # Preserve the starting kernel's BF16 dot operands and FP32 accumulation. - acc = acc * correction[:, None] + tl.dot( - prob.to(v.dtype), v, input_precision="ieee", out_dtype=tl.float32 - ) - denominator = denominator * correction + tl.sum(prob, 1) - maximum = next_max - - safe_denominator = tl.where(denominator > 0, denominator, 1.0) - value = acc / safe_denominator[:, None] - index = (bh * SPLITS + partition) * S + rows - tl.store(Part + index[:, None] * D + dims[None, :], value, rows[:, None] < S) - if SPLITS > 1: - logsum = tl.where( - denominator > 0, - maximum + tl.log2(safe_denominator), - -float("inf"), - ) - tl.store(Lse + index, logsum, rows < S) - - -@triton.jit -def _attention_segment( - q, rows, maximum, denominator, acc, - K, V, Mask, - ks: tl.constexpr, vs: tl.constexpr, mn: tl.constexpr, - LENGTH: tl.constexpr, MASK_OFFSET: tl.constexpr, - QUERY_ROWS: tl.constexpr, D: tl.constexpr, BN: tl.constexpr, -): - dims = tl.arange(0, D) - for offset in tl.range(0, (LENGTH + BN - 1) // BN, num_stages=2): - columns = offset * BN + tl.arange(0, BN) - in_segment = columns < LENGTH - k = tl.load( - K + columns[:, None] * ks + dims[None, :], - in_segment[:, None], - other=0, - ) - mask_value = tl.load( - Mask + (MASK_OFFSET + columns) * mn, in_segment, other=1.0 - ) - valid = in_segment & (mask_value == 0) - scores = tl.dot(q, tl.trans(k), input_precision="ieee", out_dtype=tl.float32) - scores = scores * (0.0625 * 1.4426950408889634) - scores = tl.where( - valid[None, :] & (rows[:, None] < QUERY_ROWS), scores, -float("inf") - ) - next_max = tl.maximum(maximum, tl.max(scores, 1)) - safe_max = tl.where(next_max == -float("inf"), 0.0, next_max) - correction = tl.exp2(maximum - safe_max) - prob = tl.exp2(scores - safe_max[:, None]) - - v = tl.load( - V + columns[:, None] * vs + dims[None, :], - in_segment[:, None], - other=0, - ) - acc = acc * correction[:, None] + tl.dot( - prob.to(v.dtype), v, input_precision="ieee", out_dtype=tl.float32 - ) - denominator = denominator * correction + tl.sum(prob, 1) - maximum = next_max - return maximum, denominator, acc - - -@triton.jit -def _segmented_attention( - Q, PK, PV, SK, SV, Mask, Out, - qb: tl.constexpr, qh: tl.constexpr, qs: tl.constexpr, - pkb: tl.constexpr, pkh: tl.constexpr, pks: tl.constexpr, - pvb: tl.constexpr, pvh: tl.constexpr, pvs: tl.constexpr, - skb: tl.constexpr, skh: tl.constexpr, sks: tl.constexpr, - svb: tl.constexpr, svh: tl.constexpr, svs: tl.constexpr, - mb: tl.constexpr, mn: tl.constexpr, - P: tl.constexpr, S: tl.constexpr, H: tl.constexpr, - HK: tl.constexpr, D: tl.constexpr, - BM: tl.constexpr, BN: tl.constexpr, - PACK_QUERY_HEADS: tl.constexpr, -): - # The frozen prefix ends on a tile boundary; both segments share one softmax. - tl.static_assert(P % BN == 0) - bh = tl.program_id(0) - if PACK_QUERY_HEADS: - # The host restricts this path to the contiguous, single-batch MQA case. - tl.static_assert(HK == 1) - tl.static_assert(qh == S * qs) - QUERY_ROWS: tl.constexpr = H * S - batch = 0 - head = 0 - kv_head = 0 - else: - QUERY_ROWS: tl.constexpr = S - batch = bh // H - head = bh % H - kv_head = head // (H // HK) - rows = tl.program_id(1) * BM + tl.arange(0, BM) - dims = tl.arange(0, D) - q = tl.load( - Q + batch * qb + head * qh + rows[:, None] * qs + dims[None, :], - rows[:, None] < QUERY_ROWS, - other=0, - ) - maximum = tl.full((BM,), -float("inf"), tl.float32) - denominator = tl.zeros((BM,), tl.float32) - acc = tl.zeros((BM, D), tl.float32) - - maximum, denominator, acc = _attention_segment( - q, rows, maximum, denominator, acc, - PK + batch * pkb + kv_head * pkh, - PV + batch * pvb + kv_head * pvh, - Mask + batch * mb, - pks, pvs, mn, P, 0, QUERY_ROWS, D, BN, - ) - maximum, denominator, acc = _attention_segment( - q, rows, maximum, denominator, acc, - SK + batch * skb + kv_head * skh, - SV + batch * svb + kv_head * svh, - Mask + batch * mb, - sks, svs, mn, S, P, QUERY_ROWS, D, BN, - ) - - safe_denominator = tl.where(denominator > 0, denominator, 1.0) - value = acc / safe_denominator[:, None] - if PACK_QUERY_HEADS: - out_rows = rows - else: - out_rows = bh * S + rows - tl.store( - Out + out_rows[:, None] * D + dims[None, :], - value, - rows[:, None] < QUERY_ROWS, - ) - - -@triton.jit -def _merge_attention( - Part, Lse, Out, - S: tl.constexpr, D: tl.constexpr, SPLITS: tl.constexpr, -): - row = tl.program_id(0) - bh = tl.program_id(1) - partitions = tl.arange(0, SPLITS) - dims = tl.arange(0, D) - indices = (bh * SPLITS + partitions) * S + row - logsum = tl.load(Lse + indices) - maximum = tl.max(logsum, 0) - # All-empty rows have zero weights, rather than exp2(-inf - -inf). - safe_max = tl.where(maximum == -float("inf"), 0.0, maximum) - weight = tl.exp2(logsum - safe_max) - partial = tl.load(Part + indices[:, None] * D + dims[None, :]) - denominator = tl.sum(weight, 0) - safe_denominator = tl.where(denominator > 0, denominator, 1.0) - value = tl.sum(partial * weight[:, None], 0) / safe_denominator - tl.store(Out + (bh * S + row) * D + dims, value) - - -def run(x0, x1, x2, x3, x4, x5): - """Return contiguous BF16 BHSD attention without modifying any input.""" - batch, heads, size, width = x0.shape - prefix = x1.shape[2] - if (batch, heads, width) != (1, 8, 256): - raise ValueError("Unsupported query shape") - if (prefix, size) == (0, 576): - splits, block_m, block_n, part_length = 1, 32, 64, 576 - elif (prefix, size) == (576, 50): - splits, block_m, block_n, part_length = 1, 16, 64, 640 - else: - raise ValueError("Unsupported prefix/query lengths") - - pack_query_heads = ( - (prefix, size) == (576, 50) - and x1.shape[1] == 1 - and x3.shape[1] == 1 - and tuple(x0.stride()) == (102400, 12800, 256, 1) - ) - - # Shape allocation deliberately does not preserve the input query strides. - output = torch.empty( - (batch, heads, size, width), device=x0.device, dtype=x0.dtype - ) - if (prefix, size) in ((0, 576), (576, 50)): - # Pipeline both layouts; the empty prefix performs zero iterations. - # Shared K/V lets a query tile cross head boundaries in the short case. - grid = ( - (1, triton.cdiv(heads * size, block_m), 1) - if pack_query_heads - else (batch * heads, triton.cdiv(size, block_m), 1) - ) - _segmented_attention[grid]( - x0, x1, x2, x3, x4, x5, output, - *x0.stride()[:3], - *x1.stride()[:3], - *x2.stride()[:3], - *x3.stride()[:3], - *x4.stride()[:3], - x5.stride(0), x5.stride(-1), - prefix, size, heads, x1.shape[1], width, - block_m, block_n, pack_query_heads, - num_warps=4, - num_stages=2, - ) - return output - - if splits == 1: - partial = output - logsum = None - else: - partial = torch.empty( - (batch, heads, splits, size, width), - device=x0.device, - dtype=torch.float32, - ) - logsum = torch.empty( - (batch, heads, splits, size), device=x0.device, dtype=torch.float32 - ) - - # Standard Triton launches use the caller's current CUDA stream. - _partial_attention[(batch * heads, triton.cdiv(size, block_m), splits)]( - x0, x1, x2, x3, x4, x5, partial, logsum, - *x0.stride()[:3], - *x1.stride()[:3], - *x2.stride()[:3], - *x3.stride()[:3], - *x4.stride()[:3], - x5.stride(0), x5.stride(-1), - prefix, size, heads, x1.shape[1], width, - splits, part_length, block_m, block_n, - num_warps=4, - num_stages=1, - ) - if splits > 1: - _merge_attention[(size, batch * heads)]( - partial, logsum, output, size, width, splits, - num_warps=4, - num_stages=1, - ) - return output From c1e24a3094e3df0e17ced493a6f87923854ddda7 Mon Sep 17 00:00:00 2001 From: owpeter Date: Tue, 13 Oct 2026 23:16:13 -0700 Subject: [PATCH 32/44] docs: consolidate kernel tuning workflow documentation Signed-off-by: owpeter --- CONTRIBUTING.md | 19 +- README.md | 1 + docs/proposals/0008-kernel-tuning.md | 786 +++++++++++++------ docs/proposals/0009-hotspot-kernel-tuning.md | 418 ---------- scripts/kernel_tuning/registration/README.md | 4 +- 5 files changed, 571 insertions(+), 657 deletions(-) delete mode 100644 docs/proposals/0009-hotspot-kernel-tuning.md diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index fb4c382..b9506f6 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -69,7 +69,14 @@ uv run --project scripts/kernel_tuning python -m scripts.kernel_tuning generate uv run --project scripts/kernel_tuning python -m scripts.kernel_tuning check results/kernel_tuning/tasks/gated_residual ``` -To tune one model, capture its existing preparation forward and calibrate numerical +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`. + +For full operator discovery instead, capture the existing preparation forward and calibrate numerical contracts in its model runtime, then generate and preflight every eligible call (see [model discovery](docs/proposals/0008-kernel-tuning.md#model-scoped-operator-discovery)): @@ -116,7 +123,7 @@ double precision through task generation and evaluation. The old catalog workflow remains available explicitly with `generate --catalog`, `tune-all --catalog`, and `capture --catalog`. Only that workflow accepts legacy -JSONL captures, catalog model groups, operator subsets and `--estimated`; it +JSONL captures, catalog model groups and `--estimated`; it does not certify model coverage. Existing JSONL must be recaptured for model tuning. Run these commands from the repository root. Humanize2 is pinned to a reviewed @@ -140,12 +147,14 @@ plan before applying it. When inheriting system packages, verify that the instal 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 [approved contract](docs/proposals/0008-kernel-tuning.md). -The [Chinese workflow guide](docs/proposals/0008-kernel-tuning.md#11-自动化调优使用说明) +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/`. -Saved operators require a separate review/integration before runtime use. +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 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/docs/proposals/0008-kernel-tuning.md b/docs/proposals/0008-kernel-tuning.md index f99ae03..5345b12 100644 --- a/docs/proposals/0008-kernel-tuning.md +++ b/docs/proposals/0008-kernel-tuning.md @@ -1,124 +1,283 @@ -# 0008 — Reproducible kernel tuning with Humanize2 +# 0008 — Agent-driven kernel optimization: profiling to whole-model benchmarks -- Status: Accepted +- Status: Accepted for the existing tuning contracts; end-to-end extensions are partially implemented. - Author: EmbodiInfer maintainers - Date: 2026-10-02 -Implementation: repository CLI, Humanize2 Flow, evaluator, source/evidence archive, -and isolated tool lock are present. CPU/fake-agent coverage and NVIDIA RTX 5090 -Triton/CUDA C++ builder tests are verified. Thor compilation and performance -still require validation on that target machine. - -For the Chinese workflow guide, required task inputs, and stopping conditions, -see [Section 11](#11-自动化调优使用说明). +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 -Add repository developer tooling for a serial, agent-driven kernel optimization -loop. Humanize2 (`hmz`) owns coding-agent execution, budgets, traces, and resume; -EmbodiInfer supplies the task contract, trusted evaluation, promotion rules, and -saved operator artifacts. Formal measurement belongs under `benchmarks/` and -orchestration under `scripts/`. This is model-neutral tooling, not an engine API. +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 discovery/catalog support does not imply profile support. ## 2. Motivation and current gap -`embodiinfer/backend/triton/norm.py:302` and -`embodiinfer/backend/triton/split_kv_attention.py:179` contain concrete kernels and -hand-chosen launch configurations. `embodiinfer/layers/attention.py:78` provides -runtime registration, but the repository has no general operator task format, -automated search workflow, or durable archive of verified candidates. A task can -describe one of these kernels or an independent operator without importing model -or engine internals into the tuning framework. +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. -## 3. Goals and non-goals +The earlier discovery/catalog workflow optimizes each observed ATen call or +existing backend kernel against its production implementation. It remains useful +for operator investigations, but its exports do not directly install into a +hardware profile. 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. -- Accept NVlabs KDA task packages and Triton or CUDA C++ multi-file solutions. -- Run on one target machine, with an explicitly selected coding agent and model. -- Freeze precision, workload, timing, and acceptance conditions before search. -- Save every attempt and export the best verified source with its evidence. -- Support bounded execution, interruptions, and explicit resume. +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 -Runtime registration, deployment, SSH, model-level performance claims, and -automatic installation of GPU dependencies are outside this change. Hardware -neutrality describes orchestration: a backend must support the selected hardware; -CUDA and the default FlashInfer adapter require a compatible NVIDIA environment. +- 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 discovery/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 -Use the [NVlabs wishlist task template](https://github.com/NVlabs/kda/tree/2186cd90dfb3b11f1ec250b4e3f321d72d7e7ef0/example) -and [minimal flow](https://github.com/NVlabs/kda/blob/ef6ce617693ef0782b3ecb9f37e39bbf10226a90/docs/agent-flow.md). -Tasks contain `README.md`, `definition.json`, `workloads.jsonl`, `baseline.py`, -and optionally `benchmark.py`. The definition embeds the mathematical reference -`run`; `baseline.py::run` supplies the initial performance baseline. The JSON -schemas remain FlashInfer Trace schemas. Project settings live in `tuning.yaml`. - -Implement a local Humanize2 Flow against -[`hmz` revision b70d4427](https://github.com/humanfia/humanize/tree/b70d4427f11308442aef20f9b2233dfbbc34a39e). -The Flow asks an existing coding agent to inspect the task, write a hypothesis -and executable plan, produce a candidate, then respond to evaluator feedback. -Agent turns perform source/JSON checks and return the candidate; GPU compilation, -numerical verification, and timing run only in the separate evaluator. A working -GPU in the evaluator does not imply that the agent's command environment grants -GPU or compiler-cache access. Plans describe validation criteria without trying -to execute that validation inside an agent turn. -It delegates all agent sessions and cancellation to Humanize2. Domain helpers -validate tasks, launch a bounded evaluator process, persist candidate evidence, -and atomically publish the best pointer. No independent coding-agent subprocess -implementation is introduced. - -The default evaluator reuses FlashInfer Bench 0.1.2 definitions, workload -materialization, and solution building. Additional gates enforce the configured -precision contract, workload completeness, graph replay correctness when selected, -and paired baseline/candidate measurements. A task-owned `benchmark.py` can -implement the documented structured evaluation interface for other hardware or -special input semantics. Exit status or printed PASS alone cannot promote a kernel. - -Each task fixes eager or CUDA Graph timing. Promotion requires the weighted mean -latency to fall by at least 3% and every workload to regress by at most 5%, against -the incumbent and initial baseline. Weights default to equal. Repeat paired -measurements before promotion. Stop at 20 attempts, two hours of active execution, -or five consecutive failures/non-improvements; all limits are configurable. - -Every timed sample starts from a cold L2 (a 256 MiB buffer is zeroed) with -freshly copied inputs. CUDA Graph samples queue the flush and the input copy -ahead of the start event and do not synchronize per sample, so the events bracket -only the replay's device execution. FlashInfer's `do_bench` synchronizes first; -its start event then precedes the host launch, adding several microseconds of -launch latency and scheduler jitter to every sample. On an RTX 5090 that -inflated microsecond kernels to about 10 µs and let three identical solutions -differ by up to 53%, far above the 3% promotion threshold. Without it, and with -1000 iterations (the generated-task default), identical solutions differed by at -most 4.5% and mostly under 2%. A single trivial kernel replays in about 4.1 µs -under this protocol, the device's floor; tasks whose baseline is already there -leave an agent nothing to win, so preflight prints every baseline latency. -Eager timing still uses FlashInfer's `time_runnable`, including launch cost. - -Runs snapshot task files and record source digests, evaluator identity, hardware -and software fingerprints, precision settings, agent settings, attempts, raw -measurements, and decisions. Resume requires matching task/tool/environment -identity. Interrupted attempts remain in the archive and consume an attempt. -Export reconstructs source files plus the official Solution JSON and evidence. -Task and artifact hashes detect accidental modifications; they are not an OS -security sandbox for coding agents or kernels running as the same user. - -Humanize2 needs Python 3.12+. Install it in a separate tooling environment; the -evaluator interpreter points to the prepared target-machine runtime (for example -Thor). EmbodiInfer's Python 3.10 floor and model dependency groups stay intact. - -An independent lightweight controller was considered and rejected because the -approved design explicitly reuses Humanize2. Runtime auto-selection was rejected -because source review and separate integration are required before deployment. +### 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. 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. +ATen-discovery/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 preserves executed operator coverage but cannot itself + establish a new fused region or an installable registry call site. Keep it as + a distinct workflow ([Appendix C](#model-scoped-operator-discovery)). +- 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 -This is developer automation and formal measurement, outside the inference -package. No engine, policy, or layer interface changes. Operator semantics belong -entirely to each task's definition and workloads. Backend and hardware details -belong to the evaluator and generated solution, not the orchestration loop. +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 -### Calibrated numerical contracts for discovered operators +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 Model discovery records execution settings separately from acceptance rules. Pure copies, indexing, exact elementwise operations, RNG transitions, and unknown @@ -159,8 +318,8 @@ 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), discovery recorder/task generation, and the replay evaluator. -There are no engine, policy, backend, dependency-pin or model-level action-parity -changes. CPU tests exercise selection, calibration, per-workload bounds, attention +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. @@ -178,48 +337,283 @@ 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. -## 7. Implementation plan - -Add `scripts/kernel_tuning/` for task validation, artifact storage, Humanize2 Flow, -and the `python -m scripts.kernel_tuning` developer entry point. Add -`benchmarks/kernel_tuning/` for measurement and a self-contained task example. -Document isolated setup in `CONTRIBUTING.md`; keep the optional tool's dependency -metadata separate from runtime requirements. Results default to the already -ignored `results/kernel_tuning/`. Nothing automatically registers saved kernels. +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, discovery and registry +capture, calibration, evaluator, batch tuning, GPU attribution, registry/profile +loader, install, model verification, and the CPU registration controller. The +pi0.5 call sites are integrated; the tuned profile directory currently contains +the loader but no committed hardware profile. Historical target-machine checks +are not new validation performed by this documentation change. + +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 use fake coding agents and evaluators to cover actual Flow execution, -resume, failed attempts, budgets, task tampering, metric completeness, precision -contracts, regressions, atomic best publication, and export. Static checks cover -all new Python. GPU evaluator tests are explicitly marked and skip without their -optional dependencies and hardware. CPU validation does not launch a real coding -agent or tune kernels. Real agent searches and GPU validation run explicitly on -a prepared target machine, with their own saved task and measurement conditions. +CPU tests under `tests/kernel_tuning/` cover task validation, Humanize2 fake-agent +execution, failed/interrupted attempts, budgets/resume, evidence tampering, +preflight barriers, replay layouts/mutations/RNG, calibration, hotspot attribution, +installation, model-gate arithmetic and registration repair. 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 -After deployment to a prepared GPU machine, each task records hardware, software, -dtype, shapes, seeds, warmup, iterations, trials, timing mode, and baseline source. -Paired trials evaluate the full declared workload set. Record failed and neutral -attempts as well as improvements. Kernel-level results do not establish model -latency or action parity; those require the existing model benchmark separately. +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; no new performance claim is made here. ## 10. Risks and limitations -Humanize2 and FlashInfer are evolving dependencies; use pinned versions and fail -explicitly on incompatible contracts. Arbitrary operators may need custom input -generation or timing through `benchmark.py`. Input coverage bounds correctness; -finite tests cannot prove an arbitrary generated kernel correct. Thermal state, -other GPU users, and clock changes can affect timing; paired remeasurement reduces -but does not remove this uncertainty. User-selected coding agents may execute -commands without interactive approval; run in a disposable workspace/account -appropriate to the task. Actual Thor compilation and performance remain target -machine validation work. The pinned Humanize2 journal uses POSIX primitives: real -searches run on the Linux target machine; local Windows checks and export are -supported. CPU tests on Windows adapt only upstream atomic journal file writing. +- 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. + +## 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 +``` + +### 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. If a preflight-only batch will later resume search, +specify its agent 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/`. -### Task-owned benchmark interface +## 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 @@ -288,7 +682,7 @@ incomplete evidence. A custom adapter remains responsible for truthful measureme and numerical checks: the controller cannot infer GPU correctness from a boolean. The default adapter supplies these checks for deterministic pure tensor operators. -### Archive and reproduction +## 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 @@ -323,55 +717,7 @@ 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. -### Generated catalog tasks and batch tuning - -Hand-writing a task package per operator does not scale to every core kernel. -`scripts/kernel_tuning/operators.py` declares each core Triton operator once: -its eager Torch reference, the production call used as the baseline, the -model-derived workload shapes, and the numerical contract. `generate --catalog` renders a -complete task package from an entry. The baseline is the production kernel -itself: the defining repository module and the repository modules it imports -are copied byte-for-byte under `vendor/` (absolute `embodiinfer.` imports are -rewritten to relative ones; package `__init__` files are never copied), and -`baseline.py` calls it as the engine does. `generated.json` records the -repository revision and source digests. Structured integer inputs such as -segment offsets are written as safetensors data inside the task. - -The catalog covers the Triton paths of pi0.5 (`ada_rms_norm`, `gated_residual`, -`gated_gelu`, `rotate_qk`, `split_kv_attention`), the Qwen2.5-VL vision tower of -the Qwen R2R low-level and panoramic policies (`rotate_half_rope`, -`segmented_attention`), and StreamVLN's Qwen2 decode path (`rms_norm`, -`add_rms_norm`, `swiglu`). ActiveVLN runs the Transformers vision tower and an -eager attention backend, so it currently calls none of these kernels. Workloads -come from committed captures (see below); the catalog's estimates are only the -fallback (`--estimated`). All generated tasks time with CUDA Graphs. - -Contracts follow the production kernel, not an aspiration. `gated_residual` -and `rotate_qk` are bit-exact against Torch. The other kernels already round -differently from eager Torch, so their contract is a tolerance with a stated -reason: `atol = 2**-10` and `rtol = k * 2**-8`, where `k` counts BF16 roundings -the production kernel and the eager reference do not share (2 when only the -final casts differ; 3 for `rms_norm` and `swiglu`; 4 for `gated_gelu` and -`add_rms_norm`). The `gpu`-marked -`test_production_baseline_meets_generated_contract` checks every production -baseline against its contract, including changed-input CUDA Graph replay, and -prints the fraction of the bound it uses. Every baseline passed on an RTX 4060 -Laptop GPU (Torch 2.6, Triton 3.2; at most 0.85 of its bound) and on an RTX -5090 (Torch 2.12, Triton 3.7; at most 0.90), and the evaluator preflight passed -for all ten tasks on both. Re-run both on each target before trusting a -contract there. - -`tune-all --catalog` generates the selected operators into a new batch directory under -`results/kernel_tuning/batches/`, measures every production baseline with the -evaluator (preflight), and then runs each operator serially in its own -`run`/`resume` subprocess. A failed operator does not stop the batch; an -interrupt does, and `tune-all --resume BATCH` continues it, resuming existing -run archives. Promoted kernels are exported to `exports/`, and -`summary.md`/`summary.json` report attempts, promotions, and the improvement -over the production baseline from the weakest paired round. Machine-specific -settings (`evaluator_python`, `device`, budgets) are passed with `--set` at -generation time, and `--hardware-notes` copies target-hardware notes into each -task as `HARDWARE.md` for the agent. +## Appendix C. Other task sources ### Model-scoped operator discovery @@ -543,55 +889,31 @@ Capture is a single-process preparation facility; it does not certify unvisited branches, other worker processes, or checkpoint-level action parity. Existing fused kernels count as one call; no new fused operators are proposed or generated. -### Captured catalog workloads - -Catalog shapes are estimates. `capture --catalog` replaces them with production traffic: -it runs an unmodified model script in the model's runtime with every catalog -kernel wrapped, converts each call to the operator's task axes, scalars, and -fixed integer inputs, and writes one JSONL row per distinct case with its call -count. Calls outside a task's semantics (another dtype, broadcast tables, -custom attention scaling) are counted as skipped with the reason. Kernels -launched while a CUDA Graph is captured are attached to that graph and counted on -every replay, so production graph configurations can be captured; compilation -should be disabled. Captures committed under `benchmarks/kernel_tuning/captures/` -replace the catalog estimates by default, using each operator's most frequent -cases weighted by call count; `--captured FILE` selects another capture and -`--estimated` restores the estimates. - -Shapes depend on the configuration and the inputs, never on weight values. -`skeleton REPO --revision REV` therefore builds a random-weight copy of a -Hugging Face checkpoint: it downloads the small files and synthesizes each -safetensors file from its header, read with an HTTP range request. Benchmark -scripts then run unchanged on it, with synthetic frames where the policy resizes -images to a fixed size anyway. Random weights do not reproduce generated text, so -decode lengths are fixed in the capture configuration. - -`benchmarks/kernel_tuning/captures/pi05-libero10.jsonl` comes from -`benchmarks/pi05-benchmark` on 200 LIBERO-10 frames (10 tasks), with the -production Triton paths enabled (`native_inference`, Triton prefix and denoise -attention). It corrected three estimates. The prefix holds two cameras plus a -48- or 64-token language bucket (560 or 576 tokens, not 968). `split_kv_attention` -also runs once per layer over the whole prefix with no cached keys (`P=0`, -`S=576`). The checkpoint's AdaRMS modulation is BF16, not FP32; the catalog now -declares BF16 (recalibrated: at most 0.79 of its bound on an RTX 5090). At B=1, -four of the five pi0.5 kernels measure about 4 µs on an RTX 5090, the device's -launch floor; only `split_kv_attention` (16-35 µs) leaves room to tune. - -The Qwen R2R captures (`qwen-r2r-low.jsonl`, `qwen-r2r-panoramic.jsonl`) use -skeleton checkpoints and synthetic R2R/RxR episodes (the trajectory release is -gated; episode lengths follow the benchmark READMEs) with `attention_backend: -triton` and CUDA Graphs on: only the graph path runs EmbodiInfer's Triton vision -forward. Each call encodes 4 history frames plus the current one: 1,980 patches -in 46 ragged windows (low-level) or 7,704 patches in 172 windows (panoramic), -not the estimated uniform 1,024/4,096. The rotary tables are BF16, which makes -the production RoPE bit-exact; the catalog now declares both. These shapes put -`segmented_attention` at 357 µs and 1,405 µs and `rotate_half_rope` at 18 µs and -61 µs on an RTX 5090. The StreamVLN capture (`streamvln-r2r-rxr.jsonl`, skeleton -weights, decode capped at the five tokens of its fast action path) sees only -single-row decode calls: `swiglu` never runs on prefill chunks, and all three -StreamVLN kernels sit at the launch floor. +### 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 discovery contracts. ```bash -PYTHONPATH=/path/to/EmbodiInfer /path/to/model-runtime/python -m scripts.kernel_tuning capture --catalog --output pi05.jsonl -- benchmark.py --config capture-config.yaml -python -m scripts.kernel_tuning tune-all --catalog --model pi05 --captured pi05.jsonl --agent 'HARNESS/MODEL:EFFORT' +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/docs/proposals/0009-hotspot-kernel-tuning.md b/docs/proposals/0009-hotspot-kernel-tuning.md deleted file mode 100644 index 214b42c..0000000 --- a/docs/proposals/0009-hotspot-kernel-tuning.md +++ /dev/null @@ -1,418 +0,0 @@ -# 0009 — Hotspot-driven Triton tuning with installable hardware profiles - -- Status: Draft -- Author: EmbodiInfer maintainers -- Date: 2026-10-09 - -## 1. Summary - -Close the loop from kernel search to model latency. Profile a model in its -production configuration, rank code regions by the time a better kernel could -recover, let a coding agent choose a few regions (including new fusion -boundaries) and write Triton implementations that must beat the model's Torch -implementation, then install the verified kernels in the repository under a -hardware profile (for example `rtx4090`). At runtime an operator registry -selects them behind an off-by-default switch and falls back to the existing -implementation whenever a kernel does not apply. Success is measured as model -end-to-end latency under a frozen model-level numerical gate, not as isolated -kernel speedups. The registry is a model-agnostic `layers/` contract; region -definitions and call sites stay in policies; tuned kernels are `backend/` code. -New work builds on this proposal instead of the per-operator workflow of -[0008](0008-kernel-tuning.md); its Humanize2 flow, evaluator timing protocol and -run archive are reused. - -## 2. Motivation and current gap - -- Tuned kernels never reach the model. 0008 exports verified sources, but - "saved operators require a separate review/integration before runtime use" - and there is no place to plug them in. -- There is no general dispatch point. Attention - (`embodiinfer/layers/attention.py:74`) and quantized linear - (`embodiinfer/layers/linear.py:253`) have registries; norm, activation and - RoPE are chosen by per-model branches that import `embodiinfer.backend.triton` - directly (`embodiinfer/policies/pi05/modeling_pi05.py:122`, `:161`, `:314`, - `:669`; `embodiinfer/models/streamvln/backbone.py:12`; - `embodiinfer/models/qwen25_vl/vision.py:47`). -- Isolated operators are the wrong unit. At B=1, four of five pi0.5 Triton - kernels and all StreamVLN kernels already run near the ~4 µs device floor - (0008, "Captured catalog workloads"). Remaining time sits in a few large - kernels and in chains of small memory-bound kernels between GEMMs; only - changing fusion boundaries reduces kernel count and intermediate traffic. -- Tuning everything is wasteful. Model discovery produced 94 pi0.5 tasks, most - without measurable model-level effect. -- The comparison target is wrong. 0008 compares a candidate with the same ATen - call or the existing Triton kernel. The goal is now for Triton to beat the - model's Torch implementation, which in production may be compiled by Inductor - (`_inference_helper`, `embodiinfer/policies/pi05/modeling_pi05.py:65`). - -## 3. Goals and non-goals - -Goals: - -- One resumable pipeline per model and hardware: profile → select → tune → - install → verify → benchmark. Only a new fusion region needs a code change, - made once and merged like any other change. -- A hotspot report that attributes device time to model regions and estimates - recoverable time; tuning only the selected regions. -- Regions may span several existing operators (new fusion boundaries). -- Baseline: the region's Torch implementation under the model's production - settings (dtype, Inductor when configured, CUDA Graphs). Existing hand-written - Triton kernels are optional starting points, not the baseline. -- Tuned kernels committed in the repository per hardware profile with - provenance, selectable at runtime, reusable by others, and later promotable to - the default Triton implementation through a normal review. - -Non-goals: - -- Making any tuned kernel the default (a separate change with evidence). -- Lower precision, quantization or changed model mathematics. A region may - reorder floating-point work, but not change dtype or matmul precision. -- Automatic merging of agent-written policy changes; they are reviewed as PRs. -- Multi-GPU and batch > 1 tuning. -- Compatibility with the 0008 discovery/catalog task formats. Those commands - stay as they are and are not extended. - -## 4. Design - -### 4.1 Operator registry (`embodiinfer/layers/ops.py`) - -An operator is named by its mathematics (`ada_rms_norm`, `gated_residual`, -`add_rms_norm`, `rotate_qk`), never by a model. A policy-local fused region is -registered by its policy under a namespaced name (`pi05.residual_ada_rms_norm`). - -Every existing Torch/Triton switch point is migrated to the registry once, as an -ordinary code change: the registry call becomes the model's default code path, -with today's implementations registered as `torch` and `triton`. Tuning runs -never edit these call sites; they only add implementations to a profile. - -```python -register_op(op, impl, fn) # "torch" reference, "triton", ... -register_lazy_op(op, impl, module, function) # import on first use -selection = KernelSelection(kernel_profile) # one per policy -selection.prepare() # resolve everything eagerly -selection.resolve(op, prefer) -> Callable # cached lookup -``` - -- Torch references live in `embodiinfer/backend/torch/ops.py`; the existing - Triton kernels are registered unchanged. -- A policy builds its selection on first use on its device and drops it on - refit, device or dtype change. `prepare()` resolves every operator before any - `torch.compile` region, so graph tracing only reads the cache. -- Torch references run uncompiled. pi0.5's per-helper Inductor wrappers were - first created while Dynamo traced the prefix encoder, which stores the raw - function, so they never ran compiled outside that graph; keeping the - references uncompiled preserves those numerics. Compiling them is a separate, - measured change. -- `prefer` keeps the policy's current choice (`torch` or `triton`, from its - existing flags). A selected profile overrides only operators it contains. -- A tuned callable is wrapped with its `applies(*args)` predicate (dtype, - contiguity, shape divisibility, size limits). When it returns false the call - goes to the preferred implementation. StreamVLN's `supports_*` checks - (`embodiinfer/models/streamvln/backbone.py:125`) follow the same pattern. -- With no profile selected, resolution returns the current callables: behavior - is bit-exact with the code before this change. - -Runtime switch: `kernel_profile: none | auto | | torch` in the policy -configuration (default `none`). `auto` picks the profile whose device name, -compute capability and Triton/Torch versions match the current process, else -`none` with one log line. An explicit name that does not match raises. `torch` -forces every registered operator to its Torch reference and is the baseline -configuration for benchmarks. - -### 4.2 Hardware profiles (`embodiinfer/backend/triton/tuned//`) - -Each profile directory holds one module per tuned operator and `profile.json`: - -- match keys: GPU name, compute capability, Torch/Triton/CUDA versions tested; -- per operator: module/function, `applies` predicate, covered workloads and - weights, source digest, kernel-level and model-level gate results, the run id - and the digest of its evidence archive; -- the model, checkpoint revision and configuration the profile was verified on. - -Sources and this small manifest are tracked in Git. Evidence archives stay -outside the tracked tree, as with existing benchmark snapshots. Loading verifies -source digests, so a hand-edited kernel is reported instead of silently used. -Tuned modules are pure tensor functions and import nothing from `models`, -`policies` or `engine`. - -### 4.3 Hotspot analysis (`scripts.kernel_tuning profile`) - -Runs the model's inference driver in its production configuration on recorded -inputs, in two passes: - -1. Timing: production settings, CUDA events, end-to-end and prefill/decode - segments as the model benchmarks measure them. -2. Attribution: `torch.profiler` (CUPTI) with a `record_function` range around - every registry call and every `nn.Module.forward`. Ranges are normalized - across repeated layers (`layers.*`). When per-node kernels inside CUDA Graphs - are not reported, the pass reruns without graphs and records that it did. - -Per region the report gives device time, share of end-to-end time, calls, kernel -count, bytes moved and FLOPs from tensor metadata, and a lower bound -`max(bytes / measured bandwidth, FLOPs / measured peak)` from a device -microbenchmark. Recoverable time is `time - max(lower bound, floor × kernels)`. -Fusion candidates are adjacent producer→consumer chains of memory-bound regions; -their bound counts the eliminated intermediate traffic once. Outputs: -`hotspots.json` and `hotspots.md`. - -### 4.4 Selection and new fusion regions - -A Humanize2 selection flow gives the agent the hotspot report and the model -source. It returns `plan.json`: at most K regions (default 5), each either an -already registered operator or a proposed new fusion region, with reason and -estimated saving. Regions under 1% of end-to-end time are rejected. A person may -edit the plan or supply one with `--plan`. - -A registered operator goes straight to tuning. A new fusion region does not -exist in the model yet, so it needs a one-time code change: the agent writes the -region's Torch reference (a composition of the existing operations), registers -it, and replaces the covered call sites with one registry call. The change is -accepted only if the model with `kernel_profile: none` is bit-exact against the -unchanged model on the evaluation inputs (§6). It then lands in the repository -as a normal reviewed change; from that point the region is an ordinary -registered operator for every hardware profile and later runs reuse it. - -Workloads for every selected operator are captured by recording the registry -operator's real arguments and call counts during a production run. - -### 4.5 Region tuning - -A region task contains the Torch reference, the baseline (that reference under -the policy's compile settings, for example Inductor with the options used by -`_inference_helper`), an optional starting kernel, the captured workloads with -call-count weights, and frozen numerical bounds (§6). The 0008 Humanize2 flow, -CUDA Graph cold-L2 timing, paired promotion rule (≥ 3% weighted, ≤ 5% per -workload) and archive are reused. A candidate must also export `applies`, which -the evaluator checks is true on every workload. Expected model saving is -`Σ count × Δt`. - -Each planning turn receives the recent candidates' verified source alongside -their measured failures, including rejected candidates that never became the -incumbent. This lets a fresh agent inspect failed configurations before choosing -a new change. Optional hardware notes are part of the frozen task context. - -### 4.6 Install, verify, benchmark - -`install` writes promoted sources and updates `profile.json`. It refuses an -operator that is not registered by the current source, so a fusion region's code -change must be merged first. The installed kernel applies only to the argument -layouts the evaluator measured (shape, dtype and strides; strides of size-one -dimensions and empty tensors are ignored) and otherwise falls back. Capture runs -eagerly, so a policy must give registry arguments the same layout with and -without `torch.compile`: on Jetson Thor the pi0.5 prefix K/V cache was a view of -the fused projection in eager mode but contiguous under Inductor, and a tuned -split-KV kernel was never selected until the cache was made contiguous in both. -`verify` runs the -model with `kernel_profile: ` against `kernel_profile: torch` on the -evaluation set and applies the model-level gate; an operator is kept only if the -whole profile passes. `benchmark` reports Torch, existing Triton and the profile -on the model's benchmark protocol, plus a per-operator ablation. The pipeline -command `optimize --model M --config C --profile P --agent A` runs all stages -with resume under `results/kernel_tuning/optimize//`. - -### 4.7 Alternatives considered - -- Add selection to the 0008 per-ATen discovery: it cannot express fusion and its - baseline is the same ATen call, not the Torch implementation. Rejected. -- Rely on Inductor (`max-autotune`) for fusion: that is the baseline being beaten. - It is reported as an extra benchmark row where it compiles. -- Keep tuned kernels outside the repository and download them: rejected, because - the team wants reviewable, shareable sources that can become defaults. -- Uncached per-call resolution: rejected for launch-path overhead; resolution - is a dictionary lookup after `prepare()` and is skipped entirely by CUDA Graph - replay. -- Apply agent policy patches without review: rejected. - -### 4.8 Agent-assisted operator registration - -`python -m scripts.kernel_tuning register` implements the registration step -between manual operator selection and tuning. Its developer-only controller is -in `scripts/kernel_tuning/registration/`; it changes no inference API or default -runtime behavior. The caller supplies one mathematical operator and exact source -expressions. The initial supported routing convention is pi0.5's existing `_op` -helper and explicit `KernelSelection` parameters. Other policies must first gain -an equivalent integration; this tool does not create model-specific engine logic. - -The [registration contract](../../scripts/kernel_tuning/registration/README.md) -defines the YAML specification and `new`, `run`, `resume`, `status`, and `finalize` -commands. `new` requires a clean checkout and probes the public registry in a -separately selected `test_python`. `--repository` may select a clean model checkout -separate from an uncommitted controller checkout. It records BASE, creates a detached baseline -and a `registry/-` worktree, and stops after AST call-site review. -Explicit `run` confirms the reviewed sites and freezes the execution settings. -Humanize2 plans and implements in the branch worktree, keeping a session for -repair turns and opening a fresh planning session after two equal check failures. -An interrupted run retains its completed stages and elapsed budget. - -The controller applies five fail-fast CPU gates: file scope (including untracked -and staged changes); annotated Torch reference, standalone lazy registration, -preserved branches, parameter forwarding and public registry resolution; Ruff; -declared existing tests and new test nodes; and baseline effectiveness. Every new -test node must execute on BASE and fail by assertion, with all parametrizations -covered. Collection errors, missing dependencies and skips are not proof. The -baseline overlay is journaled before copying and restored after failures or -interruptions. Check evidence is bound to source bytes and the registration -manifest, so changing a patch invalidates completed gates. -Baseline integrity hashes tracked symbolic links by their link text without -following their targets; selected source and test files must still be regular files. - -This is lossless routing, not kernel optimization: the reference must preserve -ATen calls, arithmetic order, dtype conversions and shapes. Generated tests must -compare original and reference expressions with `torch.equal` on CPU FP32 and -BF16, exercise a recording profile, and verify preservation branches. The §6 -whole-model target-GPU bit-exact gate remains mandatory before merge. Static -checks and finite CPU tests cannot prove arbitrary numerical equivalence or -that agent-authored tests cover every behavior. - -Passing runs export `patch.diff` and `SUMMARY.md`. Finalization prints the branch -by default; commits and pushes each need their explicit flags. No merge is -provided. Operator selection, fusion, GPU testing and tuning remain separate. -The alternative of adding registration to the GPU tuning flow was rejected: -registration needs baseline code checkouts and routing tests, whereas tuning -measures candidate kernels against a frozen numerical/performance contract. - -Validation uses disposable clones of existing commit objects, structural -counterexamples, real CPU subprocesses, and official Humanize2 fake agents for -repair/session restart/resume. It requires no model checkpoint and makes no -performance claim. Coding agents and tests run with the user's process -permissions; directory bounds, integrity checks and hidden accelerator devices -are not an OS security sandbox. Actual Humanize2 execution requires POSIX and -Python 3.12+, while runtime imports use the separately prepared test interpreter. - -## 5. Model-agnosticism verdict - -The registry, profile loader and `kernel_profile` switch are model-agnostic -`layers/` code: operators are named by mathematics and resolution uses only the -current device and versions. Tuned kernels are backend code. Region definitions, -call sites and policy-namespaced fused operators are policy-layer and optional; -with `kernel_profile: none` no path changes. The profiler uses module hooks and -registry ranges only; model-specific inputs come from the existing drivers. - -## 6. Losslessness and precision criterion - -Triton kernels reorder accumulation, so bit-exact output is not generally -achievable. This proposal requests explicit approval of calibrated tolerances, -following the frozen thresholds already used for the pi0.5 native path -(`benchmarks/pi05-benchmark/README.md`). - -Exact gates: - -- `kernel_profile: none` is bit-exact with the code before this change. -- Integrating a region with its Torch reference as default is bit-exact against - the unpatched model (`max_i ||Δa_i||_inf = 0`). -- Precision conditions are identical: dtype, `float32_matmul_precision`, TF32 - flags and attention semantics. - -Region level (during tuning): with FP64 reference `R` of the Torch reference, -Torch baseline output `B` and candidate `C` in the production dtype with epsilon -`u`, each floating output must satisfy `|C - R| <= atol + u·|R|` with -`atol = max(u·RMS(R), 2·max(|B - R| - u·|R|, 0), tiny)`, the 0008 formula with -the Torch implementation as `B`. Inputs are the captured fixtures plus seeded -random, zero and alternating-sign cancellation inputs. Integer, boolean, layout, -aliasing and mutation results are exact; non-finite values fail. In words: a -candidate may be at most about twice as wrong as Torch. - -Model level (required for install): - -- Continuous outputs (action chunks): compute an FP32 Torch reference from the - same weights and inputs, and measure each run's max-abs and RMSE error per - observation. Over the whole evaluation set, the candidate's mean max-abs and - mean RMSE may exceed those of the `kernel_profile: torch` run by at most 25%, - and its worst max-abs may be at most twice Torch's worst (floors `1e-4` and - `1e-5`). Thresholds are frozen before any candidate runs. A per-observation - bound of twice Torch's error was tried first and rejected: on Jetson Thor, - 40 pi0.5 observations, the existing Triton path failed it on 2 observations - where Torch happened to land within 1e-3 of the reference, while its mean - errors equalled Torch's and two Torch runs themselves differed per observation. -- Discrete outputs (generated tokens or discrete actions): sequences must equal - the `kernel_profile: torch` run on every evaluation input. A difference fails; - there is no automatic waiver. -- Evaluation set: the model benchmark's frozen cross-task references (pi0.5: - the 10 references in `compare.py`) plus 200 benchmark frames. - -Inductor's runtime kernel-configuration choice makes a fresh compilation of -the same Torch path produce different outputs: on Jetson Thor, two pi0.5 -processes without a shared compile cache differed by up to 4.6e-2 in normalized -actions. Parity checks therefore run with `TORCHINDUCTOR_DETERMINISTIC=1`, which -made fresh compilations bit-exact, and model-level thresholds are derived from -the same compiled baseline that is timed. - -The thresholds, reference identity and evaluation-set digest are saved in the -optimize run and summarized in `profile.json`; `scripts.kernel_tuning verify` -reproduces them. - -## 7. Implementation plan - -1. Registry and profiles: `embodiinfer/layers/ops.py`, profile loader, - `kernel_profile` switch. One-time migration of the existing switch points to - the registry as the default code path: pi0.5's `gated_residual`, - `ada_rms_norm`, `gated_gelu` and `rotate_qk` (`split_kv_attention` already - uses the attention registry and gains profile lookup), then the Qwen R2R and - StreamVLN norm, activation and RoPE call sites. Install one kernel by hand and - run `verify`/`benchmark` on an RTX 4090 to prove the loop. -2. `profile` command and hotspot report. -3. Tasks for registered operators with Torch baselines, tuning, automated - `install`. -4. Agent selection and new fusion regions (one-time code changes, §4.4). -5. Further hardware profiles. - -Status (2026-10-12): steps 1–3 are implemented for pi0.5; the loop was -exercised end to end on Jetson Thor instead of an RTX 4090. No hardware profile -is committed yet. Qwen R2R and StreamVLN call sites, step 4 and committed -profiles are not done; neither is shape-range tuning (§10). The manually selected, -single-operator registration portion of step 4 now has the CPU controller in -§4.8; automatic selection and new fusion remain future work. - -New tooling lives in `scripts/kernel_tuning/` and `benchmarks/kernel_tuning/`. -Existing 0008 commands keep working and are not extended. No public symbol is -removed. - -## 8. Test plan - -CPU: registry resolution and preference; `applies` fallback; profile matching -and mismatch handling for `auto` and explicit names; digest tamper detection; -install round trip; hotspot aggregation and fusion-chain detection on synthetic -profiler traces; plan validation; region and model gate arithmetic, including -frozen thresholds and discrete-output equality. - -GPU (marked): pi0.5 with `kernel_profile: none` bit-exact against the previous -code; tuned and fallback paths inside CUDA Graph capture; the integration -bit-exact gate; `verify` on real weights. - -## 9. Benchmark plan - -pi0.5 on one RTX 4090, B=1, LIBERO-10 1,600 frames, BF16, Inductor and full -denoise-loop CUDA Graph, the protocol in `benchmarks/pi05-benchmark/README.md`. -Rows: `kernel_profile: torch`, current Triton, the tuned profile, and Inductor -`max-autotune` where it compiles; end-to-end mean/P50/P95 and prefill/decode -split; per-operator ablation. Expected gain is bounded by the recoverable time -in the hotspot report; it disappears where regions are already at the bandwidth -or launch floor. Batch > 1 and other hardware are reported only when measured. - -With `kernel_profile: none`, 200 benchmark frames were bit-identical to the -previous code on both the native and the Torch path -(`TORCHINDUCTOR_DETERMINISTIC=1`). - -## 10. Risks and limitations - -- Agent kernels are often shape-specialized (a 2026-10-07 `split_kv_attention` - candidate asserted `P % PART_LENGTH == 0`); `applies` and fallback mitigate, - but uncovered shapes run the old path. A kernel tuned on captured B=1 shapes - with a 576-token prefix covers only those; the 48-token text bucket, B > 1 and - other checkpoints fall back. Follow-up: tasks over declared shape ranges - (buckets, batch sizes, several layouts per shape), kernels with runtime strides - and scales, and range-based `applies`. -- Region selection is manual: the hotspot report ranks time, but a person picks - the operators. Follow-up: rank by `Σ count × (t − roofline floor)`, detect - chains of memory-bound kernels between GEMMs as fusion candidates, and let an - agent turn a candidate into a registered operator by PR, with a bit-exact - `kernel_profile: none` parity test as the proof. -- GEMMs (the largest share of pi0.5 time on Thor) are not registry operators. -- Profiles are tied to Triton/Torch versions; a version upgrade disables them - under `auto` until re-verified. -- Attribution under CUDA Graphs depends on CUPTI support; the eager fallback pass - can shift shares slightly. -- Finite evaluation sets do not prove correctness for every input or task - success rate. -- Generated code adds review load; each install goes through a PR. -- Coding agents run commands without interactive approval, as in 0008. diff --git a/scripts/kernel_tuning/registration/README.md b/scripts/kernel_tuning/registration/README.md index 5ef8b8b..cffdda4 100644 --- a/scripts/kernel_tuning/registration/README.md +++ b/scripts/kernel_tuning/registration/README.md @@ -2,7 +2,7 @@ `python -m scripts.kernel_tuning register …` 把一个尚未注册的数学算子接入算子注册表 （`embodiinfer/layers/ops.py`），并把选定的调用点改为经过注册表调用。设计归属见 -[0009 §4.8](../../../docs/proposals/0009-hotspot-kernel-tuning.md#48-agent-assisted-operator-registration)。 +[0008 §4.2](../../../docs/proposals/0008-kernel-tuning.md#42-select-and-if-necessary-register-an-operator)。 工具由一个控制器驱动编码智能体，循环执行"计划 → 实现 → CPU 检查 → 反馈修复"，直到全部检查通过或预算用尽。 最终产物是一个待人工审阅的分支和 `patch.diff`。 @@ -282,7 +282,7 @@ $TOOL register new --model pi05 --op linear \ #### 合并前 CPU 检查不能证明 GPU 上逐位一致。例如 pi0.5 的 native 路径只在 CUDA 上启用，CPU 测试执行不到其中的调用点。 -合并前必须在目标 GPU 上按 [0009 §6](../../../docs/proposals/0009-hotspot-kernel-tuning.md#6-losslessness-and-precision-criterion) +合并前必须在目标 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` 的末尾也会写明这一项尚未完成。 From dbafdcae7e63113a1e0ffa51c3d9d256f39e7bf5 Mon Sep 17 00:00:00 2001 From: owpeter Date: Wed, 14 Oct 2026 05:50:22 -0700 Subject: [PATCH 33/44] feat: allow running agents without the Humanize2 Landlock fence Humanize2 fences each agent session with Linux Landlock, so tuning and registration cannot open a session on kernels built without it (for example Jetson L4T). Add --unfenced to run, tune-all and register run. Roles are declared with every scope; runs without the flag narrow the agent back to Humanize2's default permission, so fenced behavior is unchanged. The choice is frozen in the run or batch manifest, reused on resume and passed to every batch run, and a warning is printed before agents start. Signed-off-by: owpeter --- CONTRIBUTING.md | 7 ++++ docs/proposals/0008-kernel-tuning.md | 9 +++-- scripts/kernel_tuning/__main__.py | 21 +++++++++-- scripts/kernel_tuning/artifacts.py | 24 +++++++++++-- scripts/kernel_tuning/batch.py | 14 +++++++- scripts/kernel_tuning/flow.py | 9 ++--- scripts/kernel_tuning/permissions.py | 36 +++++++++++++++++++ scripts/kernel_tuning/registration/README.md | 8 +++-- scripts/kernel_tuning/registration/cli.py | 13 +++++-- scripts/kernel_tuning/registration/flow.py | 9 ++--- .../kernel_tuning/registration/workspace.py | 17 ++++++--- tests/kernel_tuning/test_batch.py | 25 +++++++++++++ tests/kernel_tuning/test_flow.py | 14 ++++++++ tests/kernel_tuning/test_registration.py | 28 +++++++++++++++ 14 files changed, 209 insertions(+), 25 deletions(-) create mode 100644 scripts/kernel_tuning/permissions.py diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index b9506f6..177e628 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -166,6 +166,13 @@ 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: diff --git a/docs/proposals/0008-kernel-tuning.md b/docs/proposals/0008-kernel-tuning.md index 5345b12..ab55af8 100644 --- a/docs/proposals/0008-kernel-tuning.md +++ b/docs/proposals/0008-kernel-tuning.md @@ -446,7 +446,10 @@ is not evidence for another; no new performance claim is made here. 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. + 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 @@ -551,7 +554,9 @@ uv run --project scripts/kernel_tuning python -m scripts.kernel_tuning tune-all ``` Use `--dry-run` to generate tasks only, or `--preflight-only` to measure baselines -without starting agents. If a preflight-only batch will later resume search, +without starting agents. On a host whose kernel lacks Landlock, add `--unfenced` +when creating the batch and run it under an isolated account +([CONTRIBUTING](../../CONTRIBUTING.md#optional-kernel-tuning-tools)). If a preflight-only batch will later resume search, specify its agent 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. diff --git a/scripts/kernel_tuning/__main__.py b/scripts/kernel_tuning/__main__.py index f28de4e..2d201f0 100644 --- a/scripts/kernel_tuning/__main__.py +++ b/scripts/kernel_tuning/__main__.py @@ -10,7 +10,15 @@ import sys from pathlib import Path -from .artifacts import REPOSITORY, RunStore, atomic_json, run_lock, runtime_identity +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 @@ -39,6 +47,7 @@ def _execute(store: RunStore, *, resume: bool) -> dict: 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"), @@ -178,11 +187,14 @@ def _tune_all(args: argparse.Namespace) -> int: 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, and --agent") + raise ContractError( + "--resume continues the recorded batch; omit selection, --set, --agent, and --unfenced" + ) root = args.resume.resolve(strict=True) else: agent = args.agent or "" @@ -192,6 +204,7 @@ def _tune_all(args: argparse.Namespace) -> int: "agent": agent, "overrides": parse_overrides(args.set), "hardware_notes": args.hardware_notes, + "unfenced": args.unfenced, } if args.catalog: operators, notes = apply_file( @@ -247,6 +260,7 @@ def main(argv: list[str] | None = None) -> int: 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") @@ -264,6 +278,7 @@ def main(argv: list[str] | None = None) -> int: 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") @@ -456,7 +471,7 @@ def main(argv: list[str] | None = None) -> int: from .batch import verify_model_task model_batch = verify_model_task(task) - store = RunStore.create(task, args.output, agent) + 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] diff --git a/scripts/kernel_tuning/artifacts.py b/scripts/kernel_tuning/artifacts.py index d54d23c..e90241e 100644 --- a/scripts/kernel_tuning/artifacts.py +++ b/scripts/kernel_tuning/artifacts.py @@ -19,6 +19,21 @@ 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: @@ -127,8 +142,12 @@ def __init__(self, root: Path) -> None: raise ContractError("Task snapshot no longer matches its manifest") @classmethod - def create(cls, task: TaskPackage, output: Path, agent: str) -> RunStore: - """Snapshot a task and baseline under a fresh, collision-resistant run ID.""" + 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: @@ -164,6 +183,7 @@ def create(cls, task: TaskPackage, output: Path, agent: str) -> RunStore: "tool_digest": digest(source_hashes), "tool_runtime": runtime_identity(), "agent": agent, + "unfenced": unfenced, "baseline_digest": digest(baseline), "environment": None, "elapsed_seconds": 0.0, diff --git a/scripts/kernel_tuning/batch.py b/scripts/kernel_tuning/batch.py index 7037df8..f6881f5 100644 --- a/scripts/kernel_tuning/batch.py +++ b/scripts/kernel_tuning/batch.py @@ -49,6 +49,7 @@ def create( 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: @@ -74,6 +75,7 @@ def create( "schema_version": 1, "created_at": datetime.now(timezone.utc).isoformat(), "agent": agent, + "unfenced": unfenced, "overrides": overrides, "items": items, }, @@ -147,7 +149,13 @@ def preflight(root: Path, names: list[str] | None = None) -> dict[str, str | Non def create_model( - capture: ModelCapture, output: Path, *, agent: str, overrides: dict[str, Any], hardware_notes: Path | None + 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 @@ -186,6 +194,7 @@ def create_model( "schema_version": 2, "created_at": datetime.now(timezone.utc).isoformat(), "agent": agent, + "unfenced": unfenced, "overrides": overrides, "model": capture.manifest["target"], "capture_identity": capture.identity, @@ -299,6 +308,8 @@ def execute(root: Path, *, launcher: Launcher = launch) -> dict[str, Any]: "--output", str(root / "runs"), ] + if batch.get("unfenced"): + command.append("--unfenced") item["status"] = "running" atomic_json(root / "batch.json", batch) try: @@ -367,6 +378,7 @@ def write_summary(root: Path) -> list[dict[str, Any]]: 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 |", "|---|---|---|---|---|---|---|---|", diff --git a/scripts/kernel_tuning/flow.py b/scripts/kernel_tuning/flow.py index 1865eed..a2eb9d1 100644 --- a/scripts/kernel_tuning/flow.py +++ b/scripts/kernel_tuning/flow.py @@ -15,7 +15,6 @@ from typing import Any from hmz.flows import ( - Agent, AgentCollection, Budget, BudgetExceeded, @@ -35,6 +34,7 @@ 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 @@ -42,7 +42,7 @@ class Agents(AgentCollection): """An existing coding agent, selected by an explicit hmz agent specification.""" - coder: Agent + coder: Coder class Envs(EnvCollection): @@ -74,6 +74,7 @@ async def kernel_tuning( 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() @@ -135,7 +136,7 @@ async def measure( session = None try: work_env = await envs["workspace"].derive_subdir(subdir=candidate) - session = await agents["coder"].spawn(env=work_env) + session = await coder.spawn(env=work_env) for stage, prompt in ( ("plan", plan_prompt(store.task)), ("implement", implementation_prompt(store.task)), @@ -146,7 +147,7 @@ async def measure( seconds = min(cfg.search.agent_timeout_seconds, remaining()) (archive / f"{stage}-prompt.txt").write_text(prompt, encoding="utf-8") reply = await asyncio.wait_for( - agents["coder"].run( + coder.run( prompt, session=session, budget=Budget(duration=timedelta(seconds=seconds), graceful=False), 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/registration/README.md b/scripts/kernel_tuning/registration/README.md index cffdda4..9faabc2 100644 --- a/scripts/kernel_tuning/registration/README.md +++ b/scripts/kernel_tuning/registration/README.md @@ -20,6 +20,8 @@ - **工具环境**：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` 的检查都在这个解释器的子进程中执行。工具环境自带 @@ -155,11 +157,11 @@ $TOOL register new --model pi05 --op linear \ ### 4.3 `register run`：启动智能体 -`run RUN --agent HARNESS/MODEL:EFFORT [--set test_python=…] [--yes]` 执行一次即视为确认调用点，然后： +`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` 和 `sites.json` 哈希，状态变为 `ready`。 +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`； @@ -171,7 +173,7 @@ $TOOL register new --model pi05 --op linear \ 智能体的 `LocalEnv` 是从中派生出的 `RUN/worktree`；框架日志归档到 `RUN/workspace/`， 实际的 epic 路径记录在 `humanize.json` 中。 -`run` 只能执行一次。之后使用 `resume`，它不接受更换 agent 或 `test_python`。 +`run` 只能执行一次。之后使用 `resume`，它不接受更换 agent、`test_python` 或 `unfenced`。 **每次尝试的阶段：** diff --git a/scripts/kernel_tuning/registration/cli.py b/scripts/kernel_tuning/registration/cli.py index d550d19..637254b 100644 --- a/scripts/kernel_tuning/registration/cli.py +++ b/scripts/kernel_tuning/registration/cli.py @@ -11,7 +11,14 @@ from pathlib import Path from typing import Any -from scripts.kernel_tuning.artifacts import REPOSITORY, atomic_json, run_lock, runtime_identity +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 @@ -41,6 +48,7 @@ def add_parser(commands: Any) -> None: 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") @@ -82,6 +90,7 @@ def execute(run: RegistrationRun, *, resume: bool) -> dict[str, Any]: 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( @@ -153,7 +162,7 @@ def dispatch(args: argparse.Namespace) -> int: run.root / "test-runtime.log", ) ) - run.prepare(args.agent, python) + 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: diff --git a/scripts/kernel_tuning/registration/flow.py b/scripts/kernel_tuning/registration/flow.py index d8e5c00..a460af8 100644 --- a/scripts/kernel_tuning/registration/flow.py +++ b/scripts/kernel_tuning/registration/flow.py @@ -11,7 +11,6 @@ from typing import Any from hmz.flows import ( - Agent, AgentCollection, Budget, BudgetExceeded, @@ -24,6 +23,7 @@ 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 @@ -33,7 +33,7 @@ class Agents(AgentCollection): """One explicitly selected coding-agent harness.""" - coder: Agent + coder: Coder class Envs(EnvCollection): @@ -66,6 +66,7 @@ async def operator_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() @@ -133,13 +134,13 @@ async def heartbeat() -> None: if remaining() <= 0: raise TimeoutError("Registration duration exhausted") if session is None: - session = await agents["coder"].spawn(env=work_env) + 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( - agents["coder"].run( + coder.run( instruction, session=session, budget=Budget(duration=timedelta(seconds=seconds), graceful=False), diff --git a/scripts/kernel_tuning/registration/workspace.py b/scripts/kernel_tuning/registration/workspace.py index df6d9fb..04b8632 100644 --- a/scripts/kernel_tuning/registration/workspace.py +++ b/scripts/kernel_tuning/registration/workspace.py @@ -290,13 +290,21 @@ def recover_base(self) -> None: self.manifest["base_overlay"] = {} self.save() - def prepare(self, agent: str, python: str | None = None) -> None: - """Freeze reviewed call sites and build context before the first agent turn.""" + 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"]): - raise ContractError("Agent/test_python is frozen; resume without overrides") + 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)) @@ -320,6 +328,7 @@ def prepare(self, agent: str, python: str | None = None) -> None: (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"), diff --git a/tests/kernel_tuning/test_batch.py b/tests/kernel_tuning/test_batch.py index b21fd6c..98bbad8 100644 --- a/tests/kernel_tuning/test_batch.py +++ b/tests/kernel_tuning/test_batch.py @@ -183,6 +183,31 @@ def test_cli_validates_agent_and_resume_arguments(tmp_path: Path, capsys) -> Non 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: diff --git a/tests/kernel_tuning/test_flow.py b/tests/kernel_tuning/test_flow.py index 916e812..aed5405 100644 --- a/tests/kernel_tuning/test_flow.py +++ b/tests/kernel_tuning/test_flow.py @@ -227,3 +227,17 @@ def failed_copy(*args, **kwargs): 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_registration.py b/tests/kernel_tuning/test_registration.py index 4a610eb..b39afa1 100644 --- a/tests/kernel_tuning/test_registration.py +++ b/tests/kernel_tuning/test_registration.py @@ -858,3 +858,31 @@ def restore(): 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") From c0db3f477af20aa6f4d06f04756eebd38e8682f5 Mon Sep 17 00:00:00 2001 From: owpeter Date: Wed, 14 Oct 2026 05:55:31 -0700 Subject: [PATCH 34/44] docs: describe --unfenced agent runs in the kernel tuning proposal Add a section on hosts without Landlock: how to check for it, where to pass --unfenced for run, tune-all and register run, how the choice is frozen across resume, the isolation an unfenced agent needs, and how a refused session shows up in run evidence. Signed-off-by: owpeter --- docs/proposals/0008-kernel-tuning.md | 73 +++++++++++++++++++++++++--- 1 file changed, 66 insertions(+), 7 deletions(-) diff --git a/docs/proposals/0008-kernel-tuning.md b/docs/proposals/0008-kernel-tuning.md index ab55af8..3bbe44e 100644 --- a/docs/proposals/0008-kernel-tuning.md +++ b/docs/proposals/0008-kernel-tuning.md @@ -372,7 +372,8 @@ whole-model parity each have distinct reference paths and evidence. Implemented components are the Humanize2 Flow/archive, discovery and registry capture, calibration, evaluator, batch tuning, GPU attribution, registry/profile -loader, install, model verification, and the CPU registration controller. The +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; the tuned profile directory currently contains the loader but no committed hardware profile. Historical target-machine checks are not new validation performed by this documentation change. @@ -402,8 +403,10 @@ must be checked by model benchmarks before becoming performance claims. CPU tests under `tests/kernel_tuning/` cover task validation, Humanize2 fake-agent execution, failed/interrupted attempts, budgets/resume, evidence tampering, preflight barriers, replay layouts/mutations/RNG, calibration, hotspot attribution, -installation, model-gate arithmetic and registration repair. Registry tests also -cover preference, fallback, profile mismatch and digest checks. +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 @@ -488,6 +491,61 @@ subsequent profiling warmup and measurement window. --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 @@ -555,10 +613,11 @@ uv run --project scripts/kernel_tuning python -m scripts.kernel_tuning tune-all 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 -([CONTRIBUTING](../../CONTRIBUTING.md#optional-kernel-tuning-tools)). If a preflight-only batch will later resume search, -specify its agent when creating it. A normal run starts agents only after every -baseline passes. Default limits apply per task: 20 attempts, two active hours, +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. From f68a4d7b7f1b2366009002e2d9a97c6d16e51cdf Mon Sep 17 00:00:00 2001 From: owpeter Date: Thu, 15 Oct 2026 05:51:17 -0700 Subject: [PATCH 35/44] fix: take capture-ops output contracts from the production call capture-ops derived each registry case's output layout from the Torch reference, while downstream operators' inputs were recorded from the implementation the model actually ran. When the two disagree (for example the Triton rotate_qk returns a contiguous query while the reference keeps the input view layout), a tuned kernel follows the reference layout and the next tuned kernel's layout gate no longer matches in production. Record outputs, aliases and dtypes from the model's own call (run once, its result returned unchanged). The reference must agree on everything except strides and offsets; differing layouts are kept as reference_outputs and the generated baseline copies the reference result into the production layout. Signed-off-by: owpeter --- docs/proposals/0008-kernel-tuning.md | 12 ++- scripts/kernel_tuning/discovery/tasks.py | 55 +++++++++++ scripts/kernel_tuning/regions.py | 95 +++++++++++++++---- tests/kernel_tuning/test_regions.py | 116 +++++++++++++++++++++++ 4 files changed, 261 insertions(+), 17 deletions(-) create mode 100644 tests/kernel_tuning/test_regions.py diff --git a/docs/proposals/0008-kernel-tuning.md b/docs/proposals/0008-kernel-tuning.md index 3bbe44e..17ba53c 100644 --- a/docs/proposals/0008-kernel-tuning.md +++ b/docs/proposals/0008-kernel-tuning.md @@ -144,7 +144,17 @@ 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. Registry tasks currently default to CUDA Graph kernel timing, with an +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. diff --git a/scripts/kernel_tuning/discovery/tasks.py b/scripts/kernel_tuning/discovery/tasks.py index 33c48ef..3465e0e 100644 --- a/scripts/kernel_tuning/discovery/tasks.py +++ b/scripts/kernel_tuning/discovery/tasks.py @@ -107,12 +107,60 @@ def _sources(op: dict[str, Any]) -> tuple[str, dict[str, str]]: 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 identity["kind"] == "registry" else {} + 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]: @@ -282,6 +330,13 @@ def _plan( "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", diff --git a/scripts/kernel_tuning/regions.py b/scripts/kernel_tuning/regions.py index 6346e42..b34dc5b 100644 --- a/scripts/kernel_tuning/regions.py +++ b/scripts/kernel_tuning/regions.py @@ -8,6 +8,10 @@ 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. @@ -33,6 +37,28 @@ 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 @@ -185,7 +211,15 @@ def _source_digests(self, module: str) -> dict[str, str]: return self._digests[module] def record(self, op: str, implementation: Any, function: Any, args: tuple, kwargs: dict) -> Any: - """Snapshot one call, derive its output contract from the Torch reference, then run it.""" + """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 ( @@ -193,7 +227,6 @@ def record(self, op: str, implementation: Any, function: Any, args: tuple, kwarg describe_inputs, flatten, restore, - storage_key, tensor_outputs, tensor_spec, ) @@ -206,8 +239,8 @@ def record(self, op: str, implementation: Any, function: Any, args: tuple, kwarg local = clone_inputs(tensors) ref_args, ref_kwargs = restore(recipe, local) result = reference(*ref_args, **ref_kwargs) - outputs = tensor_outputs(result) - if not outputs: + reference_outputs = tensor_outputs(result) + if not reference_outputs: raise ContractError("registry operators must return tensors") identity = { "kind": "registry", @@ -227,29 +260,59 @@ def record(self, op: str, implementation: Any, function: Any, args: tuple, kwarg } # 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": [ - [i for i, tensor in enumerate(local) if storage_key(tensor) == storage_key(output)] - for output in outputs - ], - "output_groups": [ - [i for i, other in enumerate(outputs) if storage_key(output) == storage_key(other)] - for output 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: - message = f"registry.{op}: {exc}" - if message not in self.errors: - self.errors.append(message) - return function(*args, **kwargs) + 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 diff --git a/tests/kernel_tuning/test_regions.py b/tests/kernel_tuning/test_regions.py new file mode 100644 index 0000000..1e958fa --- /dev/null +++ b/tests/kernel_tuning/test_regions.py @@ -0,0 +1,116 @@ +"""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 From d8cc325f8c8a05fe611bbb66c676290158d150a1 Mon Sep 17 00:00:00 2001 From: owpeter Date: Wed, 14 Oct 2026 07:15:58 -0700 Subject: [PATCH 36/44] feat: route pi05 calls through the linear registry operator Register a torch reference for linear and route the pi0.5 fused gate/up GEMM, MLP down projection and native SigLIP q/k/v, out, fc1 and fc2 projections through it. Produced by the automated registration run 20261014T134024Z-ed6d7ff626 (codex gpt-6-astra:xhigh, unfenced) and passed its C1-C5 CPU checks. Signed-off-by: owpeter --- embodiinfer/backend/torch/ops.py | 12 + embodiinfer/layers/ops.py | 3 + embodiinfer/policies/pi05/modeling_pi05.py | 49 +- tests/test_registry_linear.py | 654 +++++++++++++++++++++ 4 files changed, 708 insertions(+), 10 deletions(-) create mode 100644 tests/test_registry_linear.py diff --git a/embodiinfer/backend/torch/ops.py b/embodiinfer/backend/torch/ops.py index 3d9c866..3545182 100644 --- a/embodiinfer/backend/torch/ops.py +++ b/embodiinfer/backend/torch/ops.py @@ -12,6 +12,18 @@ 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 diff --git a/embodiinfer/layers/ops.py b/embodiinfer/layers/ops.py index fa37b13..da72454 100644 --- a/embodiinfer/layers/ops.py +++ b/embodiinfer/layers/ops.py @@ -157,6 +157,9 @@ def _select_profile(name: str) -> Any: return profile +register_lazy_op("linear", "torch", "embodiinfer.backend.torch.ops", "linear") + + for _op, _module in ( ("gated_residual", "norm"), ("ada_rms_norm", "norm"), diff --git a/embodiinfer/policies/pi05/modeling_pi05.py b/embodiinfer/policies/pi05/modeling_pi05.py index a81d12a..ef463f4 100644 --- a/embodiinfer/policies/pi05/modeling_pi05.py +++ b/embodiinfer/policies/pi05/modeling_pi05.py @@ -40,6 +40,7 @@ import torch import torch.nn.functional as F +from torch import nn from ...engine.parallel import TensorParallelContext from ...layers import KernelSelection, get_attention_backend, get_split_kv_attention_backend @@ -207,7 +208,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 +) -> 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) @@ -221,11 +224,13 @@ def _fused_linear_pair(x: torch.Tensor, first: Any, second: Any) -> tuple[torch. ) fused = packed(x) return fused.split((first.out_features, second.out_features), dim=-1) + if type(first) is not nn.Linear or type(second) is not nn.Linear: + return first(x), second(x) weight = getattr(first, "_vvla_fused_pair_weight", None) 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) @@ -305,14 +310,17 @@ def _mlp( kernels: KernelSelection | None = None, ) -> 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) activated = _op(kernels, "gated_gelu", use_triton)(gate, up) - return mlp.down_proj(activated) + if type(mlp.down_proj) is nn.Linear: + return _op(kernels, "linear", False)(activated, mlp.down_proj.weight, mlp.down_proj.bias) + else: + return mlp.down_proj(activated) @dataclass @@ -534,14 +542,35 @@ 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) + activated = layer.mlp.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 diff --git a/tests/test_registry_linear.py b/tests/test_registry_linear.py new file mode 100644 index 0000000..bcd88aa --- /dev/null +++ b/tests/test_registry_linear.py @@ -0,0 +1,654 @@ +"""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, + _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) + for swapped in (False, True): + first, second = make(7, 5, dtype, False), _track(_dense(7, 9, dtype, False)) + if swapped: + first, second = second, first + 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) + assert calls == [] + + 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_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 == [] From 4457b91b8f63f9cfc50106772e68de9088291853 Mon Sep 17 00:00:00 2001 From: owpeter Date: Wed, 14 Oct 2026 07:19:52 -0700 Subject: [PATCH 37/44] fix: keep packing every non-quantized pi05 gate/up pair into one GEMM The registered linear route skipped the packed gate/up GEMM for any projection that is not exactly nn.Linear, so tensor-parallel ColumnParallelLinear shards fell back to two GEMMs. BASE packs every non-quantized inference pair regardless of module type; restore that and test it for an overriding nn.Linear subclass and ColumnParallelLinear. Signed-off-by: owpeter --- embodiinfer/policies/pi05/modeling_pi05.py | 2 - tests/test_registry_linear.py | 44 +++++++++++++++------- 2 files changed, 31 insertions(+), 15 deletions(-) diff --git a/embodiinfer/policies/pi05/modeling_pi05.py b/embodiinfer/policies/pi05/modeling_pi05.py index ef463f4..12e680d 100644 --- a/embodiinfer/policies/pi05/modeling_pi05.py +++ b/embodiinfer/policies/pi05/modeling_pi05.py @@ -224,8 +224,6 @@ def _fused_linear_pair( ) fused = packed(x) return fused.split((first.out_features, second.out_features), dim=-1) - if type(first) is not nn.Linear or type(second) is not nn.Linear: - return first(x), second(x) weight = getattr(first, "_vvla_fused_pair_weight", None) 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() diff --git a/tests/test_registry_linear.py b/tests/test_registry_linear.py index bcd88aa..6373c4d 100644 --- a/tests/test_registry_linear.py +++ b/tests/test_registry_linear.py @@ -204,19 +204,6 @@ def _track(module): def _check_module_bypass(make, selection, calls, dtype): x = torch.randn(2, 3, 7, dtype=dtype) - for swapped in (False, True): - first, second = make(7, 5, dtype, False), _track(_dense(7, 9, dtype, False)) - if swapped: - first, second = second, first - 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) - assert calls == [] - mlp = _mlp(dtype) mlp.down_proj = make(11, 7, dtype, True) expected = _original_mlp(mlp, x, fuse=False) @@ -570,6 +557,37 @@ def row(k, n, dtype, bias): 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() From 8d1cc04dbec3091ff2835e3d700515bf407c2bec Mon Sep 17 00:00:00 2001 From: owpeter Date: Wed, 14 Oct 2026 08:13:13 -0700 Subject: [PATCH 38/44] feat: route pi05 SigLIP tanh GELU through the gelu_tanh operator Register a torch reference for gelu_tanh and route the native SigLIP MLP activation through it when the activation is exactly a tanh GELU. Produced by the automated registration run 20261014T144936Z-50adc4a729 (codex gpt-6-astra:xhigh, unfenced); it passed C1-C5 on attempt 2. Signed-off-by: owpeter --- embodiinfer/backend/torch/ops.py | 10 + embodiinfer/layers/ops.py | 1 + embodiinfer/policies/pi05/modeling_pi05.py | 30 +- tests/test_registry_gelu_tanh.py | 507 +++++++++++++++++++++ 4 files changed, 547 insertions(+), 1 deletion(-) create mode 100644 tests/test_registry_gelu_tanh.py diff --git a/embodiinfer/backend/torch/ops.py b/embodiinfer/backend/torch/ops.py index 3545182..4dc041d 100644 --- a/embodiinfer/backend/torch/ops.py +++ b/embodiinfer/backend/torch/ops.py @@ -49,6 +49,16 @@ def gated_gelu(gate: torch.Tensor, up: torch.Tensor) -> torch.Tensor: 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 :] diff --git a/embodiinfer/layers/ops.py b/embodiinfer/layers/ops.py index da72454..2b2d382 100644 --- a/embodiinfer/layers/ops.py +++ b/embodiinfer/layers/ops.py @@ -158,6 +158,7 @@ def _select_profile(name: str) -> Any: 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 ( diff --git a/embodiinfer/policies/pi05/modeling_pi05.py b/embodiinfer/policies/pi05/modeling_pi05.py index 12e680d..db482e3 100644 --- a/embodiinfer/policies/pi05/modeling_pi05.py +++ b/embodiinfer/policies/pi05/modeling_pi05.py @@ -33,6 +33,7 @@ from __future__ import annotations +import sys from collections.abc import Mapping from dataclasses import dataclass from threading import Lock @@ -99,6 +100,29 @@ def _op(kernels: KernelSelection | None, op: str, use_triton: bool): 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 + # 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, @@ -562,7 +586,11 @@ def _embed_image(self, image: torch.Tensor) -> torch.Tensor: activated = _op(self._kernels, "linear", False)(x, layer.mlp.fc1.weight, layer.mlp.fc1.bias) else: activated = layer.mlp.fc1(x) - activated = layer.mlp.activation_fn(activated) + 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 diff --git a/tests/test_registry_gelu_tanh.py b/tests/test_registry_gelu_tanh.py new file mode 100644 index 0000000..50f3712 --- /dev/null +++ b/tests/test_registry_gelu_tanh.py @@ -0,0 +1,507 @@ +"""CPU bit-exact parity and guarded SigLIP dispatch for tanh GELU.""" + +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): + def forward(self, x): + return F.gelu(x, approximate="tanh") + + 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, + _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(), + 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) From 2df0eb211f40fa6ad5f4977bc1a3b9863cd99641 Mon Sep 17 00:00:00 2001 From: owpeter Date: Wed, 14 Oct 2026 08:16:54 -0700 Subject: [PATCH 39/44] fix: route only native-partial GELUTanh activations through gelu_tanh Transformers builds GELUTanh for both gelu_pytorch_tanh and gelu_python_tanh; the latter evaluates a Python formula that differs from F.gelu(approximate="tanh") in the last bit. Require the activation to be the native functools.partial before using the operator, mirror the real class in the test double, and check the installed Transformers directly. Signed-off-by: owpeter --- embodiinfer/policies/pi05/modeling_pi05.py | 12 +++++++++ tests/test_registry_gelu_tanh.py | 29 +++++++++++++++++++++- 2 files changed, 40 insertions(+), 1 deletion(-) diff --git a/embodiinfer/policies/pi05/modeling_pi05.py b/embodiinfer/policies/pi05/modeling_pi05.py index db482e3..2059522 100644 --- a/embodiinfer/policies/pi05/modeling_pi05.py +++ b/embodiinfer/policies/pi05/modeling_pi05.py @@ -33,6 +33,7 @@ from __future__ import annotations +import functools import sys from collections.abc import Mapping from dataclasses import dataclass @@ -111,6 +112,17 @@ def _is_gelu_tanh(activation: Any) -> bool: 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__ diff --git a/tests/test_registry_gelu_tanh.py b/tests/test_registry_gelu_tanh.py index 50f3712..7b2588d 100644 --- a/tests/test_registry_gelu_tanh.py +++ b/tests/test_registry_gelu_tanh.py @@ -1,5 +1,7 @@ """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 @@ -42,8 +44,20 @@ def transformers_tanh(monkeypatch): 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 F.gelu(x, approximate="tanh") + return self.act(x) PytorchGELUTanh.__module__ = module.__name__ module.PytorchGELUTanh = PytorchGELUTanh @@ -340,6 +354,7 @@ def hook(module, args, output, hook_calls=hook_calls): custom, AlteredGELU(approximate="tanh"), AlteredTransformersGELU(), + transformers_tanh(use_gelu_tanh_python=True), overridden, hooked, impostor, @@ -505,3 +520,15 @@ def activation_fn(self): _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) From 9c9ae931f35bddc66fc55dc3f702ef01a762e7a2 Mon Sep 17 00:00:00 2001 From: owpeter Date: Thu, 15 Oct 2026 17:43:29 -0700 Subject: [PATCH 40/44] feat: add the orin_pi05 kernel profile for Jetson AGX Orin Commit the agent-tuned split_kv_attention, linear (prefix down_proj) and rotate_qk kernels installed from their evaluator-verified exports. rotate_qk composes separately tuned prefix and decode variants with disjoint verified layouts; the decode variant was tuned against the production output layout. Record the end-to-end measurement in proposal 0008: on Jetson AGX Orin with pi05_base, B=1, mixed FP32/BF16, eager prefix and a full denoise-loop CUDA Graph, mean latency is 537.8 ms against 609.8 ms for the existing Triton path and 825.5 ms for Torch. kernel_profile still defaults to none. Signed-off-by: owpeter --- docs/proposals/0008-kernel-tuning.md | 67 ++++- .../triton/tuned/orin_pi05/linear/__init__.py | 62 ++++ .../tuned/orin_pi05/linear/candidate.py | 144 ++++++++++ .../triton/tuned/orin_pi05/profile.json | 267 +++++++++++++++++ .../tuned/orin_pi05/rotate_qk/__init__.py | 27 ++ .../orin_pi05/rotate_qk/decode/__init__.py | 62 ++++ .../orin_pi05/rotate_qk/decode/candidate.py | 117 ++++++++ .../orin_pi05/rotate_qk/prefix/__init__.py | 62 ++++ .../orin_pi05/rotate_qk/prefix/candidate.py | 149 ++++++++++ .../orin_pi05/split_kv_attention/__init__.py | 62 ++++ .../orin_pi05/split_kv_attention/candidate.py | 270 ++++++++++++++++++ 11 files changed, 1285 insertions(+), 4 deletions(-) create mode 100644 embodiinfer/backend/triton/tuned/orin_pi05/linear/__init__.py create mode 100644 embodiinfer/backend/triton/tuned/orin_pi05/linear/candidate.py create mode 100644 embodiinfer/backend/triton/tuned/orin_pi05/profile.json create mode 100644 embodiinfer/backend/triton/tuned/orin_pi05/rotate_qk/__init__.py create mode 100644 embodiinfer/backend/triton/tuned/orin_pi05/rotate_qk/decode/__init__.py create mode 100644 embodiinfer/backend/triton/tuned/orin_pi05/rotate_qk/decode/candidate.py create mode 100644 embodiinfer/backend/triton/tuned/orin_pi05/rotate_qk/prefix/__init__.py create mode 100644 embodiinfer/backend/triton/tuned/orin_pi05/rotate_qk/prefix/candidate.py create mode 100644 embodiinfer/backend/triton/tuned/orin_pi05/split_kv_attention/__init__.py create mode 100644 embodiinfer/backend/triton/tuned/orin_pi05/split_kv_attention/candidate.py diff --git a/docs/proposals/0008-kernel-tuning.md b/docs/proposals/0008-kernel-tuning.md index 17ba53c..dc33811 100644 --- a/docs/proposals/0008-kernel-tuning.md +++ b/docs/proposals/0008-kernel-tuning.md @@ -384,9 +384,9 @@ Implemented components are the Humanize2 Flow/archive, discovery and 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; the tuned profile directory currently contains -the loader but no committed hardware profile. Historical target-machine checks -are not new validation performed by this documentation change. +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: @@ -440,7 +440,66 @@ 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; no new performance claim is made here. +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 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 From 5351e76a7f8c1bafd3d64ba77af0c0c7234614dd Mon Sep 17 00:00:00 2001 From: owpeter Date: Thu, 15 Oct 2026 23:20:50 -0700 Subject: [PATCH 41/44] fix: adapt kernel registry routing to main's pi05 refactor Merging main left the compiled-helper warmup referencing the RoPE and tanh-GELU helpers that now resolve through registry operators, and made `kernels` a required argument of `_fused_linear_pair` that main's callers and its packed-weight benchmark wrapper do not pass. Default `kernels` to the unprofiled selection, forward keyword arguments in the benchmark wrapper, and give the registry tests' policy stubs main's `_get_optimizations` hook. Signed-off-by: owpeter --- benchmarks/pi05-benchmark/validate_nonquant.py | 4 ++-- embodiinfer/policies/pi05/modeling_pi05.py | 13 +++++++------ tests/test_registry_gelu_tanh.py | 1 + tests/test_registry_linear.py | 1 + 4 files changed, 11 insertions(+), 8 deletions(-) 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/embodiinfer/policies/pi05/modeling_pi05.py b/embodiinfer/policies/pi05/modeling_pi05.py index fa2e7af..9d52421 100644 --- a/embodiinfer/policies/pi05/modeling_pi05.py +++ b/embodiinfer/policies/pi05/modeling_pi05.py @@ -254,7 +254,7 @@ def _packed_quantized_projection( def _fused_linear_pair( - x: torch.Tensor, first: Any, second: Any, *, kernels: KernelSelection | None + 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: @@ -375,14 +375,12 @@ def _prepare_compiled_inference_helpers() -> None: 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) @@ -872,7 +870,10 @@ def _attn_sublayer( q, k = operators.rotate(q, cos, sin), operators.rotate(k, cos, sin) else: q, k = _apply_rope( - q, k, cos, sin, + q, + k, + cos, + sin, self.compile_backend == "inductor", native_expert and q.dtype in (torch.bfloat16, torch.float16), self._kernels, diff --git a/tests/test_registry_gelu_tanh.py b/tests/test_registry_gelu_tanh.py index 7b2588d..b060b80 100644 --- a/tests/test_registry_gelu_tanh.py +++ b/tests/test_registry_gelu_tanh.py @@ -145,6 +145,7 @@ def _vision(dtype, selection, activation=None): 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)) diff --git a/tests/test_registry_linear.py b/tests/test_registry_linear.py index 6373c4d..4a8e682 100644 --- a/tests/test_registry_linear.py +++ b/tests/test_registry_linear.py @@ -163,6 +163,7 @@ def _vision(dtype, selection): 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")] From 97a1114ba914dac023320ddfc7ed38779a9c050f Mon Sep 17 00:00:00 2001 From: owpeter Date: Thu, 15 Oct 2026 23:20:51 -0700 Subject: [PATCH 42/44] test: use an unregistered example op in registration tests The registration tests register `linear` as a worked example against the current checkout, but the routed pi05 projections now register a real `linear` operator and `tests/test_registry_linear.py`. Name the example `matmul_t` so it remains a genuine registry gap. Signed-off-by: owpeter --- tests/kernel_tuning/test_registration.py | 50 ++++++++++++------------ 1 file changed, 25 insertions(+), 25 deletions(-) diff --git a/tests/kernel_tuning/test_registration.py b/tests/kernel_tuning/test_registration.py index b39afa1..38c1510 100644 --- a/tests/kernel_tuning/test_registration.py +++ b/tests/kernel_tuning/test_registration.py @@ -36,14 +36,14 @@ ) SOURCE = "embodiinfer/policies/pi05/modeling_pi05.py" -NEW_TEST = "tests/test_registry_linear.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": "linear", + "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}, } @@ -72,7 +72,7 @@ async def preflight(*args, **kwargs): "--model", "pi05", "--op", - "linear", + "matmul_t", "--site", f"{SOURCE}::_fused_qkv::F.linear(x, weight)", ] @@ -142,30 +142,30 @@ def write_candidate(run: RegistrationRun) -> 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, "linear", False)(x, weight)') + + 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 linear(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:\n """Preserve the original bias-free F.linear expression."""\n return F.linear(x, weight)\n' + '\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("linear", "torch", "embodiinfer.backend.torch.ops", "linear")\n') + 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": "linear", - "reference": f"{REFERENCE}::linear", + "op": "matmul_t", + "reference": f"{REFERENCE}::matmul_t", "sites": [ { "id": "S1", "file": SOURCE, "function": "_fused_qkv", - "call": '_op(kernels, "linear", False)(x, weight)', + "call": '_op(kernels, "matmul_t", False)(x, weight)', } ], "preserved": [], @@ -198,12 +198,12 @@ def test_spec_rejects_invalid_contracts(overrides) -> None: def test_yaml_and_cli_share_validation(tmp_path) -> None: path = tmp_path / "spec.yaml" path.write_text( - f"model: pi05\nop: linear\nsites:\n - source: {SOURCE}\n function: _fused_qkv\n expression: F.linear(x, weight)\n", + 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="linear", sites=[f"{SOURCE}::_fused_qkv::F.linear(x, weight)"] + None, model="pi05", op="matmul_t", sites=[f"{SOURCE}::_fused_qkv::F.linear(x, weight)"] ) assert yaml_spec == cli_spec with pytest.raises(ContractError): @@ -333,15 +333,15 @@ def test_structure_rejects_misregistration(registration_run, mutation, message) elif mutation == "loop": p = run.worktree / REGISTRY text = p.read_text(encoding="utf-8") - line = 'register_lazy_op("linear", "torch", "embodiinfer.backend.torch.ops", "linear")' + 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, "linear", False)(x, weight)', - 'unused = F.linear(x, weight)\n fused = _op(kernels, "linear", False)(x, weight)', + '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:") @@ -355,7 +355,7 @@ def test_structure_rejects_misregistration(registration_run, mutation, message) p = run.worktree / "REGISTRATION.json" document = read_json(p) if mutation == "triton": - document["sites"][0]["call"] = '_op(kernels, "linear", True)(x, weight)' + document["sites"][0]["call"] = '_op(kernels, "matmul_t", True)(x, weight)' else: import ast @@ -457,7 +457,7 @@ def test_real_registry_resolution_in_isolated_interpreter(registration_run) -> N pytest.importorskip("numpy") pytest.importorskip("PIL") run = registration_run - asyncio.run(check_unregistered(run.base, "linear", sys.executable, run.root / "real-preflight.log")) + 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 @@ -782,23 +782,23 @@ def resolve(self, op): source = worktree / SOURCE source.write_text( source.read_text(encoding="utf-8").replace( - "return raw(x, weight)", 'return _op(kernels, "linear", False)(x, weight)' + "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 linear(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:\n """Synthetic multiplication reference used only to test the controller."""\n return x * weight\n' + '\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("linear", "torch", "embodiinfer.backend.torch.ops", "linear")\n') + 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, "linear") - assert ops.linear(2, 3) == 6 + assert hasattr(ops, "matmul_t") + assert ops.matmul_t(2, 3) == 6 def test_route(): @@ -812,21 +812,21 @@ def resolve(self, op): return lambda x, w: x * w assert compute(2, 3, Recording()) == 6 - assert calls == ["linear"] + assert calls == ["matmul_t"] """, encoding="utf-8", ) atomic_json( worktree / "REGISTRATION.json", { - "op": "linear", - "reference": f"{REFERENCE}::linear", + "op": "matmul_t", + "reference": f"{REFERENCE}::matmul_t", "sites": [ { "id": "S1", "file": SOURCE, "function": "compute", - "call": '_op(kernels, "linear", False)(x, weight)', + "call": '_op(kernels, "matmul_t", False)(x, weight)', } ], "preserved": [{"site": "S1", "branch": "quantized"}], From 4b0d5ef67ba89170af2f9b0c18307d82e62554cf Mon Sep 17 00:00:00 2001 From: owpeter Date: Thu, 15 Oct 2026 23:20:51 -0700 Subject: [PATCH 43/44] fix: fail hotspot attribution clearly without Kineto activity types Stock Torch 2.10 Kineto events do not expose activity_type(), which the hotspot attribution needs to separate launches from device kernels. Raise an explicit error instead of an AttributeError, and test both behaviours. Signed-off-by: owpeter --- scripts/kernel_tuning/hotspots.py | 9 ++++++++- tests/kernel_tuning/test_hotspots.py | 6 +++++- 2 files changed, 13 insertions(+), 2 deletions(-) diff --git a/scripts/kernel_tuning/hotspots.py b/scripts/kernel_tuning/hotspots.py index 88e85e2..0788155 100644 --- a/scripts/kernel_tuning/hotspots.py +++ b/scripts/kernel_tuning/hotspots.py @@ -573,7 +573,14 @@ def kernels(self) -> list[Kernel]: 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 = [], [], [] - for event in profiler.profiler.kineto_results.events(): + 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): diff --git a/tests/kernel_tuning/test_hotspots.py b/tests/kernel_tuning/test_hotspots.py index 8c2766b..0f427e8 100644 --- a/tests/kernel_tuning/test_hotspots.py +++ b/tests/kernel_tuning/test_hotspots.py @@ -149,7 +149,11 @@ def test_session_profiles_the_window_and_restores_patches(fake_engine): 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 - assert session.kernels() == [] # CPU-only run: no device kernels + 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): From ca3f71ec1c2dd0b9c23ac36d14aaaa48aa07ff89 Mon Sep 17 00:00:00 2001 From: owpeter Date: Wed, 21 Oct 2026 03:37:52 -0700 Subject: [PATCH 44/44] refactor(kernel-tuning)!: remove full ATen discovery Remove the obsolete capture --model path, ATen/RNG recording and replay branches, and their tests. Keep registry capture-ops, catalog capture, shared tensor/layout helpers, and the tuning preflight and evaluator. Update regression coverage and the canonical workflow documentation. Legacy full-discovery artifacts must be recaptured with capture-ops; registry captures need recalibration when numerical-source hashes change. Signed-off-by: owpeter --- CONTRIBUTING.md | 72 +- benchmarks/kernel_tuning/calibration.py | 134 +-- benchmarks/kernel_tuning/discovery_inputs.py | 6 +- benchmarks/kernel_tuning/replay_adapter.py | 59 +- docs/proposals/0008-kernel-tuning.md | 267 ++---- scripts/kernel_tuning/__main__.py | 54 +- scripts/kernel_tuning/batch.py | 4 +- scripts/kernel_tuning/capture.py | 2 +- scripts/kernel_tuning/discovery/__init__.py | 2 +- scripts/kernel_tuning/discovery/contracts.py | 69 +- scripts/kernel_tuning/discovery/drive.py | 10 +- scripts/kernel_tuning/discovery/numerics.py | 30 +- scripts/kernel_tuning/discovery/recorder.py | 792 ------------------ scripts/kernel_tuning/discovery/rng.py | 133 --- scripts/kernel_tuning/discovery/tasks.py | 147 +--- scripts/kernel_tuning/discovery/tensors.py | 16 - scripts/kernel_tuning/prompts.py | 2 +- scripts/kernel_tuning/regions.py | 2 +- tests/kernel_tuning/conftest.py | 33 + tests/kernel_tuning/test_calibration.py | 131 +-- tests/kernel_tuning/test_cli.py | 27 + tests/kernel_tuning/test_discovery.py | 605 ------------- tests/kernel_tuning/test_drive.py | 54 ++ tests/kernel_tuning/test_model_capture.py | 235 ++++++ tests/kernel_tuning/test_regions.py | 51 ++ tests/kernel_tuning/test_replay_extensions.py | 322 ------- 26 files changed, 654 insertions(+), 2605 deletions(-) delete mode 100644 scripts/kernel_tuning/discovery/recorder.py delete mode 100644 scripts/kernel_tuning/discovery/rng.py delete mode 100644 tests/kernel_tuning/test_discovery.py create mode 100644 tests/kernel_tuning/test_drive.py create mode 100644 tests/kernel_tuning/test_model_capture.py delete mode 100644 tests/kernel_tuning/test_replay_extensions.py diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 913d1ee..b305032 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -184,55 +184,41 @@ Use `profile` to select registered operators, then `capture-ops`, `calibrate`, manual; there is no unified `optimize` command. Only registry captures produce exports accepted by `install`. -For full operator discovery instead, capture the existing preparation forward and calibrate numerical -contracts in its model runtime, then generate and preflight every eligible call -(see [model discovery](docs/proposals/0008-kernel-tuning.md#model-scoped-operator-discovery)): +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 --model pi05 --fixture-bytes 8589934592 --output results/kernel_tuning/captures/pi05 -- /absolute/preparation.py -uv run --project scripts/kernel_tuning python -m scripts.kernel_tuning generate --model pi05 --captured results/kernel_tuning/captures/pi05 --list +/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` is your existing preparation driver, constructing the -selected canonical policy through `make_policy` and running its normal forward -path. Disable compilation explicitly; select the same checkpoint, backend, -dtype, batching and decoding settings that will be tuned. Discovery covers only -paths exercised by that preparation. The capture records ATen overloads and -existing backend entrypoints, all observed shapes/layouts, and call counts. -Quantized calls and their internals are excluded; views/allocations and host work -are listed separately. No new fusion is introduced. A missing replay contract, -required fixture, or failing baseline blocks the entire model batch. `--dry-run` -only builds tasks; `--preflight-only` measures baselines without opening agents. -Model batches cannot use `--skip-preflight`, partial operator selection, or -estimated shapes. A standalone `run` cannot bypass their preflight barrier. - -The default fixture budget is 64 MiB; the example raises the upper limit to 8 GiB -without allocating it up front. Supported floating matmul/linear, reductions, -softmax and attention require complete real input fixtures, including weights. -Increase `--fixture-bytes` for the observed model; exhaustion fails explicitly. -Other exact tasks retain optional small floating fixtures and required integer -fixtures. Calibration compares only the production baseline and an FP64 reference -on recorded, random and boundary inputs. It stores per-workload/profile/output -bounds in a new hashed capture. Generation copies them to `numerics.json`; every -element must pass, and seeds/thresholds/execution settings cannot be overridden. -Calibration costs extra storage and target-device compute, before any agent runs. -Use a new capture/calibration/batch when changing settings. Old runs keep their -original contract; old captures with incomplete real inputs need recapture. - -Default-generator `randn` tasks save required before/after RNG fixtures, check -both output bytes and generator advancement, and support eager timing only. -Explicit generators and other unsupported RNG calls remain coverage gaps. -SDPA with zero/default dropout is eligible; nonzero dropout still needs an adapter. -Discovered float64 tensors use the opt-in replay definition schema and retain -double precision through task generation and evaluation. - -The old catalog workflow remains available explicitly with `generate --catalog`, -`tune-all --catalog`, and `capture --catalog`. Only that workflow accepts legacy -JSONL captures, catalog model groups and `--estimated`; it -does not certify model coverage. Existing JSONL must be recaptured for model tuning. +`/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 @@ -291,7 +277,7 @@ uv run --project scripts/kernel_tuning ruff check scripts/kernel_tuning benchmar uv run --project scripts/kernel_tuning ruff format --check scripts/kernel_tuning benchmarks/kernel_tuning tests/kernel_tuning ``` -CPU discovery/replay tests also need Torch in the selected interpreter; they do +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. diff --git a/benchmarks/kernel_tuning/calibration.py b/benchmarks/kernel_tuning/calibration.py index b76887e..576d60c 100644 --- a/benchmarks/kernel_tuning/calibration.py +++ b/benchmarks/kernel_tuning/calibration.py @@ -82,41 +82,18 @@ def captured_settings(precision: dict[str, Any], flags: dict[str, Any]) -> Itera 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"] == "aten": - _, name, overload = identity["name"].split(".") - fn = getattr(getattr(torch.ops.aten, name), overload) - else: - fn = getattr(importlib.import_module(identity["module"]), identity["entry"]) + 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]: - args, kwargs = restore(identity["arguments"], inputs) - if identity["kind"] == "aten": - _, name, overload = identity["name"].split(".") - schema = getattr(getattr(torch.ops.aten, name), overload)._schema - return { - arg.name: args[i] if i < len(args) else kwargs.get(arg.name, arg.default_value) - for i, arg in enumerate(schema.arguments) - } - if identity["kind"] == "registry": - import inspect + import inspect - function = getattr(importlib.import_module(identity["module"]), identity["entry"]) - return {**dict(zip(inspect.signature(function).parameters, args)), **kwargs} - names = { - "split_kv_attention": ( - "query", - "prefix_key", - "prefix_value", - "suffix_key", - "suffix_value", - "mask", - "scaling", - ), - "segmented_attention": ("q", "k", "v", "q_segment_offsets", "kv_segment_offsets"), - }[policy(identity)] - return {**dict(zip(names, args)), **kwargs} + 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( @@ -160,103 +137,12 @@ def validation_inputs( return inputs -def _attention( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - *, - mask: Any = None, - scale: float | None = None, - causal: bool = False, - gqa: bool = False, -) -> torch.Tensor: - if gqa and q.shape[-3] != k.shape[-3]: - k = k.repeat_interleave(q.shape[-3] // k.shape[-3], dim=-3) - v = v.repeat_interleave(q.shape[-3] // v.shape[-3], dim=-3) - scale = q.shape[-1] ** -0.5 if scale is None else scale - rows = [] - # Bound oracle scratch memory without changing each row's reduction semantics. - for start in range(0, q.shape[-2], 128): - stop = min(start + 128, q.shape[-2]) - scores = (q[..., start:stop, :] @ k.transpose(-1, -2)) * scale - if causal: - visible = ( - torch.arange(k.shape[-2], device=q.device) - <= torch.arange(start, stop, device=q.device)[:, None] - ) - scores = scores.masked_fill(~visible, float("-inf")) - if mask is not None: - local = mask[..., start:stop, :] if mask.ndim >= 2 and mask.shape[-2] != 1 else mask - scores = ( - scores.masked_fill(~local, float("-inf")) if local.dtype == torch.bool else scores + local - ) - probs = torch.softmax(scores, dim=-1) - probs = torch.where(torch.isneginf(scores).all(dim=-1, keepdim=True), 0.0, probs) - rows.append(probs @ v) - return torch.cat(rows, dim=-2) if rows else q.new_empty((*q.shape[:-1], v.shape[-1])) - - def high_precision(identity: dict[str, Any], inputs: list[torch.Tensor]) -> list[torch.Tensor]: - """Evaluate the declared mathematical operation in FP64 with unchanged attention semantics.""" - kind = policy(identity) - if kind is None: + """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] - args = _arguments(identity, promoted) - if kind == "registry": - # The Torch reference is dtype-generic: FP64 inputs give an FP64 oracle. - return production(identity, promoted) - if kind == "aten": - _, name, overload = identity["name"].split(".") - if name == "scaled_dot_product_attention": - if args["dropout_p"] != 0: - raise ContractError("Attention calibration requires dropout_p=0") - return [ - _attention( - args["query"], - args["key"], - args["value"], - mask=args["attn_mask"], - scale=args["scale"], - causal=args["is_causal"], - gqa=args.get("enable_gqa", False), - ) - ] - if "dtype" in args: - args["dtype"] = torch.float64 - if "half_to_float" in args: - args["half_to_float"] = False - if name == "rms_norm" and args.get("eps") is None: - # The default epsilon belongs to the captured input dtype, not the oracle dtype. - args["eps"] = torch.finfo(inputs[0].dtype).eps - fn = getattr(getattr(torch.ops.aten, name), overload) - return tensor_outputs(fn(**args)) - if kind == "split_kv_attention": - mask = args.get("mask") - return [ - _attention( - args["query"], - torch.cat((args["prefix_key"], args["suffix_key"]), dim=2), - torch.cat((args["prefix_value"], args["suffix_value"]), dim=2), - mask=mask == 0 if mask is not None else None, - scale=args["scaling"], - gqa=True, - ) - ] - q, k, v = args["q"], args["k"], args["v"] - qb = args["q_segment_offsets"].tolist() - kb = args.get("kv_segment_offsets") - kb = qb if kb is None else kb.tolist() - output = torch.empty_like(q) - for qs, qe, ks, ke in zip(qb[:-1], qb[1:], kb[:-1], kb[1:]): - output[qs:qe] = _attention( - q[qs:qe].transpose(0, 1), - k[ks:ke].transpose(0, 1), - v[ks:ke].transpose(0, 1), - scale=args.get("scaling"), - gqa=True, - ).transpose(0, 1) - return [output] + return production(identity, promoted) def derive_bound(baseline: torch.Tensor, reference: torch.Tensor) -> dict[str, Any]: diff --git a/benchmarks/kernel_tuning/discovery_inputs.py b/benchmarks/kernel_tuning/discovery_inputs.py index aa8a95e..f5c0b65 100644 --- a/benchmarks/kernel_tuning/discovery_inputs.py +++ b/benchmarks/kernel_tuning/discovery_inputs.py @@ -1,7 +1,7 @@ -"""Export benchmark data as end-to-end inputs for model operator discovery. +"""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 discovery sees the same images, states, +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 \\ @@ -146,7 +146,7 @@ def streamvln(args: argparse.Namespace) -> list[list[RawPolicyRequest]]: def main(argv: list[str] | None = None) -> None: - """Write one model's discovery inputs.""" + """Write one model's recorded inference inputs.""" parser = argparse.ArgumentParser( description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter ) diff --git a/benchmarks/kernel_tuning/replay_adapter.py b/benchmarks/kernel_tuning/replay_adapter.py index 3557fdf..806f97a 100644 --- a/benchmarks/kernel_tuning/replay_adapter.py +++ b/benchmarks/kernel_tuning/replay_adapter.py @@ -1,8 +1,7 @@ -"""Evaluate discovered calls with captured layouts, aliases, and mutable state.""" +"""Evaluate registry calls with captured layouts, aliases, and mutable state.""" from __future__ import annotations -from contextlib import nullcontext from typing import Any from scripts.kernel_tuning.contracts import ContractError, Precision, TaskPackage, read_json @@ -32,7 +31,6 @@ def __init__(self, task: TaskPackage) -> None: torch.use_deterministic_algorithms(flags["deterministic_algorithms"]) self._precision_flags = self._flags() self.case: dict[str, Any] = {} - self._rng: Any = None 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 @@ -42,23 +40,13 @@ def __init__(self, task: TaskPackage) -> None: "Calibration and evaluator environments differ; recalibrate on the target" ) if self.cfg.profile: - raise ContractError("Discovered tasks need a replay-aware NCU adapter before profile=true") - if self.replay.get("rng") and self.cfg.timing.mode != "eager": - raise ContractError("RNG replay currently requires eager timing") + 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 evaluate(self, solutions: dict[str, dict[str, Any]]) -> dict[str, Any]: - """Evaluate stateful RNG calls without leaking their state into the caller.""" - from scripts.kernel_tuning.discovery.rng import preserve - - rng = self.replay.get("rng") - with preserve(rng["devices"]) if rng else nullcontext(): - return super().evaluate(solutions) - def _build_reference(self, registry: Any) -> Any: from flashinfer_bench.data import Solution @@ -66,22 +54,11 @@ def _build_reference(self, registry: Any) -> Any: # snapshotted relative imports. Both reference and baseline are immutable. return registry.build(self.definition, Solution.model_validate(self.task.baseline())) - def _seed(self, seed: int) -> None: - # RngReplay resets only the declared generators. The ordinary adapter's - # global seed would also mutate unrelated CUDA devices. - if not self.replay.get("rng"): - super()._seed(seed) - 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] - from scripts.kernel_tuning.discovery.rng import RngReplay - - self._rng = ( - RngReplay(self.replay["rng"], self.case, self.task.root, seed) if self.replay.get("rng") else None - ) if getattr(self, "numerical", None) is not None: from scripts.kernel_tuning.discovery.numerics import profile_key @@ -138,7 +115,7 @@ def _check_outputs(self, inputs: list[Any], outputs: list[Any]) -> None: def _call(self, runnable: Any, inputs: list[Any]) -> tuple[list[Any], list[Any]]: if runnable.metadata.destination_passing_style: - raise ContractError("Discovered tasks require return-value style") + raise ContractError("Registry tasks require return-value style") outputs = self._outputs(runnable(*inputs)) self._check_outputs(inputs, outputs) return list(inputs), outputs @@ -146,8 +123,6 @@ def _call(self, runnable: Any, inputs: list[Any]) -> tuple[list[Any], list[Any]] def _verify_call( self, runnable: Any, reference: Any, inputs: list[Any], seed: int ) -> tuple[float, float]: - from scripts.kernel_tuning.discovery.rng import preserve - if getattr(self, "numerical", None) is not None: from scripts.kernel_tuning.discovery.numerics import PROFILES, profile_key @@ -164,33 +139,20 @@ def _verify_call( current = self._verify_transition(runnable, reference, values, seed) error = tuple(max(a, b) for a, b in zip(error, current)) return error - rng = getattr(self, "_rng", None) - with preserve(rng.devices) if rng else nullcontext(): - return self._verify_transition(runnable, reference, inputs, seed) + 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 - from scripts.kernel_tuning.discovery.rng import snapshot - rng = getattr(self, "_rng", None) expected_inputs = self._clone(inputs) self._seed(seed) - if rng: - rng.reset() _, expected = self._call(reference, expected_inputs) - if rng: - expected_rng = snapshot(rng.devices) - rng.check_reference(expected_rng) expected = [value.clone() for value in expected] self._seed(seed) local = self._clone(inputs) - if rng: - rng.reset() _, outputs = self._call(runnable, local) - if rng: - rng.check(snapshot(rng.devices), expected_rng) torch.cuda.synchronize(self.device) self._check_flags() # Check every byte of backing storage, including aliases and untouched @@ -200,10 +162,6 @@ def _verify_transition( 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]: - # Unspecified bytes (unused dropout RNG state) carry no result to validate. - undefined = set(self.replay.get("undefined_outputs", ())) - actual = [value for i, value in enumerate(actual) if i not in undefined] - expected = [value for i, value in enumerate(expected) if i not in undefined] if getattr(self, "numerical", None) is None: return compare_outputs(actual, expected, self.cfg.precision) from .calibration import compare_calibrated, high_precision @@ -328,13 +286,6 @@ def _verify_calibrated_graph( return error def _time(self, runnable: Any, inputs: list[Any]) -> float: - from scripts.kernel_tuning.discovery.rng import preserve - - rng = getattr(self, "_rng", None) - with preserve(rng.devices) if rng else nullcontext(): - return self._time_samples(runnable, inputs) - - def _time_samples(self, runnable: Any, inputs: list[Any]) -> float: import torch local = self._clone(inputs) @@ -353,8 +304,6 @@ def execute() -> Any: def prepare() -> None: for dst, src in pairs: dst.copy_(src) - if self._rng: - self._rng.reset() # Reset copies touch L2 too; evict them before the timed call. self._flush.zero_() diff --git a/docs/proposals/0008-kernel-tuning.md b/docs/proposals/0008-kernel-tuning.md index dc33811..ba3a157 100644 --- a/docs/proposals/0008-kernel-tuning.md +++ b/docs/proposals/0008-kernel-tuning.md @@ -42,7 +42,7 @@ There is **no `optimize`, `select`, or `benchmark` subcommand** in 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 discovery/catalog support does not imply profile support. +outstanding; their benchmark-input/catalog support does not imply profile support. ## 2. Motivation and current gap @@ -51,12 +51,11 @@ 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 earlier discovery/catalog workflow optimizes each observed ATen call or -existing backend kernel against its production implementation. It remains useful -for operator investigations, but its exports do not directly install into a -hardware profile. 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 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 @@ -71,7 +70,7 @@ new fusion integration, shape-range coverage and orchestration across stages. 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 discovery/catalog and custom-task workflows usable, with +- Keep the existing catalog and custom-task workflows usable, with their original contracts and artifact formats. Robot/simulator communication, task success metrics, deployment orchestration, @@ -203,7 +202,7 @@ archives and reproduction. `install` accepts an exported registry task from `capture-ops`, with a Python/Triton solution and an operator already registered in the current tree. -ATen-discovery/catalog exports and CUDA C++ solutions are not installable through +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`. @@ -254,9 +253,9 @@ Profiler timings and weighted kernel gains cannot substitute for this run. - A separate agent controller would duplicate Humanize2's sessions, budgets and journals; reuse its Flow instead. -- Full ATen discovery preserves executed operator coverage but cannot itself - establish a new fused region or an installable registry call site. Keep it as - a distinct workflow ([Appendix C](#model-scoped-operator-discovery)). +- 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. @@ -289,12 +288,10 @@ candidate fails. ### Operator gate -Model discovery records execution settings separately from acceptance rules. -Pure copies, indexing, exact elementwise operations, RNG transitions, and unknown -semantics retain bit-exact validation. An explicit registry selects floating -matrix multiplication, linear, reductions, softmax, and supported attention/norm -calls for tolerance validation; dtype alone never relaxes an unknown operator. -Existing hand-written/catalog contracts and archived runs keep their contracts. +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 @@ -312,7 +309,7 @@ Every candidate element must satisfy abs(C-R) <= atol+rtol*abs(R). The one-epsil 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 -and RNG constraints remain exact. Nonfinite reference/baseline/output values +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. @@ -327,7 +324,7 @@ 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), discovery recorder/task generation, and the replay evaluator. +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; @@ -380,7 +377,7 @@ whole-model parity each have distinct reference paths and evidence. ## 7. Implementation status and remaining design -Implemented components are the Humanize2 Flow/archive, discovery and registry +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 @@ -412,7 +409,7 @@ must be checked by model benchmarks before becoming performance claims. CPU tests under `tests/kernel_tuning/` cover task validation, Humanize2 fake-agent execution, failed/interrupted attempts, budgets/resume, evidence tampering, -preflight barriers, replay layouts/mutations/RNG, calibration, hotspot attribution, +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 @@ -852,175 +849,61 @@ are not established by CPU/fake-agent tests. ## Appendix C. Other task sources -### Model-scoped operator discovery - -The model workflow selects exactly one canonical policy ID before preparation. -During the existing preparation execution, a scoped ATen observer and wrappers -around existing backend entrypoints collect complete input/output metadata. -Existing compound kernels remain single operators; this change introduces no -new fusion, quantization, runtime registration, or engine API. - -The observed scope is end-to-end inference, not a list of policy methods. It -opens at the model-agnostic entries every deployment path uses: `EngineCore` -construction and execution, the `GenerationBackend` rollout calls, the policy's -`collate`, and its serving adapter's construction and `infer`/`infer_batch`. -Everything below an entry is recorded, including decoder-specific recurrent -decoding, RL sampling, request preprocessing and action postprocessing. Scope is -tracked per thread and the observer is installed on whichever thread enters, so -the async engine's executor and batched-serving workers are covered. Model -weights read by arithmetic outside every entry (a driver calling policy -internals directly) and a run without any entry are capture errors, so a -bypassed path cannot silently produce a smaller operator set. - -The capture freezes the model type, supplied checkpoint/configuration descriptors, Torch version, -numerical settings, operator overloads, tensor shapes/strides, argument trees, -aliasing and mutation contracts, and every observed workload (no top-eight -truncation or estimated-shape fallback). Quantized calls and their internal -floating-point operations are excluded by invocation scope, including dynamically -created quantized projections and first-forward packing. Metadata-only operations -and host-side work are reported separately from device computation. Discovery -covers the executed preparation paths, not unvisited branches. - -Uncompiled preparation is required; compilation is rejected rather than silently -changing the model's attention or execution path. CUDA Graph metadata is collected -during warmup/capture and calls are counted on replay. Bounded fixtures may be saved -outside graph capture; cases requiring unavailable fixtures are coverage gaps. -Unknown external operators, unrepresentable arguments/state, and unsupported -contracts also remain explicit gaps. Every eligible workload must have a task and -all production baselines must pass preflight before any model batch starts an -agent. No partial-model or skip-preflight escape hatch is provided. - -ATen tasks replay the exact captured overload as their frozen bit-exact reference -and production baseline. Backend tasks snapshot their existing implementation. -The existing Humanize2 search, evaluator, paired promotion rules, archive and -export remain in use. Layout/state-aware task adapters extend evaluation without -weakening the default pure-operator checks. Inputs, current stream, numerical -settings, changed-input graph replay, and declared mutations are validated. - -This is repository tooling under scripts/kernel_tuning and measurement under -benchmarks/kernel_tuning. A static operator whitelist was rejected because it -cannot discover model calls; full-model graph compilation was rejected because -it changes operator boundaries and is outside the no-new-fusion scope. Existing -catalog tasks remain available explicitly as a legacy workflow. Legacy captures -cannot establish model coverage and must be recaptured for model tuning. - -CPU regression coverage must exercise uncatalogued operators, model identity, -all shapes/layouts, quantized scopes (including dynamic modules), unsupported -cases, immutable capture artifacts, and the all-baselines preflight barrier. -GPU/checkpoint validation remains explicit in the matching model environment; -CPU/fake-agent passes do not establish GPU correctness or performance. - -#### Commands and artifacts - -Run preparation in the selected model's own runtime. The bundled driver -`scripts.kernel_tuning.discovery.drive` constructs the policy and engine with the -serving transports' arguments and replays recorded inputs: serving requests go -through the policy's adapter as HTTP requests would, engine observations through -`collate` and `EngineCore.execute`, with one session per recurrent episode. -`benchmarks/kernel_tuning/discovery_inputs.py` exports such inputs from the pi05, -Qwen navigation and StreamVLN benchmark data; any script that drives one of the -entries above may be used instead. Compilation must be disabled explicitly (for -example `--no-cuda-graph` for StreamVLN, whose graph mode compiles its prefill). -The capture does not launch an extra forward or change the attention backend. -Canonical policy IDs are those accepted by the factory, for example `pi05`, -`streamvln`, or `qwen2.5-vl-3b-r2r-low-level`; catalog group names are not model IDs. - -```bash -/absolute/runtime/python -m benchmarks.kernel_tuning.discovery_inputs pi05 --config PI05_BENCHMARK.yaml --output pi05-inputs.pt --serving-config pi05-serving.json -/absolute/runtime/python -m scripts.kernel_tuning capture --model pi05 --fixture-bytes 8589934592 --output results/kernel_tuning/captures/pi05 -- -m scripts.kernel_tuning.discovery.drive --policy pi05 --checkpoint CHECKPOINT --adapter-config pi05-serving.json --inputs pi05-inputs.pt -uv run --project scripts/kernel_tuning python -m scripts.kernel_tuning generate --model pi05 --captured results/kernel_tuning/captures/pi05 --list -/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 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 -``` - -`--dry-run` builds and validates all tasks without evaluation. `--preflight-only` -evaluates all baselines without agents; if that batch is later resumed, supply -its agent at creation time. Model preflight failures prevent every agent from -starting. After all baselines pass, an individual search failure is recorded and -other prepared tasks may continue. Resume rechecks capture/task/tool identities -and reruns model preflight. The single-task `run`/`resume` paths also check the -parent model batch. There is no model-mode `--force`, `--estimated`, operator -filter, multiple-model selection, or `--skip-preflight` option. - -Capture artifacts are `manifest.json` (actual policy type, builder configuration, -revision, Torch version, numerical settings, completeness and hashes), -`operators.json` (overload/backend identity, signature, every workload and counts), -`coverage.json` (ready/excluded_quantized/metadata_only/host_only/blocked counts), -and bounded `fixtures/`. A model batch snapshots this directory under `capture/`, -creates a task per call signature under `tasks//`, and retains -`batch.json`, preflight records, run archives, exports and summaries. Each task -includes `replay.json` with storage/layout/mutation contracts and model identity; -its baseline calls the exact ATen overload or snapshotted backend implementation. -Different shapes and strides remain workloads of the same signature; overload, -scalar/static arguments, tensor rank/dtype and output structure distinguish tasks. - -The default 64 MiB fixture budget is explicit (`--fixture-bytes`); the example -raises the maximum to 8 GiB. Integer/index storage and all inputs of supported -numerically calibrated operators require real fixtures, even for large weights. -For other exact tasks, floating storage up to 1 MiB is captured when budget -permits and other floating storage uses seeded generation. First observed -fixtures represent each structural case; calibration tests those real values, -fresh random values, zeros and cancellation inputs with masks/indices preserved. -This is finite validation, not proof for every input. Required fixture exhaustion -is a coverage gap. `calibrate` adds `calibration.json` to a new capture; generated -tasks carry each operator's frozen bounds as `numerics.json`. See the numerical -policy in Section 6 for the formula, environment binding and limitations. - -Replay preserves shared storage, strides and offsets. It checks output values, -layouts, input/output and output/output aliases, declared mutations and unchanged -input storage. Timing resets backing storage before every sample, outside timed -events, in both eager and graph modes. The default is eager; `timing.mode=cuda_graph` -also requires successful changed-input A/B/A replay. NCU profiling currently -requires a dedicated replay-aware adapter. Numerical overrides cannot change -the selected exact contract or the frozen calibrated tolerance contract. - -Default-generator `aten.randn.default` and `aten.randn.generator` with -`generator=None` have an eager replay contract. Capture stores the opaque CPU -and target CUDA generator states immediately before and after the existing call, -within the same required-fixture budget. Seed 0 replays the actual captured -position; other validation seeds use independent default-generator states. -Validation compares output bytes and the complete post-call generator states, -rejecting extra draws or rewinds even when output values match. State resets -happen before each warmup/sample, outside timing events, and evaluation restores -the caller's Torch RNG states on success or failure. Explicit generator objects, -other random overloads, and RNG tasks using CUDA Graph timing remain gaps. - -SDPA's dispatcher RNG tag is conditional: `dropout_p=0` (including its schema -default) does not require an RNG replay contract. Discovery resolves positional, -keyword and omitted dropout arguments without changing the invocation, backend, -mask, scale or attention semantics. Nonzero dropout remains a coverage gap until -its own state contract is implemented; other nondeterministic tags still apply. -The same rule covers the backend overloads (flash, efficient, cuDNN, CPU) that -model code reaches under `no_grad`. Without dropout their `philox_seed` and -`philox_offset` outputs are left uninitialized, so replay records them as -`undefined_outputs` and checks every other output exactly; no other operator may -declare undefined outputs. - -Replay tasks declare `definition_schema=embodiinfer-replay-v1`. Their local -`ReplayDefinition` extends the pinned FlashInfer tensor schema with `float64`, -retains upstream structural validation, and uses the same builders with native -double-precision tensors. Captured fixtures, input materialization and output -checks preserve float64 without a cast or relaxed tolerance. Ordinary FlashInfer -definitions and their accepted dtype set are unchanged. - -Existing RoPE cache writes and greedy-sampling workspace mutations have explicit -contracts; graph GQA and prefill entries remain individual existing operators. -`GreedyWorkspace` has a typed tensor-state recipe; arbitrary Python objects still -require an adapter. Dtypes unsupported by the replay schema remain task -construction gaps rather than silently converted inputs. - -Known quantization scopes are `QuantizedLinear` (FP8/INT8/NVFP4), conversion and -packing helpers, their backend functions, and explicit native quantization -dispatcher operators. Normal integer indices are not a quantization signal. -Unknown dispatcher extensions and raw Triton launches are reported as gaps; -random operators outside the default-generator randn contract, opaque Python -arguments, and calls without tensor outputs also require an adapter. Custom -C++/CUDA calls that bypass both the dispatcher and -these entrypoints need explicit instrumentation before claiming coverage. -Capture is a single-process preparation facility; it does not certify unvisited -branches, other worker processes, or checkpoint-level action parity. Existing -fused kernels count as one call; no new fused operators are proposed or generated. +### 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 @@ -1036,7 +919,7 @@ The catalog covers pi0.5 norm/activation/RoPE/split-KV kernels, Qwen2.5-VL visio 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 discovery contracts. +not silently replaced by calibrated registry contracts. ```bash uv run --project scripts/kernel_tuning python -m scripts.kernel_tuning generate \ diff --git a/scripts/kernel_tuning/__main__.py b/scripts/kernel_tuning/__main__.py index 2d201f0..4d91594 100644 --- a/scripts/kernel_tuning/__main__.py +++ b/scripts/kernel_tuning/__main__.py @@ -125,7 +125,7 @@ def _model_capture(args: argparse.Namespace) -> ModelCapture: "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 --model first") + 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)) @@ -266,13 +266,15 @@ def main(argv: list[str] | None = None) -> int: 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 every discovered model task; no execution") + 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 discovered operators and coverage gaps") + 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") @@ -284,17 +286,11 @@ def main(argv: list[str] | None = None) -> int: 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 production kernel shapes while a model script runs (model runtime)" + "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", help="Record legacy catalog-only JSONL") - capture.add_argument("--model", help="Canonical policy name constructed by the preparation script") - capture.add_argument( - "--fixture-bytes", type=int, default=64 * 1024 * 1024, help="Maximum stored input fixture bytes" - ) - capture.add_argument( - "--output", type=Path, required=True, help="New model capture directory (JSONL for --catalog)" - ) + 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)" @@ -357,31 +353,15 @@ def main(argv: list[str] | None = None) -> int: return dispatch(args) if args.command == "capture": - if args.catalog: - from .capture import run as run_capture - from .operators import select - - if args.model: - raise ContractError("Legacy catalog capture does not bind a model; omit --model") - 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']}") - else: - from .discovery.recorder import run as run_model_capture - - if not args.model or args.operators: - raise ContractError( - "Capture requires one --model and no operator filter; use --catalog for legacy capture" - ) - result = run_model_capture( - command or args.script, args.output, args.model, fixture_bytes=args.fixture_bytes - ) - print(json.dumps(result, indent=2, ensure_ascii=False)) + 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 diff --git a/scripts/kernel_tuning/batch.py b/scripts/kernel_tuning/batch.py index f6881f5..16bfb81 100644 --- a/scripts/kernel_tuning/batch.py +++ b/scripts/kernel_tuning/batch.py @@ -94,7 +94,7 @@ def preflight(root: Path, names: list[str] | None = None) -> dict[str, str | Non _verify_model(root, batch) if batch.get("model"): if names is not None: - raise ContractError("Model preflight must include every discovered task") + 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] = {} @@ -254,7 +254,7 @@ def verify_model_task(task: TaskPackage) -> Path: root = task.root.parent.parent if not (root / "batch.json").is_file(): raise ContractError( - "Start discovered tasks with tune-all; standalone run cannot bypass model preflight" + "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: diff --git a/scripts/kernel_tuning/capture.py b/scripts/kernel_tuning/capture.py index 05115fb..3ef2f49 100644 --- a/scripts/kernel_tuning/capture.py +++ b/scripts/kernel_tuning/capture.py @@ -1,6 +1,6 @@ """Record the shapes of production kernel calls while a model actually runs. -``python -m scripts.kernel_tuning capture --output calls.jsonl -- SCRIPT [ARGS]`` +``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 diff --git a/scripts/kernel_tuning/discovery/__init__.py b/scripts/kernel_tuning/discovery/__init__.py index 2176a78..a5238ea 100644 --- a/scripts/kernel_tuning/discovery/__init__.py +++ b/scripts/kernel_tuning/discovery/__init__.py @@ -1,4 +1,4 @@ -"""Model-scoped task discovery; Torch is imported only in the capture runtime.""" +"""Shared registry capture, replay and driver utilities; the historical module paths are retained.""" from .contracts import ModelCapture diff --git a/scripts/kernel_tuning/discovery/contracts.py b/scripts/kernel_tuning/discovery/contracts.py index 0030a36..ce0b5e3 100644 --- a/scripts/kernel_tuning/discovery/contracts.py +++ b/scripts/kernel_tuning/discovery/contracts.py @@ -2,7 +2,6 @@ from __future__ import annotations -import re from dataclasses import dataclass from pathlib import Path from typing import Any @@ -10,39 +9,10 @@ from ..contracts import ContractError, digest, read_json, relative_path, tree_hashes SCHEMA = "embodiinfer-model-capture-v1" -DISPOSITIONS = {"ready", "excluded_quantized", "metadata_only", "host_only", "blocked"} -RNG_OVERLOADS = {"aten.randn.default", "aten.randn.generator"} +DISPOSITIONS = {"ready", "blocked"} REPLAY_SCHEMA = "embodiinfer-replay-v1" -def validate_rng(contract: dict[str, Any] | None, case: dict[str, Any], files: dict[str, str]) -> None: - """Bind both RNG-state transitions to immutable files for the declared devices.""" - if contract is None: - if "rng" in case: - raise ContractError("Undeclared RNG state in workload") - return - if not isinstance(contract, dict): - raise ContractError("Invalid default-generator RNG contract") - devices = contract.get("devices") - if ( - set(contract) != {"kind", "devices"} - or contract["kind"] != "torch_default" - or not isinstance(devices, list) - or not 1 <= len(devices) <= 2 - or devices[0] != "cpu" - or (len(devices) == 2 and not re.fullmatch(r"cuda:\d+", str(devices[1]))) - ): - raise ContractError("Invalid default-generator RNG contract") - states = case.get("rng", {}) - if not isinstance(states, dict) or set(states) != {"before", "after"}: - raise ContractError("RNG replay requires before and after states") - for phase in states.values(): - if not isinstance(phase, dict) or set(phase) != set(devices): - raise ContractError("RNG state devices differ from the contract") - if any(relative_path(path) not in files for path in phase.values()): - raise ContractError("RNG fixture missing from the capture/task") - - def validate_replay( replay: dict[str, Any], definition: dict[str, Any], traces: list[dict[str, Any]], hashes: dict[str, str] ) -> None: @@ -61,27 +31,19 @@ def validate_replay( and schema != REPLAY_SCHEMA ): raise ContractError("float64 requires the explicit replay definition schema") - if replay.get("rng") is not None and replay.get("operator") not in RNG_OVERLOADS: - raise ContractError("Unsupported RNG replay operator") + 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") - undefined = replay.get("undefined_outputs", []) - if ( - not isinstance(undefined, list) - or (undefined and not str(replay.get("operator")).startswith("aten._scaled_dot_product_")) - or len(set(undefined)) != len(undefined) - or len(undefined) >= len(definition["outputs"]) - or any(type(index) is not int or not 0 <= index < len(definition["outputs"]) for index in undefined) - ): - raise ContractError("Only unused attention dropout state may be excluded from output checks") for trace in traces: work = trace["workload"] case = replay["cases"][work["uuid"]] - validate_rng(replay.get("rng"), case, hashes) + 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]): @@ -113,7 +75,7 @@ def validate_replay( @dataclass(frozen=True) class ModelCapture: - """A completed preparation run bound to one actual policy and all observed cases.""" + """A registry capture bound to one policy and every case of its selected operators.""" root: Path manifest: dict[str, Any] @@ -133,6 +95,10 @@ def load(cls, root: Path, model: str | None = None, *, require_ready: bool = Tru 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") @@ -151,6 +117,12 @@ def load(cls, root: Path, model: str | None = None, *, require_ready: bool = Tru 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): @@ -163,18 +135,15 @@ def load(cls, root: Path, model: str | None = None, *, require_ready: bool = Tru 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 op["status"] == "ready": - rng = op["identity"].get("rng") - if rng is not None and op["name"] not in RNG_OVERLOADS: - raise ContractError("Unsupported RNG capture operator") - validate_rng(rng, case, hashes) + 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 non-quantized operator preparation gaps.""" + """Reject incomplete runs and any selected-operator capture gaps.""" gaps = list(self.manifest.get("errors", [])) gaps += [ f"{op['name']}: {op.get('reason', 'unsupported')}" diff --git a/scripts/kernel_tuning/discovery/drive.py b/scripts/kernel_tuning/discovery/drive.py index dfef885..d9cc9a6 100644 --- a/scripts/kernel_tuning/discovery/drive.py +++ b/scripts/kernel_tuning/discovery/drive.py @@ -1,16 +1,16 @@ """Run recorded inputs through one model's production inference path. -This is the preparation driver for model operator discovery:: +This shared inference driver serves profiling, registry capture and verification:: - python -m scripts.kernel_tuning capture --model POLICY --output CAPTURE -- \\ - -m scripts.kernel_tuning.discovery.drive --policy POLICY --inputs INPUTS.pt [serving options] + 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 capture records whatever this -inference executes; the driver itself chooses no operators or code paths. +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 diff --git a/scripts/kernel_tuning/discovery/numerics.py b/scripts/kernel_tuning/discovery/numerics.py index 5d76ec0..a9d3dba 100644 --- a/scripts/kernel_tuning/discovery/numerics.py +++ b/scripts/kernel_tuning/discovery/numerics.py @@ -12,31 +12,6 @@ VERSION = "baseline-fp64-v1" SEEDS = (0, 1, 2) PROFILES = (("recorded", 0), ("random", 1), ("random", 2), ("zeros", 0), ("cancellation", 0)) -# Explicit semantic adapters, not substring matching on arbitrary operator names. -_ATEN = { - "mm", - "bmm", - "matmul", - "mv", - "dot", - "addmm", - "baddbmm", - "linear", - "sum", - "mean", - "_softmax", - "softmax", - "_log_softmax", - "log_softmax", - "native_layer_norm", - "layer_norm", - "rms_norm", - "scaled_dot_product_attention", -} -_BACKENDS = { - "embodiinfer.backend.triton.split_kv_attention.split_kv_attention": "split_kv_attention", - "embodiinfer.backend.triton.segmented_attention.segmented_attention": "segmented_attention", -} def implementation_digest() -> str: @@ -66,10 +41,7 @@ def policy(identity: dict[str, Any]) -> str | None: # Registry operators define their mathematics by a Torch reference, # which also evaluates it in FP64 (embodiinfer.backend.torch.ops). return "registry" - name = identity["name"] - if identity["kind"] == "aten" and name.split(".")[1] in _ATEN: - return "aten" - return _BACKENDS.get(name) + return None def profile_key(profile: str, seed: int) -> str: diff --git a/scripts/kernel_tuning/discovery/recorder.py b/scripts/kernel_tuning/discovery/recorder.py deleted file mode 100644 index 7ae88d9..0000000 --- a/scripts/kernel_tuning/discovery/recorder.py +++ /dev/null @@ -1,792 +0,0 @@ -"""Observe one model's existing execution, without replacing its numerical path.""" - -from __future__ import annotations - -import functools -import hashlib -import importlib -import inspect -import runpy -import sys -import threading -from collections import Counter -from collections.abc import Iterator -from contextlib import contextmanager -from dataclasses import fields, is_dataclass -from pathlib import Path -from typing import Any - -import torch -from torch.utils._python_dispatch import TorchDispatchMode -from torch.utils._pytree import tree_leaves - -from ..artifacts import REPOSITORY, atomic_json -from ..contracts import ContractError, digest, tree_hashes -from ..generate import _git_revision, vendor_sources -from .contracts import SCHEMA -from .numerics import policy -from .rng import contract_for, requires_rng_adapter, undefined_outputs -from .rng import snapshot as rng_snapshot -from .tensors import describe_inputs, extent, flatten, storage_key, tensor_outputs, tensor_spec - -_VIEWS = { - "alias", - "detach", - "view", - "_unsafe_view", - "transpose", - "permute", - "as_strided", - "expand", - "squeeze", - "unsqueeze", - "slice", - "select", - "split", - "split_with_sizes", - "unbind", - "t", - "as_strided_", - "transpose_", - "squeeze_", - "unsqueeze_", - "detach_", - "resize_", - "resize_as_", - "set_", -} -_ALLOCATIONS = {"empty", "empty_like", "empty_strided", "new_empty", "new_empty_strided"} -_HOST = {"item", "_local_scalar_dense", "is_nonzero"} -# Weight loading, casting and device placement are not inference arithmetic. -_COPIES = {"_to_copy", "copy_", "clone", "contiguous", "lift_fresh", "_copy_from", "_copy_from_and_resize"} -# Production inference enters through the model-agnostic engine contract. Every -# operator running below one of these calls, on any thread, belongs to the model. -_ENTRIES = ( - ("embodiinfer.engine.core", "EngineCore", ("__init__", "execute", "execute_pipelined", "_pipeline_step")), - ( - "embodiinfer.engine.rollout.generation_backend", - "GenerationBackend", - ("generate", "generate_with_logprob", "sample_group", "best_of_n"), - ), -) -_QUANTIZED_ATEN = { - "_scaled_mm", - "_weight_int8pack_mm", - "_weight_int4pack_mm", - "_convert_weight_to_int4pack", - "quantize_per_tensor", - "quantize_per_channel", - "dequantize", -} -# These contracts describe existing public kernels, not an operator discovery whitelist. -_PURE_KERNELS = { - "rms_norm", - "add_rms_norm", - "ada_rms_norm", - "gated_residual", - "swiglu", - "gated_gelu", - "rotate_qk", - "rotate_half_rope", - "segmented_attention", - "split_kv_attention", - "flash_prefill_attention", - "flash_prefill_attention_graph", - "gqa_decode_graph", -} -_KERNEL_WRITES = { - "fused_rope_cache": ("key_cache", "value_cache"), - "fused_rope_cache_graph": ("key_cache", "value_cache"), - "fused_prefill_rope_cache_graph": ("key_cache", "value_cache"), - "fused_greedy": ("workspace",), - "fused_lm_head_greedy": ("workspace",), - "mark_penalized": ("workspace",), -} - - -def _portable(value: Any) -> Any: - if value is None or isinstance(value, (str, int, float, bool)): - return value - if isinstance(value, (torch.dtype, torch.device, Path)): - return str(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] - if is_dataclass(value) and not isinstance(value, type): - return {field.name: _portable(getattr(value, field.name)) for field in fields(value)} - return {"type": f"{type(value).__module__}.{type(value).__qualname__}"} - - -class _Observer(TorchDispatchMode): - def __init__(self, session: CaptureSession): - super().__init__() - self.session = session - - def __torch_dispatch__(self, func: Any, types: Any, args: tuple = (), kwargs: dict | None = None) -> Any: - session = self.session - kwargs = kwargs or {} - if session.suppressed: - return func(*args, **kwargs) - if not session.depth: - session.outside(func, args, kwargs) - return func(*args, **kwargs) - name = str(func) - short = name.split(".")[1] if name.startswith("aten.") else name - reason = None - with session.silence(): - rng = contract_for(name, kwargs) - if not name.startswith("aten."): - reason = "External dispatcher operator needs a reference adapter" - if rng is None and requires_rng_adapter(func, args, kwargs): - reason = "Random/nondeterministic operator needs an RNG replay adapter" - writes = [] - for index, argument in enumerate(func._schema.arguments): - if argument.alias_info is not None and argument.alias_info.is_write: - value = args[index] if index < len(args) else kwargs.get(argument.name) - writes.append(value) - quantized = name.split(".")[0] in {"quantized", "quantized_decomposed"} or short in _QUANTIZED_ATEN - origin = {"kind": "aten", "name": name, "schema": str(func._schema)} - if undefined := undefined_outputs(func, args, kwargs): - origin["undefined_outputs"] = undefined - return session.call( - origin, - func, - args, - kwargs, - writes=writes, - reason=reason, - rng=rng if rng is not None and rng["devices"][-1].split(":")[0] == session.device_type else None, - classification="excluded_quantized" - if quantized - else "metadata_only" - if short in _VIEWS | _ALLOCATIONS - else "host_only" - if short in _HOST - else None, - ) - - -class CaptureSession: - """Collect one bound model's calls; failures remain explicit coverage gaps. - - Use ``observe`` around existing preparation execution for custom drivers. The - CLI additionally binds a target constructed by the public policy factory. - ``device_type='cpu'`` is useful for deterministic observer unit tests only. - """ - - def __init__( - self, model: str, output: Path, *, fixture_bytes: int = 64 * 1024 * 1024, device_type: str = "cuda" - ) -> None: - if not model or fixture_bytes < 0: - raise ContractError("Select one model and a nonnegative fixture byte budget") - self.model, self.root = model, output.resolve() - if self.root.exists() and any(self.root.iterdir()): - raise ContractError("Capture destination must be empty; use a new directory") - self.root.mkdir(parents=True, exist_ok=True) - self.fixture_bytes, self.used_bytes, self.device_type = fixture_bytes, 0, device_type - self.target: dict[str, Any] = {} - self.instance: Any = None - # Dispatch modes are thread-local, so inference scope is tracked per thread. - self._local = threading.local() - self.entries = 0 - self.bypassed: Counter = Counter() - self._weights: tuple[int, frozenset[int]] = (0, frozenset()) - self.operators: dict[str, dict[str, Any]] = {} - self.errors: list[str] = [] - self.replacements: list[tuple[Any, str, Any, bool]] = [] - self._fixtures: dict[str, list[dict[str, Any]]] = {} - self.graphs: dict[int, Counter] = {} - self.graph_stack: list[Counter] = [] - self._numerics_bound = False - self._set_numerics() - - def _set_numerics(self) -> None: - 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(), - } - - def _check_numerics(self) -> None: - previous = (self.precision, self.flags) - self._set_numerics() - if self._numerics_bound and previous != (self.precision, self.flags): - self.errors.append("Numerical settings changed during model preparation") - self._numerics_bound = True - - @property - def depth(self) -> int: - return getattr(self._local, "depth", 0) - - @depth.setter - def depth(self, value: int) -> None: - self._local.depth = value - - @property - def suppressed(self) -> int: - return getattr(self._local, "suppressed", 0) - - @suppressed.setter - def suppressed(self, value: int) -> None: - self._local.suppressed = value - - @contextmanager - def inference(self) -> Iterator[None]: - """Observe every operator below one inference call on the current thread.""" - observer = None - if not getattr(self._local, "observing", False): - observer = _Observer(self) - observer.__enter__() - self._local.observing = True - if not self.depth: - self.entries += 1 - self.depth += 1 - try: - yield - finally: - self.depth -= 1 - if observer is not None: - self._local.observing = False - observer.__exit__(None, None, None) - - @contextmanager - def watch(self) -> Iterator[None]: - """Watch the driver thread outside inference to reject bypassed model arithmetic.""" - self._local.observing = True - try: - with _Observer(self): - yield - finally: - self._local.observing = False - - def _entry(self, function: Any) -> Any: - @functools.wraps(function) - def wrapper(*args: Any, **kwargs: Any) -> Any: - with self.inference(): - return function(*args, **kwargs) - - return wrapper - - def _weight_storages(self) -> frozenset[int]: - tensors = [*self.instance.parameters(), *self.instance.buffers()] - if not tensors: - return frozenset() - probe = tensors[0].untyped_storage().data_ptr() - # Moving or casting the model replaces storages; refresh only then. - if probe != self._weights[0]: - self._weights = (probe, frozenset(t.untyped_storage().data_ptr() for t in tensors) - {0}) - return self._weights[1] - - def outside(self, func: Any, args: tuple, kwargs: dict) -> None: - """Record model-weight arithmetic that bypassed every inference entry.""" - if not isinstance(self.instance, torch.nn.Module): - return - name = str(func) - short = name.split(".")[1] if name.startswith("aten.") else name - if short in _VIEWS | _ALLOCATIONS | _HOST | _COPIES: - return - with self.silence(): - weights = self._weight_storages() - if any( - isinstance(t, torch.Tensor) and t.untyped_storage().data_ptr() in weights - for t in tree_leaves((args, kwargs)) - ): - self.bypassed[name] += 1 - - @contextmanager - def silence(self) -> Iterator[None]: - """Prevent instrumentation and atomic/quantized internals from becoming tasks.""" - self.suppressed += 1 - try: - yield - finally: - self.suppressed -= 1 - - def bind(self, instance: Any, config: dict[str, Any]) -> None: - """Bind to an actual object rather than trusting an unverified model label.""" - if self.instance is not None and self.instance is not instance: - raise ContractError("Prepare exactly one target policy instance per capture") - modules = instance.modules() if isinstance(instance, torch.nn.Module) else (instance,) - if config.get("compile_backend", "none") != "none" or any( - hasattr(module, "_orig_mod") for module in modules - ): - raise ContractError("Model discovery requires explicit uncompiled preparation") - self.instance = instance - self.target = { - "model": self.model, - "type": f"{type(instance).__module__}.{type(instance).__qualname__}", - "config": _portable(config), - } - - @contextmanager - def observe(self, instance: Any, **config: Any) -> Iterator[None]: - """Observe the existing preparation call without running it a second time.""" - self.bind(instance, config) - with self.inference(): - yield - - def _snapshot( - self, key: str, tensors: list[torch.Tensor], specs: list[dict[str, Any]], *, real_inputs: bool = False - ) -> list[dict[str, Any]]: - if key in self._fixtures: - return self._fixtures[key] - if torch.cuda.is_available() and torch.cuda.is_current_stream_capturing(): - if any(not tensor.is_floating_point() for tensor in tensors): - raise ContractError("Valid non-floating fixtures must be observed during eager warmup") - return 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) - byte_count = size * tensor.element_size() - required = real_inputs or not tensor.is_floating_point() - if required and byte_count > self.fixture_bytes: - raise ContractError( - "Required input fixture exceeds the capture byte budget; increase --fixture-bytes" - ) - # Calibrated arithmetic needs real weights too. Other floating tasks - # retain small fixtures when space permits; indices are never guessed. - save = required or byte_count <= 1024 * 1024 - if save and not required and self.used_bytes + byte_count > self.fixture_bytes: - save = False - if save: - raw = ( - tensor.detach() - .as_strided((size,), (1,), 0) - .contiguous() - .view(torch.uint8) - .cpu() - .numpy() - .tobytes() - ) - path = f"fixtures/{hashlib.sha256(raw).hexdigest()}.bin" - file = self.root / path - if not file.exists(): - if self.used_bytes + byte_count > self.fixture_bytes: - raise ContractError( - "Required real input fixture exceeds the capture byte budget; increase --fixture-bytes" - ) - file.parent.mkdir(exist_ok=True) - file.write_bytes(raw) - self.used_bytes += len(raw) - for i in indices: - specs[i]["fixture"] = path - self._fixtures[key] = specs - return specs - - def _rng_snapshot(self, devices: list[str]) -> dict[str, str]: - if torch.cuda.is_available() and torch.cuda.is_current_stream_capturing(): - raise ContractError("RNG fixtures require eager preparation outside CUDA Graph capture") - result = {} - for device, state in rng_snapshot(devices).items(): - raw = bytes(state.tolist()) - path = f"fixtures/{hashlib.sha256(raw).hexdigest()}.bin" - file = self.root / path - if not file.exists(): - if self.used_bytes + len(raw) > self.fixture_bytes: - raise ContractError("Required RNG fixture exceeds the capture byte budget") - file.parent.mkdir(exist_ok=True) - file.write_bytes(raw) - self.used_bytes += len(raw) - result[device] = path - return result - - def _add(self, identity: dict[str, Any], case: dict[str, Any], status: str, reason: str | None) -> None: - op_id = digest(identity)[:20] - op = self.operators.setdefault( - op_id, - { - "id": op_id, - "name": identity["name"], - "identity": identity, - "status": status, - "reason": reason, - "workloads": [], - }, - ) - if status == "blocked": - op.update(status=status, reason=reason) - case_id = digest({k: v for k, v in case.items() if k != "count"})[:20] - existing = next((w for w in op["workloads"] if w["id"] == case_id), None) - if existing is None: - existing = {**case, "id": case_id, "count": 0} - op["workloads"].append(existing) - if self.graph_stack: - self.graph_stack[-1][op_id, case_id] += 1 - # Retain capture-only cases too; replay traffic is counted separately. - existing["count"] += 1 - - def call( - self, - origin: dict[str, Any], - function: Any, - args: tuple, - kwargs: dict, - *, - writes: list[Any] | None = None, - reason: str | None = None, - rng: dict[str, Any] | None = None, - classification: str | None = None, - ) -> Any: - """Record before/after metadata while invoking the original callable exactly once.""" - with self.silence(): - self._check_numerics() - tensors: list[torch.Tensor] = [] - identity = dict(origin) - case: dict[str, Any] = {"inputs": [], "outputs": []} - status = classification or "ready" - explicit = reason - try: - recipe = flatten((args, kwargs), tensors) - if any(t.is_quantized for t in tensors): - classification = status = "excluded_quantized" - specs = describe_inputs(tensors) - identity.update( - arguments=recipe, - input_dtypes=[s["dtype"] for s in specs], - input_ranks=[len(s["shape"]) for s in specs], - ) - case["inputs"] = specs - if not any(t.device.type == self.device_type for t in tensors): - status = classification or "host_only" - mutated: list[torch.Tensor] = [] - for value in writes or []: - flatten(value, mutated) - identity["mutates"] = [ - i for i, tensor in enumerate(tensors) if any(tensor is other for other in mutated) - ] - if reason: - raise ContractError(reason) - if status == "ready": - case["inputs"] = self._snapshot( - digest([identity, specs]), tensors, specs, real_inputs=policy(identity) is not None - ) - if rng is not None: - identity["rng"] = rng - case["rng"] = {"before": self._rng_snapshot(rng["devices"])} - except (ContractError, RuntimeError, TypeError, ValueError) as exc: - status, reason = ( - ("excluded_quantized" if classification == "excluded_quantized" else "blocked"), - str(exc), - ) - if explicit is None and classification in {"metadata_only", "host_only"}: - status = classification # Not a task, so no replay recipe is needed. - try: - result = function(*args, **kwargs) - except Exception as exc: - self.errors.append(f"{origin['name']} failed during preparation: {type(exc).__name__}: {exc}") - raise - if ( - status == "blocked" - and explicit is None - and rng is None - and not any( - isinstance(t, torch.Tensor) and t.device.type == self.device_type - for t in tree_leaves((args, kwargs, result)) - ) - ): - # Host preprocessing with an unrepresentable argument is not a device task. - status = "host_only" - try: - if rng is not None and status != "blocked": - case["rng"]["after"] = self._rng_snapshot(rng["devices"]) - if ( - status == "host_only" - and classification is None - and any(t.device.type == self.device_type for t in tensor_outputs(result)) - ): - status = "ready" - # Host inputs were not snapshotted before the call; capture them now. - if identity["mutates"]: - raise ContractError("Host-to-device calls mutating inputs need a dedicated adapter") - case["inputs"] = self._snapshot(digest([identity, specs]), tensors, specs) - if status not in {"metadata_only", "host_only", "excluded_quantized", "blocked"}: - outputs = tensor_outputs(result) - identity["returns"] = flatten(result, []) - # Mutated tensors are part of the task's observable outputs too. - outputs += [tensors[i] for i in identity["mutates"]] - if not outputs: - raise ContractError("Calls without tensor outputs require a dedicated task adapter") - case["outputs"] = [tensor_spec(t) for t in outputs] - identity["output_dtypes"] = [s["dtype"] for s in case["outputs"]] - identity["output_ranks"] = [len(s["shape"]) for s in case["outputs"]] - case["output_aliases"] = [ - [i for i, tensor in enumerate(tensors) if storage_key(tensor) == storage_key(output)] - for output in outputs - ] - case["output_groups"] = [ - [i for i, other in enumerate(outputs) if storage_key(output) == storage_key(other)] - for output in outputs - ] - else: - case["outputs"] = [tensor_spec(t) for t in tensor_outputs(result)] - except (ContractError, RuntimeError, TypeError, ValueError) as exc: - if rng is not None or status not in {"metadata_only", "host_only", "excluded_quantized"}: - status, reason = "blocked", str(exc) - self._add(identity, case, status, reason) - return result - - def _replace(self, holder: Any, name: str, replacement: Any) -> None: - self.replacements.append((holder, name, getattr(holder, name), name in vars(holder))) - setattr(holder, name, replacement) - - def _replace_function(self, original: Any, wrapper: Any) -> None: - for module in list(sys.modules.values()): - if module is None: - continue - for name, value in list(vars(module).items()): - if value is original: - self._replace(module, name, wrapper) - - def _atomic(self, function: Any, *, quantized: bool = False) -> Any: - source_digests: dict[str, str] | None = None - - @functools.wraps(function) - def wrapper(*args: Any, **kwargs: Any) -> Any: - nonlocal source_digests - if not self.depth or self.suppressed: - return function(*args, **kwargs) - origin = { - "kind": "backend", - "name": f"{function.__module__}.{function.__name__}", - "module": function.__module__, - "entry": function.__name__, - } - if quantized: - # A module instance is provenance, not a replayable tensor input. - call_args = args[1:] if args and isinstance(args[0], torch.nn.Module) else args - return self.call( - origin, - lambda *a, **kw: function(*args, **kwargs), - call_args, - kwargs, - classification="excluded_quantized", - ) - bound = inspect.signature(function).bind(*args, **kwargs) - bound.apply_defaults() - source = Path(inspect.getfile(function)) - origin["source_digest"] = hashlib.sha256(source.read_text(encoding="utf-8").encode()).hexdigest() - names = _KERNEL_WRITES.get(function.__name__, ()) - reason = ( - None - if function.__name__ in _PURE_KERNELS or names - else "Backend entry requires an explicit mutation/reference contract" - ) - try: - if source_digests is None: - source_digests = { - path: source_digest - for path, (_, source_digest) in vendor_sources(function.__module__).items() - } - origin["source_digests"] = source_digests - except ContractError as exc: - reason = str(exc) - return self.call( - origin, function, args, kwargs, writes=[bound.arguments[n] for n in names], reason=reason - ) - - return wrapper - - def install(self) -> None: - """Install process-local quantization, backend, factory, and graph observation hooks.""" - from embodiinfer.models import linear - from embodiinfer.policies import factory - - if self.model not in factory.available_policies(): - raise ContractError(f"Unknown policy {self.model!r}; choose from {factory.available_policies()}") - self._replace( - linear.QuantizedLinear, "forward", self._atomic(linear.QuantizedLinear.forward, quantized=True) - ) - for name in ("quantize_linear", "pack_quantized_linears"): - original = getattr(linear, name) - self._replace_function(original, self._atomic(original, quantized=True)) - for method in ("fp8", "int8", "nvfp4"): - module = importlib.import_module(f"embodiinfer.backend.torch.{method}") - for name, function in list(vars(module).items()): - if ( - inspect.isfunction(function) - and function.__module__ == module.__name__ - and not name.startswith("_") - ): - self._replace_function(function, self._atomic(function, quantized=True)) - for path in sorted((REPOSITORY / "embodiinfer/backend/triton").glob("*.py")): - if path.stem in {"__init__", "capability"}: - continue - module = importlib.import_module(f"embodiinfer.backend.triton.{path.stem}") - for name, function in list(vars(module).items()): - if ( - inspect.isfunction(function) - and function.__module__ == module.__name__ - and not name.startswith(("_", "supports_", "warmup_")) - and not name.endswith(("_available", "_capability")) - ): - self._replace_function(function, self._atomic(function, quantized=path.stem == "fp8")) - original_factory = factory.make_policy - - @functools.wraps(original_factory) - def construct(name: str, **kwargs: Any) -> Any: - if name != self.model: - raise ContractError(f"Preparation constructed {name!r}, expected {self.model!r}") - if kwargs.get("compile_backend", "none") != "none": - raise ContractError("Disable torch.compile explicitly before model discovery") - policy = original_factory(name, **kwargs) - self.bind(policy, kwargs) - # Request preprocessing and serving adapters belong to end-to-end inference. - if callable(getattr(policy, "collate", None)): - self._replace(policy, "collate", self._entry(policy.collate)) - build = getattr(policy, "build_serving_adapter", None) - if callable(build): - - @functools.wraps(build) - def serving(*args: Any, **kw: Any) -> Any: - with self.inference(): - adapter = build(*args, **kw) - for method in ("infer", "infer_batch"): - if callable(getattr(adapter, method, None)): - self._replace(adapter, method, self._entry(getattr(adapter, method))) - return adapter - - self._replace(policy, "build_serving_adapter", serving) - return policy - - self._replace_function(original_factory, construct) - for module, owner, methods in _ENTRIES: - holder = getattr(importlib.import_module(module), owner) - for method in methods: - self._replace(holder, method, self._entry(getattr(holder, method))) - - def compile_guard(*args: Any, **kwargs: Any) -> Any: - raise ContractError("torch.compile is unsupported during model preparation; disable compilation") - - self._replace_function(torch.compile, compile_guard) - try: - from triton.runtime import JITFunction - except ImportError: - pass - else: - original_run = JITFunction.run - - def triton_run(kernel: Any, *args: Any, **kwargs: Any) -> Any: - if not self.depth or self.suppressed: - return original_run(kernel, *args, **kwargs) - origin = {"kind": "unregistered_triton", "name": f"{kernel.__module__}.{kernel.__name__}"} - return self.call( - origin, - lambda *a, **kw: original_run(kernel, *a, **kw), - args, - kwargs, - reason="Raw Triton entry needs a public callable and replay contract", - ) - - self._replace(JITFunction, "run", triton_run) - graph = torch.cuda.CUDAGraph - begin, end, replay = graph.capture_begin, graph.capture_end, graph.replay - - def capture_begin(instance: Any, *args: Any, **kwargs: Any) -> Any: - result = begin(instance, *args, **kwargs) - self.graph_stack.append(Counter()) - return result - - def capture_end(instance: Any, *args: Any, **kwargs: Any) -> Any: - result = end(instance, *args, **kwargs) - self.graphs[id(instance)] = self.graph_stack.pop() - return result - - def graph_replay(instance: Any, *args: Any, **kwargs: Any) -> Any: - if self.depth and not self.suppressed and id(instance) not in self.graphs: - self.errors.append("Model replayed a CUDA graph not observed during this preparation") - result = replay(instance, *args, **kwargs) - for (op_id, case_id), count in self.graphs.get(id(instance), {}).items(): - case = next(w for w in self.operators[op_id]["workloads"] if w["id"] == case_id) - case["count"] += count - return result - - for name, fn in ( - ("capture_begin", capture_begin), - ("capture_end", capture_end), - ("replay", graph_replay), - ): - self._replace(graph, name, fn) - - def close(self, *, complete: bool) -> None: - """Restore all hooks and publish a hashed success or failure capture.""" - for holder, name, original, owned in reversed(self.replacements): - if owned: - setattr(holder, name, original) - else: - delattr(holder, name) - self.replacements.clear() - if not self.target: - self.errors.append("No target policy was bound through make_policy or observe") - elif not self.entries: - self.errors.append( - "No end-to-end inference ran; drive the model through EngineCore, " - "GenerationBackend or its serving adapter" - ) - for name, count in sorted(self.bypassed.items()): - self.errors.append( - f"{name} read model weights outside an inference entry ({count} calls); " - "drive the model through EngineCore, GenerationBackend or its serving adapter" - ) - operators = list(self.operators.values()) - atomic_json(self.root / "operators.json", operators) - coverage = dict(Counter(op["status"] for op in operators)) - atomic_json(self.root / "coverage.json", {"operators": coverage, "errors": self.errors}) - atomic_json( - self.root / "manifest.json", - { - "schema": SCHEMA, - "target": self.target, - "complete": complete, - "errors": self.errors, - "repository_revision": _git_revision(), - "torch_version": str(torch.__version__), - "precision": self.precision, - "flags": self.flags, - "device_type": self.device_type, - "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, *, fixture_bytes: int = 64 * 1024 * 1024 -) -> dict[str, Any]: - """Execute an existing model preparation script with model-scoped observation.""" - if not command: - raise ContractError("Supply the preparation script after --") - session = CaptureSession(model, output, fixture_bytes=fixture_bytes) - argv, path = sys.argv, sys.path[:] - complete = False - try: - session.install() - with session.watch(): - 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 - session.close(complete=complete) - from .contracts import ModelCapture - - capture = ModelCapture.load(output, model) - return {"model": model, "identity": capture.identity, "operators": len(capture.operators)} diff --git a/scripts/kernel_tuning/discovery/rng.py b/scripts/kernel_tuning/discovery/rng.py deleted file mode 100644 index f0488d3..0000000 --- a/scripts/kernel_tuning/discovery/rng.py +++ /dev/null @@ -1,133 +0,0 @@ -"""Argument-aware RNG classification and exact default-generator randn replay.""" - -from __future__ import annotations - -from collections.abc import Iterator -from contextlib import contextmanager -from pathlib import Path -from typing import Any - -import torch - -from ..contracts import ContractError -from .contracts import RNG_OVERLOADS - - -def _zero_dropout_attention(function: Any, args: tuple[Any, ...], kwargs: dict[str, Any]) -> bool: - # The composite call decomposes to a backend overload (flash, efficient, - # cuDNN, CPU) before Python dispatch; none draws random state without dropout. - name = str(function) - if name != "aten.scaled_dot_product_attention.default" and not name.startswith( - "aten._scaled_dot_product_" - ): - return False - for index, argument in enumerate(function._schema.arguments): - if argument.name == "dropout_p": - dropout = kwargs.get(argument.name, args[index] if index < len(args) else argument.default_value) - return type(dropout) in (int, float) and dropout == 0 - return False - - -def requires_rng_adapter(function: Any, args: tuple[Any, ...], kwargs: dict[str, Any]) -> bool: - """Treat seeded tags as conditional for SDPA with explicitly zero/default dropout.""" - tags = {tag for tag in getattr(function, "tags", ()) if "nondeterministic" in str(tag)} - if _zero_dropout_attention(function, args, kwargs): - tags.discard(torch.Tag.nondeterministic_seeded) - return bool(tags) - - -def undefined_outputs(function: Any, args: tuple[Any, ...], kwargs: dict[str, Any]) -> list[int]: - """Outputs whose bytes are unspecified: dropout RNG state left uninitialized without dropout.""" - if not _zero_dropout_attention(function, args, kwargs): - return [] - return [ - i - for i, value in enumerate(function._schema.returns) - if value.name in {"philox_seed", "philox_offset"} - ] - - -def contract_for(name: str, kwargs: dict[str, Any]) -> dict[str, Any] | None: - """Recognize randn with the default CPU/CUDA generator; other RNG calls stay gaps.""" - if name not in RNG_OVERLOADS or kwargs.get("generator") is not None: - return None - device = torch.device(kwargs.get("device") or torch.get_default_device()) - if device.type not in {"cpu", "cuda"}: - return None - devices = ["cpu"] - if device.type == "cuda": - index = torch.cuda.current_device() if device.index is None else device.index - devices.append(f"cuda:{index}") - return {"kind": "torch_default", "devices": devices} - - -def snapshot(devices: list[str]) -> dict[str, torch.Tensor]: - """Read opaque generator states without consuming random numbers.""" - return { - device: (torch.get_rng_state() if device == "cpu" else torch.cuda.get_rng_state(device)).clone() - for device in devices - } - - -def restore(states: dict[str, torch.Tensor]) -> None: - """Restore every declared default generator from its opaque byte state.""" - for device, state in states.items(): - if device == "cpu": - torch.set_rng_state(state) - else: - torch.cuda.set_rng_state(state, device) - - -@contextmanager -def preserve(devices: list[str]) -> Iterator[None]: - """Restore the caller's RNG state on both successful and failed evaluation.""" - saved = snapshot(devices) - try: - yield - finally: - restore(saved) - - -class RngReplay: - """Replay the captured transition at seed 0 and independent states at other seeds.""" - - def __init__(self, contract: dict[str, Any], case: dict[str, Any], root: Path, seed: int) -> None: - self.devices = contract["devices"] - states = { - phase: { - device: torch.tensor(list((root / path).read_bytes()), dtype=torch.uint8) - for device, path in paths.items() - } - for phase, paths in case["rng"].items() - } - current = snapshot(self.devices) - if any( - state.shape != current[device].shape - for phase in states.values() - for device, state in phase.items() - ): - raise ContractError("RNG fixture state size differs from the current Torch runtime") - self.before = ( - states["before"] - if seed == 0 - else { - device: torch.Generator(device=device).manual_seed(seed).get_state() - for device in self.devices - } - ) - self.captured_after = states["after"] if seed == 0 else None - - def reset(self) -> None: - """Restore the selected workload state before a call, outside timed events.""" - restore(self.before) - - def check_reference(self, actual: dict[str, torch.Tensor]) -> None: - """Require the frozen overload to reproduce its observed seed-0 transition.""" - if self.captured_after is not None: - self.check(actual, self.captured_after) - - @staticmethod - def check(actual: dict[str, torch.Tensor], expected: dict[str, torch.Tensor]) -> None: - """Reject changed generator advancement even when output values are correct.""" - if set(actual) != set(expected) or any(not torch.equal(actual[d], expected[d]) for d in expected): - raise ContractError("RNG state after the call differs from the reference") diff --git a/scripts/kernel_tuning/discovery/tasks.py b/scripts/kernel_tuning/discovery/tasks.py index 3465e0e..6e05da2 100644 --- a/scripts/kernel_tuning/discovery/tasks.py +++ b/scripts/kernel_tuning/discovery/tasks.py @@ -2,13 +2,11 @@ from __future__ import annotations -import hashlib import json -import re from pathlib import Path from typing import Any -from ..artifacts import REPOSITORY, atomic_json +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 @@ -47,8 +45,6 @@ def _expression(recipe: dict[str, Any]) -> str: 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()) + "}" - if kind == "greedy_workspace": - return "GreedyWorkspace(**" + _expression({"dict": value}) + ")" raise ContractError(f"Unrecognized argument recipe {kind}") @@ -69,45 +65,21 @@ def _return_expressions( def _sources(op: dict[str, Any]) -> tuple[str, dict[str, str]]: identity = op["identity"] files = {} - if identity["kind"] == "aten": - if not re.fullmatch(r"aten\.\w+\.\w+", identity["name"]): - raise ContractError("Invalid captured ATen overload") - imports, callable_name = "import torch\n", f"torch.ops.{identity['name']}" - elif identity["kind"] == "backend": - module, entry = identity["module"], identity["entry"] - if not re.fullmatch(r"embodiinfer\.backend\.\w+(?:\.\w+)*", module) or not entry.isidentifier(): - raise ContractError("Unsupported backend source identity") - source = REPOSITORY.joinpath(*module.split(".")).with_suffix(".py") - if ( - hashlib.sha256(source.read_text(encoding="utf-8").encode()).hexdigest() - != identity["source_digest"] - ): - raise ContractError(f"Backend source changed since capture: {identity['name']}") - vendored = vendor_sources(module) - if {path: source_digest for path, (_, source_digest) in vendored.items()} != identity[ - "source_digests" - ]: - raise ContractError(f"Backend dependency changed since capture: {identity['name']}") - files = {path: text for path, (text, _) in vendored.items()} - imports, callable_name = f"import torch\nfrom .vendor.{module} import {entry}\n", entry - if module == "embodiinfer.backend.triton.sampling": - imports += f"from .vendor.{module} import GreedyWorkspace\n" - elif identity["kind"] == "registry": - 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 - else: - raise ContractError(f"No task builder for {identity['kind']}") + 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 identity["kind"] == "registry" else {} + layouts = _production_layouts(op) if layouts: imports += _RELAYOUT.format(layouts=repr(layouts)) code = imports + f"\ndef run({', '.join(names)}):\n" @@ -164,8 +136,7 @@ def _production_layouts(op: dict[str, Any]) -> dict[tuple[Any, ...], tuple[Any, def _plan( capture: ModelCapture, op: dict[str, Any], overrides: dict[str, Any], hardware_notes: Path | None ) -> dict[str, str | bytes]: - registry = op["identity"]["kind"] == "registry" - name = (f"registry_{op['identity']['op']}_" if registry else "discovered_") + op["id"] + name = f"registry_{op['identity']['op']}_{op['id']}" unsupported = { spec["dtype"] for case in op["workloads"] for spec in case["inputs"] + case["outputs"] } - _DTYPES @@ -221,9 +192,6 @@ def _plan( for tensor in case["inputs"]: if "fixture" in tensor: sources[tensor["fixture"]] = (capture.root / tensor["fixture"]).read_bytes() - for phase in case.get("rng", {}).values(): - for path in phase.values(): - sources[path] = (capture.root / path).read_bytes() numerical = None precision = {"mode": "bit_exact", **execution_precision(capture.manifest["precision"])} if policy(identity) is not None: @@ -251,16 +219,12 @@ def _plan( 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." - if registry - else "Single observed operator; no new fusion.", + 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" if registry else "eager", + timing="cuda_graph", ) settings = tuning_settings(descriptor, overrides) - if identity.get("rng") and settings["timing"]["mode"] != "eager": - raise ContractError("RNG replay currently requires eager timing") 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): @@ -290,10 +254,6 @@ def _plan( "capture_identity": capture.identity, "target": capture.manifest["target"], } - if "rng" in identity: - replay["rng"] = identity["rng"] - if "undefined_outputs" in identity: - replay["undefined_outputs"] = identity["undefined_outputs"] if numerical is not None: replay["numerical_identity"] = identity sources.update( @@ -304,47 +264,38 @@ def _plan( "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", - "README.md": f"# {op['name']}\n\nTarget: `{capture.manifest['target']['model']}`.\nCapture: `{capture.identity}`.\n\nOptimize only this captured operator, preserving all output layouts, aliases and declared input mutations in replay.json. Return tensor outputs in definition order, including the mutated input buffers. Do not fuse with neighboring operators, introduce quantization, or change numerical settings. Use return-value style (destination_passing_style=false). All recorded workloads are mandatory.\n", - "generated.json": json.dumps( - { - "generator": "scripts.kernel_tuning.discovery", - "capture_identity": capture.identity, - "operator_id": op["id"], - } - ), } ) - if registry: - 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"], - } + 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: @@ -355,14 +306,6 @@ def _plan( "cancellation inputs must all pass. Reassociation is permitted within these " "bounds; changing dtype, masks, accumulation settings or the contract is not.\n" ) - if "rng" in identity: - sources["README.md"] += ( - "\nThis call consumes the default PyTorch generator. The evaluator restores the " - "opaque before-state before every call and checks the exact after-state as well as " - "the tensor outputs. Seed 0 replays the recorded transition; other seeds test fresh " - "generator states. Preserve both CPU and target-CUDA RNG state advancement. RNG " - "reset is outside timed events; CUDA Graph timing is not supported for this task.\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 " diff --git a/scripts/kernel_tuning/discovery/tensors.py b/scripts/kernel_tuning/discovery/tensors.py index f778829..5c39b09 100644 --- a/scripts/kernel_tuning/discovery/tensors.py +++ b/scripts/kernel_tuning/discovery/tensors.py @@ -3,7 +3,6 @@ from __future__ import annotations import math -from dataclasses import fields from pathlib import Path from typing import Any @@ -35,15 +34,6 @@ def flatten(value: Any, tensors: list[torch.Tensor]) -> Any: index = len(tensors) tensors.append(value) return {"tensor": index} - if ( - type(value).__module__ == "embodiinfer.backend.triton.sampling" - and type(value).__name__ == "GreedyWorkspace" - ): - return { - "greedy_workspace": { - field.name: flatten(getattr(value, field.name), tensors) for field in fields(value) - } - } 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): @@ -73,10 +63,6 @@ def restore(recipe: Any, tensors: list[torch.Tensor]) -> Any: return tuple(result) if kind == "tuple" else result if kind == "dict": return {key: restore(item, tensors) for key, item in value.items()} - if kind == "greedy_workspace": - from embodiinfer.backend.triton.sampling import GreedyWorkspace - - return GreedyWorkspace(**{key: restore(item, tensors) for key, item in value.items()}) raise ContractError(f"Unknown argument recipe: {kind}") @@ -156,8 +142,6 @@ def tensor_outputs(result: Any) -> list[torch.Tensor]: recipe = flatten(result, tensors) def check(node: Any) -> None: - if "greedy_workspace" in node: - raise ContractError("Workspace outputs require explicit tensor returns") 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"): diff --git a/scripts/kernel_tuning/prompts.py b/scripts/kernel_tuning/prompts.py index eac6024..16b057e 100644 --- a/scripts/kernel_tuning/prompts.py +++ b/scripts/kernel_tuning/prompts.py @@ -63,7 +63,7 @@ def implementation_prompt(task: TaskPackage) -> str: ], } replay = ( - "This discovered 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" + "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 "" ) diff --git a/scripts/kernel_tuning/regions.py b/scripts/kernel_tuning/regions.py index b34dc5b..a59341b 100644 --- a/scripts/kernel_tuning/regions.py +++ b/scripts/kernel_tuning/regions.py @@ -4,7 +4,7 @@ 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 discovery format, so ``calibrate``, ``generate`` and ``tune-all`` apply +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. diff --git a/tests/kernel_tuning/conftest.py b/tests/kernel_tuning/conftest.py index e5c8542..c8776c4 100644 --- a/tests/kernel_tuning/conftest.py +++ b/tests/kernel_tuning/conftest.py @@ -162,3 +162,36 @@ def candidate(store: RunStore) -> dict[str, Any]: {"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_calibration.py b/tests/kernel_tuning/test_calibration.py index 1bc426c..52c7454 100644 --- a/tests/kernel_tuning/test_calibration.py +++ b/tests/kernel_tuning/test_calibration.py @@ -25,59 +25,31 @@ 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.recorder import CaptureSession # noqa: E402 from scripts.kernel_tuning.discovery.tasks import render_model # noqa: E402 -from scripts.kernel_tuning.discovery.tensors import flatten # noqa: E402 - - -def record(root: Path, fn, *, budget: int = 1024 * 1024) -> ModelCapture: - session = CaptureSession("test", root / "raw", device_type="cpu", fixture_bytes=budget) - try: - with torch.inference_mode(), session.observe(object()): - fn() - finally: - session.close(complete=True) - return ModelCapture.load(session.root, require_ready=False) - - -def identity(name: str, *args, **kwargs) -> tuple[dict, list]: - tensors = [] - recipe = flatten((args, kwargs), tensors) - return { - "kind": "aten", - "name": name, - "arguments": recipe, - "input_dtypes": [str(t.dtype).removeprefix("torch.") for t in tensors], - "output_dtypes": ["float32"], - "mutates": [], - }, tensors @pytest.fixture -def calibrated(tmp_path: Path) -> ModelCapture: +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 = record(tmp_path, lambda: (a @ b, c @ b)) + raw = registry_capture(tmp_path, [("linear", (a, b.T), {}), ("linear", (c, b.T), {})]) return calibrate(raw, tmp_path / "calibrated") -def test_semantics_do_not_relax_copies_rng_mutation_integer_or_fp64() -> None: - x = torch.ones(3, 3) - for name in ( - "aten.mm.default", - "aten.linear.default", - "aten.sum.default", - "aten.scaled_dot_product_attention.default", - ): - spec, _ = identity(name, x) - assert policy(spec) == "aten" - for name in ("aten.clone.default", "aten.copy_.default", "aten.relu.default", "aten.index.Tensor"): - spec, _ = identity(name, x) - assert policy(spec) is None - spec, _ = identity("aten.mm.default", x) +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"]}, @@ -86,9 +58,11 @@ def test_semantics_do_not_relax_copies_rng_mutation_integer_or_fp64() -> None: assert policy({**spec, **extra}) is None -def test_generation_requires_calibration_and_capture_keeps_only_execution_switches(tmp_path: Path) -> None: +def test_generation_requires_calibration_and_capture_keeps_only_execution_switches( + tmp_path: Path, registry_capture +) -> None: x = torch.ones(3, 3) - capture = record(tmp_path, lambda: x @ x) + 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") @@ -139,11 +113,16 @@ def test_task_loader_rejects_calibration_edits(calibrated: ModelCapture, tmp_pat TaskPackage.load(task.root) -def test_real_inputs_cannot_silently_fall_back_to_random_when_budget_exhausted(tmp_path: Path) -> None: +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 = record(tmp_path, lambda: x @ x, budget=0) - assert capture.operators[0]["status"] == "blocked" - assert "fixture" in capture.operators[0]["reason"] + 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") @@ -182,35 +161,6 @@ def test_calibration_restores_global_precision_on_failure() -> None: assert environment("cpu") == before -@pytest.mark.parametrize( - "causal,gqa,masked", [(False, False, True), (True, False, False), (False, True, True)] -) -def test_fp64_attention_matches_math_reference(causal: bool, gqa: bool, masked: bool) -> None: - generator = torch.Generator().manual_seed(2) - q = torch.randn(1, 4, 7, 8, generator=generator) - k = torch.randn(1, 2 if gqa else 4, 9, 8, generator=generator) - v = torch.randn(1, 2 if gqa else 4, 9, 8, generator=generator) - mask = torch.ones(7, 9, dtype=torch.bool) if masked else None - if mask is not None: - mask[0] = False - mask[1:, -2:] = False - spec, tensors = identity( - "aten.scaled_dot_product_attention.default", - q, - k, - v, - attn_mask=mask, - dropout_p=0.0, - is_causal=causal, - scale=0.25, - enable_gqa=gqa, - ) - expected = torch.nn.functional.scaled_dot_product_attention( - q.double(), k.double(), v.double(), attn_mask=mask, is_causal=causal, scale=0.25, enable_gqa=gqa - ) - torch.testing.assert_close(high_precision(spec, tensors)[0], expected, atol=1e-12, rtol=1e-12) - - def test_replay_uses_calibrated_bounds_and_all_profiles( calibrated: ModelCapture, tmp_path: Path, monkeypatch ) -> None: @@ -264,33 +214,32 @@ def replay(): assert asdict(task.settings.precision)["mode"] == "tolerance" -def test_random_profiles_preserve_masks_and_change_real_values(tmp_path: Path) -> None: +def test_random_profiles_preserve_masks_and_change_real_values(tmp_path: Path, registry_capture) -> None: q = torch.ones(1, 1, 3, 4) - mask = torch.tensor([[True, False, True]]).expand(3, 3) - capture = record(tmp_path, lambda: torch.ops.aten.scaled_dot_product_attention.default(q, q, q, mask)) - # Current Torch observes the public SDPA overload atomically. - op = next(op for op in capture.operators if op["name"] == "aten.scaled_dot_product_attention.default") + 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_rms_norm_default_epsilon_preserves_input_dtype_semantics() -> None: - value = torch.full((2, 4), 0.001, dtype=torch.bfloat16) - spec, tensors = identity("aten.rms_norm.default", value, [4]) - actual = high_precision(spec, tensors)[0] - expected = torch.nn.functional.rms_norm(value.double(), [4], eps=torch.finfo(value.dtype).eps) - torch.testing.assert_close(actual, expected, rtol=0, atol=0) - - -def test_calibrate_cli_writes_a_new_capture_before_generation(tmp_path: Path, capsys) -> None: +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 = record(tmp_path, lambda: value @ weight) + 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 diff --git a/tests/kernel_tuning/test_cli.py b/tests/kernel_tuning/test_cli.py index 71b8146..e7e9ef7 100644 --- a/tests/kernel_tuning/test_cli.py +++ b/tests/kernel_tuning/test_cli.py @@ -6,6 +6,7 @@ import sys from pathlib import Path +import pytest from scripts.kernel_tuning.__main__ import main @@ -30,3 +31,29 @@ def test_bad_yaml_is_a_clear_static_validation_error(task_dir: Path, capsys) -> (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_discovery.py b/tests/kernel_tuning/test_discovery.py deleted file mode 100644 index 7233ec1..0000000 --- a/tests/kernel_tuning/test_discovery.py +++ /dev/null @@ -1,605 +0,0 @@ -"""Model discovery, replay contracts and the all-task barrier, using CPU tensors.""" - -from __future__ import annotations - -import contextlib -from pathlib import Path - -import pytest - -torch = pytest.importorskip("torch") - -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, read_json # noqa: E402 -from scripts.kernel_tuning.discovery import ModelCapture # noqa: E402 -from scripts.kernel_tuning.discovery.recorder import CaptureSession # 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 - - -@pytest.fixture(autouse=True) -def cpu_capture(monkeypatch): - """Keep observer tests independent of CUDA initialization and numerical globals.""" - monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) - old = torch.get_float32_matmul_precision() - torch.set_float32_matmul_precision("highest") - yield - torch.set_float32_matmul_precision(old) - - -def capture_calls(tmp_path: Path, fn, *, fixture_bytes: int = 1024 * 1024) -> ModelCapture: - session = CaptureSession( - "mock_flow_vla", tmp_path / "capture", device_type="cpu", fixture_bytes=fixture_bytes - ) - try: - with session.observe(object(), checkpoint="test-revision"): - fn() - finally: - session.close(complete=True) - return ModelCapture.load(session.root, require_ready=False) - - -def generated_function(task: TaskPackage): - scope = {} - exec((task.root / "baseline.py").read_text(encoding="utf-8"), scope) - return scope["run"] - - -def test_discovers_uncatalogued_ops_all_shapes_and_call_counts(tmp_path: Path) -> None: - from benchmarks.kernel_tuning.calibration import calibrate - - values = [torch.ones(n, 4) for n in range(1, 12)] - weight = torch.ones(4, 3) - - def forward(): - for value in [*values, values[0]]: - (value @ weight).relu() - - capture = capture_calls(tmp_path, forward) - ready = {op["name"]: op for op in capture.operators if op["status"] == "ready"} - assert set(ready) == {"aten.mm.default", "aten.relu.default"} - assert len(ready["aten.mm.default"]["workloads"]) == 11 - assert sum(case["count"] for case in ready["aten.mm.default"]["workloads"]) == 12 - tasks = render_model(calibrate(capture, tmp_path / "calibrated"), tmp_path / "tasks") - assert len(tasks) == 2 - assert all(len(task.workloads) == 11 for task in tasks) - for task in tasks: - replay = read_json(task.root / "replay.json") - assert task.settings.precision.mode == ("tolerance" if "mm." in replay["operator"] else "bit_exact") - for case in replay["cases"].values(): - inputs = materialize(case["inputs"], task.root, 0) - result = generated_function(task)(*inputs) - assert list(result.shape) == case["outputs"][0]["shape"] - - -def test_noncontiguous_aliases_mutations_and_metadata_are_preserved(tmp_path: Path) -> None: - backing = torch.arange(24, dtype=torch.float32).reshape(4, 6) - left, right = backing[:, 1:3], backing[:, 3:5] - capture = capture_calls(tmp_path, lambda: (left.add_(right), left.transpose(0, 1))) - assert any(op["status"] == "metadata_only" for op in capture.operators) - (task,) = render_model(capture, tmp_path / "tasks") - replay = read_json(task.root / "replay.json") - (case,) = replay["cases"].values() - assert replay["mutates"] == [0] - inputs = materialize(case["inputs"], task.root, 0) - assert storage_key(inputs[0]) == storage_key(inputs[1]) - assert inputs[0].stride() == (6, 1) - assert inputs[0].storage_offset() == 1 - cloned = clone_inputs(inputs) - assert storage_key(cloned[0]) == storage_key(cloned[1]) - assert storage_key(cloned[0]) != storage_key(inputs[0]) - expected = inputs[0] + inputs[1] - output, mutated = generated_function(task)(*cloned) - assert output is mutated and output is cloned[0] - assert torch.equal(output, expected) - - -def test_required_integer_fixtures_are_valid_and_never_treated_as_quantized(tmp_path: Path) -> None: - weight = torch.arange(20, dtype=torch.float32).reshape(5, 4) - indices = torch.tensor([4, 1, 0], dtype=torch.int64) - capture = capture_calls(tmp_path, lambda: torch.nn.functional.embedding(indices, weight)) - assert not any(op["status"] == "excluded_quantized" for op in capture.operators) - (task,) = render_model(capture, tmp_path / "tasks") - (case,) = read_json(task.root / "replay.json")["cases"].values() - args = materialize(case["inputs"], task.root, 0) - assert torch.equal(generated_function(task)(*args), weight[indices]) - - -def test_tensor_factories_without_tensor_inputs_get_tasks(tmp_path: Path) -> None: - capture = capture_calls(tmp_path, lambda: torch.arange(9, device="cpu")) - (task,) = render_model(capture, tmp_path / "tasks") - assert task.definition["inputs"] == {} - assert task.settings.device == "cpu" - assert torch.equal(generated_function(task)(), torch.arange(9)) - - -def test_scalar_reads_are_host_work_not_coverage_gaps(tmp_path: Path) -> None: - from scripts.kernel_tuning.discovery.recorder import _Observer - - value = torch.arange(4, dtype=torch.float32).sum() - session = CaptureSession("mock_flow_vla", tmp_path / "capture", device_type="cpu") - try: - with session.observe(object(), checkpoint="test-revision"): - # Torch 2.12 delivers Tensor.item() to the observer as aten.item; older - # releases decompose it first, so deliver the overload directly. - result = _Observer(session).__torch_dispatch__(torch.ops.aten.item.default, (), (value,)) - assert result == 6.0 - finally: - session.close(complete=True) - capture = ModelCapture.load(session.root, require_ready=False) - statuses = {op["name"]: op["status"] for op in capture.operators} - assert statuses["aten.item.default"] == "host_only" - assert "blocked" not in statuses.values() - - -def test_host_to_device_copies_keep_required_input_fixtures(tmp_path: Path) -> None: - # "meta" stands in for the target device: host inputs, target-device output. - session = CaptureSession("mock_flow_vla", tmp_path / "capture", device_type="meta") - indices = torch.tensor([4, 1, 0], dtype=torch.int64) - try: - with session.observe(object(), checkpoint="test-revision"): - indices.to("meta") - finally: - session.close(complete=True) - capture = ModelCapture.load(session.root, require_ready=False) - (op,) = [op for op in capture.operators if op["status"] == "ready"] - (case,) = op["workloads"] - assert case["inputs"][0]["device"] == "cpu" - assert case["inputs"][0].get("fixture") - - -@pytest.mark.parametrize("device", ["cpu", "meta"]) -def test_unrepresentable_host_preprocessing_is_not_a_device_gap(tmp_path: Path, device: str) -> None: - # "meta" stands in for the target device; serving adapters check requests on the host. - session = CaptureSession("mock_flow_vla", tmp_path / "capture", device_type="meta") - try: - with session.observe(object(), checkpoint="test-revision"): - torch.ones(2, device=device).ne(float("inf")) # torch.isfinite's decomposition - finally: - session.close(complete=True) - capture = ModelCapture.load(session.root, require_ready=False) - (op,) = [op for op in capture.operators if op["name"] == "aten.ne.Scalar"] - # A device call with the same argument still needs a replay recipe. - assert op["status"] == ("host_only" if device == "cpu" else "blocked") - - -def test_quantized_scope_includes_dynamic_unregistered_module_and_float_fallback(tmp_path: Path) -> None: - from embodiinfer.layers.linear import INT8Config - from embodiinfer.models import linear - - dense = torch.nn.Linear(4, 3, bias=False) - quantized = linear.INT8Linear.from_linear(dense, INT8Config(backend="torch")) - owner = torch.nn.Module() - object.__setattr__(owner, "projection", quantized) - assert list(owner.children()) == [] - value = torch.ones(2, 4) - original = linear.QuantizedLinear.forward - session = CaptureSession("mock_flow_vla", tmp_path / "capture", device_type="cpu") - try: - session.install() - with session.observe(owner): - output = owner.projection(value) - output.relu() - finally: - session.close(complete=True) - assert linear.QuantizedLinear.forward is original - capture = ModelCapture.load(session.root) - assert [op["name"] for op in capture.operators if op["status"] == "ready"] == ["aten.relu.default"] - assert sum(op["status"] == "excluded_quantized" for op in capture.operators) == 1 - assert not any("mm." in op["name"] for op in capture.operators) - - -@pytest.mark.parametrize("failure", ["random", "fixture", "external"]) -def test_any_coverage_gap_blocks_all_task_generation(tmp_path: Path, failure: str) -> None: - value = torch.ones(4) - index = torch.tensor([1, 0]) - - def forward(): - value.relu() - if failure == "random": - torch.rand_like(value) - elif failure == "fixture": - value[index] - else: - custom(value) - - library = torch.library.Library("discovery_test", "FRAGMENT") - if failure == "external": - library.define("opaque(Tensor x) -> Tensor") - library.impl("opaque", lambda x: x + 1, "CPU") - custom = torch.ops.discovery_test.opaque.default - capture = capture_calls(tmp_path, forward, fixture_bytes=0) - assert any(op["status"] == "blocked" for op in capture.operators) - with pytest.raises(ContractError, match="incomplete"): - render_model(capture, tmp_path / "tasks") - assert not (tmp_path / "tasks").exists() - - -def test_capture_model_and_immutable_fixtures_are_verified(tmp_path: Path) -> None: - value = torch.ones(4) - capture = capture_calls(tmp_path, value.relu) - with pytest.raises(ContractError, match="differs from requested"): - ModelCapture.load(capture.root, "pi05") - fixture = next((capture.root / "fixtures").iterdir()) - fixture.write_bytes(b"changed") - with pytest.raises(ContractError, match="content changed"): - render_model(capture, tmp_path / "tasks") - - -def test_single_actual_target_and_uncompiled_execution_are_required(tmp_path: Path) -> None: - session = CaptureSession("mock_flow_vla", tmp_path / "capture", device_type="cpu") - try: - session.install() - from embodiinfer.policies.factory import make_policy - - with pytest.raises(ContractError, match="expected"): - make_policy("pi05") - with pytest.raises(ContractError, match="compilation"): - torch.compile(lambda x: x) - session.bind(object(), {}) - with pytest.raises(ContractError, match="one target"): - session.bind(object(), {}) - finally: - session.close(complete=False) - - -def test_model_batch_never_starts_agent_until_every_baseline_passes(tmp_path: Path, monkeypatch) -> None: - from test_batch import measurement - - value = torch.ones(4) - capture = capture_calls(tmp_path, lambda: value.relu().square()) - 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 "pow" in read_json(store.task.root / "replay.json")["operator"]: - 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()) - assert any(error for error in 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 - # A subsequently changed task invalidates the whole model, including resume. - 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_model_cli_rejects_partial_selection_and_legacy_capture(tmp_path: Path, capsys) -> None: - assert main(["generate", "--output", str(tmp_path / "tasks")]) == 2 - assert "exactly one --model" in capsys.readouterr().err - value = torch.ones(4) - capture = capture_calls(tmp_path, value.relu) - args = ["--model", "mock_flow_vla", "--captured", str(capture.root)] - assert main(["generate", *args, "--list"]) == 0 - assert "aten.relu.default" in capsys.readouterr().out - assert main(["tune-all", *args, "--dry-run", "--output", str(tmp_path / "batches")]) == 0 - assert main(["generate", *args, "relu"]) == 2 - assert main(["tune-all", *args, "--agent", "fake/model", "--skip-preflight"]) == 2 - assert main(["generate", "--model", "pi05", "--model", "streamvln"]) == 2 - legacy = tmp_path / "legacy.jsonl" - legacy.write_text("{}", encoding="utf-8") - assert main(["generate", "--model", "pi05", "--captured", str(legacy)]) == 2 - (task,) = render_model(capture, tmp_path / "tasks") - assert main(["run", str(task.root), "--agent", "fake/model", "--output", str(tmp_path / "runs")]) == 2 - assert not (tmp_path / "runs").exists() - - -_ENGINE_DRIVER = """ -from concurrent.futures import ThreadPoolExecutor - -import torch - -from embodiinfer.engine import EngineConfig, EngineCore, GenerationBackend -from embodiinfer.policies.factory import make_policy -from embodiinfer.types import Observation - -torch.manual_seed(0) -policy = make_policy("mock_flow_vla", preset="tiny") -cfg = policy.config -observation = 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,)), -) -x = torch.ones(3) -x.digamma() # driver arithmetic outside the model is not a model operator -core = EngineCore(policy, EngineConfig(device="cpu", use_cuda_graph=False, num_steps=2)) -MODE -""" -_ENGINE_MODES = { - # The async engine and batched serving execute the engine on a worker thread. - "thread": "with ThreadPoolExecutor(1) as pool:\n" - " pool.submit(core.execute, policy.collate([observation], ['r0'])).result()\n", - # The RL rollout path samples through decoder methods the engine never calls. - "rollout": "GenerationBackend(core).generate_with_logprob([observation], num_steps=2)\n", - # Benchmark-style drivers that call policy internals bypass the engine. - "bypass": "batch = policy.collate([observation], ['r0']).to('cpu', torch.float32)\n" - "policy.encode_prefix(batch)\n", -} - - -def _engine_capture(tmp_path: Path, monkeypatch, mode: str) -> ModelCapture: - from scripts.kernel_tuning.discovery import recorder - - # The public CLI always targets CUDA; this test substitutes the observer's - # device classification while exercising the real driver and engine hooks. - monkeypatch.setattr( - recorder, "CaptureSession", lambda *args, **kwargs: CaptureSession(*args, device_type="cpu", **kwargs) - ) - script = tmp_path / f"{mode}.py" - script.write_text(_ENGINE_DRIVER.replace("MODE", _ENGINE_MODES[mode]), encoding="utf-8") - with contextlib.suppress(ContractError): # Gaps are asserted from the published capture. - recorder.run([str(script)], tmp_path / mode, "mock_flow_vla") - return ModelCapture.load(tmp_path / mode, require_ready=False) - - -def _reference_operators(mode: str) -> set[str]: - """Every dispatcher operator of the same inference, collected without the recorder.""" - from torch.utils._python_dispatch import TorchDispatchMode - - from embodiinfer.engine import EngineConfig, EngineCore, GenerationBackend - from embodiinfer.policies.factory import make_policy - from embodiinfer.types import Observation - - names: set[str] = set() - - class Collect(TorchDispatchMode): - def __torch_dispatch__(self, func, types, args=(), kwargs=None): - names.add(str(func)) - return func(*args, **(kwargs or {})) - - torch.manual_seed(0) - policy = make_policy("mock_flow_vla", preset="tiny") - cfg = policy.config - observation = 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,)), - ) - with Collect(): - core = EngineCore(policy, EngineConfig(device="cpu", use_cuda_graph=False, num_steps=2)) - if mode == "thread": - core.execute(policy.collate([observation], ["r0"])) - else: - GenerationBackend(core).generate_with_logprob([observation], num_steps=2) - return names - - -@pytest.mark.parametrize("mode", ["thread", "rollout"]) -def test_engine_inference_is_captured_end_to_end(tmp_path: Path, monkeypatch, mode: str) -> None: - capture = _engine_capture(tmp_path, monkeypatch, mode) - capture.require_ready() - assert capture.manifest["target"]["config"] == {"preset": "tiny"} - captured = {op["name"] for op in capture.operators} - # No policy method list: the whole inference matches an independent collector. - assert captured == _reference_operators(mode) - assert "aten.digamma.default" not in captured - - -def test_drive_replays_recorded_observations_through_the_engine(tmp_path: Path, monkeypatch) -> None: - import json - - from scripts.kernel_tuning.discovery import recorder - from scripts.kernel_tuning.discovery.drive import save_inputs - - from embodiinfer.policies.mock import preset_config - from embodiinfer.types import Observation - - cfg = preset_config("tiny") - - def observation() -> 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,)), - ) - - save_inputs(tmp_path / "inputs.pt", "observations", [[observation(), observation()], [observation()]]) - (tmp_path / "serving.json").write_text( - json.dumps({"policy_kwargs": {"preset": "tiny"}}), encoding="utf-8" - ) - monkeypatch.setattr( - recorder, "CaptureSession", lambda *args, **kwargs: CaptureSession(*args, device_type="cpu", **kwargs) - ) - command = ["-m", "scripts.kernel_tuning.discovery.drive", "--policy", "mock_flow_vla", "--device", "cpu"] - command += ["--adapter-config", str(tmp_path / "serving.json"), "--inputs", str(tmp_path / "inputs.pt")] - recorder.run(command, tmp_path / "capture", "mock_flow_vla") - capture = ModelCapture.load(tmp_path / "capture", "mock_flow_vla") - assert capture.manifest["target"]["config"] == {"preset": "tiny"} - addmm = next(op for op in capture.operators if op["name"] == "aten.addmm.default") - # Three requests, each with a prefix and the default denoise loop, were observed. - assert sum(case["count"] for case in addmm["workloads"]) > 3 - - -def test_model_arithmetic_outside_the_engine_is_a_coverage_gap(tmp_path: Path, monkeypatch) -> None: - capture = _engine_capture(tmp_path, monkeypatch, "bypass") - errors = capture.manifest["errors"] - assert any("read model weights outside an inference entry" in error for error in errors) - with pytest.raises(ContractError, match="outside an inference entry"): - capture.require_ready() - - -def test_capture_without_inference_is_rejected(tmp_path: Path, monkeypatch) -> None: - from scripts.kernel_tuning.discovery import recorder - - from embodiinfer.policies import factory - - monkeypatch.setitem(factory._REGISTRY, "mock_flow_vla", lambda **kwargs: torch.nn.Linear(2, 2)) - monkeypatch.setattr( - recorder, "CaptureSession", lambda *args, **kwargs: CaptureSession(*args, device_type="cpu", **kwargs) - ) - script = tmp_path / "prepare.py" - script.write_text( - "from embodiinfer.policies.factory import make_policy\nmake_policy('mock_flow_vla')\n", - encoding="utf-8", - ) - with pytest.raises(ContractError): - recorder.run([str(script)], tmp_path / "capture", "mock_flow_vla") - manifest = read_json(tmp_path / "capture" / "manifest.json") - assert any("No end-to-end inference ran" in error for error in manifest["errors"]) - - -def test_replay_adapter_checks_output_alias_and_undeclared_mutation_on_cpu(tmp_path: Path) -> None: - from benchmarks.kernel_tuning.replay_adapter import ReplayAdapter - - value = torch.ones(4) - capture = capture_calls(tmp_path, lambda: value.add_(2)) - (task,) = render_model(capture, tmp_path / "tasks") - replay = read_json(task.root / "replay.json") - (case,) = replay["cases"].values() - # Exercise the contract logic without constructing the GPU/compiler adapter. - adapter = ReplayAdapter.__new__(ReplayAdapter) - adapter.case, adapter.replay, adapter.cfg = case, replay, task.settings - inputs = materialize(case["inputs"], task.root, 0) - expected = clone_inputs(inputs) - outputs = generated_function(task)(*expected) - adapter._check_outputs(expected, list(outputs)) - assert adapter._check_state(expected, expected, inputs) == (0.0, 0.0) - with pytest.raises(ContractError, match="aliases"): - adapter._check_outputs(expected, [x.clone() for x in outputs]) - adapter.replay = {**replay, "mutates": []} - with pytest.raises(ContractError, match="bytes"): - adapter._check_state(expected, expected, inputs) - - -def test_missing_workload_and_unbound_output_axes_are_rejected(tmp_path: Path) -> None: - from benchmarks.kernel_tuning.calibration import calibrate - - value = torch.ones(4, 2) - capture = capture_calls(tmp_path, lambda: value.sum(1)) - (task,) = render_model(calibrate(capture, tmp_path / "calibrated"), tmp_path / "tasks") - 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) - - -@pytest.mark.parametrize("known", [True, False]) -def test_existing_backend_is_atomic_and_requires_a_replay_contract( - tmp_path: Path, monkeypatch, known: bool -) -> None: - from scripts.kernel_tuning.discovery import recorder - - from embodiinfer.backend.triton import norm - - if not known: - monkeypatch.setattr(recorder, "_PURE_KERNELS", recorder._PURE_KERNELS - {"gated_residual"}) - - values = [torch.ones(1, 2, 4) for _ in range(3)] - session = CaptureSession("mock_flow_vla", tmp_path / "capture", device_type="cpu") - original = norm.gated_residual - try: - session.install() - with session.observe(object()): - result = norm.gated_residual(*values) - finally: - session.close(complete=True) - assert norm.gated_residual is original - assert torch.equal(result, torch.full_like(result, 2)) - capture = ModelCapture.load(session.root, require_ready=False) - (op,) = capture.operators - assert op["identity"]["kind"] == "backend" - assert op["name"].endswith(".gated_residual") - if not known: - with pytest.raises(ContractError, match="explicit mutation/reference contract"): - capture.require_ready() - return - (task,) = render_model(capture, tmp_path / "tasks") - assert "vendor/embodiinfer/backend/triton/norm.py" in task.hashes - assert "from .vendor.embodiinfer.backend.triton.norm" in task.definition["reference"] - - -def test_changed_numerical_settings_block_capture(tmp_path: Path) -> None: - value = torch.ones(4) - - def forward(): - value.relu() - torch.set_float32_matmul_precision("high") - value.square() - - capture = capture_calls(tmp_path, forward) - with pytest.raises(ContractError, match="Numerical settings changed"): - capture.require_ready() - - -def test_greedy_workspace_has_an_explicit_tensor_state_recipe() -> None: - from scripts.kernel_tuning.discovery.tasks import _expression - from scripts.kernel_tuning.discovery.tensors import flatten, restore - - from embodiinfer.backend.triton.sampling import GreedyWorkspace - - workspace = GreedyWorkspace.allocate(16, torch.device("cpu")) - tensors = [] - recipe = flatten(workspace, tensors) - cloned = clone_inputs(tensors) - restored = restore(recipe, cloned) - assert restored.vocab_size == 16 - assert restored.token is cloned[recipe["greedy_workspace"]["token"]["tensor"]] - scope = {"GreedyWorkspace": GreedyWorkspace, **{f"x{i}": tensor for i, tensor in enumerate(cloned)}} - constructed = eval(_expression(recipe), scope) - assert constructed.penalized is restored.penalized - assert constructed.token is restored.token - - -@pytest.mark.gpu -@pytest.mark.parametrize("timing", ["eager", "cuda_graph"]) -def test_discovered_mutable_baseline_passes_cuda_replay(tmp_path: Path, timing: str) -> None: - """Target-runtime smoke test: real builders, alias checks, resets and graph A/B/A.""" - if not torch.cuda.is_available(): - pytest.skip("requires CUDA") - pytest.importorskip("flashinfer_bench") - pytest.importorskip("triton") - from benchmarks.kernel_tuning.replay_adapter import ReplayAdapter - from scripts.kernel_tuning.contracts import validate_measurement - - value = torch.arange(16, dtype=torch.float32, device="cuda:0").reshape(2, 8)[:, 1:5] - session = CaptureSession("mock_flow_vla", tmp_path / "capture") - try: - with session.observe(object()): - value.add_(1) - finally: - session.close(complete=True) - (task,) = render_model( - ModelCapture.load(session.root), - tmp_path / "tasks", - overrides={"timing.mode": timing, "timing.warmup": 2, "timing.iterations": 5, "timing.trials": 1}, - ) - adapter = ReplayAdapter(task) - validate_measurement( - task, adapter.evaluate({role: task.baseline() for role in ("baseline", "incumbent", "candidate")}) - ) 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_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_regions.py b/tests/kernel_tuning/test_regions.py index 1e958fa..25006fc 100644 --- a/tests/kernel_tuning/test_regions.py +++ b/tests/kernel_tuning/test_regions.py @@ -114,3 +114,54 @@ def test_baseline_code_is_unchanged_without_layout_differences(tmp_path): (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_replay_extensions.py b/tests/kernel_tuning/test_replay_extensions.py deleted file mode 100644 index 1b007f5..0000000 --- a/tests/kernel_tuning/test_replay_extensions.py +++ /dev/null @@ -1,322 +0,0 @@ -"""Lossless random-state and float64 contracts for discovered operator tasks.""" - -from __future__ import annotations - -from pathlib import Path - -import pytest - -torch = pytest.importorskip("torch") - -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.recorder import CaptureSession # noqa: E402 -from scripts.kernel_tuning.discovery.tasks import render_model # noqa: E402 -from scripts.kernel_tuning.discovery.tensors import materialize # noqa: E402 - - -@pytest.fixture(autouse=True) -def isolated_numerics(monkeypatch): - """CPU tests must neither initialize CUDA nor leak numerical/RNG settings.""" - monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) - precision, state = torch.get_float32_matmul_precision(), torch.get_rng_state() - torch.set_float32_matmul_precision("highest") - yield - torch.set_float32_matmul_precision(precision) - torch.set_rng_state(state) - - -def record(tmp_path: Path, fn, *, fixture_bytes: int = 1024 * 1024) -> ModelCapture: - """Run one CPU operation through the production observer and immutable loader.""" - session = CaptureSession( - "mock_flow_vla", tmp_path / "capture", device_type="cpu", fixture_bytes=fixture_bytes - ) - try: - with session.observe(object()): - fn() - finally: - session.close(complete=True) - return ModelCapture.load(session.root, require_ready=False) - - -def baseline(task): - """Load the generated single-file ATen baseline without a GPU builder.""" - scope = {} - exec((task.root / "baseline.py").read_text(encoding="utf-8"), scope) - return scope["run"] - - -def test_default_randn_records_real_state_and_generates_a_task(tmp_path: Path) -> None: - torch.manual_seed(123) - torch.randn(7) # Include a non-initial generator position/cache. - before = torch.get_rng_state().clone() - observed = [] - capture = record(tmp_path, lambda: observed.append(torch.randn(3, 5, generator=None))) - after = torch.get_rng_state().clone() - capture.require_ready() - (task,) = render_model(capture, tmp_path / "tasks") - replay = read_json(task.root / "replay.json") - assert replay["rng"] == {"kind": "torch_default", "devices": ["cpu"]} - (case,) = replay["cases"].values() - for phase, state in (("before", before), ("after", after)): - assert (task.root / case["rng"][phase]["cpu"]).read_bytes() == bytes(state.tolist()) - torch.set_rng_state(before) - assert torch.equal(baseline(task)(), observed[0]) - assert torch.equal(torch.get_rng_state(), after) - - -def test_float64_task_preserves_dtype_and_values(tmp_path: Path) -> None: - value = torch.tensor([1.0 + 2**-30, 1.0 - 2**-30], dtype=torch.float64) - capture = record(tmp_path, value.sin) - (task,) = render_model(capture, tmp_path / "tasks") - assert task.definition["inputs"]["x0"]["dtype"] == "float64" - assert task.definition["outputs"]["y0"]["dtype"] == "float64" - replay = read_json(task.root / "replay.json") - assert replay["definition_schema"] == "embodiinfer-replay-v1" - (case,) = replay["cases"].values() - inputs = materialize(case["inputs"], task.root, 0) - actual = baseline(task)(*inputs) - assert actual.dtype == torch.float64 - assert torch.equal(actual, value.sin()) - assert not torch.equal(actual, value.float().sin().double()) - - -def test_rng_state_is_required_bounded_and_immutable(tmp_path: Path) -> None: - capture = record(tmp_path / "bounded", lambda: torch.randn(2), fixture_bytes=0) - with pytest.raises(ContractError, match="RNG fixture.*budget"): - capture.require_ready() - assert len(capture.operators) == 1 # Instrumentation does not capture its own state operations. - capture = record( - tmp_path / "after_bounded", lambda: torch.randn(2), fixture_bytes=torch.get_rng_state().numel() - ) - with pytest.raises(ContractError, match="RNG fixture.*budget"): - capture.require_ready() - capture = record(tmp_path / "valid", lambda: torch.randn(2)) - (task,) = render_model(capture, tmp_path / "tasks") - replay = read_json(task.root / "replay.json") - (case,) = replay["cases"].values() - (task.root / case["rng"]["after"]["cpu"]).unlink() - with pytest.raises(ContractError, match="RNG fixture missing"): - TaskPackage.load(task.root) - with pytest.raises(ContractError, match="eager timing"): - render_model(capture, tmp_path / "graphs", overrides={"timing.mode": "cuda_graph"}) - - -def test_unsupported_rng_and_explicit_generators_still_block(tmp_path: Path) -> None: - for name, fn in ( - ("explicit", lambda: torch.randn(2, generator=torch.Generator().manual_seed(7))), - ("uniform", lambda: torch.rand(2)), - ): - capture = record(tmp_path / name, fn) - with pytest.raises(ContractError, match="incomplete"): - capture.require_ready() - - -@pytest.mark.parametrize("style", ["omitted", "positional", "keyword"]) -def test_zero_dropout_sdpa_is_not_a_random_call(tmp_path: Path, style: str) -> None: - from benchmarks.kernel_tuning.calibration import calibrate - - q = torch.randn(1, 2, 8, 16) - state = torch.get_rng_state().clone() - outputs = [] - - def forward(): - op = torch.ops.aten.scaled_dot_product_attention.default - for _ in range(2): - if style == "omitted": - outputs.append(op(q, q, q)) - elif style == "positional": - outputs.append(op(q, q, q, None, 0.0)) - else: - outputs.append(op(q, q, q, dropout_p=0.0)) - - with torch.inference_mode(): - capture = record(tmp_path, forward) - capture.require_ready() - assert len(capture.operators) == 1 - assert capture.operators[0]["name"] == "aten.scaled_dot_product_attention.default" - assert sum(c["count"] for c in capture.operators[0]["workloads"]) == 2 - assert torch.equal(state, torch.get_rng_state()) - assert torch.equal(outputs[0], outputs[1]) - (task,) = render_model(calibrate(capture, tmp_path / "calibrated"), tmp_path / "tasks") - assert "rng" not in read_json(task.root / "replay.json") - - -def test_zero_dropout_sdpa_backend_overload_is_not_a_random_call(tmp_path: Path) -> None: - # Under the engine's no_grad, the functional API reaches the observer as a backend overload. - q = torch.randn(1, 2, 8, 16) - with torch.no_grad(): - capture = record(tmp_path, lambda: torch.nn.functional.scaled_dot_product_attention(q, q, q)) - assert [op["name"] for op in capture.operators] == [ - "aten._scaled_dot_product_flash_attention_for_cpu.default" - ] - capture.require_ready() - - -def test_unused_attention_dropout_state_is_not_compared(tmp_path: Path) -> None: - from types import SimpleNamespace - - from benchmarks.kernel_tuning.replay_adapter import ReplayAdapter - from scripts.kernel_tuning.discovery.rng import undefined_outputs - - op = torch.ops.aten._scaled_dot_product_efficient_attention.default - q = torch.empty(1, 2, 8, 16) - # CUDA leaves philox seed/offset uninitialized when dropout is zero. - assert undefined_outputs(op, (q, q, q, None, False), {}) == [2, 3] - assert undefined_outputs(op, (q, q, q, None, False, 0.1), {}) == [] - assert undefined_outputs(torch.ops.aten.mm.default, (q, q), {}) == [] - - adapter = ReplayAdapter.__new__(ReplayAdapter) - adapter.cfg = SimpleNamespace(precision=SimpleNamespace(mode="bit_exact")) - adapter.replay = {"undefined_outputs": [1]} - value = torch.ones(3) - assert adapter._compare([value, torch.tensor(16)], [value, torch.tensor(0)], []) == (0.0, 0.0) - with pytest.raises(ContractError, match="bit-exact"): - adapter._compare([value + 1, torch.tensor(0)], [value, torch.tensor(0)], []) - - -@pytest.mark.parametrize("style", ["positional", "keyword"]) -def test_nonzero_dropout_sdpa_remains_a_coverage_gap(tmp_path: Path, style: str) -> None: - q = torch.randn(1, 2, 8, 16) - - def forward(): - op = torch.ops.aten.scaled_dot_product_attention.default - return op(q, q, q, None, 0.25) if style == "positional" else op(q, q, q, dropout_p=0.25) - - with torch.inference_mode(): - capture = record(tmp_path, forward) - with pytest.raises(ContractError, match="RNG replay adapter"): - capture.require_ready() - - -def test_rng_outputs_and_advancement_are_checked_and_caller_state_restored( - tmp_path: Path, monkeypatch -) -> None: - from types import SimpleNamespace - - from benchmarks.kernel_tuning.replay_adapter import ReplayAdapter - - capture = record(tmp_path, lambda: torch.randn(15, generator=None)) - (task,) = render_model(capture, tmp_path / "tasks") - adapter = ReplayAdapter.__new__(ReplayAdapter) - adapter.task, adapter.cfg = task, task.settings - adapter.definition = SimpleNamespace( - outputs=task.definition["outputs"], torch_output_dtypes=[torch.float32] - ) - adapter.replay = read_json(task.root / "replay.json") - adapter.device = "cpu" - adapter._check_flags = lambda: None - monkeypatch.setattr(torch.cuda, "synchronize", lambda *args: None) - - def unexpected_global_seed(*args, **kwargs): - pytest.fail("RNG replay must not seed unrelated generators") - - monkeypatch.setattr(torch, "manual_seed", unexpected_global_seed) - - class Runnable: - metadata = SimpleNamespace(destination_passing_style=False) - - def __init__(self, fn): - self.fn = fn - - def __call__(self, *args): - return self.fn(*args) - - fn = baseline(task) - reference = Runnable(fn) - - def extra_draw(): - output = fn() - torch.randn(1) - return output - - def rewind(): - state = torch.get_rng_state() - output = fn() - torch.set_rng_state(state) - return output - - def crash(): - torch.randn(1) - raise RuntimeError("candidate failed") - - for seed in (0, 1, 2): - inputs = adapter._inputs(SimpleNamespace(uuid=task.workload_ids[0]), seed) - saved = torch.get_rng_state().clone() - assert adapter._verify_call(reference, reference, inputs, seed) == (0.0, 0.0) - assert torch.equal(torch.get_rng_state(), saved) - for bad in (extra_draw, rewind): - with pytest.raises(ContractError, match="RNG state"): - adapter._verify_call(Runnable(bad), reference, inputs, seed) - assert torch.equal(torch.get_rng_state(), saved) - with pytest.raises(RuntimeError, match="candidate failed"): - adapter._verify_call(Runnable(crash), reference, inputs, seed) - assert torch.equal(torch.get_rng_state(), saved) - - -def test_float64_extension_requires_replay_and_retains_upstream_validation(tmp_path: Path) -> None: - pytest.importorskip("flashinfer_bench") - from benchmarks.kernel_tuning.replay_schema import ReplayDefinition - from flashinfer_bench.data import Definition - from pydantic import ValidationError - - capture = record(tmp_path, lambda: torch.linspace(0, 1, 8, dtype=torch.float64)) - (task,) = render_model(capture, tmp_path / "tasks") - definition = ReplayDefinition.model_validate(task.definition) - assert isinstance(definition, Definition) - assert definition.torch_output_dtypes == [torch.float64] - assert definition.outputs["y0"].dtype == "float64" - assert definition.model_dump(mode="json")["outputs"]["y0"]["dtype"] == "float64" - with pytest.raises(ValidationError): - Definition.model_validate(task.definition) # No process-wide dtype monkeypatch. - wrong = {**task.definition, "outputs": {"y0": {"shape": ["unknown_axis"], "dtype": "float64"}}} - with pytest.raises(ValidationError): - ReplayDefinition.model_validate(wrong) - replay = read_json(task.root / "replay.json") - del replay["definition_schema"] - import json - - (task.root / "replay.json").write_text(json.dumps(replay), encoding="utf-8") - with pytest.raises(ContractError, match="explicit replay definition schema"): - TaskPackage.load(task.root) - - -@pytest.mark.gpu -@pytest.mark.parametrize("operation", ["randn", "float64", "sdpa"]) -def test_cuda_rng_and_float64_baselines_pass_real_evaluator(tmp_path: Path, operation: str) -> None: - """Exercise real Python builders, byte comparison, RNG transitions and eager resets.""" - 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 benchmarks.kernel_tuning.replay_adapter import ReplayAdapter - from scripts.kernel_tuning.contracts import validate_measurement - - state = torch.cuda.get_rng_state().clone() - query = torch.ones(1, 2, 8, 16, device="cuda:0", dtype=torch.bfloat16) - session = CaptureSession("mock_flow_vla", tmp_path / "capture") - try: - with torch.inference_mode(), session.observe(object()): - if operation == "randn": - torch.randn(1, 50, 32, device="cuda:0", dtype=torch.bfloat16, generator=None) - elif operation == "sdpa": - torch.nn.functional.scaled_dot_product_attention(query, query, query, dropout_p=0.0) - else: - times = torch.linspace(0, 1, 16, device="cuda:0", dtype=torch.float64) - (1000.0**times).reciprocal().sin().to(torch.bfloat16) - finally: - session.close(complete=True) - torch.cuda.set_rng_state(state) - tasks = render_model( - calibrate(ModelCapture.load(session.root), tmp_path / "calibrated"), - tmp_path / "tasks", - overrides={"timing.warmup": 1, "timing.iterations": 2, "timing.trials": 1}, - ) - for task in tasks: - adapter = ReplayAdapter(task) - before = torch.cuda.get_rng_state().clone() - report = adapter.evaluate({role: task.baseline() for role in ("baseline", "incumbent", "candidate")}) - validate_measurement(task, report) - if operation == "randn": - assert torch.equal(before, torch.cuda.get_rng_state())