diff --git a/examples/inference/optimizations/README.md b/examples/inference/optimizations/README.md index 70e3bc9fe2..cd139ac3a7 100644 --- a/examples/inference/optimizations/README.md +++ b/examples/inference/optimizations/README.md @@ -52,3 +52,49 @@ Result files use `schema_version: 1` and record wall time, generation time, peak memory, environment identity, optional frames path, and failure reasons. `candidate` mode writes `candidate_result.json`, so it cannot overwrite a previously validated `optimized_result.json`. + +## Generic graph dispatch + +Once a kernel has been packaged as an artifact bundle, FastVideo can run it +without any model-specific code. Point the runtime at a **trusted** directory +of bundles: + +```bash +FASTVIDEO_OPTIMIZATION_ARTIFACT_DIR=/path/to/artifacts \ +FASTVIDEO_OPTIMIZATION_ARTIFACT_MODEL_ID=Wan-AI/Wan2.1-T2V-1.3B-Diffusers \ +FASTVIDEO_OPTIMIZATION_ARTIFACT_DIAGNOSTICS=/tmp/dispatch.json \ +python examples/inference/optimizations/generation_launcher.py \ + --workload /path/to/motionkernel/workloads/wan_t2v_1.3b_480p.yaml \ + --mode candidate \ + --output-dir /tmp/wan_candidate +``` + +Dispatch attaches to repeated block stacks -- children of an `nn.ModuleList` +that share a class -- exactly as capture does, so no architecture is named +anywhere. Per stack and per observed input signature it runs the first call +natively (which is what reveals the output signature), recomputes the module's +graph fingerprint with the capture tracer, and selects a bundle whose +fingerprint, tensor signatures and declared environment all match. The chosen +entry point is called as `candidate(module, *args, **kwargs)`. + +Everything is fail-safe. A missing match, an unloadable bundle, an untraceable +module or an exception raised by the candidate falls back to native execution +and records a structured reason; the candidate is not retried afterwards. + +| Variable | Default | Purpose | +| --- | --- | --- | +| `FASTVIDEO_OPTIMIZATION_ARTIFACT_DIR` | `""` | Trusted artifact root. **Unset means nothing is wrapped at all** and generation is byte-for-byte identical to a build without this feature. | +| `FASTVIDEO_OPTIMIZATION_ARTIFACT_TRACER` | `symbolic` | Tracer used to recompute the fingerprint. `symbolic` does not re-execute the module with real inputs; `export` and `dynamo` do. | +| `FASTVIDEO_OPTIMIZATION_ARTIFACT_MODEL_ID` | profile model id | Model identity matched against each bundle. | +| `FASTVIDEO_OPTIMIZATION_ARTIFACT_MODEL_REVISION` | `*` | Revision matched against each bundle. | +| `FASTVIDEO_OPTIMIZATION_ARTIFACT_DISTRIBUTED_MODE` | auto | Sharding mode. Auto-detection only resolves a single-rank run; a multi-rank run must declare its mode or no artifact is selected. | +| `FASTVIDEO_OPTIMIZATION_ARTIFACT_MAX_SCOPES` | `64` | Upper bound on dispatched block stacks. | +| `FASTVIDEO_OPTIMIZATION_ARTIFACT_MAX_SHAPES` | `8` | Upper bound on resolved input signatures per stack. | +| `FASTVIDEO_OPTIMIZATION_ARTIFACT_DIAGNOSTICS` | `""` | Optional path for the structured dispatch/fallback report. | + +The diagnostics report is metadata only: it records each scope, shape key, +decision reason, artifact id, rejection codes and call counts, plus the +registry and runtime identity. It never contains tensor or prompt data. + +The bundle format, its packager and the matching rules are documented in +MotionKernel's `docs/ARTIFACT_BUNDLE.md`. diff --git a/fastvideo/envs.py b/fastvideo/envs.py index 3d68898a55..57ccb297c7 100644 --- a/fastvideo/envs.py +++ b/fastvideo/envs.py @@ -45,6 +45,14 @@ FASTVIDEO_OPTIMIZATION_PROFILE_FX_TRACER: str = "auto" FASTVIDEO_OPTIMIZATION_PROFILE_FX_MAX_SCOPES: int = 64 FASTVIDEO_OPTIMIZATION_PROFILE_FX_MAX_SHAPES: int = 8 + FASTVIDEO_OPTIMIZATION_ARTIFACT_DIR: str = "" + FASTVIDEO_OPTIMIZATION_ARTIFACT_TRACER: str = "symbolic" + FASTVIDEO_OPTIMIZATION_ARTIFACT_MAX_SCOPES: int = 64 + FASTVIDEO_OPTIMIZATION_ARTIFACT_MAX_SHAPES: int = 8 + FASTVIDEO_OPTIMIZATION_ARTIFACT_MODEL_ID: str = "" + FASTVIDEO_OPTIMIZATION_ARTIFACT_MODEL_REVISION: str = "*" + FASTVIDEO_OPTIMIZATION_ARTIFACT_DISTRIBUTED_MODE: str = "" + FASTVIDEO_OPTIMIZATION_ARTIFACT_DIAGNOSTICS: str = "" FASTVIDEO_TRACE_ACTIVATIONS: bool = False FASTVIDEO_TRACE_LAYERS: str = "" FASTVIDEO_TRACE_STATS: str = "abs_mean,sum" @@ -305,6 +313,39 @@ def maybe_convert_int(value: str | None) -> int | None: "FASTVIDEO_OPTIMIZATION_PROFILE_FX_MAX_SHAPES": lambda: int(os.getenv("FASTVIDEO_OPTIMIZATION_PROFILE_FX_MAX_SHAPES", "8")), + # Trusted directory holding packaged optimization artifacts. Executable + # candidate code is loaded only from here, and only after every declared + # file matches the hash recorded in its manifest. Leaving this unset + # disables graph dispatch entirely: no forward is wrapped and generation + # behaves exactly as it does without the feature. + "FASTVIDEO_OPTIMIZATION_ARTIFACT_DIR": + lambda: os.getenv("FASTVIDEO_OPTIMIZATION_ARTIFACT_DIR", ""), + # Tracer used to recompute a module's graph fingerprint at dispatch time. + # ``symbolic`` is the default because, unlike ``export`` and ``dynamo``, it + # does not re-execute the module with real inputs. + "FASTVIDEO_OPTIMIZATION_ARTIFACT_TRACER": + lambda: os.getenv("FASTVIDEO_OPTIMIZATION_ARTIFACT_TRACER", "symbolic"), + # Upper bound on dispatched block stacks. + "FASTVIDEO_OPTIMIZATION_ARTIFACT_MAX_SCOPES": + lambda: int(os.getenv("FASTVIDEO_OPTIMIZATION_ARTIFACT_MAX_SCOPES", "64")), + # Upper bound on distinct input signatures resolved per stack. + "FASTVIDEO_OPTIMIZATION_ARTIFACT_MAX_SHAPES": + lambda: int(os.getenv("FASTVIDEO_OPTIMIZATION_ARTIFACT_MAX_SHAPES", "8")), + # Model identity matched against each artifact's declared compatibility. + # Falls back to the optimization profile's model id when unset. + "FASTVIDEO_OPTIMIZATION_ARTIFACT_MODEL_ID": + lambda: os.getenv("FASTVIDEO_OPTIMIZATION_ARTIFACT_MODEL_ID", ""), + "FASTVIDEO_OPTIMIZATION_ARTIFACT_MODEL_REVISION": + lambda: os.getenv("FASTVIDEO_OPTIMIZATION_ARTIFACT_MODEL_REVISION", "*"), + # Sharding mode declared to the matcher. Empty means auto-detect, which + # only resolves a single-rank run; a multi-rank run must declare its mode + # explicitly or no artifact is selected. + "FASTVIDEO_OPTIMIZATION_ARTIFACT_DISTRIBUTED_MODE": + lambda: os.getenv("FASTVIDEO_OPTIMIZATION_ARTIFACT_DISTRIBUTED_MODE", ""), + # Optional path for the structured dispatch/fallback report. + "FASTVIDEO_OPTIMIZATION_ARTIFACT_DIAGNOSTICS": + lambda: os.getenv("FASTVIDEO_OPTIMIZATION_ARTIFACT_DIAGNOSTICS", ""), + # Enable activation trace hooks if set. "FASTVIDEO_TRACE_ACTIVATIONS": lambda: bool(os.getenv("FASTVIDEO_TRACE_ACTIVATIONS", "0") != "0"), diff --git a/fastvideo/optimization/__init__.py b/fastvideo/optimization/__init__.py index d8938c52aa..7b0cea5bfa 100644 --- a/fastvideo/optimization/__init__.py +++ b/fastvideo/optimization/__init__.py @@ -1,6 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 -"""Model-independent optimization discovery helpers.""" +"""Model-independent optimization discovery and dispatch helpers.""" +from fastvideo.optimization.dispatch import (attach_graph_dispatch, detach_graph_dispatch) from fastvideo.optimization.profiler import optimization_profile -__all__ = ["optimization_profile"] +__all__ = [ + "attach_graph_dispatch", + "detach_graph_dispatch", + "optimization_profile", +] diff --git a/fastvideo/optimization/artifact.py b/fastvideo/optimization/artifact.py new file mode 100644 index 0000000000..5236dbfe0f --- /dev/null +++ b/fastvideo/optimization/artifact.py @@ -0,0 +1,744 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Trusted artifact bundle loading for generic graph dispatch. + +A bundle is a directory holding an ``artifact.json`` manifest plus the files it +pins by SHA-256. The manifest is produced by MotionKernel, which owns the +normative schema; this module is the consumer half and deliberately re-derives +the checks rather than importing the producer, exactly as +:mod:`fastvideo.optimization.fx_capture` re-derives the graph fingerprint. + +Three rules hold everywhere in this module: + +* Executable code is imported only from inside the explicitly configured + trusted root. Nothing here reads a path from the manifest and follows it. +* Every declared file is hashed and compared before the entry point is + imported. +* Anything unexpected -- an unknown schema version, a missing field, a changed + byte -- is a hard rejection with a structured reason, never a best-effort + load. +""" + +from __future__ import annotations + +import hashlib +import importlib.util +import json +import math +import re +import sys +from dataclasses import dataclass +from pathlib import Path +from typing import Any +from collections.abc import Callable + +from fastvideo.logger import init_logger + +logger = init_logger(__name__) + +#: Manifest format understood here. A bundle declaring anything else is +#: rejected rather than parsed optimistically. +SUPPORTED_ARTIFACT_SCHEMA_VERSION = 1 + +MANIFEST_FILENAME = "artifact.json" + +#: Wildcard accepted by the string-valued compatibility fields. +ANY = "*" + +#: Synthetic package the bundles are imported under, so they never shadow a +#: real module for the rest of the process. +_MODULE_NAMESPACE = "fastvideo._artifacts" + +_IGNORED_DIRECTORIES = frozenset({"__pycache__"}) +_READ_CHUNK_BYTES = 1 << 20 + +_FINGERPRINT_PATTERN = re.compile(r"^[0-9a-f]{32}$") +_SHA256_PATTERN = re.compile(r"^[0-9a-f]{64}$") +_SYMBOL_PATTERN = re.compile(r"^[A-Za-z_][A-Za-z0-9_]{0,127}$") +_RELATIVE_FILE_PATTERN = re.compile(r"^[A-Za-z0-9_][A-Za-z0-9._/-]{0,255}$") + +# Structured rejection reasons. These are logged and exported verbatim, so they +# are part of the contract with the producer. +REASON_FINGERPRINT_MISMATCH = "fingerprint_mismatch" +REASON_INPUT_SIGNATURE_MISMATCH = "input_signature_mismatch" +REASON_OUTPUT_SIGNATURE_MISMATCH = "output_signature_mismatch" +REASON_MODEL_MISMATCH = "model_mismatch" +REASON_REVISION_MISMATCH = "model_revision_mismatch" +REASON_ARCHITECTURE_MISMATCH = "gpu_architecture_mismatch" +REASON_TORCH_VERSION = "torch_version_unsupported" +REASON_CUDA_VERSION = "cuda_version_unsupported" +REASON_TRITON_VERSION = "triton_version_unsupported" +REASON_EXECUTION_MODE = "execution_mode_unsupported" +REASON_DISTRIBUTED_MODE = "distributed_mode_unsupported" +REASON_NOT_PROMOTED = "not_promoted" +REASON_EVIDENCE_INCOMPLETE = "evidence_incomplete" + + +class ArtifactError(ValueError): + """Raised when a bundle is malformed, altered, or unsafe to load.""" + + +def _fail(source: object, location: str, message: str) -> ArtifactError: + return ArtifactError(f"artifact bundle {source!r}: {location}: {message}") + + +def _mapping(value: Any, source: object, location: str) -> dict[str, Any]: + if not isinstance(value, dict) or not value: + raise _fail(source, location, "must be a non-empty object") + for key in value: + if not isinstance(key, str) or not key: + raise _fail(source, location, "keys must be non-empty strings") + return value + + +def _sequence(value: Any, source: object, location: str) -> list[Any]: + if not isinstance(value, list): + raise _fail(source, location, "must be a list") + return value + + +def _text(value: Any, source: object, location: str) -> str: + if not isinstance(value, str) or not value.strip(): + raise _fail(source, location, "must be a non-empty string") + return value + + +def _bool(value: Any, source: object, location: str) -> bool: + if not isinstance(value, bool): + raise _fail(source, location, "must be a bool") + return value + + +def _pattern(value: Any, pattern: re.Pattern[str], source: object, location: str, description: str) -> str: + text = _text(value, source, location) + if not pattern.fullmatch(text): + raise _fail(source, location, f"must be {description}") + return text + + +def _relative_path(value: Any, source: object, location: str) -> str: + text = _pattern(value, _RELATIVE_FILE_PATTERN, source, location, "a relative POSIX path") + if any(part in ("", ".", "..") for part in text.split("/")): + raise _fail(source, location, "must not contain empty or relative segments") + return text + + +def parse_version(text: str | None) -> tuple[int, ...] | None: + """Parse a leading dotted-numeric version, ignoring local/pre-release tags. + + ``"2.8.0+cu128"`` and ``"2.8.0a0"`` both parse to ``(2, 8, 0)``. ``None`` is + returned when there is no numeric prefix, which callers treat as "cannot + compare" rather than "compatible". + """ + if text is None: + return None + parts: list[int] = [] + for chunk in str(text).strip().split("."): + digits = "" + for character in chunk: + if not character.isdigit(): + break + digits += character + if not digits: + break + parts.append(int(digits)) + return tuple(parts) or None + + +def _compare_versions(left: tuple[int, ...], right: tuple[int, ...]) -> int: + width = max(len(left), len(right)) + padded_left = left + (0, ) * (width - len(left)) + padded_right = right + (0, ) * (width - len(right)) + if padded_left < padded_right: + return -1 + return 0 if padded_left == padded_right else 1 + + +@dataclass(frozen=True) +class VersionRange: + """An inclusive-minimum, exclusive-maximum version window.""" + + minimum: str | None = None + maximum_exclusive: str | None = None + + @classmethod + def from_dict(cls, raw: Any, *, source: object, location: str) -> VersionRange: + if raw is None: + return cls() + if not isinstance(raw, dict): + raise _fail(source, location, "must be an object") + unknown = sorted(set(raw) - {"min", "max_exclusive"}) + if unknown: + raise _fail(source, location, f"unknown field(s) {unknown}") + minimum = raw.get("min") + maximum = raw.get("max_exclusive") + if minimum is not None: + minimum = _text(minimum, source, f"{location}.min") + if maximum is not None: + maximum = _text(maximum, source, f"{location}.max_exclusive") + low = parse_version(minimum) + high = parse_version(maximum) + if minimum is not None and low is None: + raise _fail(source, f"{location}.min", "must be a dotted version") + if maximum is not None and high is None: + raise _fail( + source, + f"{location}.max_exclusive", + "must be a dotted version", + ) + if low is not None and high is not None and _compare_versions(low, high) >= 0: + raise _fail(source, location, "min must be lower than max_exclusive") + return cls(minimum=minimum, maximum_exclusive=maximum) + + @property + def unbounded(self) -> bool: + return self.minimum is None and self.maximum_exclusive is None + + def contains(self, version: str | None) -> bool: + """Whether ``version`` satisfies this range. + + A bound that cannot be evaluated -- because the runtime version is + missing or unparsable -- is never assumed to hold. + """ + if self.unbounded: + return True + observed = parse_version(version) + if observed is None: + return False + if self.minimum is not None: + low = parse_version(self.minimum) + if low is None or _compare_versions(observed, low) < 0: + return False + if self.maximum_exclusive is not None: + high = parse_version(self.maximum_exclusive) + if high is None or _compare_versions(observed, high) >= 0: + return False + return True + + def describe(self) -> str: + if self.unbounded: + return "any" + low = self.minimum if self.minimum is not None else "any" + high = self.maximum_exclusive if self.maximum_exclusive is not None else "any" + return f">={low},<{high}" + + +def signature_key(meta: dict[str, Any]) -> tuple[Any, ...]: + """The comparable identity of one tensor signature. + + Built from the capture module's tensor metadata so a live invocation and a + packaged artifact are compared on exactly the same fields. ``name`` is + excluded: argument naming is a runtime detail, not part of the layout. + """ + return ( + tuple(int(dim) for dim in meta.get("shape", ())), + tuple(int(step) for step in meta.get("stride", ())), + str(meta.get("dtype", "")), + str(meta.get("device_type", "")), + bool(meta.get("requires_grad", False)), + ) + + +def _signature_keys(raw: Any, source: object, location: str) -> tuple[tuple[Any, ...], ...]: + items = _sequence(raw, source, location) + if not items: + raise _fail(source, location, "must be a non-empty list") + keys = [] + for index, item in enumerate(items): + entry = _mapping(item, source, f"{location}[{index}]") + unknown = sorted(set(entry) - {"name", "shape", "stride", "dtype", "device_type", "requires_grad"}) + if unknown: + raise _fail( + source, + f"{location}[{index}]", + f"unknown field(s) {unknown}", + ) + shape = _sequence(entry.get("shape"), source, f"{location}[{index}].shape") + stride = _sequence(entry.get("stride"), source, f"{location}[{index}].stride") + for values, name in ((shape, "shape"), (stride, "stride")): + for position, value in enumerate(values): + if isinstance(value, bool) or not isinstance(value, int): + raise _fail( + source, + f"{location}[{index}].{name}[{position}]", + "must be an integer", + ) + if name == "shape" and value < 0: + raise _fail( + source, + f"{location}[{index}].shape[{position}]", + "must be non-negative", + ) + if len(shape) != len(stride): + raise _fail( + source, + f"{location}[{index}].stride", + "must have the same length as shape", + ) + _text(entry.get("dtype"), source, f"{location}[{index}].dtype") + _text(entry.get("device_type"), source, f"{location}[{index}].device_type") + _bool( + entry.get("requires_grad", False), + source, + f"{location}[{index}].requires_grad", + ) + keys.append(signature_key(entry)) + return tuple(keys) + + +@dataclass(frozen=True) +class ArtifactFile: + """One bundled file, pinned by size and content hash.""" + + path: str + sha256: str + size: int + + +@dataclass(frozen=True) +class ArtifactManifest: + """The subset of a bundle manifest this runtime acts on.""" + + artifact_id: str + graph_fingerprint: str + operation_name: str + parent_module: str + input_keys: tuple[tuple[Any, ...], ...] + output_keys: tuple[tuple[Any, ...], ...] + entry_file: str + entry_symbol: str + files: tuple[ArtifactFile, ...] + model_id: str + model_revision: str + gpu_architectures: tuple[str, ...] + torch_range: VersionRange + cuda_range: VersionRange + triton_range: VersionRange + execution_modes: tuple[str, ...] + distributed_modes: tuple[str, ...] + promotion_decision: str + evidence_passed: bool + speedup: float + directory: Path + + @classmethod + def from_dict(cls, raw_value: Any, *, directory: Path) -> ArtifactManifest: + source = str(directory) + raw = _mapping(raw_value, source, "top level") + version = raw.get("schema_version") + if isinstance(version, bool) or not isinstance(version, int): + raise _fail(source, "schema_version", "must be an integer") + if version != SUPPORTED_ARTIFACT_SCHEMA_VERSION: + raise _fail( + source, + "schema_version", + f"unsupported version {version}; expected " + f"{SUPPORTED_ARTIFACT_SCHEMA_VERSION}", + ) + + operation = _mapping(raw.get("operation"), source, "operation") + signature = _mapping(raw.get("signature"), source, "signature") + entry_point = _mapping(raw.get("entry_point"), source, "entry_point") + compatibility = _mapping(raw.get("compatibility"), source, "compatibility") + evidence = _mapping(raw.get("evidence"), source, "evidence") + promotion = _mapping(raw.get("promotion"), source, "promotion") + + files = [] + for index, item in enumerate(_sequence(raw.get("files"), source, "files")): + entry = _mapping(item, source, f"files[{index}]") + size = entry.get("bytes") + if isinstance(size, bool) or not isinstance(size, int) or size < 0: + raise _fail(source, f"files[{index}].bytes", "must be a non-negative integer") + files.append( + ArtifactFile( + path=_relative_path(entry.get("path"), source, f"files[{index}].path"), + sha256=_pattern( + entry.get("sha256"), + _SHA256_PATTERN, + source, + f"files[{index}].sha256", + "64 lowercase hex characters", + ), + size=size, + )) + if not files: + raise _fail(source, "files", "must be a non-empty list") + paths = [item.path for item in files] + if len(paths) != len(set(paths)): + raise _fail(source, "files", "contains duplicate paths") + + entry_file = _relative_path(entry_point.get("file"), source, "entry_point.file") + if entry_file not in paths: + raise _fail(source, "entry_point.file", f"{entry_file!r} is not a declared file") + if not entry_file.endswith(".py"): + raise _fail(source, "entry_point.file", "must be a .py file") + + benchmark = _mapping(evidence.get("benchmark"), source, "evidence.benchmark") + generation = _mapping(evidence.get("generation"), source, "evidence.generation") + speedup = benchmark.get("speedup") + if (isinstance(speedup, bool) or not isinstance(speedup, int | float) or not math.isfinite(float(speedup)) + or float(speedup) < 0): + raise _fail( + source, + "evidence.benchmark.speedup", + "must be a finite non-negative number", + ) + + architectures = _sequence( + compatibility.get("gpu_architectures"), + source, + "compatibility.gpu_architectures", + ) + execution_modes = _sequence( + compatibility.get("execution_modes"), + source, + "compatibility.execution_modes", + ) + distributed_modes = _sequence( + compatibility.get("distributed_modes"), + source, + "compatibility.distributed_modes", + ) + for items, location in ( + (architectures, "compatibility.gpu_architectures"), + (execution_modes, "compatibility.execution_modes"), + (distributed_modes, "compatibility.distributed_modes"), + ): + if not items: + raise _fail(source, location, "must be a non-empty list") + + return cls( + artifact_id=_text(raw.get("artifact_id"), source, "artifact_id"), + graph_fingerprint=_pattern( + operation.get("graph_fingerprint"), + _FINGERPRINT_PATTERN, + source, + "operation.graph_fingerprint", + "32 lowercase hex characters", + ), + operation_name=_text(operation.get("name"), source, "operation.name"), + parent_module=_text(operation.get("parent_module"), source, "operation.parent_module"), + input_keys=_signature_keys(signature.get("inputs"), source, "signature.inputs"), + output_keys=_signature_keys(signature.get("outputs"), source, "signature.outputs"), + entry_file=entry_file, + entry_symbol=_pattern( + entry_point.get("symbol"), + _SYMBOL_PATTERN, + source, + "entry_point.symbol", + "a Python identifier", + ), + files=tuple(files), + model_id=_text(compatibility.get("model_id"), source, "compatibility.model_id"), + model_revision=_text( + compatibility.get("model_revision"), + source, + "compatibility.model_revision", + ), + gpu_architectures=tuple( + _text(item, source, f"compatibility.gpu_architectures[{index}]") + for index, item in enumerate(architectures)), + torch_range=VersionRange.from_dict( + compatibility.get("torch"), + source=source, + location="compatibility.torch", + ), + cuda_range=VersionRange.from_dict( + compatibility.get("cuda"), + source=source, + location="compatibility.cuda", + ), + triton_range=VersionRange.from_dict( + compatibility.get("triton"), + source=source, + location="compatibility.triton", + ), + execution_modes=tuple( + _text(item, source, f"compatibility.execution_modes[{index}]") + for index, item in enumerate(execution_modes)), + distributed_modes=tuple( + _text(item, source, f"compatibility.distributed_modes[{index}]") + for index, item in enumerate(distributed_modes)), + promotion_decision=_text(promotion.get("decision"), source, "promotion.decision"), + evidence_passed=(_bool(benchmark.get("passed"), source, "evidence.benchmark.passed") + and _bool(generation.get("passed"), source, "evidence.generation.passed")), + speedup=float(speedup), + directory=directory, + ) + + +@dataclass(frozen=True) +class RuntimeProfile: + """The environment this process is executing in.""" + + model_id: str + model_revision: str + gpu_architecture: str + torch_version: str + cuda_version: str | None = None + triton_version: str | None = None + execution_mode: str = "inference" + distributed_mode: str = "single" + + @classmethod + def detect( + cls, + *, + model_id: str, + model_revision: str = ANY, + execution_mode: str = "inference", + distributed_mode: str = "single", + ) -> RuntimeProfile: + """Read torch/CUDA/Triton identity from the running process.""" + import torch + + architecture = "cpu" + cuda_version = None + if torch.cuda.is_available(): + index = torch.cuda.current_device() + major, minor = torch.cuda.get_device_capability(index) + architecture = f"sm{major}{minor}" + cuda_version = getattr(torch.version, "cuda", None) + try: + import triton + + triton_version = getattr(triton, "__version__", None) + except Exception: # noqa: BLE001 - Triton is optional on every platform + triton_version = None + return cls( + model_id=model_id, + model_revision=model_revision, + gpu_architecture=architecture, + torch_version=str(getattr(torch, "__version__", "")), + cuda_version=cuda_version, + triton_version=triton_version, + execution_mode=execution_mode, + distributed_mode=distributed_mode, + ) + + +def _wildcard_equal(declared: str, observed: str) -> bool: + return declared in (ANY, observed) + + +def check_compatibility( + manifest: ArtifactManifest, + *, + graph_fingerprint: str, + input_keys: tuple[tuple[Any, ...], ...], + output_keys: tuple[tuple[Any, ...], ...], + runtime: RuntimeProfile, +) -> str | None: + """Return the reason ``manifest`` cannot serve this call, or ``None``. + + Ordering is cheapest-first: graph identity, then tensor layout, then the + declared environment window. + """ + if manifest.graph_fingerprint != graph_fingerprint: + return REASON_FINGERPRINT_MISMATCH + if manifest.input_keys != input_keys: + return REASON_INPUT_SIGNATURE_MISMATCH + if manifest.output_keys != output_keys: + return REASON_OUTPUT_SIGNATURE_MISMATCH + if not _wildcard_equal(manifest.model_id, runtime.model_id): + return REASON_MODEL_MISMATCH + if not _wildcard_equal(manifest.model_revision, runtime.model_revision): + return REASON_REVISION_MISMATCH + if not any(_wildcard_equal(item, runtime.gpu_architecture) for item in manifest.gpu_architectures): + return REASON_ARCHITECTURE_MISMATCH + if not manifest.torch_range.contains(runtime.torch_version): + return REASON_TORCH_VERSION + if not manifest.cuda_range.contains(runtime.cuda_version): + return REASON_CUDA_VERSION + if not manifest.triton_range.contains(runtime.triton_version): + return REASON_TRITON_VERSION + if runtime.execution_mode not in manifest.execution_modes: + return REASON_EXECUTION_MODE + if runtime.distributed_mode not in manifest.distributed_modes: + return REASON_DISTRIBUTED_MODE + if manifest.promotion_decision != "promoted": + return REASON_NOT_PROMOTED + if not manifest.evidence_passed: + return REASON_EVIDENCE_INCOMPLETE + return None + + +def file_sha256(path: Path) -> str: + """Content hash of one file, streamed so large bundles stay cheap.""" + digest = hashlib.sha256() + with open(path, "rb") as handle: + while True: + chunk = handle.read(_READ_CHUNK_BYTES) + if not chunk: + break + digest.update(chunk) + return digest.hexdigest() + + +def _bundle_contents(directory: Path) -> set[str]: + present = set() + for path in directory.rglob("*"): + if path.is_dir(): + continue + relative = path.relative_to(directory) + if any(part in _IGNORED_DIRECTORIES for part in relative.parts): + continue + present.add(relative.as_posix()) + return present + + +def verify_bundle(directory: Path) -> ArtifactManifest: + """Parse a bundle and confirm every declared file is byte-for-byte intact. + + Undeclared files are refused as well: an attacker who can drop an extra + module next to a signed entry point could otherwise have it imported. + """ + manifest_path = directory / MANIFEST_FILENAME + if not manifest_path.is_file(): + raise _fail(str(directory), MANIFEST_FILENAME, "not found") + try: + raw = json.loads(manifest_path.read_text(encoding="utf-8")) + except json.JSONDecodeError as exc: + raise _fail(str(directory), MANIFEST_FILENAME, f"invalid JSON: {exc}") from exc + manifest = ArtifactManifest.from_dict(raw, directory=directory) + + declared = {item.path for item in manifest.files} + undeclared = sorted(_bundle_contents(directory) - declared - {MANIFEST_FILENAME}) + if undeclared: + raise _fail(str(directory), "files", f"undeclared file(s) {undeclared}") + + for entry in sorted(manifest.files, key=lambda item: item.path): + path = directory / entry.path + if not path.is_file(): + raise _fail(str(directory), "files", f"{entry.path!r} is missing") + size = path.stat().st_size + if size != entry.size: + raise _fail( + str(directory), + "files", + f"{entry.path!r} is {size} bytes, manifest records {entry.size}", + ) + actual = file_sha256(path) + if actual != entry.sha256: + raise _fail( + str(directory), + "files", + f"{entry.path!r} hash {actual} does not match manifest {entry.sha256}", + ) + return manifest + + +def _resolve_inside(root: Path, candidate: Path) -> Path: + """Resolve ``candidate`` and require it to stay under ``root``.""" + resolved_root = root.resolve(strict=True) + resolved = candidate.resolve(strict=True) + if resolved != resolved_root and resolved_root not in resolved.parents: + raise ArtifactError(f"artifact bundle {str(candidate)!r}: resolves outside the trusted " + f"root {str(resolved_root)!r}") + return resolved + + +def load_entry_point(manifest: ArtifactManifest, *, trusted_root: Path) -> Callable[..., Any]: + """Re-verify a bundle and import its candidate callable. + + The bundle is verified again here even if the caller already validated it: + the gap between validation and import is precisely where a swapped file + would land. + """ + directory = _resolve_inside(trusted_root, manifest.directory) + verified = verify_bundle(directory) + if verified.artifact_id != manifest.artifact_id: + raise _fail( + str(directory), + "artifact_id", + f"changed from {manifest.artifact_id!r} to {verified.artifact_id!r} " + "since validation", + ) + + entry_file = _resolve_inside(directory, directory / verified.entry_file) + readable_id = re.sub(r"[^A-Za-z0-9_]", "_", verified.artifact_id) + unique_id = hashlib.sha256(verified.artifact_id.encode("utf-8")).hexdigest()[:16] + module_name = f"{_MODULE_NAMESPACE}.{readable_id}_{unique_id}" + spec = importlib.util.spec_from_file_location(module_name, entry_file) + if spec is None or spec.loader is None: + raise _fail(str(directory), "entry_point", f"cannot load {verified.entry_file!r}") + module = importlib.util.module_from_spec(spec) + # Registered before execution so the module can look itself up, and removed + # again on failure so a half-initialized module is never reachable. + sys.modules[module_name] = module + try: + spec.loader.exec_module(module) + except Exception as exc: # noqa: BLE001 - untrusted code, any failure is a rejection + sys.modules.pop(module_name, None) + raise _fail( + str(directory), + "entry_point", + f"importing {verified.entry_file!r} raised {type(exc).__name__}", + ) from exc + + candidate = getattr(module, verified.entry_symbol, None) + if candidate is None or not callable(candidate): + sys.modules.pop(module_name, None) + raise _fail( + str(directory), + "entry_point", + f"{verified.entry_symbol!r} is not a callable in {verified.entry_file!r}", + ) + return candidate + + +def discover_bundles(root: Path) -> list[Path]: + """List bundle directories under ``root``: the root itself plus children.""" + if not root.is_dir(): + return [] + found = [] + if (root / MANIFEST_FILENAME).is_file(): + found.append(root) + for child in sorted(root.iterdir()): + if child.is_dir() and (child / MANIFEST_FILENAME).is_file(): + found.append(child) + return found + + +class ArtifactRegistry: + """Every valid bundle found in the trusted artifact directory. + + Loading is eager and one-shot: bundles are parsed and hashed when the + registry is built, so a corrupt artifact surfaces before generation starts + rather than mid-run. Bundles that fail validation are recorded in + :attr:`errors` and skipped; they never disable the ones that passed. + """ + + def __init__(self, root: Path | None) -> None: + self.root = root + self.manifests: list[ArtifactManifest] = [] + self.errors: list[str] = [] + if root is None: + return + if not root.is_dir(): + self.errors.append(f"artifact directory {str(root)!r}: not a directory") + return + for path in discover_bundles(root): + try: + self.manifests.append(verify_bundle(path)) + except ArtifactError as exc: + self.errors.append(str(exc)) + + @property + def enabled(self) -> bool: + return bool(self.manifests) + + def candidates_for(self, input_keys: tuple[tuple[Any, ...], ...]) -> list[ArtifactManifest]: + """Bundles whose input layout matches, before any graph is traced. + + This is the cheap pre-filter that keeps dispatch from tracing a module + when the store holds nothing shaped like the live call. + """ + return [item for item in self.manifests if item.input_keys == input_keys] + + def summary(self) -> dict[str, Any]: + return { + "root": str(self.root) if self.root is not None else None, + "loaded": len(self.manifests), + "artifact_ids": sorted(item.artifact_id for item in self.manifests), + "errors": list(self.errors), + } diff --git a/fastvideo/optimization/dispatch.py b/fastvideo/optimization/dispatch.py new file mode 100644 index 0000000000..dd38eb84f8 --- /dev/null +++ b/fastvideo/optimization/dispatch.py @@ -0,0 +1,511 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Generic graph dispatch: run a packaged kernel, or fall back to native. + +Dispatch attaches to the same model-independent structure the capture pipeline +uses -- stacks of identically-typed children under an ``nn.ModuleList`` -- and +decides, per stack and per observed input signature, whether a trusted artifact +may replace the native forward. Nothing in this module knows the name of a +model, a block class, or an architecture; a new model is supported by +publishing an artifact, never by editing code here. + +The decision sequence for one scope and input signature is: + +1. The first call always runs natively. It is what reveals the output + signature, which is part of the artifact's identity. +2. The registry is pre-filtered on the input layout. If nothing in the store + is shaped like this call, no graph is ever traced. +3. The module is traced once to recompute its graph fingerprint, using the + capture module so the value matches what the producer recorded. +4. Compatibility is checked, the winning bundle is re-verified and imported, + and the callable is cached for every later call with that signature. + +Every step is guarded. A failure at any point demotes the signature to native +execution permanently and records a structured reason; it never propagates to +the caller. With no artifact directory configured, nothing is attached at all +and the model runs byte-for-byte as it does today. + +Candidate calling convention +---------------------------- +A bundle's entry point is called as ``candidate(module, *args, **kwargs)``: +the native module comes first, then the exact arguments its ``forward`` was +given. Passing the module is what lets one artifact serve every block in a +stack -- the kernel reads the parameters it needs from the module it was handed +instead of the loader having to know which parameters exist. +""" + +from __future__ import annotations + +import json +from collections import defaultdict +from contextlib import contextmanager +from dataclasses import dataclass +from pathlib import Path +from typing import Any +from collections.abc import Callable + +from torch import nn + +from fastvideo import envs +from fastvideo.logger import init_logger +from fastvideo.optimization.artifact import ( + ArtifactManifest, + ArtifactRegistry, + RuntimeProfile, + check_compatibility, + load_entry_point, + signature_key, +) +from fastvideo.optimization.fx_capture import default_capture_targets +from fastvideo.optimization.identity import ( + graph_identity, + input_signatures, + output_signatures, + shape_key_for, +) + +logger = init_logger(__name__) + +#: Diagnostic report format. Bumped independently of the capture and profiler +#: schemas so a consumer can tell the three apart. +DISPATCH_SCHEMA_VERSION = 1 + +# Structured fallback reasons. Consumers group on these, so they are stable. +REASON_SELECTED = "artifact_selected" +REASON_NO_SIGNATURE_MATCH = "no_artifact_for_input_signature" +REASON_NO_COMPATIBLE_ARTIFACT = "no_compatible_artifact" +REASON_NO_TENSOR_INPUTS = "no_tensor_inputs" +_REASON_IDENTITY_PREFIX = "graph_identity_unavailable" +_REASON_LOAD_PREFIX = "artifact_load_failed" +_REASON_RUNTIME_PREFIX = "candidate_runtime_error" +_REASON_DECISION_PREFIX = "decision_failed" + + +@dataclass +class _Decision: + """What to do for one scope and input signature, and why.""" + + candidate: Callable[..., Any] | None + reason: str + artifact_id: str | None = None + rejections: tuple[str, ...] = () + #: Calls made *after* this decision was resolved. The one call that + #: produced it ran natively and is not counted here. + calls: int = 0 + candidate_calls: int = 0 + runtime_fallbacks: int = 0 + + def as_dict(self, scope: str, shape_key: str) -> dict[str, Any]: + return { + "scope": scope, + "shape_key": shape_key, + "reason": self.reason, + "artifact_id": self.artifact_id, + "rejections": list(self.rejections), + "calls": self.calls, + "candidate_calls": self.candidate_calls, + "runtime_fallbacks": self.runtime_fallbacks, + "active": self.candidate is not None, + } + + +@dataclass +class _Wrapper: + """One patched module, remembered so ``detach`` can restore it exactly.""" + + scope: str + module: nn.Module + native_forward: Callable[..., Any] + had_own_forward: bool = False + + +class GraphDispatchSession: + """Route repeated module stacks through trusted artifacts when possible.""" + + def __init__( + self, + registry: ArtifactRegistry, + runtime: RuntimeProfile, + *, + tracer: str = "symbolic", + max_scopes: int = 64, + max_shape_variants: int = 8, + ) -> None: + self.registry = registry + self.runtime = runtime + self.tracer = tracer + self.max_scopes = max_scopes + self.max_shape_variants = max_shape_variants + self._wrappers: list[_Wrapper] = [] + self._scopes: set[str] = set() + self._decisions: dict[tuple[str, str], _Decision] = {} + self._dropped_variants: dict[str, int] = defaultdict(int) + self._dropped_scopes = 0 + self._errors: list[str] = [] + + # -- attach --------------------------------------------------------- + + def attach_modules(self, modules: dict[str, Any] | None) -> int: + """Wrap every repeated block stack in a pipeline's module mapping.""" + if not modules or not self.registry.enabled: + return 0 + wrapped = 0 + for name, module in modules.items(): + if isinstance(module, nn.Module): + wrapped += self.attach(module, prefix=str(name)) + return wrapped + + def attach(self, root: nn.Module, *, prefix: str = "") -> int: + """Wrap every repeated block stack under ``root``.""" + try: + targets = default_capture_targets(root) + except Exception as exc: # noqa: BLE001 - dispatch never breaks generation + self._record_exception("target_selection_failed", exc) + return 0 + wrapped = 0 + for scope, module in targets: + qualified = f"{prefix}.{scope}" if prefix else scope + if qualified not in self._scopes and len(self._scopes) >= self.max_scopes: + self._dropped_scopes += 1 + continue + try: + self._install(qualified, module) + self._scopes.add(qualified) + wrapped += 1 + except Exception as exc: # noqa: BLE001 + self._record_exception("install_failed", exc, scope=qualified) + return wrapped + + def _install(self, scope: str, module: nn.Module) -> None: + """Replace ``module.forward`` with the dispatching wrapper.""" + native_forward = module.forward + wrapper = _Wrapper( + scope=scope, + module=module, + native_forward=native_forward, + had_own_forward="forward" in vars(module), + ) + + def dispatching_forward(*args: Any, **kwargs: Any) -> Any: + return self._dispatch(wrapper, args, kwargs) + + module.forward = dispatching_forward # type: ignore[method-assign] + self._wrappers.append(wrapper) + + def detach(self) -> None: + """Restore every native forward. Safe to call more than once.""" + for wrapper in reversed(self._wrappers): + try: + if wrapper.had_own_forward: + wrapper.module.forward = wrapper.native_forward # type: ignore[method-assign] + else: + # The wrapper shadowed the class-level method; deleting the + # instance attribute restores the original binding exactly. + del wrapper.module.forward # type: ignore[misc] + except Exception as exc: # noqa: BLE001 + self._record_exception("detach_failed", exc, scope=wrapper.scope) + self._wrappers.clear() + + # -- dispatch ------------------------------------------------------- + + def _dispatch(self, wrapper: _Wrapper, args: tuple[Any, ...], kwargs: dict[str, Any]) -> Any: + native = wrapper.native_forward + try: + input_metas = input_signatures(args, kwargs) + shape_key = shape_key_for(input_metas) + except Exception as exc: # noqa: BLE001 + self._record_exception("shape_key_failed", exc, scope=wrapper.scope) + return native(*args, **kwargs) + + key = (wrapper.scope, shape_key) + decision = self._decisions.get(key) + if decision is None: + # First call for this signature: run native, learn the output + # layout, then decide once for every later call. + output = native(*args, **kwargs) + self._decisions[key] = self._decide(wrapper, args, kwargs, output, input_metas) + return output + + decision.calls += 1 + if decision.candidate is None: + return native(*args, **kwargs) + try: + result = decision.candidate(wrapper.module, *args, **kwargs) + except Exception as exc: # noqa: BLE001 - untrusted candidate code + # Demote permanently: a candidate that raised once is not trusted + # to be retried thousands of times over the rest of the run. + decision.runtime_fallbacks += 1 + decision.candidate = None + decision.reason = f"{_REASON_RUNTIME_PREFIX}:{type(exc).__name__}" + logger.warning( + "Artifact %s failed at runtime for %s; falling back to native " + "execution for the rest of this run", + decision.artifact_id, + wrapper.scope, + exc_info=True, + ) + return native(*args, **kwargs) + decision.candidate_calls += 1 + return result + + def _decide( + self, + wrapper: _Wrapper, + args: tuple[Any, ...], + kwargs: dict[str, Any], + output: Any, + input_metas: list[dict[str, Any]], + ) -> _Decision: + """Resolve one signature to an artifact or to a native fallback.""" + try: + return self._resolve(wrapper, args, kwargs, output, input_metas) + except Exception as exc: # noqa: BLE001 - resolution is best effort + self._record_exception("decide_failed", exc, scope=wrapper.scope) + return _Decision(None, f"{_REASON_DECISION_PREFIX}:{type(exc).__name__}") + + def _resolve( + self, + wrapper: _Wrapper, + args: tuple[Any, ...], + kwargs: dict[str, Any], + output: Any, + input_metas: list[dict[str, Any]], + ) -> _Decision: + scope = wrapper.scope + if self._variant_count(scope) >= self.max_shape_variants: + self._dropped_variants[scope] += 1 + return _Decision(None, "shape_variant_budget_exhausted") + + if not input_metas: + return _Decision(None, REASON_NO_TENSOR_INPUTS) + input_keys = tuple(signature_key(meta) for meta in input_metas) + candidates = self.registry.candidates_for(input_keys) + if not candidates: + return _Decision(None, REASON_NO_SIGNATURE_MATCH) + + try: + with self._native_forward_for_identity(wrapper): + region = graph_identity( + wrapper.module, + args, + kwargs, + output, + scope=scope, + tracer=self.tracer, + ) + except Exception as exc: # noqa: BLE001 - untraceable modules stay native + reason = f"{_REASON_IDENTITY_PREFIX}:{type(exc).__name__}" + logger.info("Graph identity unavailable for %s: %s", scope, reason) + return _Decision(None, reason) + + fingerprint = str(region.get("fingerprint", "")) + output_keys = tuple(signature_key(meta) for meta in output_signatures(output)) + matched: list[ArtifactManifest] = [] + rejections: list[str] = [] + for manifest in candidates: + compatibility_reason = check_compatibility( + manifest, + graph_fingerprint=fingerprint, + input_keys=input_keys, + output_keys=output_keys, + runtime=self.runtime, + ) + if compatibility_reason is None: + matched.append(manifest) + else: + rejections.append(f"{manifest.artifact_id}:{compatibility_reason}") + if not matched: + return _Decision( + None, + REASON_NO_COMPATIBLE_ARTIFACT, + rejections=tuple(sorted(rejections)), + ) + + best = max(matched, key=lambda item: (item.speedup, item.artifact_id)) + rejections.extend(f"{item.artifact_id}:not_selected" for item in matched if item is not best) + root = self.registry.root + if root is None: + return _Decision(None, f"{_REASON_LOAD_PREFIX}:NoTrustedRoot") + try: + candidate = load_entry_point(best, trusted_root=root) + except Exception as exc: # noqa: BLE001 - a bad bundle must not stop generation + logger.warning("Artifact %s could not be loaded; staying native", best.artifact_id, exc_info=True) + return _Decision( + None, + f"{_REASON_LOAD_PREFIX}:{type(exc).__name__}", + artifact_id=best.artifact_id, + rejections=tuple(sorted(rejections)), + ) + logger.info( + "Dispatching %s to artifact %s (fingerprint %s)", + scope, + best.artifact_id, + fingerprint, + ) + return _Decision( + candidate, + REASON_SELECTED, + artifact_id=best.artifact_id, + rejections=tuple(sorted(rejections)), + ) + + @contextmanager + def _native_forward_for_identity(self, wrapper: _Wrapper): + """Expose the real forward while export/Dynamo recompute identity. + + Symbolic FX traces the class method, but export and Dynamo call the + instance forward. Leaving the dispatch wrapper installed would recurse + back into this session during identity capture. + """ + installed_forward = wrapper.module.forward + wrapper.module.forward = wrapper.native_forward # type: ignore[method-assign] + try: + yield + finally: + wrapper.module.forward = installed_forward # type: ignore[method-assign] + + def _variant_count(self, scope: str) -> int: + return sum(1 for existing_scope, _ in self._decisions if existing_scope == scope) + + # -- diagnostics ---------------------------------------------------- + + def _record_error(self, message: str) -> None: + if len(self._errors) < 64: + self._errors.append(message) + + def _record_exception(self, code: str, exc: Exception, *, scope: str | None = None) -> None: + # Exception text can carry argument reprs; only the type is recorded. + location = f"[{scope}]" if scope else "" + self._record_error(f"{code}{location}:{type(exc).__name__}") + + def diagnostics(self) -> dict[str, Any]: + """A metadata-only report of every dispatch decision made.""" + decisions = [ + decision.as_dict(scope, shape_key) for (scope, shape_key), decision in sorted(self._decisions.items()) + ] + reason_counts: dict[str, int] = defaultdict(int) + for decision in decisions: + reason_counts[str(decision["reason"])] += 1 + return { + "dispatch": { + "dispatch_schema_version": DISPATCH_SCHEMA_VERSION, + "registry": self.registry.summary(), + "runtime": { + "model_id": self.runtime.model_id, + "model_revision": self.runtime.model_revision, + "gpu_architecture": self.runtime.gpu_architecture, + "torch_version": self.runtime.torch_version, + "cuda_version": self.runtime.cuda_version, + "triton_version": self.runtime.triton_version, + "execution_mode": self.runtime.execution_mode, + "distributed_mode": self.runtime.distributed_mode, + }, + "tracer": self.tracer, + "scopes": sorted(self._scopes), + "dropped_scopes": self._dropped_scopes, + "dropped_shape_variants": dict(sorted(self._dropped_variants.items())), + "errors": list(self._errors), + "reason_counts": dict(sorted(reason_counts.items())), + }, + "decisions": decisions, + } + + def write_diagnostics(self, path: str | Path) -> Path | None: + """Write the diagnostic report, returning ``None`` if that failed.""" + output = Path(path).expanduser() + try: + output.parent.mkdir(parents=True, exist_ok=True) + temporary = output.with_suffix(output.suffix + ".tmp") + temporary.write_text( + json.dumps(self.diagnostics(), indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + temporary.replace(output) + except Exception: # noqa: BLE001 - diagnostics must never break a run + logger.warning("Could not write dispatch diagnostics to %s", output, exc_info=True) + return None + return output + + +def _execution_mode(modules: dict[str, Any] | None) -> str: + """Report ``training`` when any hooked module is in training mode.""" + for module in (modules or {}).values(): + if isinstance(module, nn.Module) and any(child.training for child in module.modules()): + return "training" + return "inference" + + +def _distributed_mode() -> str: + """Report the sharding mode this process runs in. + + A multi-rank run whose mode is not declared explicitly reports + ``unspecified``, which matches no artifact: an unsharded kernel silently + applied to a sharded module would be a correctness bug, so the ambiguous + case fails closed. + """ + declared = envs.FASTVIDEO_OPTIMIZATION_ARTIFACT_DISTRIBUTED_MODE + if declared: + return declared + try: + import torch.distributed as distributed + + if distributed.is_available() and distributed.is_initialized(): + return "single" if distributed.get_world_size() == 1 else "unspecified" + except Exception: # noqa: BLE001 - absence of distributed means single process + return "single" + return "single" + + +def attach_graph_dispatch(modules: dict[str, Any] | None) -> GraphDispatchSession | None: + """Attach generic dispatch, or return ``None`` when it is not configured. + + Returning ``None`` is the zero-effect path: no forward is wrapped, no graph + is traced, and no artifact code is read. This is what makes an unset + ``FASTVIDEO_OPTIMIZATION_ARTIFACT_DIR`` behaviorally identical to a build + without this feature. + """ + configured = envs.FASTVIDEO_OPTIMIZATION_ARTIFACT_DIR + if not configured: + return None + try: + root = Path(configured).expanduser() + registry = ArtifactRegistry(root) + for error in registry.errors: + logger.warning("Skipping artifact bundle: %s", error) + if not registry.enabled: + logger.info("No usable artifacts in %s; running natively", root) + return None + runtime = RuntimeProfile.detect( + model_id=(envs.FASTVIDEO_OPTIMIZATION_ARTIFACT_MODEL_ID or envs.FASTVIDEO_OPTIMIZATION_PROFILE_MODEL_ID), + model_revision=envs.FASTVIDEO_OPTIMIZATION_ARTIFACT_MODEL_REVISION, + execution_mode=_execution_mode(modules), + distributed_mode=_distributed_mode(), + ) + session = GraphDispatchSession( + registry, + runtime, + tracer=envs.FASTVIDEO_OPTIMIZATION_ARTIFACT_TRACER, + max_scopes=envs.FASTVIDEO_OPTIMIZATION_ARTIFACT_MAX_SCOPES, + max_shape_variants=envs.FASTVIDEO_OPTIMIZATION_ARTIFACT_MAX_SHAPES, + ) + wrapped = session.attach_modules(modules) + except Exception: # noqa: BLE001 - dispatch setup never breaks generation + logger.warning("Graph dispatch could not be started; running natively", exc_info=True) + return None + if not wrapped: + logger.info("Graph dispatch found no repeated module stacks; running natively") + return None + logger.info("Graph dispatch attached to %d module(s) from %s", wrapped, registry.root) + return session + + +def detach_graph_dispatch(session: GraphDispatchSession | None) -> None: + """Restore every native forward and write diagnostics if requested.""" + if session is None: + return + try: + session.detach() + except Exception: # noqa: BLE001 + logger.warning("Graph dispatch detach failed", exc_info=True) + output = envs.FASTVIDEO_OPTIMIZATION_ARTIFACT_DIAGNOSTICS + if output: + session.write_diagnostics(output) diff --git a/fastvideo/optimization/fx_capture.py b/fastvideo/optimization/fx_capture.py index 3c619e92e0..8b09524783 100644 --- a/fastvideo/optimization/fx_capture.py +++ b/fastvideo/optimization/fx_capture.py @@ -159,22 +159,18 @@ def _tensor_leaves( leaves: list[tuple[str, torch.Tensor]] = [] if isinstance(value, tuple | list): for index, item in enumerate(value): - leaves.extend( - _tensor_leaves( - item, - prefix=f"{prefix}_{index}", - seen=seen, - ) - ) + leaves.extend(_tensor_leaves( + item, + prefix=f"{prefix}_{index}", + seen=seen, + )) elif isinstance(value, Mapping): for key, item in value.items(): - leaves.extend( - _tensor_leaves( - item, - prefix=f"{prefix}_{_safe_name(str(key))}", - seen=seen, - ) - ) + leaves.extend(_tensor_leaves( + item, + prefix=f"{prefix}_{_safe_name(str(key))}", + seen=seen, + )) elif is_dataclass(value) and not isinstance(value, type): for item in fields(value): leaves.extend( @@ -182,8 +178,7 @@ def _tensor_leaves( getattr(value, item.name), prefix=f"{prefix}_{_safe_name(item.name)}", seen=seen, - ) - ) + )) return leaves @@ -194,25 +189,18 @@ def _input_metas( leaves: list[tuple[str, torch.Tensor]] = [] seen: set[int] = set() for index, value in enumerate(args): - leaves.extend( - _tensor_leaves(value, prefix=f"input_{index}", seen=seen) - ) + leaves.extend(_tensor_leaves(value, prefix=f"input_{index}", seen=seen)) for key, value in kwargs.items(): - leaves.extend( - _tensor_leaves( - value, - prefix=f"kwarg_{_safe_name(str(key))}", - seen=seen, - ) - ) + leaves.extend(_tensor_leaves( + value, + prefix=f"kwarg_{_safe_name(str(key))}", + seen=seen, + )) return [_tensor_meta(name, tensor) for name, tensor in leaves] def _output_metas(output: Any) -> list[dict[str, Any]]: - return [ - _tensor_meta(name, tensor) - for name, tensor in _tensor_leaves(output, prefix="output") - ] + return [_tensor_meta(name, tensor) for name, tensor in _tensor_leaves(output, prefix="output")] def _input_shape_key(args: tuple[Any, ...], kwargs: dict[str, Any]) -> str: @@ -255,25 +243,19 @@ def _sanitize_constants(raw: dict[str, Any]) -> dict[str, Any]: def _capture_failure_reason(mode: str, exc: Exception) -> str: """Classify failures without exporting exception text or source snippets.""" text = str(exc).lower() - if any( - marker in text - for marker in ( + if any(marker in text for marker in ( "data-dependent", "data dependent", "guardondatadependentsymnode", ".item()", - ) - ): + )): code = "data_dependent_control_flow" - elif any( - marker in text - for marker in ( + elif any(marker in text for marker in ( "control flow", "proxy object", "symbolically traced variables", "cannot be iterated", - ) - ): + )): code = "dynamic_python_control_flow" elif "alias" in text: code = "unknown_aliasing" @@ -324,11 +306,7 @@ def _capture_ready_module( return forward_hooks = getattr(manager, "forward_hooks", {}) - unsupported_hooks = [ - str(name) - for name in forward_hooks - if str(name) != "LayerwiseOffloadHook" - ] + unsupported_hooks = [str(name) for name in forward_hooks if str(name) != "LayerwiseOffloadHook"] if unsupported_hooks: raise RuntimeError("unsupported_module_hooks") @@ -342,15 +320,8 @@ def _capture_ready_module( if offload_hook is not None and state is None: raise RuntimeError("missing_offload_state") - parameter_data = { - name: parameter.data - for name, parameter in module.named_parameters() - } - gpu_parameters = ( - dict(getattr(state, "gpu_named_parameters", {})) - if state is not None - else {} - ) + parameter_data = {name: parameter.data for name, parameter in module.named_parameters()} + gpu_parameters = (dict(getattr(state, "gpu_named_parameters", {})) if state is not None else {}) try: if state is not None: state.wait_and_replace_params() @@ -369,9 +340,7 @@ def _capture_ready_module( @contextmanager -def _capture_forward_context( - observed_context: tuple[Any, Any] | None, -) -> Iterator[None]: +def _capture_forward_context(observed_context: tuple[Any, Any] | None, ) -> Iterator[None]: """Re-establish the bounded runtime context needed by attention layers.""" if observed_context is None: yield @@ -381,9 +350,9 @@ def _capture_forward_context( current_timestep, attention_metadata = observed_context with set_forward_context( - current_timestep=current_timestep, - attn_metadata=attention_metadata, - forward_batch=None, + current_timestep=current_timestep, + attn_metadata=attention_metadata, + forward_batch=None, ): yield @@ -487,27 +456,21 @@ def record_constants(node: Any, key: str) -> None: if _is_safe_constant(value): constants[f"{node.name}.arg{arg_index}"] = value elif value is not None and not tuple(node_names(value)): - notes.append( - f"{key}: unsafe positional constant arg{arg_index}" - ) + notes.append(f"{key}: unsafe positional constant arg{arg_index}") for kwarg, value in (node.kwargs or {}).items(): if hasattr(value, "op"): continue if _is_safe_constant(value): constants[f"{node.name}.{kwarg}"] = value elif value is not None and not tuple(node_names(value)): - notes.append( - f"{key}: non-scalar constant {kwarg!r} dropped" - ) + notes.append(f"{key}: non-scalar constant {kwarg!r} dropped") for node in graph.nodes: if node.op in {"placeholder", "output"}: continue if node.op == "get_attr": target = str(node.target) - notes.append( - f"get_attr:{target}: lifted attribute value not exported" - ) + notes.append(f"get_attr:{target}: lifted attribute value not exported") continue if node.op == "call_function": key = _op_key(node.target) @@ -589,14 +552,12 @@ def _ir_argument( return {"device": "runtime"} if isinstance(value, tuple | list): kind = "tuple" if isinstance(value, tuple) else "list" - return { - kind: [_ir_argument(item, node_expressions) for item in value] - } + return {kind: [_ir_argument(item, node_expressions) for item in value]} # Static SymInts are shape metadata and safe to materialize. if type(value).__module__.startswith("torch") and type(value).__name__ in { - "SymInt", - "SymBool", - "SymFloat", + "SymInt", + "SymBool", + "SymFloat", }: try: scalar = int(value) if type(value).__name__ != "SymFloat" else float(value) @@ -625,9 +586,7 @@ def _extract_executable_ir(exported: Any, graph: Any) -> dict[str, Any]: inputs: list[dict[str, Any]] = [] runtime_index = 0 lifted_index = 0 - for index, (node, input_spec) in enumerate( - zip(placeholders, input_specs, strict=True) - ): + for index, (node, input_spec) in enumerate(zip(placeholders, input_specs, strict=True)): meta = _ir_tensor_meta((node.meta or {}).get("val")) if meta is None: constant = _ir_argument((node.meta or {}).get("val"), {}) @@ -661,9 +620,7 @@ def _extract_executable_ir(exported: Any, graph: Any) -> dict[str, Any]: item: dict[str, Any] = { "id": node_id, "target": _ir_target(node), - "args": [ - _ir_argument(value, node_expressions) for value in node.args - ], + "args": [_ir_argument(value, node_expressions) for value in node.args], "kwargs": { _safe_name(str(key)): _ir_argument(value, node_expressions) for key, value in (node.kwargs or {}).items() @@ -678,11 +635,7 @@ def _extract_executable_ir(exported: Any, graph: Any) -> dict[str, Any]: if len(output_nodes) != 1 or not output_nodes[0].args: raise RuntimeError("invalid_graph_output") encoded_output = _ir_argument(output_nodes[0].args[0], node_expressions) - outputs = ( - encoded_output["tuple"] - if set(encoded_output) == {"tuple"} - else [encoded_output] - ) + outputs = (encoded_output["tuple"] if set(encoded_output) == {"tuple"} else [encoded_output]) return { "schema_version": 1, "inputs": inputs, @@ -898,11 +851,9 @@ def _mode_order(self) -> tuple[str, ...]: if self.tracer in {"auto", "fallback"}: return _CAPTURE_MODES if self.tracer in _CAPTURE_MODES: - return (self.tracer,) - raise ValueError( - f"unsupported tracer {self.tracer!r}; " - "use auto, symbolic, export, or dynamo" - ) + return (self.tracer, ) + raise ValueError(f"unsupported tracer {self.tracer!r}; " + "use auto, symbolic, export, or dynamo") def _trace( self, @@ -915,11 +866,9 @@ def _trace( if args is None or kwargs is None: raise RuntimeError("example arguments unavailable") with _capture_ready_module(module, args, kwargs) as ( - capture_args, - capture_kwargs, - ), _capture_forward_context( - variant.observed_context - ), _traceable_module_forwards(module): + capture_args, + capture_kwargs, + ), _capture_forward_context(variant.observed_context), _traceable_module_forwards(module): if mode == "symbolic": return torch.fx.symbolic_trace(module) if mode == "export": @@ -937,11 +886,7 @@ def _trace( graph_module = getattr(exported, "graph_module", None) if graph_module is None and isinstance(exported, tuple): graph_module = exported[0] - return ( - graph_module - if graph_module is not None - else exported - ) + return (graph_module if graph_module is not None else exported) raise ValueError(f"unsupported capture mode {mode!r}") def _regions_for(self, record: _Scope) -> list[dict[str, Any]]: @@ -962,13 +907,9 @@ def _regions_for(self, record: _Scope) -> list[dict[str, Any]]: graph = getattr(traced, "graph", None) if graph is None: raise RuntimeError("trace result has no FX graph") - operations, dependencies, constants, notes = _extract_graph( - graph - ) + operations, dependencies, constants, notes = _extract_graph(graph) if not operations: - raise RuntimeError( - "trace result has no tensor operations" - ) + raise RuntimeError("trace result has no tensor operations") capture_mode = mode if mode == "export": try: @@ -979,20 +920,18 @@ def _regions_for(self, record: _Scope) -> list[dict[str, Any]]: except Exception as exc: # noqa: BLE001 known_reason = str(exc) if known_reason not in { - "graph_signature_input_mismatch", - "invalid_graph_output", - "missing_graph_signature", - "unlifted_graph_attribute", - "unsupported_graph_node_kind", - "unsupported_non_tensor_graph_input", + "graph_signature_input_mismatch", + "invalid_graph_output", + "missing_graph_signature", + "unlifted_graph_attribute", + "unsupported_graph_node_kind", + "unsupported_non_tensor_graph_input", }: known_reason = type(exc).__name__ self._graph_breaks.append({ "scope": record.scope, - "reason": ( - "executable_ir_unavailable:" - f"{known_reason}" - ), + "reason": ("executable_ir_unavailable:" + f"{known_reason}"), "count": max(variant.calls, 1), }) break @@ -1054,9 +993,7 @@ def _regions_for(self, record: _Scope) -> list[dict[str, Any]]: "module_class": record.class_name, "tracer": self.tracer, "capture_mode": capture_mode, - "capture_attempts": list(attempts)[ - : list(attempts).index(capture_mode) + 1 - ], + "capture_attempts": list(attempts)[:list(attempts).index(capture_mode) + 1], "capture_failures": failures, }, } @@ -1087,32 +1024,28 @@ def finalize(self) -> dict[str, Any]: unsupported = _coalesce(self._unsupported, ("op_name", "reason", "scope")) payload = { "capture": { - "capture_schema_version": CAPTURE_SCHEMA_VERSION, - "tracer": self.tracer, - "scopes": sorted(self._scopes), + "capture_schema_version": + CAPTURE_SCHEMA_VERSION, + "tracer": + self.tracer, + "scopes": + sorted(self._scopes), "scope_calls": { record.scope: record.calls for record in self._scopes.values() }, - "dropped_scopes": self._dropped_scopes, - "dropped_shape_variants": sum(record.dropped_variants for record in self._scopes.values()), - "errors": list(self._errors), - "capture_mode_breakdown": dict( - sorted( - ( - mode, - sum( - 1 - for region in regions - if region.get("attributes", {}).get( - "capture_mode" - ) - == mode - ), - ) - for mode in _CAPTURE_MODES - ) - ), + "dropped_scopes": + self._dropped_scopes, + "dropped_shape_variants": + sum(record.dropped_variants for record in self._scopes.values()), + "errors": + list(self._errors), + "capture_mode_breakdown": + dict( + sorted(( + mode, + sum(1 for region in regions if region.get("attributes", {}).get("capture_mode") == mode), + ) for mode in _CAPTURE_MODES)), }, "regions": regions, "graph_breaks": breaks, @@ -1122,6 +1055,56 @@ def finalize(self) -> dict[str, Any]: return payload +# Public dispatch-facing capture API. Keeping these operations here makes +# exported/runtime fingerprints share one implementation without forcing the +# dispatcher to import underscore-private capture internals. +def invocation_input_signatures( + args: tuple[Any, ...], + kwargs: dict[str, Any], +) -> list[dict[str, Any]]: + """Return metadata-only tensor signatures for a module invocation.""" + return _input_metas(args, kwargs) + + +def invocation_output_signatures(output: Any) -> list[dict[str, Any]]: + """Return metadata-only tensor signatures for a module output.""" + return _output_metas(output) + + +def invocation_shape_key(metas: list[dict[str, Any]]) -> str: + """Return the canonical capture shape key for an invocation.""" + return _shape_key(metas) + + +def capture_invocation_identity( + module: nn.Module, + args: tuple[Any, ...], + kwargs: dict[str, Any], + output: Any, + *, + scope: str, + tracer: str = "symbolic", +) -> dict[str, Any]: + """Capture one observed invocation and return its region identity. + + Failures are raised because dispatch needs one definitive decision rather + than a full discovery report containing graph-break records. + """ + session = FXCaptureSession(tracer=tracer, max_scopes=1, max_shape_variants=1) + record = _Scope( + scope=scope, + class_name=type(module).__name__, + module=module, + ) + session._scopes[scope] = record + session._observe(record, args, kwargs, output) + regions = session._regions_for(record) + if not regions: + reasons = [str(item.get("reason", "")) for item in session._graph_breaks] + raise RuntimeError(reasons[0] if reasons else "trace_produced_no_region") + return regions[0] + + def _coalesce(records: list[dict[str, Any]], keys: tuple[str, ...]) -> list[dict[str, Any]]: """Merge duplicate records, summing their ``count`` fields.""" merged: dict[tuple[Any, ...], dict[str, Any]] = {} diff --git a/fastvideo/optimization/identity.py b/fastvideo/optimization/identity.py new file mode 100644 index 0000000000..85c5545771 --- /dev/null +++ b/fastvideo/optimization/identity.py @@ -0,0 +1,30 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Graph identity for a single observed module invocation. + +Dispatch has to answer one question about a live call: *which captured region +is this?* The answer must be the same value the producer recorded, or no +artifact would ever match. + +That is why this module reuses :mod:`fastvideo.optimization.fx_capture`'s own +helpers rather than re-deriving tensor metadata, shape keys or fingerprints. +Those helpers are the canonical implementation; a second implementation here +would be one refactor away from silently disagreeing with the exports it is +supposed to match, and a fingerprint that disagrees is indistinguishable from +"no artifact available". +""" + +from __future__ import annotations + +from fastvideo.optimization.fx_capture import ( + capture_invocation_identity as graph_identity, + invocation_input_signatures as input_signatures, + invocation_output_signatures as output_signatures, + invocation_shape_key as shape_key_for, +) + +__all__ = [ + "graph_identity", + "input_signatures", + "output_signatures", + "shape_key_for", +] diff --git a/fastvideo/pipelines/composed_pipeline_base.py b/fastvideo/pipelines/composed_pipeline_base.py index 4aa925ac10..78acd18537 100644 --- a/fastvideo/pipelines/composed_pipeline_base.py +++ b/fastvideo/pipelines/composed_pipeline_base.py @@ -19,7 +19,7 @@ from fastvideo.hooks.activation_trace import attach_activation_trace, detach_activation_trace from fastvideo.logger import init_logger from fastvideo.profiler import get_or_create_profiler -from fastvideo.optimization import optimization_profile +from fastvideo.optimization import (attach_graph_dispatch, detach_graph_dispatch, optimization_profile) from fastvideo.models.loader.component_loader import PipelineComponentLoader from fastvideo.pipelines.pipeline_batch_info import ForwardBatch from fastvideo.pipelines.stages import PipelineStage @@ -65,6 +65,7 @@ def __init__(self, self._stages: list[PipelineStage] = [] self._stage_name_mapping: dict[str, PipelineStage] = {} self._trace_mgr = None + self._dispatch_session = None self._optimization_profile_calls = 0 if required_config_modules is not None: @@ -237,6 +238,9 @@ def post_init(self) -> None: logger.info("Torch Compile enabled for audio VAE") self._trace_mgr = attach_activation_trace(self.modules.get("transformer")) + # Generic graph dispatch. Returns None unless a trusted artifact + # directory is configured, in which case nothing is wrapped at all. + self._dispatch_session = attach_graph_dispatch(self.modules) if not self.fastvideo_args.training_mode: logger.info("Creating pipeline stages...") @@ -526,6 +530,8 @@ def train(self) -> None: def close(self) -> None: detach_activation_trace(getattr(self, "_trace_mgr", None)) self._trace_mgr = None + detach_graph_dispatch(getattr(self, "_dispatch_session", None)) + self._dispatch_session = None def __del__(self): self.close() diff --git a/fastvideo/tests/optimization/test_dispatch.py b/fastvideo/tests/optimization/test_dispatch.py new file mode 100644 index 0000000000..3f14a28101 --- /dev/null +++ b/fastvideo/tests/optimization/test_dispatch.py @@ -0,0 +1,813 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CPU tests for generic artifact dispatch and native fallback. + +Every kernel and module here is a CPU fake: no CUDA, no Triton, and no model +code. What is exercised is the contract -- identity matching, hash +verification, fallback behavior and diagnostics -- not any particular model. +""" + +from __future__ import annotations + +import copy +import hashlib +import json +from pathlib import Path + +import pytest +import torch +from torch import nn + +from fastvideo import envs +from fastvideo.optimization.artifact import (ANY, MANIFEST_FILENAME, ArtifactRegistry, RuntimeProfile, load_entry_point, + verify_bundle) +from fastvideo.optimization.dispatch import (REASON_NO_COMPATIBLE_ARTIFACT, REASON_NO_SIGNATURE_MATCH, REASON_SELECTED, + GraphDispatchSession, attach_graph_dispatch, detach_graph_dispatch) +from fastvideo.optimization.identity import (graph_identity, input_signatures, output_signatures) + +WIDTH = 4 +MARKER = 7.0 + +# A CPU fake kernel. It reads what it needs from the module it is handed, which +# is what lets one artifact serve every block in a stack. +FAKE_KERNEL = '''"""CPU fake kernel used by the dispatch tests.""" +import torch + + +def fused_block(module, hidden): + return torch.full_like(hidden, float(module.marker)) +''' + +RAISING_KERNEL = '''"""CPU fake kernel that fails at call time.""" + + +def fused_block(module, hidden): + raise RuntimeError("candidate exploded") +''' + +IMPORT_FAILURE_KERNEL = '''"""CPU fake kernel that fails at import time.""" +raise RuntimeError("import side effect") + + +def fused_block(module, hidden): + return hidden +''' + + +class _Block(nn.Module): + """One traceable block, standing in for a transformer block.""" + + def __init__(self, width: int) -> None: + super().__init__() + self.proj = nn.Linear(width, width) + self.marker = MARKER + + def forward(self, hidden: torch.Tensor) -> torch.Tensor: + return torch.nn.functional.silu(self.proj(hidden)) + hidden + + +class _UntraceableBlock(nn.Module): + """A block whose data-dependent branch defeats symbolic tracing.""" + + def __init__(self, width: int) -> None: + super().__init__() + self.proj = nn.Linear(width, width) + self.marker = MARKER + + def forward(self, hidden: torch.Tensor) -> torch.Tensor: + if bool(hidden.sum() > 0): + return self.proj(hidden) + return self.proj(hidden) * 2 + + +class _Transformer(nn.Module): + """A repeated block stack: the structure dispatch keys on.""" + + def __init__(self, width: int = WIDTH, depth: int = 2, block_cls=_Block) -> None: + super().__init__() + self.blocks = nn.ModuleList([block_cls(width) for _ in range(depth)]) + + def forward(self, hidden: torch.Tensor) -> torch.Tensor: + for block in self.blocks: + hidden = block(hidden) + return hidden + + +def _pipeline_modules(block_cls=_Block) -> dict[str, nn.Module]: + torch.manual_seed(0) + return {"transformer": _Transformer(block_cls=block_cls).eval()} + + +def _runtime(**overrides) -> RuntimeProfile: + profile = { + "model_id": "fake/model", + "model_revision": "main", + "gpu_architecture": "cpu", + "torch_version": torch.__version__, + "cuda_version": None, + "triton_version": None, + "execution_mode": "inference", + "distributed_mode": "single", + } + profile.update(overrides) + return RuntimeProfile(**profile) + + +def _observed_identity( + module: nn.Module, + hidden: torch.Tensor, + *, + scope: str = "transformer.blocks", + tracer: str = "symbolic", +): + """Trace one call the way dispatch does, to learn the artifact identity.""" + with torch.no_grad(): + output = module(hidden) + region = graph_identity(module, (hidden, ), {}, output, scope=scope, tracer=tracer) + return ( + region["fingerprint"], + input_signatures((hidden, ), {}), + output_signatures(output), + ) + + +def _sections(fingerprint, inputs, outputs, **overrides) -> dict: + sections = { + "schema_version": 1, + "artifact_id": "fake-fused-block", + "created_at": "2026-07-31T00:00:00+00:00", + "producer": { + "name": "motionkernel", + "version": "1.0.0" + }, + "operation": { + "name": "generated_blocks_fused", + "graph_fingerprint": fingerprint, + "parent_module": "transformer.blocks", + "operations": ["aten::linear", "aten::silu", "aten::add"], + }, + "signature": { + "inputs": copy.deepcopy(inputs), + "outputs": copy.deepcopy(outputs), + }, + "entry_point": { + "file": "kernel.py", + "symbol": "fused_block" + }, + "compatibility": { + "model_id": "fake/model", + "model_revision": ANY, + "gpu_architectures": ["cpu"], + "torch": {}, + "cuda": {}, + "triton": {}, + "execution_modes": ["inference"], + "distributed_modes": ["single"], + }, + "evidence": { + "benchmark": { + "harness": "motionkernel-bench", + "device": "cpu", + "samples": 20, + "baseline_us": 100.0, + "candidate_us": 50.0, + "speedup": 2.0, + "max_abs_error": 0.0, + "max_rel_error": 0.0, + "atol": 1e-5, + "rtol": 1e-5, + "passed": True, + "result_ref": "", + }, + "generation": { + "workload_id": "fake-workload", + "steps": 2, + "metric": "max_abs_latent_diff", + "value": 0.0, + "threshold": 1e-3, + "passed": True, + "baseline_ref": "", + "candidate_ref": "", + }, + }, + "promotion": { + "decision": "promoted", + "reason": "2x with full-generation parity", + "decided_at": "2026-07-31T00:00:00+00:00", + "campaign": { + "campaign_id": "campaign-1", + "source": "cpu-fake", + "target_name": "blocks_fused", + }, + }, + } + sections.update(overrides) + return sections + + +def _write_bundle(root: Path, sections: dict, *, kernel_source: str = FAKE_KERNEL, name: str = "fused") -> Path: + """Write a bundle the way the producer would, hashes included.""" + directory = root / name + directory.mkdir(parents=True, exist_ok=True) + kernel = directory / "kernel.py" + kernel.write_text(kernel_source, encoding="utf-8") + document = dict(sections) + document["files"] = [{ + "path": "kernel.py", + "sha256": hashlib.sha256(kernel.read_bytes()).hexdigest(), + "bytes": kernel.stat().st_size, + }] + (directory / MANIFEST_FILENAME).write_text(json.dumps(document, indent=2, sort_keys=True), encoding="utf-8") + return directory + + +@pytest.fixture() +def hidden() -> torch.Tensor: + torch.manual_seed(1) + return torch.randn(2, WIDTH) + + +@pytest.fixture() +def store(tmp_path: Path) -> Path: + directory = tmp_path / "artifacts" + directory.mkdir() + return directory + + +def _session(store: Path, *, tracer: str = "symbolic", **overrides) -> GraphDispatchSession: + return GraphDispatchSession( + ArtifactRegistry(store), + _runtime(**overrides), + tracer=tracer, + ) + + +def _decisions(session: GraphDispatchSession) -> list[dict]: + return session.diagnostics()["decisions"] + + +def _reasons(session: GraphDispatchSession) -> list[str]: + return [item["reason"] for item in _decisions(session)] + + +# -- zero-effect baseline ----------------------------------------------------- + + +def test_no_artifact_directory_behaves_exactly_like_native(monkeypatch, hidden): + monkeypatch.setattr(envs, "FASTVIDEO_OPTIMIZATION_ARTIFACT_DIR", "") + modules = _pipeline_modules() + transformer = modules["transformer"] + before = [type(block).forward for block in transformer.blocks] + with torch.no_grad(): + expected = transformer(hidden) + + session = attach_graph_dispatch(modules) + + assert session is None + # No instance attribute shadows the class method: the module graph is + # untouched, so behavior is identical to a build without this feature. + assert all("forward" not in vars(block) for block in transformer.blocks) + assert [type(block).forward for block in transformer.blocks] == before + with torch.no_grad(): + assert torch.equal(transformer(hidden), expected) + + +def test_empty_artifact_directory_stays_native(store, monkeypatch, hidden): + monkeypatch.setattr(envs, "FASTVIDEO_OPTIMIZATION_ARTIFACT_DIR", str(store)) + modules = _pipeline_modules() + + session = attach_graph_dispatch(modules) + + assert session is None + + +def test_detach_restores_the_native_forward(store, hidden): + modules = _pipeline_modules() + transformer = modules["transformer"] + with torch.no_grad(): + expected = transformer(hidden) + fingerprint, inputs, outputs = _observed_identity(transformer.blocks[0], hidden) + _write_bundle(store, _sections(fingerprint, inputs, outputs)) + session = _session(store) + + assert session.attach_modules(modules) == 2 + assert all("forward" in vars(block) for block in transformer.blocks) + + session.detach() + + assert all("forward" not in vars(block) for block in transformer.blocks) + with torch.no_grad(): + assert torch.equal(transformer(hidden), expected) + + +# -- selection ---------------------------------------------------------------- + + +def test_compatible_artifact_is_selected(store, hidden): + modules = _pipeline_modules() + transformer = modules["transformer"] + fingerprint, inputs, outputs = _observed_identity(transformer.blocks[0], hidden) + _write_bundle(store, _sections(fingerprint, inputs, outputs)) + session = _session(store) + session.attach_modules(modules) + + with torch.no_grad(): + first = transformer(hidden) + second = transformer(hidden) + + session.detach() + # One decision covers the whole stack, so only the very first block call + # runs natively; every later block -- including the second block of the + # first pass -- goes through the candidate, which fills the tensor with the + # marker it reads off the module it was handed. + assert torch.allclose(first, torch.full_like(first, MARKER)) + assert torch.allclose(second, torch.full_like(second, MARKER)) + assert _reasons(session) == [REASON_SELECTED] + decision = _decisions(session)[0] + assert decision["artifact_id"] == "fake-fused-block" + assert decision["candidate_calls"] == 3 + assert decision["runtime_fallbacks"] == 0 + + +def test_export_identity_traces_native_forward_without_dispatch_recursion(store, hidden): + modules = _pipeline_modules() + transformer = modules["transformer"] + fingerprint, inputs, outputs = _observed_identity( + transformer.blocks[0], hidden, tracer="export" + ) + _write_bundle(store, _sections(fingerprint, inputs, outputs)) + session = _session(store, tracer="export") + session.attach_modules(modules) + + with torch.no_grad(): + first = transformer(hidden) + second = transformer(hidden) + + session.detach() + assert torch.allclose(first, torch.full_like(first, MARKER)) + assert torch.allclose(second, torch.full_like(second, MARKER)) + assert _reasons(session) == [REASON_SELECTED] + + +def test_fastest_compatible_artifact_wins(store, hidden): + modules = _pipeline_modules() + transformer = modules["transformer"] + fingerprint, inputs, outputs = _observed_identity(transformer.blocks[0], hidden) + _write_bundle(store, _sections(fingerprint, inputs, outputs), name="slow") + fast = _sections(fingerprint, inputs, outputs) + fast["artifact_id"] = "fake-fused-block-fast" + fast["evidence"]["benchmark"]["speedup"] = 4.0 + _write_bundle(store, fast, name="fast") + session = _session(store) + session.attach_modules(modules) + + with torch.no_grad(): + transformer(hidden) + transformer(hidden) + + session.detach() + decision = _decisions(session)[0] + assert decision["artifact_id"] == "fake-fused-block-fast" + assert decision["rejections"] == ["fake-fused-block:not_selected"] + + +# -- rejections --------------------------------------------------------------- + + +def test_fingerprint_mismatch_falls_back_to_native(store, hidden): + modules = _pipeline_modules() + transformer = modules["transformer"] + _, inputs, outputs = _observed_identity(transformer.blocks[0], hidden) + _write_bundle(store, _sections("0" * 32, inputs, outputs)) + session = _session(store) + session.attach_modules(modules) + + with torch.no_grad(): + expected = _pipeline_modules()["transformer"](hidden) + transformer(hidden) + actual = transformer(hidden) + + session.detach() + assert torch.equal(actual, expected) + assert _reasons(session) == [REASON_NO_COMPATIBLE_ARTIFACT] + assert _decisions(session)[0]["rejections"] == ["fake-fused-block:fingerprint_mismatch"] + + +@pytest.mark.parametrize( + ("mutate", "field"), + [ + (lambda item: item.update({"dtype": "bfloat16"}), "dtype"), + (lambda item: item.update({"shape": [8, WIDTH]}), "shape"), + (lambda item: item.update({"device_type": "cuda"}), "device_type"), + ], +) +def test_signature_mismatches_never_reach_a_trace(store, hidden, mutate, field): + modules = _pipeline_modules() + transformer = modules["transformer"] + fingerprint, inputs, outputs = _observed_identity(transformer.blocks[0], hidden) + mutated = copy.deepcopy(inputs) + mutate(mutated[0]) + _write_bundle(store, _sections(fingerprint, mutated, outputs)) + session = _session(store) + session.attach_modules(modules) + + with torch.no_grad(): + expected = _pipeline_modules()["transformer"](hidden) + transformer(hidden) + actual = transformer(hidden) + + session.detach() + assert torch.equal(actual, expected) + # The input pre-filter rejects these before any graph is traced. + assert _reasons(session) == [REASON_NO_SIGNATURE_MATCH] + + +def test_output_signature_mismatch_falls_back_to_native(store, hidden): + modules = _pipeline_modules() + transformer = modules["transformer"] + fingerprint, inputs, outputs = _observed_identity(transformer.blocks[0], hidden) + mutated = copy.deepcopy(outputs) + mutated[0]["shape"] = [2, WIDTH * 2] + mutated[0]["stride"] = [WIDTH * 2, 1] + _write_bundle(store, _sections(fingerprint, inputs, mutated)) + session = _session(store) + session.attach_modules(modules) + + with torch.no_grad(): + transformer(hidden) + transformer(hidden) + + session.detach() + assert _decisions(session)[0]["rejections"] == ["fake-fused-block:output_signature_mismatch"] + + +@pytest.mark.parametrize( + ("override", "expected"), + [ + ({ + "gpu_architectures": ["sm90"] + }, "gpu_architecture_mismatch"), + ({ + "torch": { + "min": "99.0.0" + } + }, "torch_version_unsupported"), + ({ + "cuda": { + "min": "12.0" + } + }, "cuda_version_unsupported"), + ({ + "triton": { + "min": "3.0.0" + } + }, "triton_version_unsupported"), + ({ + "model_id": "other/model" + }, "model_mismatch"), + ({ + "execution_modes": ["training"] + }, "execution_mode_unsupported"), + ({ + "distributed_modes": ["tensor_parallel"] + }, "distributed_mode_unsupported"), + ], +) +def test_environment_mismatches_fall_back_to_native(store, hidden, override, expected): + modules = _pipeline_modules() + transformer = modules["transformer"] + fingerprint, inputs, outputs = _observed_identity(transformer.blocks[0], hidden) + sections = _sections(fingerprint, inputs, outputs) + sections["compatibility"].update(override) + _write_bundle(store, sections) + session = _session(store) + session.attach_modules(modules) + + with torch.no_grad(): + native = _pipeline_modules()["transformer"](hidden) + transformer(hidden) + actual = transformer(hidden) + + session.detach() + assert torch.equal(actual, native) + assert _decisions(session)[0]["rejections"] == [f"fake-fused-block:{expected}"] + + +def test_unpromoted_artifact_is_never_dispatched(store, hidden): + modules = _pipeline_modules() + transformer = modules["transformer"] + fingerprint, inputs, outputs = _observed_identity(transformer.blocks[0], hidden) + sections = _sections(fingerprint, inputs, outputs) + sections["promotion"]["decision"] = "quarantined" + _write_bundle(store, sections) + session = _session(store) + session.attach_modules(modules) + + with torch.no_grad(): + transformer(hidden) + transformer(hidden) + + session.detach() + assert _decisions(session)[0]["rejections"] == ["fake-fused-block:not_promoted"] + + +# -- tampering and load failures ---------------------------------------------- + + +def test_tampered_kernel_is_rejected_before_import(store, hidden): + modules = _pipeline_modules() + transformer = modules["transformer"] + fingerprint, inputs, outputs = _observed_identity(transformer.blocks[0], hidden) + directory = _write_bundle(store, _sections(fingerprint, inputs, outputs)) + # Swap the kernel for one that would be obvious if it ever executed. + (directory / "kernel.py").write_text('raise SystemExit("payload executed")\n', encoding="utf-8") + + registry = ArtifactRegistry(store) + + assert registry.manifests == [] + assert len(registry.errors) == 1 + assert "does not match manifest" in registry.errors[0] or "bytes" in registry.errors[0] + + session = GraphDispatchSession(registry, _runtime()) + assert session.attach_modules(modules) == 0 + with torch.no_grad(): + assert torch.equal(transformer(hidden), _pipeline_modules()["transformer"](hidden)) + + +def test_undeclared_file_is_rejected(store, hidden): + modules = _pipeline_modules() + fingerprint, inputs, outputs = _observed_identity(modules["transformer"].blocks[0], hidden) + directory = _write_bundle(store, _sections(fingerprint, inputs, outputs)) + (directory / "extra.py").write_text("SECRET = 1\n", encoding="utf-8") + + registry = ArtifactRegistry(store) + + assert registry.manifests == [] + assert "undeclared file" in registry.errors[0] + + +def test_unknown_schema_version_is_rejected(store, hidden): + modules = _pipeline_modules() + fingerprint, inputs, outputs = _observed_identity(modules["transformer"].blocks[0], hidden) + sections = _sections(fingerprint, inputs, outputs) + sections["schema_version"] = 99 + _write_bundle(store, sections) + + registry = ArtifactRegistry(store) + + assert registry.manifests == [] + assert "unsupported version" in registry.errors[0] + + +@pytest.mark.parametrize( + ("mutate", "match"), + [ + ( + lambda document: document["signature"]["inputs"][0].update( + requires_grad="false" + ), + "requires_grad", + ), + ( + lambda document: document["evidence"]["benchmark"].update( + speedup=float("nan") + ), + "finite non-negative", + ), + ( + lambda document: document["compatibility"].update( + torch={"min": "3.0", "max_exclusive": "2.0"} + ), + "min must be lower", + ), + ], +) +def test_malformed_acted_on_manifest_fields_fail_closed( + store, hidden, mutate, match +): + modules = _pipeline_modules() + fingerprint, inputs, outputs = _observed_identity( + modules["transformer"].blocks[0], hidden + ) + document = _sections(fingerprint, inputs, outputs) + mutate(document) + _write_bundle(store, document) + + registry = ArtifactRegistry(store) + + assert registry.manifests == [] + assert match in registry.errors[0] + + +def test_artifact_module_names_do_not_collide_after_id_sanitizing(store, hidden): + modules = _pipeline_modules() + fingerprint, inputs, outputs = _observed_identity( + modules["transformer"].blocks[0], hidden + ) + dashed = _sections(fingerprint, inputs, outputs, artifact_id="my-kernel-v1") + underscored = _sections( + fingerprint, inputs, outputs, artifact_id="my_kernel_v1" + ) + dashed_dir = _write_bundle(store, dashed, name="dashed") + underscored_dir = _write_bundle(store, underscored, name="underscored") + + dashed_candidate = load_entry_point( + verify_bundle(dashed_dir), trusted_root=store + ) + underscored_candidate = load_entry_point( + verify_bundle(underscored_dir), trusted_root=store + ) + + assert dashed_candidate.__module__ != underscored_candidate.__module__ + + +def test_failing_import_falls_back_to_native(store, hidden): + modules = _pipeline_modules() + transformer = modules["transformer"] + fingerprint, inputs, outputs = _observed_identity(transformer.blocks[0], hidden) + _write_bundle(store, _sections(fingerprint, inputs, outputs), kernel_source=IMPORT_FAILURE_KERNEL) + session = _session(store) + session.attach_modules(modules) + + with torch.no_grad(): + native = _pipeline_modules()["transformer"](hidden) + transformer(hidden) + actual = transformer(hidden) + + session.detach() + assert torch.equal(actual, native) + assert _reasons(session) == ["artifact_load_failed:ArtifactError"] + + +def test_bundle_outside_the_trusted_root_is_never_loaded(store, tmp_path, hidden): + modules = _pipeline_modules() + fingerprint, inputs, outputs = _observed_identity(modules["transformer"].blocks[0], hidden) + outside = _write_bundle(tmp_path / "elsewhere", _sections(fingerprint, inputs, outputs)) + + # The registry only ever looks inside its own root. + assert ArtifactRegistry(store).manifests == [] + assert verify_bundle(outside).artifact_id == "fake-fused-block" + + +# -- runtime failure ---------------------------------------------------------- + + +def test_candidate_exception_falls_back_to_native_output(store, hidden): + modules = _pipeline_modules() + transformer = modules["transformer"] + fingerprint, inputs, outputs = _observed_identity(transformer.blocks[0], hidden) + _write_bundle(store, _sections(fingerprint, inputs, outputs), kernel_source=RAISING_KERNEL) + session = _session(store) + session.attach_modules(modules) + + with torch.no_grad(): + native = _pipeline_modules()["transformer"](hidden) + transformer(hidden) + actual = transformer(hidden) + again = transformer(hidden) + + session.detach() + assert torch.equal(actual, native) + assert torch.equal(again, native) + decision = _decisions(session)[0] + assert decision["reason"] == "candidate_runtime_error:RuntimeError" + assert decision["active"] is False + # Demoted after the first failure, not retried on every later call. + assert decision["runtime_fallbacks"] == 1 + + +# -- tracing failures --------------------------------------------------------- + + +def test_untraceable_module_stays_native(store, hidden): + modules = _pipeline_modules(block_cls=_UntraceableBlock) + transformer = modules["transformer"] + with torch.no_grad(): + reference = transformer(hidden) + # Build a bundle whose input signature matches, so the pre-filter passes + # and the tracer is genuinely attempted. + inputs = input_signatures((hidden, ), {}) + outputs = output_signatures(reference) + _write_bundle(store, _sections("1" * 32, inputs, outputs)) + session = _session(store) + session.attach_modules(modules) + + with torch.no_grad(): + transformer(hidden) + actual = transformer(hidden) + + session.detach() + assert torch.equal(actual, reference) + assert _reasons(session)[0].startswith("graph_identity_unavailable") + + +# -- diagnostics -------------------------------------------------------------- + + +def test_diagnostics_are_metadata_only(store, tmp_path, hidden): + modules = _pipeline_modules() + transformer = modules["transformer"] + fingerprint, inputs, outputs = _observed_identity(transformer.blocks[0], hidden) + _write_bundle(store, _sections(fingerprint, inputs, outputs)) + session = _session(store) + session.attach_modules(modules) + with torch.no_grad(): + transformer(hidden) + transformer(hidden) + session.detach() + + output = session.write_diagnostics(tmp_path / "dispatch.json") + + assert output is not None + report = json.loads(output.read_text(encoding="utf-8")) + assert report["dispatch"]["dispatch_schema_version"] == 1 + assert report["dispatch"]["registry"]["artifact_ids"] == ["fake-fused-block"] + assert report["dispatch"]["reason_counts"] == {REASON_SELECTED: 1} + # Nothing tensor-shaped or prompt-shaped may appear anywhere in the report. + serialized = json.dumps(report).lower() + for forbidden in ("prompt", "tensor_values", "weights", "password", "secret"): + assert forbidden not in serialized + + +def test_attach_graph_dispatch_uses_the_configured_directory(store, monkeypatch, hidden): + modules = _pipeline_modules() + transformer = modules["transformer"] + fingerprint, inputs, outputs = _observed_identity(transformer.blocks[0], hidden) + _write_bundle(store, _sections(fingerprint, inputs, outputs)) + monkeypatch.setattr(envs, "FASTVIDEO_OPTIMIZATION_ARTIFACT_DIR", str(store)) + monkeypatch.setattr(envs, "FASTVIDEO_OPTIMIZATION_ARTIFACT_MODEL_ID", "fake/model") + monkeypatch.setattr(envs, "FASTVIDEO_OPTIMIZATION_ARTIFACT_MODEL_REVISION", ANY) + monkeypatch.setattr(envs, "FASTVIDEO_OPTIMIZATION_ARTIFACT_TRACER", "symbolic") + monkeypatch.setattr(envs, "FASTVIDEO_OPTIMIZATION_ARTIFACT_MAX_SCOPES", 64) + monkeypatch.setattr(envs, "FASTVIDEO_OPTIMIZATION_ARTIFACT_MAX_SHAPES", 8) + monkeypatch.setattr(envs, "FASTVIDEO_OPTIMIZATION_ARTIFACT_DISTRIBUTED_MODE", "") + monkeypatch.setattr(envs, "FASTVIDEO_OPTIMIZATION_ARTIFACT_DIAGNOSTICS", "") + + session = attach_graph_dispatch(modules) + + assert session is not None + try: + with torch.no_grad(): + transformer(hidden) + actual = transformer(hidden) + assert torch.allclose(actual, torch.full_like(actual, MARKER)) + finally: + detach_graph_dispatch(session) + assert all("forward" not in vars(block) for block in transformer.blocks) + + +def test_training_mode_is_detected_and_rejects_inference_only_artifacts(store, monkeypatch, hidden): + modules = _pipeline_modules() + transformer = modules["transformer"] + fingerprint, inputs, outputs = _observed_identity(transformer.blocks[0], hidden) + _write_bundle(store, _sections(fingerprint, inputs, outputs)) + transformer.train() + monkeypatch.setattr(envs, "FASTVIDEO_OPTIMIZATION_ARTIFACT_DIR", str(store)) + monkeypatch.setattr(envs, "FASTVIDEO_OPTIMIZATION_ARTIFACT_MODEL_ID", "fake/model") + monkeypatch.setattr(envs, "FASTVIDEO_OPTIMIZATION_ARTIFACT_MODEL_REVISION", ANY) + monkeypatch.setattr(envs, "FASTVIDEO_OPTIMIZATION_ARTIFACT_TRACER", "symbolic") + monkeypatch.setattr(envs, "FASTVIDEO_OPTIMIZATION_ARTIFACT_MAX_SCOPES", 64) + monkeypatch.setattr(envs, "FASTVIDEO_OPTIMIZATION_ARTIFACT_MAX_SHAPES", 8) + monkeypatch.setattr(envs, "FASTVIDEO_OPTIMIZATION_ARTIFACT_DISTRIBUTED_MODE", "") + monkeypatch.setattr(envs, "FASTVIDEO_OPTIMIZATION_ARTIFACT_DIAGNOSTICS", "") + + session = attach_graph_dispatch(modules) + + assert session is not None + try: + with torch.no_grad(): + transformer(hidden) + actual = transformer(hidden) + assert not torch.allclose(actual, torch.full_like(actual, MARKER)) + assert _decisions(session)[0]["rejections"] == ["fake-fused-block:execution_mode_unsupported"] + finally: + detach_graph_dispatch(session) + + +def test_nested_training_module_is_detected_fail_closed(store, monkeypatch, hidden): + modules = _pipeline_modules() + transformer = modules["transformer"] + fingerprint, inputs, outputs = _observed_identity(transformer.blocks[0], hidden) + _write_bundle(store, _sections(fingerprint, inputs, outputs)) + transformer.eval() + transformer.blocks[0].train() + monkeypatch.setattr(envs, "FASTVIDEO_OPTIMIZATION_ARTIFACT_DIR", str(store)) + monkeypatch.setattr(envs, "FASTVIDEO_OPTIMIZATION_ARTIFACT_MODEL_ID", "fake/model") + monkeypatch.setattr(envs, "FASTVIDEO_OPTIMIZATION_ARTIFACT_MODEL_REVISION", ANY) + monkeypatch.setattr(envs, "FASTVIDEO_OPTIMIZATION_ARTIFACT_TRACER", "symbolic") + monkeypatch.setattr(envs, "FASTVIDEO_OPTIMIZATION_ARTIFACT_MAX_SCOPES", 64) + monkeypatch.setattr(envs, "FASTVIDEO_OPTIMIZATION_ARTIFACT_MAX_SHAPES", 8) + monkeypatch.setattr(envs, "FASTVIDEO_OPTIMIZATION_ARTIFACT_DISTRIBUTED_MODE", "") + monkeypatch.setattr(envs, "FASTVIDEO_OPTIMIZATION_ARTIFACT_DIAGNOSTICS", "") + + session = attach_graph_dispatch(modules) + + assert session is not None + try: + with torch.no_grad(): + transformer(hidden) + transformer(hidden) + assert _decisions(session)[0]["rejections"] == [ + "fake-fused-block:execution_mode_unsupported" + ] + finally: + detach_graph_dispatch(session)