diff --git a/docs/build.md b/docs/build.md index b4ecb73c7..4c6edbfc0 100644 --- a/docs/build.md +++ b/docs/build.md @@ -27,9 +27,60 @@ entry is `python -m pip install` with CMake options passed through | `INFINI_OPS_BUILD_DOCS` | Enable the Doxygen documentation target. | `OFF` | | `INFINI_RT_ROOT` | InfiniRT install prefix containing `include/` and `lib/`. | `$INFINI_RT_ROOT` | | `INFINI_OPS_SMOKE_BUILD` | Build only the smoke-test operator subset. | `OFF` | -| `INFINI_OPS_OPS` | Comma- or semicolon-separated operator allowlist. | empty | +| `INFINI_OPS_OPS` | Comma- or semicolon-separated operator allowlist, or a path to an `ops.json` implementation selection. | empty | | `INFINI_OPS_TORCH_OPS` | Comma- or semicolon-separated ATen operator allowlist. | empty | +An `ops.json` file selects operators and implementation slots with a top-level +operator mapping: + +```json +{ + "add": { + "implementations": "all" + }, + "argmax": { + "implementations": [8] + }, + "top_k_top_p_sampling_from_logits": { + "implementations": [16] + } +} +``` + +`"all"` keeps every available implementation for the operator. An integer +array keeps exactly those slots. Slots range from 0 through 31. The selection +is a set, not a priority order; the default dispatch selects the smallest +active slot. The selection controls generated wrappers, generated slot-8 ATen +implementations, and linked +provider resolution. Unselected linked providers do not require their external +libraries to be installed. + +Pass the file explicitly with +`-DINFINI_OPS_OPS=/path/to/ops.json`. For compatibility, +`${PROJECT_SOURCE_DIR}/ops.json` is read automatically when present. Relative +implementation header paths in legacy configurations are resolved from +`${PROJECT_SOURCE_DIR}`. An explicit inline `INFINI_OPS_OPS` allowlist takes +precedence over an implicit `${PROJECT_SOURCE_DIR}/ops.json`. When +`INFINI_OPS_TORCH_OPS` and an explicit JSON selection are both set, generated +ATen ops use their intersection. The string and string-array values supported +by the current generator remain available for checked-in implementation +headers. Structured descriptors preserve an explicit backend name, including +for implementations outside the standard backend directory layout. Generated +implementation header paths are not supported: + +```json +{ + "add": "src/native/cpu/ops/add/add.h", + "gemm": ["src/native/cpu/ops/gemm/gemm.h"], + "custom_add": [ + { + "path": "custom/add.h", + "backend": "custom" + } + ] +} +``` + Only one GPU backend should be enabled in a build. CPU may be enabled with the selected accelerator backend. diff --git a/docs/linked-operators.md b/docs/linked-operators.md index c06d71bc8..9be34be20 100644 --- a/docs/linked-operators.md +++ b/docs/linked-operators.md @@ -1,9 +1,9 @@ # Linked Operators The linked backend calls operators provided by an installed third-party shared -library. It supports exact exported C++ symbols and registered PyTorch -Dispatcher operators when a platform package does not provide source code or a -stable C API. +library. It supports exact exported symbols, TVM FFI entry points, and +registered PyTorch Dispatcher operators when a platform package does not +provide source code or a stable C API. ## Source Layout @@ -15,9 +15,11 @@ src/linked/// ops// .yaml .h - .cc + .{cc,cu} ``` +CUDA providers may use `.cu` instead of `.cc`. + The platform library file contains DSO discovery information: ```yaml @@ -25,6 +27,13 @@ python_distribution_package: vllm library_glob: vllm/_C*.so ``` +A library may also provide `include_glob` when its transport needs installed +headers. Each glob must resolve to exactly one path in the Python distribution. +A library may set `python_distribution_version` to a PEP 440 specifier. The +resolver verifies the installed distribution version before looking up its DSO. +A DSO that depends on another declared platform library lists that dependency +under the implementation's optional `link_libraries` key. + Files for an operator implementation use the provider name as their common stem. Multiple implementations for the same operator and device use distinct file stems and implementation slots. @@ -54,15 +63,19 @@ partial Dispatcher contract or a binding that mixes both forms. ## Adapter Boundary -Keep ABI behavior in `.cc`, not in YAML. Shared operator -templates own reusable tensor conversion, stream guards, layout staging, and -copy-back behavior. Provider sources own exact typed function declarations, +Keep ABI behavior in `.cc` or `.cu`, not in YAML. Shared +operator templates own reusable tensor conversion, stream guards, layout +staging, and copy-back behavior. Provider sources own exact typed declarations, synthesized arguments, and provider-specific return handling. For the `torch` transport, an implementation backend inherits its device `C10` specialization for device identity and external-stream handling, then defines its provider-specific `Call` ABI. +For the `tvm_ffi` transport, provider sources call exported TVM FFI entry +points directly. The resolver supplies installed TVM FFI headers and links the +provider DSO together with every library named in `link_libraries`. + ## Configuration At configure time, `scripts/resolve_linked_ops.py` locates the installed Python @@ -86,16 +99,23 @@ cmake -S . -B build \ -DINFINI_OPS_OPS=silu_and_mul ``` +To resolve only selected linked implementation slots, pass an `ops.json` file +through `INFINI_OPS_OPS`. The resolver reads each linked provider's slot from +its sibling C++ header before locating external libraries, so an unselected +provider does not add a package or shared-library dependency. See +[Build configuration](build.md) for the file format. + The `torch` transport uses the installed PyTorch C++ headers and libraries for `at::Tensor`, but it does not enable the standard `src/torch` operator backend. Provider and PyTorch C++ ABIs must match. Configuration fails before compilation when the distribution, shared library, or an exact required symbol is missing. InfiniOps does not bundle the provider library. Its resolved directory and the -PyTorch runtime directories are recorded in the installed binary's RPATH, so a -linked build is tied to that Python environment. Reconfigure and rebuild after -moving or replacing the provider environment. In-place changes to a resolved -provider DSO are tracked as CMake configure and link dependencies. +directories of its linked dependencies are recorded in the installed binary's +RPATH, so a linked build is tied to that Python environment. PyTorch runtime +directories are recorded for the `torch` transport as well. Reconfigure and +rebuild after moving or replacing the provider environment. In-place changes +to a resolved provider DSO are tracked as CMake configure and link dependencies. ## Implementation Slots diff --git a/pyproject.toml b/pyproject.toml index 288f9d5be..2dd4ba2db 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,5 +1,5 @@ [build-system] -requires = ["scikit-build-core", "pybind11", "libclang", "pyyaml"] +requires = ["scikit-build-core", "pybind11", "libclang", "packaging", "pyyaml"] build-backend = "scikit_build_core.build" [project] diff --git a/scripts/generate_torch_ops.py b/scripts/generate_torch_ops.py index 0c90cedad..6a93050c6 100644 --- a/scripts/generate_torch_ops.py +++ b/scripts/generate_torch_ops.py @@ -36,6 +36,9 @@ import yaml _SCRIPTS_DIR = pathlib.Path(__file__).resolve().parent +sys.path.insert(0, str(_SCRIPTS_DIR)) +import ops_config # noqa: E402 + _REPO_ROOT = _SCRIPTS_DIR.parent _OPS_YAML_PATH = _SCRIPTS_DIR / "torch_ops.yaml" _BASE_DIR = _REPO_ROOT / "src" / "base" @@ -1644,6 +1647,30 @@ def _emit(name: str, ops: list[Op], *, emit_base: bool) -> set[pathlib.Path]: return emitted_paths +def _select_op_names(cli_ops, default_ops, config): + if config is None: + return cli_ops or default_ops + + aten_names_by_public_name = collections.defaultdict(list) + for op_name in default_ops: + aten_names_by_public_name[_public_op_name(op_name)].append(op_name) + + selected_public_names = ops_config.torch_op_names( + config, aten_names_by_public_name, _PYTORCH_SLOT + ) + selected = [ + aten_name + for public_name in selected_public_names + for aten_name in aten_names_by_public_name.get(public_name, (public_name,)) + ] + + if cli_ops: + allowed = set(cli_ops) + selected = [op_name for op_name in selected if op_name in allowed] + + return selected + + def main() -> int: parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) parser.add_argument( @@ -1651,6 +1678,11 @@ def main() -> int: nargs="*", help="Override the op allowlist. If omitted, reads `scripts/torch_ops.yaml`.", ) + parser.add_argument( + "--ops-config", + type=pathlib.Path, + help="Path to an `ops.json` operator and implementation selection.", + ) parser.add_argument( "--pytorch-version", default=os.environ.get("INFINI_OPS_PYTORCH_VERSION", _DEFAULT_PYTORCH_VERSION), @@ -1666,7 +1698,9 @@ def main() -> int: global _CLANG_FORMAT _CLANG_FORMAT = _find_clang_format() - op_names = args.ops or yaml.safe_load(_OPS_YAML_PATH.read_text()) + default_ops = yaml.safe_load(_OPS_YAML_PATH.read_text()) + config = ops_config.load_ops_config(args.ops_config) if args.ops_config else None + op_names = _select_op_names(args.ops, default_ops, config) aten_entries = _load_aten_entries(args.pytorch_version) skipped: list[tuple[str, str]] = [] diff --git a/scripts/generate_wrappers.py b/scripts/generate_wrappers.py index 74818e9f8..0a188b2a0 100644 --- a/scripts/generate_wrappers.py +++ b/scripts/generate_wrappers.py @@ -2,14 +2,17 @@ import concurrent.futures import dataclasses import functools -import json import os import pathlib import re import shutil import subprocess +import sys import textwrap +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent)) +import ops_config # noqa: E402 + try: import clang.cindex from clang.cindex import CursorKind @@ -1848,6 +1851,52 @@ def _filter_ops(ops, op_allowlist, *, strict=False): return {op_name: ops[op_name] for op_name in op_allowlist if op_name in ops} +def _select_ops_from_config(ops, config, config_path): + selected = {} + + for op_name, selection in config.items(): + headers = selection["headers"] + + if headers is not None: + selected[op_name] = [ + _implementation_from_json(header) for header in headers + ] + continue + + if op_name not in ops: + raise ValueError( + f"{config_path}: operator {op_name!r} is not available for " + "the active devices" + ) + + slots = selection["implementations"] + + if slots is None: + selected[op_name] = ops[op_name] + continue + + headers_by_slot = {} + + for implementation in ops[op_name]: + slot = ops_config.implementation_slot(implementation.path) + headers_by_slot.setdefault(slot, []).append(implementation) + + missing = [slot for slot in slots if slot not in headers_by_slot] + + if missing: + formatted = ", ".join(str(slot) for slot in missing) + raise ValueError( + f"{config_path}: operator {op_name!r} has no active " + f"implementation at slot(s) {formatted}" + ) + + selected[op_name] = [ + implementation for slot in slots for implementation in headers_by_slot[slot] + ] + + return selected + + def _get_all_ops( devices, with_torch=False, @@ -2084,6 +2133,11 @@ def _dispatch_gen_batch_size(): type=str, help="Operator allowlist to generate. Accepts names separated by spaces or commas.", ) + parser.add_argument( + "--ops-config", + type=pathlib.Path, + help="Path to an `ops.json` operator and implementation selection.", + ) parser.add_argument( "--strict-ops", action="store_true", @@ -2101,25 +2155,18 @@ def _dispatch_gen_batch_size(): for directory in (_BINDINGS_DIR, _GENERATED_SRC_DIR, _INCLUDE_DIR): directory.mkdir(parents=True, exist_ok=True) - ops_json = pathlib.Path("ops.json") + config_path = args.ops_config + ops = _get_all_ops( + args.devices, + with_torch=args.with_torch, + with_ninetoothed=args.with_ninetoothed, + with_linked=args.with_linked, + with_triton=args.with_triton, + ) - if ops_json.exists(): - raw_ops = json.loads(ops_json.read_text()) - ops = { - op_name: [ - _implementation_from_json(implementation) - for implementation in implementations - ] - for op_name, implementations in raw_ops.items() - } - else: - ops = _get_all_ops( - args.devices, - with_torch=args.with_torch, - with_ninetoothed=args.with_ninetoothed, - with_linked=args.with_linked, - with_triton=args.with_triton, - ) + if config_path is not None: + config = ops_config.load_ops_config(config_path) + ops = _select_ops_from_config(ops, config, config_path) ops = _filter_ops( ops, diff --git a/scripts/ops_config.py b/scripts/ops_config.py new file mode 100644 index 000000000..b1e764fb8 --- /dev/null +++ b/scripts/ops_config.py @@ -0,0 +1,214 @@ +import json +import pathlib +import re + + +_OPERATOR_SPECIALIZATION_RE = re.compile( + r"\bclass\s+Operator<\s*[^,>]+\s*,\s*[^,>]+\s*" + r"(?:,\s*(\d+)\s*)?>" +) + + +class OpsConfigError(ValueError): + pass + + +class _StrictJsonObject(dict): + pass + + +def _strict_object(pairs): + value = _StrictJsonObject() + + for key, item in pairs: + if key in value: + raise OpsConfigError(f"duplicate key {key!r}") + value[key] = item + + return value + + +def load_ops_config(path): + path = pathlib.Path(path) + + try: + config = json.loads( + path.read_text(encoding="utf-8"), object_pairs_hook=_strict_object + ) + except (OSError, json.JSONDecodeError, OpsConfigError) as error: + raise OpsConfigError(f"failed to read {path}: {error}") from error + + if not isinstance(config, dict): + raise OpsConfigError(f"{path} must contain a JSON object") + + normalized = {} + + for op_name, value in config.items(): + if not isinstance(op_name, str) or not op_name.strip(): + raise OpsConfigError(f"{path}: operator names must be non-empty strings") + if op_name != op_name.strip(): + raise OpsConfigError( + f"{path}: operator name {op_name!r} contains surrounding whitespace" + ) + + normalized[op_name] = _normalize_selection(path, op_name, value) + + return normalized + + +def _normalize_selection(path, op_name, value): + if isinstance(value, str): + if not value.strip(): + raise OpsConfigError( + f"{path}: {op_name!r} implementation path must not be empty" + ) + + return {"headers": [value], "implementations": None} + + if isinstance(value, list): + if not value: + raise OpsConfigError( + f"{path}: {op_name!r} implementation paths must be a non-empty array" + ) + headers = [_normalize_header(path, op_name, item) for item in value] + identities = [ + header if isinstance(header, str) else (header["path"], header["backend"]) + for header in headers + ] + if len(identities) != len(set(identities)): + raise OpsConfigError( + f"{path}: {op_name!r} implementation paths contain duplicates" + ) + return {"headers": headers, "implementations": None} + + if not isinstance(value, dict): + raise OpsConfigError( + f"{path}: {op_name!r} must map to implementation path(s) or an object" + ) + + unknown_keys = sorted(set(value) - {"implementations"}) + if unknown_keys: + raise OpsConfigError( + f"{path}: {op_name!r} contains unknown keys: {', '.join(unknown_keys)}" + ) + if "implementations" not in value: + raise OpsConfigError( + f"{path}: {op_name!r} is missing required key 'implementations'" + ) + + implementations = value["implementations"] + + if implementations == "all": + implementations = None + elif isinstance(implementations, list): + if not implementations: + raise OpsConfigError( + f"{path}: {op_name!r} implementations must not be empty" + ) + if any(type(slot) is not int or not 0 <= slot < 32 for slot in implementations): + raise OpsConfigError( + f"{path}: {op_name!r} implementations must contain integers " + "between 0 and 31" + ) + if len(implementations) != len(set(implementations)): + raise OpsConfigError( + f"{path}: {op_name!r} implementations contain duplicates" + ) + implementations = tuple(implementations) + else: + raise OpsConfigError( + f"{path}: {op_name!r} implementations must be 'all' or an array" + ) + + return {"headers": None, "implementations": implementations} + + +def _normalize_header(path, op_name, value): + if isinstance(value, str): + if value.strip(): + return value + raise OpsConfigError( + f"{path}: {op_name!r} implementation path must not be empty" + ) + + if not isinstance(value, dict): + raise OpsConfigError( + f"{path}: {op_name!r} implementation entries must be paths or " + "structured descriptors" + ) + + unknown_keys = sorted(set(value) - {"path", "backend"}) + if unknown_keys: + raise OpsConfigError( + f"{path}: {op_name!r} implementation descriptor contains unknown " + f"keys: {', '.join(unknown_keys)}" + ) + + for key in ("path", "backend"): + if ( + key not in value + or not isinstance(value[key], str) + or not value[key].strip() + ): + raise OpsConfigError( + f"{path}: {op_name!r} implementation descriptor requires a " + f"non-empty string {key!r}" + ) + + return {"path": value["path"], "backend": value["backend"]} + + +def implementation_path(header): + return header if isinstance(header, str) else header["path"] + + +def selected_op_names(config): + return list(config) + + +def selected_slots(config, op_name): + selection = config.get(op_name) + + if selection is None or selection["headers"] is not None: + return None + + return selection["implementations"] + + +def implementation_slot(path): + path = pathlib.Path(path) + + try: + text = path.read_text(encoding="utf-8") + except OSError as error: + raise OpsConfigError(f"failed to read {path}: {error}") from error + + slots = { + int(match.group(1)) if match.group(1) is not None else 0 + for match in _OPERATOR_SPECIALIZATION_RE.finditer(text) + } + + if len(slots) != 1: + formatted = ", ".join(str(slot) for slot in sorted(slots)) or "none" + raise OpsConfigError( + f"{path} must declare exactly one implementation slot; found {formatted}" + ) + + return slots.pop() + + +def torch_op_names(config, default_ops=(), slot=8): + selected = [] + + for op_name, selection in config.items(): + if selection["headers"] is not None: + continue + + implementations = selection["implementations"] + + if (implementations is None and op_name in default_ops) or ( + implementations is not None and slot in implementations + ): + selected.append(op_name) + + return selected diff --git a/scripts/resolve_linked_ops.py b/scripts/resolve_linked_ops.py index 03f63be9e..4c20608d3 100644 --- a/scripts/resolve_linked_ops.py +++ b/scripts/resolve_linked_ops.py @@ -10,23 +10,31 @@ import sys import urllib.parse import urllib.request +from packaging.specifiers import InvalidSpecifier, SpecifierSet +from packaging.version import InvalidVersion, Version import yaml +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent)) +import ops_config # noqa: E402 + _PROJECT_DIR = pathlib.Path(__file__).resolve().parents[1] _DEFAULT_SOURCE_ROOT = _PROJECT_DIR / "src" / "linked" _DEFAULT_OUTPUT_DIR = _PROJECT_DIR / "generated" / "linked" _LIBRARY_KEYS = { "python_distribution_package", + "python_distribution_version", "library_glob", + "include_glob", } _BINDING_KEYS = { "library", + "link_libraries", "required_symbols", "operator_schema", "dispatch_key", } -_SUPPORTED_TRANSPORTS = {"torch"} +_SUPPORTED_TRANSPORTS = {"torch", "tvm_ffi"} class ResolutionError(RuntimeError): @@ -66,6 +74,8 @@ class LibraryConfig: path: pathlib.Path python_distribution_package: str library_glob: str + include_glob: str | None = None + python_distribution_version: str | None = None @dataclasses.dataclass(frozen=True) @@ -80,6 +90,7 @@ class BindingConfig: required_symbols: tuple[str, ...] operator_schema: str | None dispatch_key: str | None + link_libraries: tuple[str, ...] = () def _load_yaml_mapping(path, expected_keys, required_keys=None): @@ -139,10 +150,16 @@ def _find_platform_dirs(source_root, device): return platform_dirs -def _load_libraries(platform_dir, device, transport): +def _load_libraries(platform_dir, device, transport, selected_libraries=None): libraries = {} for path in sorted(platform_dir.glob("*.yaml")): - data = _load_yaml_mapping(path, _LIBRARY_KEYS) + if selected_libraries is not None and path.stem not in selected_libraries: + continue + data = _load_yaml_mapping( + path, + _LIBRARY_KEYS, + {"python_distribution_package", "library_glob"}, + ) libraries[path.stem] = LibraryConfig( device=device, transport=transport, @@ -152,23 +169,100 @@ def _load_libraries(platform_dir, device, transport): data, "python_distribution_package", path ), library_glob=_require_relative_glob(data, "library_glob", path), + python_distribution_version=( + _require_string(data, "python_distribution_version", path) + if "python_distribution_version" in data + else None + ), + include_glob=( + _require_relative_glob(data, "include_glob", path) + if "include_glob" in data + else None + ), ) return libraries -def _load_bindings(platform_dir, device, transport, selected_ops): +def _binding_is_selected(name, header, selected_ops, config): + if selected_ops is not None and name not in selected_ops: + return False + + if config is not None: + selection = config.get(name) + + if selection is None: + return False + + headers = selection["headers"] + + if headers is not None: + selected_headers = { + ( + _PROJECT_DIR / ops_config.implementation_path(selected_header) + ).resolve() + for selected_header in headers + } + + return header.resolve() in selected_headers + + slots = selection["implementations"] + + return slots is None or ops_config.implementation_slot(header) in slots + + return selected_ops is None or name in selected_ops + + +def _load_bindings(platform_dir, device, transport, selected_ops, config): bindings = [] for path in sorted((platform_dir / "ops").glob("*/*.yaml")): name = path.parent.name - if selected_ops is not None and name not in selected_ops: + header = path.with_suffix(".h") + + if config is not None and name not in config: + continue + if config is None and selected_ops is not None and name not in selected_ops: + continue + if not header.is_file(): + raise ResolutionError(f"{path}: missing sibling {header.name}") + + try: + selected = _binding_is_selected(name, header, selected_ops, config) + except ops_config.OpsConfigError as error: + raise ResolutionError(str(error)) from error + + if not selected: continue + source = path.with_suffix(".cc") + cuda_source = path.with_suffix(".cu") + if not source.is_file(): + if not cuda_source.is_file(): + raise ResolutionError( + f"{path}: missing sibling {source.name} or {cuda_source.name}" + ) + source = cuda_source + elif cuda_source.is_file(): + raise ResolutionError( + f"{path}: both {source.name} and {cuda_source.name} are present" + ) + data = _load_yaml_mapping(path, _BINDING_KEYS, {"library"}) symbols = data.get("required_symbols") operator_schema = data.get("operator_schema") dispatch_key = data.get("dispatch_key") + link_libraries = data.get("link_libraries", []) + if not isinstance(link_libraries, list) or any( + not isinstance(library, str) or not library.strip() + for library in link_libraries + ): + raise ResolutionError( + f"{path}: link_libraries must be a list of non-empty strings" + ) + link_libraries = tuple(library.strip() for library in link_libraries) + if len(link_libraries) != len(set(link_libraries)): + raise ResolutionError(f"{path}: link_libraries contains duplicates") if (symbols is None) == (operator_schema is None): raise ResolutionError( f"{path} must define exactly one of required_symbols or operator_schema" @@ -199,13 +293,6 @@ def _load_bindings(platform_dir, device, transport, selected_ops): raise ResolutionError(f"{path}: operator_schema requires dispatch_key") dispatch_key = _require_string(data, "dispatch_key", path) - header = path.with_suffix(".h") - source = path.with_suffix(".cc") - if not header.is_file(): - raise ResolutionError(f"{path}: missing sibling {header.name}") - if not source.is_file(): - raise ResolutionError(f"{path}: missing sibling {source.name}") - bindings.append( BindingConfig( device=device, @@ -218,6 +305,7 @@ def _load_bindings(platform_dir, device, transport, selected_ops): required_symbols=symbols, operator_schema=operator_schema, dispatch_key=dispatch_key, + link_libraries=link_libraries, ) ) @@ -270,7 +358,7 @@ def _locate_editable_distribution_root(distribution): return root if root.is_dir() else None -def _locate_distribution_library(config): +def _load_distribution(config): try: distribution = importlib.metadata.distribution( config.python_distribution_package @@ -281,6 +369,26 @@ def _locate_distribution_library(config): f"{config.python_distribution_package!r} required by " f"{config.path} is not installed" ) from error + if config.python_distribution_version is not None: + try: + constraint = SpecifierSet(config.python_distribution_version) + version = Version(distribution.version) + except (InvalidSpecifier, InvalidVersion) as error: + raise ResolutionError( + f"{config.path}: invalid Python distribution version constraint" + ) from error + if version not in constraint: + raise ResolutionError( + f"{config.path}: {config.python_distribution_package!r} version " + f"{distribution.version!r} does not satisfy " + f"{config.python_distribution_version!r}" + ) + + return distribution + + +def _locate_distribution_library(config): + distribution = _load_distribution(config) matches = [] distribution_root = pathlib.Path(distribution.locate_file("")).resolve() @@ -320,6 +428,37 @@ def _locate_distribution_library(config): return matches[0] +def _locate_distribution_include(config): + + if config.include_glob is None: + return None + + distribution = _load_distribution(config) + + roots = [pathlib.Path(distribution.locate_file("")).resolve()] + editable_root = _locate_editable_distribution_root(distribution) + if editable_root is not None: + roots.append(editable_root) + + matches = [] + for root in roots: + for candidate in root.glob(config.include_glob): + candidate = candidate.resolve() + if candidate.is_dir() and candidate.is_relative_to(root): + matches.append(candidate) + + matches = sorted(set(matches)) + if len(matches) != 1: + formatted = ", ".join(str(path) for path in matches) or "none" + raise ResolutionError( + f"{config.path}: include_glob {config.include_glob!r} matched " + f"{len(matches)} directories in " + f"{config.python_distribution_package!r}: {formatted}" + ) + + return matches[0] + + def _run_symbol_tool(command, library_path): try: result = subprocess.run( @@ -486,6 +625,16 @@ def _render_cmake_manifest(payload): "INFINI_OPS_LINKED_SOURCES": [ operator["source"] for operator in payload["operators"] ], + "INFINI_OPS_LINKED_TORCH_SOURCES": [ + operator["source"] + for operator in payload["operators"] + if operator["transport"] == "torch" + ], + "INFINI_OPS_LINKED_TVM_FFI_SOURCES": [ + operator["source"] + for operator in payload["operators"] + if operator["transport"] == "tvm_ffi" + ], "INFINI_OPS_LINKED_LIBRARIES": [ library["path"] for library in payload["libraries"] ], @@ -495,6 +644,11 @@ def _render_cmake_manifest(payload): "INFINI_OPS_LINKED_RUNTIME_DIRS": [ library["runtime_dir"] for library in payload["libraries"] ], + "INFINI_OPS_LINKED_INCLUDE_DIRS": [ + library["include_dir"] + for library in payload["libraries"] + if "include_dir" in library + ], "INFINI_OPS_LINKED_TRANSPORTS": [ library["transport"] for library in payload["libraries"] ], @@ -533,6 +687,7 @@ def _normalize_values(values): def resolve_linked_ops( devices, ops=None, + config_path=None, *, source_root=_DEFAULT_SOURCE_ROOT, output_dir=_DEFAULT_OUTPUT_DIR, @@ -546,16 +701,38 @@ def resolve_linked_ops( selected_ops = _normalize_values(ops) selected_op_set = set(selected_ops) if selected_ops is not None else None + try: + selection_config = ( + ops_config.load_ops_config(config_path) if config_path is not None else None + ) + except ops_config.OpsConfigError as error: + raise ResolutionError(str(error)) from error + bindings = [] library_configs = {} for device in devices: for transport, platform_dir in _find_platform_dirs(source_root, device): - libraries = _load_libraries(platform_dir, device, transport) - for name, config in libraries.items(): - library_configs[(transport, device, name)] = config - bindings.extend( - _load_bindings(platform_dir, device, transport, selected_op_set) + platform_bindings = _load_bindings( + platform_dir, + device, + transport, + selected_op_set, + selection_config, + ) + bindings.extend(platform_bindings) + selected_libraries = { + library + for binding in platform_bindings + for library in (binding.library, *binding.link_libraries) + } + libraries = _load_libraries( + platform_dir, + device, + transport, + selected_libraries, ) + for name, library_config in libraries.items(): + library_configs[(transport, device, name)] = library_config bindings.sort( key=lambda binding: ( @@ -569,17 +746,19 @@ def resolve_linked_ops( inspected_symbols = {} dispatcher_contracts = [] for binding in bindings: + for library_name in (binding.library, *binding.link_libraries): + dependency_key = (binding.transport, binding.device, library_name) + library_config = library_configs.get(dependency_key) + if library_config is None: + raise ResolutionError( + f"{binding.path}: unknown library {library_name!r} for " + f"device {binding.device}" + ) + if dependency_key not in resolved_libraries: + resolved_libraries[dependency_key] = _locate_distribution_library( + library_config + ) key = (binding.transport, binding.device, binding.library) - library_config = library_configs.get(key) - if library_config is None: - raise ResolutionError( - f"{binding.path}: unknown library {binding.library!r} for " - f"device {binding.device}" - ) - - if key not in resolved_libraries: - library_path = _locate_distribution_library(library_config) - resolved_libraries[key] = library_path library_path = resolved_libraries[key] if binding.required_symbols: @@ -628,6 +807,7 @@ def resolve_linked_ops( for key in sorted(resolved_libraries): config = library_configs[key] library_path = resolved_libraries[key] + include_dir = _locate_distribution_include(config) libraries.append( { "device": config.device, @@ -639,6 +819,8 @@ def resolve_linked_ops( "transport": config.transport, } ) + if include_dir is not None: + libraries[-1]["include_dir"] = str(include_dir) operators = [] for binding in bindings: @@ -650,6 +832,8 @@ def resolve_linked_ops( "name": binding.name, "source": str(binding.source), } + if binding.link_libraries: + operator["link_libraries"] = list(binding.link_libraries) if binding.required_symbols: operator["required_symbols"] = list(binding.required_symbols) else: @@ -677,6 +861,7 @@ def _parse_args(): ) parser.add_argument("--devices", nargs="+", required=True) parser.add_argument("--ops", nargs="*") + parser.add_argument("--ops-config", type=pathlib.Path) parser.add_argument("--source-root", default=_DEFAULT_SOURCE_ROOT) parser.add_argument("--output-dir", default=_DEFAULT_OUTPUT_DIR) parser.add_argument("--nm", default=os.environ.get("CMAKE_NM", "nm")) @@ -685,15 +870,28 @@ def _parse_args(): return parser.parse_args() +def _selection_from_environment(ops, config_path): + if ops is not None or config_path is not None: + return ops, config_path + + value = os.environ.get("INFINI_OPS_OPS") + + if value is None: + return ops, config_path + if pathlib.Path(value).suffix.lower() == ".json": + return None, pathlib.Path(value) + + return [value], None + + def main(): args = _parse_args() - ops = args.ops - if ops is None and "INFINI_OPS_OPS" in os.environ: - ops = [os.environ["INFINI_OPS_OPS"]] + ops, config_path = _selection_from_environment(args.ops, args.ops_config) try: resolve_linked_ops( args.devices, ops, + config_path, source_root=args.source_root, output_dir=args.output_dir, nm=args.nm, diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index d88058efb..0e52d7811 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -505,7 +505,7 @@ if(WITH_ASCEND) endif() set(INFINI_OPS_OPS "" CACHE STRING - "Semicolon- or comma-separated operator allowlist for generated wrappers and bindings") + "Operator allowlist or path to an ops.json implementation selection") set(INFINI_OPS_SMOKE_BUILD OFF CACHE BOOL "Build only the smoke-test operator subset") set(_infini_ops_smoke_ops @@ -532,17 +532,47 @@ if(INFINI_OPS_SMOKE_BUILD) endif() endif() -if(INFINI_OPS_OPS) +set(_infini_ops_ops_config "") +if(INFINI_OPS_OPS MATCHES "\\.json$") + get_filename_component(_infini_ops_ops_config "${INFINI_OPS_OPS}" + ABSOLUTE BASE_DIR "${PROJECT_SOURCE_DIR}") + if(NOT EXISTS "${_infini_ops_ops_config}") + message(FATAL_ERROR + "Operator selection `${_infini_ops_ops_config}` does not exist.") + endif() +elseif(NOT INFINI_OPS_OPS AND EXISTS "${PROJECT_SOURCE_DIR}/ops.json") + set(_infini_ops_ops_config "${PROJECT_SOURCE_DIR}/ops.json") +endif() + +if(_infini_ops_ops_config) + file(GLOB_RECURSE _infini_ops_implementation_headers CONFIGURE_DEPENDS + "${PROJECT_SOURCE_DIR}/src/native/*.h" + "${PROJECT_SOURCE_DIR}/src/torch/*.h" + "${PROJECT_SOURCE_DIR}/src/ninetoothed/*.h" + "${PROJECT_SOURCE_DIR}/src/linked/*.h" + "${PROJECT_SOURCE_DIR}/src/triton/*.h") + set_property(DIRECTORY APPEND PROPERTY CMAKE_CONFIGURE_DEPENDS + "${_infini_ops_ops_config}" + "${PROJECT_SOURCE_DIR}/scripts/ops_config.py" + ${_infini_ops_implementation_headers}) + message(STATUS "Operator selection: ${_infini_ops_ops_config}") +endif() + +if(INFINI_OPS_OPS AND NOT INFINI_OPS_OPS MATCHES "\\.json$") string(REPLACE "," ";" _infini_ops_op_allowlist "${INFINI_OPS_OPS}") message(STATUS "Wrapper op allowlist: ${_infini_ops_op_allowlist}") endif() set(INFINI_OPS_LINKED_SOURCES "") +set(INFINI_OPS_LINKED_TORCH_SOURCES "") +set(INFINI_OPS_LINKED_TVM_FFI_SOURCES "") set(INFINI_OPS_LINKED_LIBRARIES "") set(INFINI_OPS_LINKED_FORCE_LOAD_LIBRARIES "") set(INFINI_OPS_LINKED_RUNTIME_DIRS "") +set(INFINI_OPS_LINKED_INCLUDE_DIRS "") set(INFINI_OPS_LINKED_TRANSPORTS "") set(_infini_ops_linked_uses_torch FALSE) +set(_infini_ops_linked_uses_tvm_ffi FALSE) if(WITH_LINKED) if(NOT DEVICE_LIST) @@ -552,6 +582,7 @@ if(WITH_LINKED) file(GLOB_RECURSE _linked_resolution_inputs CONFIGURE_DEPENDS "${PROJECT_SOURCE_DIR}/src/linked/*.cc" + "${PROJECT_SOURCE_DIR}/src/linked/*.cu" "${PROJECT_SOURCE_DIR}/src/linked/*.h" "${PROJECT_SOURCE_DIR}/src/linked/*.yaml") set_property(DIRECTORY APPEND PROPERTY CMAKE_CONFIGURE_DEPENDS @@ -581,7 +612,11 @@ if(WITH_LINKED) --nm "${_linked_nm}" --readelf "${_linked_readelf}" --cxxfilt "${_linked_cxxfilt}") - if(INFINI_OPS_OPS) + if(_infini_ops_ops_config) + list(APPEND _linked_resolver_args + --ops-config "${_infini_ops_ops_config}") + endif() + if(_infini_ops_op_allowlist) list(APPEND _linked_resolver_args --ops ${_infini_ops_op_allowlist}) endif() @@ -601,6 +636,8 @@ if(WITH_LINKED) foreach(_linked_transport IN LISTS INFINI_OPS_LINKED_TRANSPORTS) if(_linked_transport STREQUAL "torch") set(_infini_ops_linked_uses_torch TRUE) + elseif(_linked_transport STREQUAL "tvm_ffi") + set(_infini_ops_linked_uses_tvm_ffi TRUE) else() message(FATAL_ERROR "Unsupported linked operator transport `${_linked_transport}`.") @@ -665,6 +702,10 @@ if(WITH_TORCH) # which we then glob below alongside any hand-written torch sources. find_package(Python COMPONENTS Interpreter REQUIRED) + set_property(DIRECTORY APPEND PROPERTY CMAKE_CONFIGURE_DEPENDS + "${PROJECT_SOURCE_DIR}/scripts/generate_torch_ops.py" + "${PROJECT_SOURCE_DIR}/scripts/torch_ops.yaml") + # Pin codegen to the locally installed torch version so vendor # forks (Cambricon's `torch_mlu` 2.1.0, etc.) get a schema whose # `at::_out` overloads match the headers they ship. Without @@ -688,6 +729,10 @@ if(WITH_TORCH) set(_torch_codegen_args ${PROJECT_SOURCE_DIR}/scripts/generate_torch_ops.py --pytorch-version ${_torch_version_tag}) + if(_infini_ops_ops_config) + list(APPEND _torch_codegen_args + --ops-config "${_infini_ops_ops_config}") + endif() if(INFINI_OPS_TORCH_OPS) string(REPLACE "," ";" _torch_op_allowlist "${INFINI_OPS_TORCH_OPS}") list(APPEND _torch_codegen_args --ops ${_torch_op_allowlist}) @@ -717,9 +762,15 @@ if(WITH_TORCH) endif() if(_infini_ops_linked_uses_torch) - list(APPEND TORCH_SOURCES ${INFINI_OPS_LINKED_SOURCES}) + list(APPEND TORCH_SOURCES ${INFINI_OPS_LINKED_TORCH_SOURCES}) endif() + +if(_infini_ops_linked_uses_tvm_ffi) + target_sources(infiniops PRIVATE ${INFINI_OPS_LINKED_TVM_FFI_SOURCES}) + target_include_directories(infiniops PRIVATE + ${INFINI_OPS_LINKED_INCLUDE_DIRS}) +endif() if(WITH_CAMBRICON AND TORCH_SOURCES) execute_process( COMMAND "${_TORCH_PYTHON}" -c @@ -988,7 +1039,11 @@ if(GENERATE_OPERATOR_CALL_INSTANTIATIONS OR GENERATE_PYTHON_BINDINGS) # failures. set(GENERATOR_ARGS --devices ${DEVICE_LIST}) - if(INFINI_OPS_OPS) + if(_infini_ops_ops_config) + list(APPEND GENERATOR_ARGS + --ops-config "${_infini_ops_ops_config}") + endif() + if(_infini_ops_op_allowlist) list(APPEND GENERATOR_ARGS --ops ${_infini_ops_op_allowlist}) endif() if(WITH_TORCH) diff --git a/src/linked/tvm_ffi/nvidia/flashinfer_sampling.yaml b/src/linked/tvm_ffi/nvidia/flashinfer_sampling.yaml new file mode 100644 index 000000000..2d801e4c0 --- /dev/null +++ b/src/linked/tvm_ffi/nvidia/flashinfer_sampling.yaml @@ -0,0 +1,3 @@ +python_distribution_package: flashinfer-jit-cache +python_distribution_version: ">=0.6.7,<0.7" +library_glob: flashinfer_jit_cache/jit_cache/sampling/sampling.so diff --git a/src/linked/tvm_ffi/nvidia/ops/top_k_top_p_sampling_from_logits/flashinfer.cu b/src/linked/tvm_ffi/nvidia/ops/top_k_top_p_sampling_from_logits/flashinfer.cu new file mode 100644 index 000000000..4dfb06fe7 --- /dev/null +++ b/src/linked/tvm_ffi/nvidia/ops/top_k_top_p_sampling_from_logits/flashinfer.cu @@ -0,0 +1,637 @@ +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +#include "data_type.h" +#include "dispatcher.h" +#include "linked/tvm_ffi/nvidia/ops/top_k_top_p_sampling_from_logits/flashinfer.h" +#include "native/cpu/caster_.h" +#include "native/cuda/nvidia/caster.cuh" +#include "native/cuda/nvidia/runtime_.h" + +extern "C" { +int __tvm_ffi_softmax(void*, const TVMFFIAny*, int32_t, TVMFFIAny*); +int __tvm_ffi_top_k_mask_logits(void*, const TVMFFIAny*, int32_t, TVMFFIAny*); +int __tvm_ffi_top_p_sampling_from_probs(void*, const TVMFFIAny*, int32_t, + TVMFFIAny*); +int __tvm_ffi_top_k_top_p_sampling_from_probs(void*, const TVMFFIAny*, int32_t, + TVMFFIAny*); +} + +namespace infini::ops { +namespace { + +using OptionalTensorView = tvm::ffi::Optional; + +constexpr std::size_t kScratchBytes = 1024 * 1024; +constexpr std::size_t kAlignment = 256; +constexpr unsigned int kThreads = 256; + +std::size_t Align(std::size_t value) { + return (value + kAlignment - 1) & ~(kAlignment - 1); +} + +std::size_t AddWorkspaceRegion(std::size_t* offset, std::size_t size) { + *offset = Align(*offset); + const auto result = *offset; + *offset += size; + return result; +} + +struct WorkspaceLayout { + explicit WorkspaceLayout(std::size_t matrix_elements, + std::size_t batch_size) { + matrix_a = AddWorkspaceRegion(&size, matrix_elements * sizeof(float)); + matrix_b = AddWorkspaceRegion(&size, matrix_elements * sizeof(float)); + top_k = AddWorkspaceRegion(&size, batch_size * sizeof(int64_t)); + top_p = AddWorkspaceRegion(&size, batch_size * sizeof(float)); + valid = AddWorkspaceRegion(&size, batch_size * sizeof(uint8_t)); + indices = AddWorkspaceRegion(&size, batch_size * sizeof(int64_t)); + scratch = AddWorkspaceRegion(&size, kScratchBytes); + size = Align(size); + } + + std::size_t matrix_a{0}; + std::size_t matrix_b{0}; + std::size_t top_k{0}; + std::size_t top_p{0}; + std::size_t valid{0}; + std::size_t indices{0}; + std::size_t scratch{0}; + std::size_t size{0}; +}; + +class DeviceGuard { + public: + explicit DeviceGuard(int device_index) { + auto status = cudaGetDevice(&previous_device_); + assert(status == cudaSuccess && + "`FlashInferSampling` failed to query the current CUDA device"); + if (previous_device_ != device_index) { + status = cudaSetDevice(device_index); + assert(status == cudaSuccess && + "`FlashInferSampling` failed to select the input CUDA device"); + restore_ = true; + } + } + + ~DeviceGuard() { + if (!restore_) return; + const auto status = cudaSetDevice(previous_device_); + assert(status == cudaSuccess && + "`FlashInferSampling` failed to restore the CUDA device"); + } + + private: + int previous_device_{0}; + bool restore_{false}; +}; + +class StreamGuard { + public: + StreamGuard(int device_index, cudaStream_t stream) + : device_index_{device_index} { + const auto status = + TVMFFIEnvSetStream(kDLCUDA, device_index, stream, &previous_stream_); + assert(status == 0 && + "`FlashInferSampling` failed to set the TVM FFI CUDA stream"); + } + + ~StreamGuard() { + const auto status = + TVMFFIEnvSetStream(kDLCUDA, device_index_, previous_stream_, nullptr); + assert(status == 0 && + "`FlashInferSampling` failed to restore the TVM FFI CUDA stream"); + } + + private: + int device_index_{0}; + TVMFFIStreamHandle previous_stream_{nullptr}; +}; + +class EventRecorder { + public: + EventRecorder(cudaEvent_t event, cudaStream_t stream, bool* recorded) + : event_{event}, stream_{stream}, recorded_{recorded} {} + + ~EventRecorder() { + if (event_ == nullptr) return; + const auto status = cudaEventRecord(event_, stream_); + assert(status == cudaSuccess && + "`FlashInferSampling` failed to record workspace completion"); + *recorded_ = true; + } + + EventRecorder(const EventRecorder&) = delete; + EventRecorder& operator=(const EventRecorder&) = delete; + + private: + cudaEvent_t event_{nullptr}; + cudaStream_t stream_{nullptr}; + bool* recorded_{nullptr}; +}; + +DLDataType Dtype(DataType dtype) { + switch (dtype) { + case DataType::kInt32: + return {kDLInt, 32, 1}; + case DataType::kInt64: + return {kDLInt, 64, 1}; + case DataType::kFloat32: + return {kDLFloat, 32, 1}; + default: + assert(false && "`FlashInferSampling` received an unsupported dtype"); + return {kDLUInt, 8, 1}; + } +} + +DLTensor MakeTensor(void* data, int device_index, int32_t ndim, int64_t* shape, + DLDataType dtype) { + return {data, {kDLCUDA, device_index}, ndim, dtype, shape, nullptr, 0}; +} + +template +__global__ void CastLogits(float* dst, const Src* src, std::size_t count) { + for (auto index = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + index < count; + index += static_cast(blockDim.x) * gridDim.x) { + dst[index] = Caster::Cast(src[index]); + } +} + +template +__global__ void GatherCastLogits(float* dst, const Src* src, + const Index* indices, std::size_t rows, + std::size_t source_rows, + std::size_t vocab_size) { + const auto count = rows * vocab_size; + for (auto index = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + index < count; + index += static_cast(blockDim.x) * gridDim.x) { + const auto row = index / vocab_size; + const auto column = index % vocab_size; + const auto source_index = indices[row]; + assert(source_index >= 0 && + static_cast(source_index) < source_rows); + (void)source_rows; + const auto source_row = static_cast(source_index); + dst[index] = Caster::Cast( + src[source_row * vocab_size + column]); + } +} + +unsigned int Blocks(std::size_t count) { + const auto blocks = (count + kThreads - 1) / kThreads; + return static_cast(std::min(blocks, 65535)); +} + +void CallSoftmax(DLTensor* scratch, DLTensor* logits, DLTensor* output) { + tvm::ffi::Function::InvokeExternC( + nullptr, __tvm_ffi_softmax, tvm::ffi::TensorView(scratch), + tvm::ffi::TensorView(logits), tvm::ffi::TensorView(output), + OptionalTensorView{}, 1.0, false); +} + +void CallTopKMask(DLTensor* logits, DLTensor* output, DLTensor* top_k, + DLTensor* scratch) { + tvm::ffi::Function::InvokeExternC( + nullptr, __tvm_ffi_top_k_mask_logits, tvm::ffi::TensorView(logits), + tvm::ffi::TensorView(output), + OptionalTensorView{tvm::ffi::TensorView(top_k)}, int64_t{0}, + tvm::ffi::TensorView(scratch)); +} + +OptionalTensorView OptionalView(DLTensor* tensor) { + return tensor == nullptr ? OptionalTensorView{} + : OptionalTensorView{tvm::ffi::TensorView(tensor)}; +} + +void CallTopP(DLTensor* probs, DLTensor* output, DLTensor* valid, + DLTensor* indices, DLTensor* top_p, bool deterministic, + uint64_t seed, uint64_t offset) { + tvm::ffi::Function::InvokeExternC( + nullptr, __tvm_ffi_top_p_sampling_from_probs, tvm::ffi::TensorView(probs), + tvm::ffi::TensorView(output), tvm::ffi::TensorView(valid), + OptionalView(indices), OptionalTensorView{tvm::ffi::TensorView(top_p)}, + 1.0, deterministic, OptionalTensorView{}, seed, OptionalTensorView{}, + offset); +} + +void CallJoint(DLTensor* probs, DLTensor* output, DLTensor* valid, + DLTensor* indices, DLTensor* top_k, DLTensor* top_p, + bool deterministic, uint64_t seed, uint64_t offset) { + tvm::ffi::Function::InvokeExternC( + nullptr, __tvm_ffi_top_k_top_p_sampling_from_probs, + tvm::ffi::TensorView(probs), tvm::ffi::TensorView(output), + tvm::ffi::TensorView(valid), OptionalView(indices), + OptionalTensorView{tvm::ffi::TensorView(top_k)}, 0.0, + OptionalTensorView{tvm::ffi::TensorView(top_p)}, 1.0, deterministic, + OptionalTensorView{}, seed, OptionalTensorView{}, offset); +} + +int64_t ReadTopK(const Tensor tensor, Tensor::Size row) { + const auto offset = row * tensor.stride(0); + return tensor.dtype() == DataType::kInt32 + ? static_cast(tensor.data())[offset] + : static_cast(tensor.data())[offset]; +} + +float ReadTopP(const Tensor tensor, Tensor::Size row) { + const auto offset = row * tensor.stride(0); + switch (tensor.dtype()) { + case DataType::kFloat16: + return Caster::Cast( + static_cast(tensor.data())[offset]); + case DataType::kBFloat16: + return Caster::Cast( + static_cast(tensor.data())[offset]); + case DataType::kFloat32: + return static_cast(tensor.data())[offset]; + case DataType::kFloat64: + return static_cast( + static_cast(tensor.data())[offset]); + default: + assert(false && "`FlashInferSampling` received invalid top-p dtype"); + return 1.0f; + } +} + +void Validate(const Tensor logits, const Tensor top_k, const Tensor top_p, + const std::optional& indices, + const std::string& filter_apply_order, bool check_nan, + Tensor out) { + assert(logits.device().type() == Device::Type::kNvidia && + out.device() == logits.device() && logits.IsContiguous() && + out.IsContiguous() && + "`FlashInferSampling` requires contiguous NVIDIA logits and output"); + assert(top_k.device().type() == Device::Type::kCpu && + top_p.device().type() == Device::Type::kCpu && + "`FlashInferSampling` requires host top-k and top-p tensors"); + assert((out.dtype() == DataType::kInt32 || out.dtype() == DataType::kInt64) && + "`FlashInferSampling` requires int32 or int64 output"); + assert(!check_nan && "`FlashInferSampling` does not support check_nan"); + if (indices) { + assert((indices->device() == logits.device() || + indices->device().type() == Device::Type::kCpu) && + indices->IsContiguous() && indices->dtype() == out.dtype() && + "`FlashInferSampling` requires contiguous CPU or NVIDIA indices " + "matching output"); + } +} + +} // namespace + +Operator::Operator( + const Tensor logits, const Tensor top_k, const Tensor top_p, + const std::optional indices, const std::string filter_apply_order, + const bool deterministic, const bool check_nan, + const std::optional seed, const std::optional offset, + Tensor out) + : TopKTopPSamplingFromLogits(logits, top_k, top_p, indices, + filter_apply_order, deterministic, check_nan, + seed, offset, out), + workspace_size_{ + WorkspaceLayout(static_cast(out.size(0)) * + static_cast(logits.size(1)), + static_cast(out.size(0))) + .size}, + logits_batch_size_{logits.size(0)}, + device_index_{logits.device().index()}, + top_k_dtype_{top_k.dtype()}, + top_p_dtype_{top_p.dtype()}, + out_dtype_{out.dtype()}, + indices_dtype_{indices ? std::optional{indices->dtype()} : std::nullopt}, + indices_device_{indices ? std::optional{indices->device()} + : std::nullopt}, + filter_apply_order_{filter_apply_order}, + deterministic_{deterministic} { + Validate(logits, top_k, top_p, indices, filter_apply_order, check_nan, out); + assert(vocab_size_ > 0 && + vocab_size_ <= + static_cast(std::numeric_limits::max()) && + "`FlashInferSampling` requires a nonempty int32-sized vocabulary"); + if (batch_size_ == 0) return; + DeviceGuard guard{device_index_}; + auto status = cudaMalloc(&default_workspace_, workspace_size_); + assert(status == cudaSuccess && + "`FlashInferSampling` failed to allocate default workspace"); + for (auto& slot : staging_slots_) { + status = cudaMallocHost( + &slot.top_p, static_cast(batch_size_) * sizeof(float)); + assert(status == cudaSuccess && + "`FlashInferSampling` failed to allocate top-p staging"); + status = cudaMallocHost( + &slot.top_k, static_cast(batch_size_) * sizeof(int64_t)); + assert(status == cudaSuccess && + "`FlashInferSampling` failed to allocate top-k staging"); + status = cudaMallocHost( + &slot.indices, static_cast(batch_size_) * sizeof(int64_t)); + assert(status == cudaSuccess && + "`FlashInferSampling` failed to allocate indices staging"); + cudaEvent_t event; + status = cudaEventCreateWithFlags(&event, cudaEventDisableTiming); + assert(status == cudaSuccess && + "`FlashInferSampling` failed to create staging event"); + slot.event = event; + } + cudaEvent_t event; + status = cudaEventCreateWithFlags(&event, cudaEventDisableTiming); + assert(status == cudaSuccess && + "`FlashInferSampling` failed to create workspace event"); + default_workspace_event_ = event; +} + +Operator::~Operator() { + if (default_workspace_ == nullptr) return; + DeviceGuard guard{device_index_}; + [[maybe_unused]] auto status = cudaSuccess; + if (default_workspace_event_recorded_) { + status = cudaEventSynchronize( + static_cast(default_workspace_event_)); + assert(status == cudaSuccess && + "`FlashInferSampling` failed to await default workspace"); + } + status = cudaEventDestroy(static_cast(default_workspace_event_)); + assert(status == cudaSuccess && + "`FlashInferSampling` failed to destroy workspace event"); + for (auto& slot : staging_slots_) { + if (slot.event_recorded) { + status = cudaEventSynchronize(static_cast(slot.event)); + assert(status == cudaSuccess && + "`FlashInferSampling` failed to await staging"); + } + status = cudaEventDestroy(static_cast(slot.event)); + assert(status == cudaSuccess && + "`FlashInferSampling` failed to destroy staging event"); + status = cudaFreeHost(slot.indices); + assert(status == cudaSuccess && + "`FlashInferSampling` failed to free indices staging"); + status = cudaFreeHost(slot.top_k); + assert(status == cudaSuccess && + "`FlashInferSampling` failed to free top-k staging"); + status = cudaFreeHost(slot.top_p); + assert(status == cudaSuccess && + "`FlashInferSampling` failed to free top-p staging"); + } + status = cudaFree(default_workspace_); + assert(status == cudaSuccess && + "`FlashInferSampling` failed to free default workspace"); +} + +std::size_t Operator::workspace_size_in_bytes() const { + return workspace_size_; +} + +void Operator::operator()(const Tensor logits, const Tensor top_k, + const Tensor top_p, + const std::optional indices, + const std::string filter_apply_order, + const bool deterministic, const bool check_nan, + const std::optional seed, + const std::optional offset, + Tensor out) const { + assert( + logits.ndim() == 2 && logits.size(0) == logits_batch_size_ && + logits.size(1) == vocab_size_ && logits.dtype() == dtype_ && + logits.device().type() == Device::Type::kNvidia && + logits.device().index() == device_index_ && top_k.ndim() == 1 && + top_k.size(0) == batch_size_ && top_k.dtype() == top_k_dtype_ && + top_k.device().type() == Device::Type::kCpu && top_p.ndim() == 1 && + top_p.size(0) == batch_size_ && top_p.dtype() == top_p_dtype_ && + top_p.device().type() == Device::Type::kCpu && out.ndim() == 1 && + out.size(0) == batch_size_ && out.dtype() == out_dtype_ && + out.device() == logits.device() && + indices.has_value() == indices_dtype_.has_value() && + filter_apply_order == filter_apply_order_ && + deterministic == deterministic_ && + "`FlashInferSampling` call metadata changed after descriptor creation"); + if (indices) { + assert(indices->ndim() == 1 && indices->size(0) == batch_size_ && + indices->dtype() == *indices_dtype_ && + indices->device() == *indices_device_ && + "`FlashInferSampling` indices metadata changed after descriptor " + "creation"); + } + assert(!offset || *offset >= 0); + Validate(logits, top_k, top_p, indices, filter_apply_order, check_nan, out); + if (batch_size_ == 0) return; + + DeviceGuard device_guard{device_index_}; + const auto stream = static_cast(stream_); + StreamGuard stream_guard{device_index_, stream}; + std::lock_guard lock{mutex_}; + + auto slot_index = next_staging_slot_; + auto* slot = &staging_slots_[slot_index]; + auto status = cudaSuccess; + if (slot->event_recorded) { + status = cudaEventQuery(static_cast(slot->event)); + if (status == cudaErrorNotReady) { + const auto other_index = (slot_index + 1) % staging_slots_.size(); + auto* other = &staging_slots_[other_index]; + auto other_status = + other->event_recorded + ? cudaEventQuery(static_cast(other->event)) + : cudaSuccess; + if (other_status == cudaSuccess) { + slot_index = other_index; + slot = other; + } else { + assert(other_status == cudaErrorNotReady && + "`FlashInferSampling` failed to query staging event"); + status = cudaEventSynchronize(static_cast(slot->event)); + assert(status == cudaSuccess && + "`FlashInferSampling` failed to await staging slot"); + } + } else { + assert(status == cudaSuccess && + "`FlashInferSampling` failed to query staging event"); + } + } + next_staging_slot_ = (slot_index + 1) % staging_slots_.size(); + auto* workspace = + static_cast(workspace_ ? workspace_ : default_workspace_); + if (!workspace_ && default_workspace_event_recorded_) { + status = cudaStreamWaitEvent( + stream, static_cast(default_workspace_event_)); + assert(status == cudaSuccess && + "`FlashInferSampling` failed to sequence default workspace"); + } + const auto available = + workspace_ ? workspace_size_in_bytes_ : workspace_size_; + const auto matrix_elements = static_cast(batch_size_) * + static_cast(logits.size(1)); + EventRecorder workspace_recorder{ + workspace_ ? nullptr : static_cast(default_workspace_event_), + stream, &default_workspace_event_recorded_}; + + const WorkspaceLayout layout{matrix_elements, + static_cast(batch_size_)}; + assert(workspace != nullptr && available >= layout.size && + "`FlashInferSampling` received insufficient workspace"); + (void)available; + auto* matrix_a = reinterpret_cast(workspace + layout.matrix_a); + auto* matrix_b = reinterpret_cast(workspace + layout.matrix_b); + auto* top_k_device = workspace + layout.top_k; + auto* top_p_device = reinterpret_cast(workspace + layout.top_p); + auto* valid_device = workspace + layout.valid; + auto* scratch_device = workspace + layout.scratch; + auto* indices_device = workspace + layout.indices; + + for (Tensor::Size row = 0; row < batch_size_; ++row) { + const auto value = ReadTopP(top_p, row); + slot->top_p[static_cast(row)] = + value > 0.0f && value < 1.0f ? value : 1.0f; + } + status = + cudaMemcpyAsync(top_p_device, slot->top_p, + static_cast(batch_size_) * sizeof(float), + cudaMemcpyHostToDevice, stream); + assert(status == cudaSuccess && + "`FlashInferSampling` failed to stage top-p values"); + + const bool top_k_is_int64 = + filter_apply_order == "joint" && out.dtype() == DataType::kInt64; + if (top_k_is_int64) { + for (Tensor::Size row = 0; row < batch_size_; ++row) { + const auto value = ReadTopK(top_k, row); + slot->top_k[static_cast(row)] = + value > 0 && value <= static_cast(vocab_size_) + ? value + : static_cast(vocab_size_); + } + status = + cudaMemcpyAsync(top_k_device, slot->top_k, + static_cast(batch_size_) * sizeof(int64_t), + cudaMemcpyHostToDevice, stream); + } else { + auto* top_k_int32 = reinterpret_cast(slot->top_k); + for (Tensor::Size row = 0; row < batch_size_; ++row) { + const auto value = ReadTopK(top_k, row); + top_k_int32[static_cast(row)] = + value > 0 && value <= static_cast(vocab_size_) + ? static_cast(value) + : static_cast(vocab_size_); + } + status = + cudaMemcpyAsync(top_k_device, top_k_int32, + static_cast(batch_size_) * sizeof(int32_t), + cudaMemcpyHostToDevice, stream); + } + assert(status == cudaSuccess && + "`FlashInferSampling` failed to stage top-k values"); + + const void* staged_indices_data = indices ? indices->data() : nullptr; + if (indices && indices->device().type() == Device::Type::kCpu) { + if (indices->dtype() == DataType::kInt32) { + for (Tensor::Size row = 0; row < batch_size_; ++row) { + const auto value = static_cast(indices->data())[row]; + assert(value >= 0 && value < logits_batch_size_ && + "`FlashInferSampling` received an out-of-range index"); + } + } else { + for (Tensor::Size row = 0; row < batch_size_; ++row) { + const auto value = static_cast(indices->data())[row]; + assert(value >= 0 && value < logits_batch_size_ && + "`FlashInferSampling` received an out-of-range index"); + } + } + const auto bytes = static_cast(batch_size_) * + kDataTypeToSize.at(indices->dtype()); + std::memcpy(slot->indices, indices->data(), bytes); + status = cudaMemcpyAsync(indices_device, slot->indices, bytes, + cudaMemcpyHostToDevice, stream); + assert(status == cudaSuccess && + "`FlashInferSampling` failed to stage indices"); + staged_indices_data = indices_device; + } + + status = cudaEventRecord(static_cast(slot->event), stream); + assert(status == cudaSuccess && + "`FlashInferSampling` failed to record staging completion"); + slot->event_recorded = true; + + DispatchFunc( + logits.dtype(), + [&](auto tag) { + using T = typename decltype(tag)::type; + if (!indices) { + CastLogits<<>>( + matrix_a, static_cast(logits.data()), matrix_elements); + } else if (indices->dtype() == DataType::kInt32) { + GatherCastLogits<<>>( + matrix_a, static_cast(logits.data()), + static_cast(staged_indices_data), + static_cast(batch_size_), + static_cast(logits_batch_size_), + static_cast(vocab_size_)); + } else { + GatherCastLogits<<>>( + matrix_a, static_cast(logits.data()), + static_cast(staged_indices_data), + static_cast(batch_size_), + static_cast(logits_batch_size_), + static_cast(vocab_size_)); + } + }, + "`FlashInferSampling` logits cast"); + + int64_t matrix_shape[2]{static_cast(batch_size_), + static_cast(vocab_size_)}; + int64_t batch_shape[1]{static_cast(batch_size_)}; + int64_t scratch_shape[1]{static_cast(kScratchBytes)}; + auto matrix_a_tensor = MakeTensor(matrix_a, device_index_, 2, matrix_shape, + Dtype(DataType::kFloat32)); + auto matrix_b_tensor = MakeTensor(matrix_b, device_index_, 2, matrix_shape, + Dtype(DataType::kFloat32)); + auto top_k_tensor = + MakeTensor(top_k_device, device_index_, 1, batch_shape, + Dtype(top_k_is_int64 ? DataType::kInt64 : DataType::kInt32)); + auto top_p_tensor = MakeTensor(top_p_device, device_index_, 1, batch_shape, + Dtype(DataType::kFloat32)); + auto valid_tensor = + MakeTensor(valid_device, device_index_, 1, batch_shape, {kDLBool, 8, 1}); + auto scratch_tensor = MakeTensor(scratch_device, device_index_, 1, + scratch_shape, {kDLUInt, 8, 1}); + auto output_tensor = + MakeTensor(out.data(), device_index_, 1, batch_shape, Dtype(out.dtype())); + const auto actual_seed = static_cast( + seed.value_or(static_cast(std::random_device{}()))); + const auto actual_offset = static_cast(offset.value_or(0)); + if (filter_apply_order == "top_k_first") { + status = cudaMemsetAsync(scratch_device, 0, kScratchBytes, stream); + assert(status == cudaSuccess && + "`FlashInferSampling` failed to initialize row-state workspace"); + CallTopKMask(&matrix_a_tensor, &matrix_b_tensor, &top_k_tensor, + &scratch_tensor); + CallSoftmax(&scratch_tensor, &matrix_b_tensor, &matrix_a_tensor); + CallTopP(&matrix_a_tensor, &output_tensor, &valid_tensor, nullptr, + &top_p_tensor, deterministic, actual_seed, actual_offset); + } else { + CallSoftmax(&scratch_tensor, &matrix_a_tensor, &matrix_b_tensor); + CallJoint(&matrix_b_tensor, &output_tensor, &valid_tensor, nullptr, + &top_k_tensor, &top_p_tensor, deterministic, actual_seed, + actual_offset); + } + + status = cudaGetLastError(); + assert(status == cudaSuccess && + "`FlashInferSampling` CUDA kernel launch failed"); +} + +} // namespace infini::ops diff --git a/src/linked/tvm_ffi/nvidia/ops/top_k_top_p_sampling_from_logits/flashinfer.h b/src/linked/tvm_ffi/nvidia/ops/top_k_top_p_sampling_from_logits/flashinfer.h new file mode 100644 index 000000000..955307aff --- /dev/null +++ b/src/linked/tvm_ffi/nvidia/ops/top_k_top_p_sampling_from_logits/flashinfer.h @@ -0,0 +1,81 @@ +#ifndef INFINI_OPS_LINKED_TVM_FFI_NVIDIA_OPS_TOP_K_TOP_P_SAMPLING_FROM_LOGITS_FLASHINFER_H_ +#define INFINI_OPS_LINKED_TVM_FFI_NVIDIA_OPS_TOP_K_TOP_P_SAMPLING_FROM_LOGITS_FLASHINFER_H_ + +#include +#include +#include +#include +#include +#include + +#include "base/top_k_top_p_sampling_from_logits.h" + +namespace infini::ops { + +template <> +class Operator + : public TopKTopPSamplingFromLogits { + public: + Operator(const Tensor logits, const Tensor top_k, const Tensor top_p, + const std::optional indices, + const std::string filter_apply_order, const bool deterministic, + const bool check_nan, const std::optional seed, + const std::optional offset, Tensor out); + + ~Operator() override; + + std::size_t workspace_size_in_bytes() const override; + + void operator()(const Tensor logits, const Tensor top_k, const Tensor top_p, + const std::optional indices, + const std::string filter_apply_order, + const bool deterministic, const bool check_nan, + const std::optional seed, + const std::optional offset, + Tensor out) const override; + + private: + struct StagingSlot { + float* top_p{nullptr}; + int64_t* top_k{nullptr}; + void* indices{nullptr}; + void* event{nullptr}; + bool event_recorded{false}; + }; + + std::size_t workspace_size_{0}; + + Tensor::Size logits_batch_size_{0}; + + int device_index_{0}; + + DataType top_k_dtype_; + + DataType top_p_dtype_; + + DataType out_dtype_; + + std::optional indices_dtype_; + + std::optional indices_device_; + + std::string filter_apply_order_; + + bool deterministic_{false}; + + mutable std::array staging_slots_; + + mutable std::size_t next_staging_slot_{0}; + + void* default_workspace_event_{nullptr}; + + mutable bool default_workspace_event_recorded_{false}; + + mutable std::mutex mutex_; + + void* default_workspace_{nullptr}; +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_LINKED_TVM_FFI_NVIDIA_OPS_TOP_K_TOP_P_SAMPLING_FROM_LOGITS_FLASHINFER_H_ diff --git a/src/linked/tvm_ffi/nvidia/ops/top_k_top_p_sampling_from_logits/flashinfer.yaml b/src/linked/tvm_ffi/nvidia/ops/top_k_top_p_sampling_from_logits/flashinfer.yaml new file mode 100644 index 000000000..e3a4f14c0 --- /dev/null +++ b/src/linked/tvm_ffi/nvidia/ops/top_k_top_p_sampling_from_logits/flashinfer.yaml @@ -0,0 +1,8 @@ +library: flashinfer_sampling +link_libraries: + - tvm_ffi +required_symbols: + - __tvm_ffi_softmax + - __tvm_ffi_top_k_mask_logits + - __tvm_ffi_top_p_sampling_from_probs + - __tvm_ffi_top_k_top_p_sampling_from_probs diff --git a/src/linked/tvm_ffi/nvidia/tvm_ffi.yaml b/src/linked/tvm_ffi/nvidia/tvm_ffi.yaml new file mode 100644 index 000000000..c386e37cf --- /dev/null +++ b/src/linked/tvm_ffi/nvidia/tvm_ffi.yaml @@ -0,0 +1,4 @@ +python_distribution_package: apache-tvm-ffi +python_distribution_version: "==0.1.10" +library_glob: tvm_ffi/lib/libtvm_ffi.so +include_glob: tvm_ffi/include diff --git a/tests/test_generate_torch_ops.py b/tests/test_generate_torch_ops.py index 43dffc4bc..6dbed1f71 100644 --- a/tests/test_generate_torch_ops.py +++ b/tests/test_generate_torch_ops.py @@ -21,6 +21,41 @@ def _load_generator_module(): return module +def test_select_op_names_honors_slot_8_and_legacy_headers(): + module = _load_generator_module() + config = { + "default_all": {"headers": None, "implementations": None}, + "default_native": {"headers": None, "implementations": (0,)}, + "explicit_torch": {"headers": None, "implementations": (8,)}, + "legacy": {"headers": ["legacy.h"], "implementations": None}, + } + + assert module._select_op_names( + None, + ["default_all", "default_native"], + config, + ) == ["default_all", "explicit_torch"] + assert module._select_op_names( + ["explicit_torch"], + ["default_all", "default_native"], + config, + ) == ["explicit_torch"] + + +def test_select_op_names_maps_public_names_to_aten_names(): + module = _load_generator_module() + config = { + "div": {"headers": None, "implementations": (8,)}, + "internal_log_softmax": {"headers": None, "implementations": (8,)}, + } + + assert module._select_op_names( + None, + ["div", "div_", "_log_softmax"], + config, + ) == ["div", "div_", "_log_softmax"] + + def test_load_aten_entries_uses_packaged_torchgen(monkeypatch): module = _load_generator_module() entries = [{"func": "relu.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!)"}] diff --git a/tests/test_generate_wrappers.py b/tests/test_generate_wrappers.py index 6025f44c6..0f34483a0 100644 --- a/tests/test_generate_wrappers.py +++ b/tests/test_generate_wrappers.py @@ -600,6 +600,34 @@ def test_torch_system_compiler_receives_host_range_profile_definition(): ) in cmake +def test_explicit_op_allowlist_precedes_implicit_ops_json(): + cmake = (pathlib.Path(__file__).parents[1] / "src" / "CMakeLists.txt").read_text( + encoding="utf-8" + ) + wrapper = ( + pathlib.Path(__file__).parents[1] / "scripts" / "generate_wrappers.py" + ).read_text(encoding="utf-8") + + assert ( + 'elseif(NOT INFINI_OPS_OPS AND EXISTS "${PROJECT_SOURCE_DIR}/ops.json")' + in cmake + ) + assert 'if(INFINI_OPS_OPS AND NOT INFINI_OPS_OPS MATCHES "\\\\.json$")' in cmake + assert "implicit_config_path" not in wrapper + + +def test_operator_selection_tracks_implementation_header_changes(): + cmake = (pathlib.Path(__file__).parents[1] / "src" / "CMakeLists.txt").read_text( + encoding="utf-8" + ) + + for root in ("native", "torch", "ninetoothed", "linked", "triton"): + assert f'"${{PROJECT_SOURCE_DIR}}/src/{root}/*.h"' in cmake + assert "${_infini_ops_implementation_headers}" in cmake + assert '"${PROJECT_SOURCE_DIR}/scripts/generate_torch_ops.py"' in cmake + assert '"${PROJECT_SOURCE_DIR}/scripts/torch_ops.yaml"' in cmake + + def test_generated_dispatch_calls_start_with_dispatch_profile_scope( tmp_path, monkeypatch ): @@ -704,6 +732,62 @@ def test_filter_ops_strict_rejects_unavailable_ops(): raise AssertionError("strict unknown ops should fail") +def test_select_ops_from_config_filters_implementation_slots(tmp_path): + module = _load_generator_module() + slot_0 = tmp_path / "slot_0.h" + slot_16 = tmp_path / "slot_16.h" + slot_0.write_text("class Operator : public Add {};") + slot_16.write_text("class Operator : public Add {};") + implementation_0 = module._Implementation(slot_0, "native") + implementation_16 = module._Implementation(slot_16, "native") + config = { + "add": {"headers": None, "implementations": (16,)}, + } + + assert module._select_ops_from_config( + {"add": [implementation_0, implementation_16]}, config, "ops.json" + ) == {"add": [implementation_16]} + + +def test_select_ops_from_config_rejects_missing_slot(tmp_path): + module = _load_generator_module() + slot_0 = tmp_path / "slot_0.h" + slot_0.write_text("class Operator : public Add {};") + implementation_0 = module._Implementation(slot_0, "native") + config = { + "add": {"headers": None, "implementations": (16,)}, + } + + try: + module._select_ops_from_config({"add": [implementation_0]}, config, "ops.json") + except ValueError as exc: + assert "slot(s) 16" in str(exc) + else: + raise AssertionError("unavailable implementation slot should fail") + + +def test_select_ops_from_config_preserves_path_and_backend_descriptors(): + module = _load_generator_module() + config = { + "add": { + "headers": [ + "src/native/cpu/ops/add/add.h", + {"path": "custom/add.h", "backend": "custom"}, + ], + "implementations": None, + }, + } + + assert module._select_ops_from_config({}, config, "ops.json") == { + "add": [ + module._Implementation( + pathlib.Path("src/native/cpu/ops/add/add.h"), "native" + ), + module._Implementation(pathlib.Path("custom/add.h"), "custom"), + ] + } + + def test_linked_implementations_require_explicit_scan_flag(monkeypatch, tmp_path): module = _load_generator_module() src_dir = tmp_path / "moore" / "src" diff --git a/tests/test_ops_config.py b/tests/test_ops_config.py new file mode 100644 index 000000000..550cf8f67 --- /dev/null +++ b/tests/test_ops_config.py @@ -0,0 +1,114 @@ +import pathlib +import sys + +import pytest + + +_SCRIPTS_DIR = pathlib.Path(__file__).resolve().parents[1] / "scripts" +sys.path.insert(0, str(_SCRIPTS_DIR)) + +import ops_config # noqa: E402 + + +def _load(tmp_path, text): + path = tmp_path / "ops.json" + path.write_text(text) + + return ops_config.load_ops_config(path) + + +def test_loads_legacy_paths_and_slot_selections(tmp_path): + config = _load( + tmp_path, + """{ + "add": "src/native/cpu/ops/add/add.h", + "gemm": [ + "first.h", + {"path": "custom/second.h", "backend": "custom"} + ], + "relu": {"implementations": "all"}, + "sampling": {"implementations": [0, 16]} +} +""", + ) + + assert config == { + "add": {"headers": ["src/native/cpu/ops/add/add.h"], "implementations": None}, + "gemm": { + "headers": ["first.h", {"path": "custom/second.h", "backend": "custom"}], + "implementations": None, + }, + "relu": {"headers": None, "implementations": None}, + "sampling": {"headers": None, "implementations": (0, 16)}, + } + assert ops_config.selected_op_names(config) == ["add", "gemm", "relu", "sampling"] + assert ops_config.selected_slots(config, "sampling") == (0, 16) + assert ops_config.torch_op_names(config, ["relu", "sampling"]) == ["relu"] + config["sampling"]["implementations"] = (8,) + assert ops_config.torch_op_names(config, ["relu"]) == ["relu", "sampling"] + + +@pytest.mark.parametrize( + ("declaration", "slot"), + ( + ("class Operator : public Add {};", 0), + ( + "class Operator : public Add {};", + 16, + ), + ), +) +def test_reads_implementation_slot(tmp_path, declaration, slot): + header = tmp_path / "implementation.h" + header.write_text(declaration) + + assert ops_config.implementation_slot(header) == slot + + +def test_reads_one_slot_declared_for_multiple_devices(tmp_path): + header = tmp_path / "implementation.h" + header.write_text( + "class Operator : public Add {};\n" + "class Operator : public Add {};\n" + ) + + assert ops_config.implementation_slot(header) == 16 + + +def test_rejects_multiple_slots_in_one_header(tmp_path): + header = tmp_path / "implementation.h" + header.write_text( + "class Operator : public Add {};\n" + "class Operator : public Add {};\n" + ) + + with pytest.raises( + ops_config.OpsConfigError, + match="exactly one implementation slot; found 16, 17", + ): + ops_config.implementation_slot(header) + + +@pytest.mark.parametrize( + ("text", "message"), + ( + ('{"add": {}, "add": {"implementations": "all"}}', "duplicate key"), + ('{"add": {}}', "missing required key"), + ('{"add": {"slots": [0]}}', "unknown keys"), + ('{"add": {"implementations": []}}', "must not be empty"), + ('{"add": {"implementations": [0, 0]}}', "contain duplicates"), + ('{"add": {"implementations": [32]}}', "between 0 and 31"), + ('{"add": {"implementations": [true]}}', "between 0 and 31"), + ('{"add": ["first.h", 1]}', "paths or structured descriptors"), + ('{"add": [{"path": "x.h"}]}', "non-empty string 'backend'"), + ('{"add": [{"backend": "custom"}]}', "non-empty string 'path'"), + ( + '{"add": [{"path": "x.h", "backend": "custom", "slot": 1}]}', + "unknown keys", + ), + ), +) +def test_rejects_invalid_config(tmp_path, text, message): + with pytest.raises(ops_config.OpsConfigError, match=message): + _load(tmp_path, text) diff --git a/tests/test_resolve_linked_ops.py b/tests/test_resolve_linked_ops.py index 72a7b45d0..7fddc19cf 100644 --- a/tests/test_resolve_linked_ops.py +++ b/tests/test_resolve_linked_ops.py @@ -226,6 +226,90 @@ def test_resolve_supports_multiple_implementations_for_one_operator( ] +def test_resolve_filters_linked_dependencies_by_slot(monkeypatch, tmp_path): + module = _load_resolver_module() + source_root = tmp_path / "linked" + platform, op_dir = _write_linked_config(source_root) + (op_dir / "vllm.h").write_text( + "class Operator {};\n" + ) + (platform / "unused.yaml").write_text( + "python_distribution_package: unused\nlibrary_glob: unused/_C*.so\n" + ) + (op_dir / "unused.yaml").write_text( + "library: unused\nrequired_symbols:\n - unused()\n" + ) + (op_dir / "unused.h").write_text( + "class Operator {};\n" + ) + (op_dir / "unused.cc").write_text("// definition\n") + config_path = tmp_path / "ops.json" + config_path.write_text('{"silu_and_mul": {"implementations": [16]}}\n') + library_path = tmp_path / "vllm" / "_C.so" + library_path.parent.mkdir() + library_path.touch() + located = [] + + def locate(config): + located.append(config.name) + return library_path + + monkeypatch.setattr(module, "_locate_distribution_library", locate) + exported = {"silu_and_mul(at::Tensor&, at::Tensor&)"} + monkeypatch.setattr( + module, + "_inspect_dynamic_symbols", + lambda *args: (exported, exported), + ) + + payload = module.resolve_linked_ops( + ["metax"], + config_path=config_path, + source_root=source_root, + output_dir=tmp_path / "generated", + ) + + assert located == ["vllm"] + assert [op["implementation"] for op in payload["operators"]] == ["vllm"] + + +def test_resolve_matches_structured_implementation_descriptor(monkeypatch, tmp_path): + module = _load_resolver_module() + source_root = tmp_path / "linked" + _, op_dir = _write_linked_config(source_root) + config_path = tmp_path / "ops.json" + config_path.write_text( + '{"silu_and_mul": [{"path": "' + + (op_dir / "vllm.h").relative_to(tmp_path).as_posix() + + '", "backend": "linked"}]}\n' + ) + library_path = tmp_path / "vllm" / "_C.so" + library_path.parent.mkdir() + library_path.touch() + + monkeypatch.setattr(module, "_PROJECT_DIR", tmp_path) + monkeypatch.setattr( + module, + "_locate_distribution_library", + lambda _config: library_path, + ) + exported = {"silu_and_mul(at::Tensor&, at::Tensor&)"} + monkeypatch.setattr( + module, + "_inspect_dynamic_symbols", + lambda *args: (exported, exported), + ) + + payload = module.resolve_linked_ops( + ["metax"], + config_path=config_path, + source_root=source_root, + output_dir=tmp_path / "generated", + ) + + assert [op["implementation"] for op in payload["operators"]] == ["vllm"] + + @pytest.mark.parametrize( ("binding", "message"), ( @@ -682,6 +766,65 @@ def read_text(self, filename): assert module._locate_distribution_library(config) == library.resolve() +@pytest.mark.parametrize( + "version, succeeds", + (("0.6.6", False), ("0.6.7.post3", True), ("0.6.16", True)), +) +def test_load_distribution_enforces_version_constraint( + monkeypatch, tmp_path, version, succeeds +): + module = _load_resolver_module() + + class FakeDistribution: + pass + + distribution = FakeDistribution() + distribution.version = version + monkeypatch.setattr( + module.importlib.metadata, "distribution", lambda name: distribution + ) + config = module.LibraryConfig( + device="nvidia", + name="sampling", + path=tmp_path / "sampling.yaml", + transport="tvm_ffi", + python_distribution_package="flashinfer-jit-cache", + python_distribution_version=">=0.6.7,<0.7", + library_glob="sampling.so", + ) + + if succeeds: + assert module._load_distribution(config) is distribution + else: + with pytest.raises(module.ResolutionError, match="does not satisfy"): + module._load_distribution(config) + + +def test_load_distribution_rejects_invalid_version_constraint(monkeypatch, tmp_path): + module = _load_resolver_module() + + class FakeDistribution: + version = "0.6.7.post3" + + monkeypatch.setattr( + module.importlib.metadata, + "distribution", + lambda name: FakeDistribution(), + ) + config = module.LibraryConfig( + device="nvidia", + name="sampling", + path=tmp_path / "sampling.yaml", + transport="tvm_ffi", + python_distribution_package="flashinfer-jit-cache", + python_distribution_version="not-a-specifier", + library_glob="sampling.so", + ) + + with pytest.raises(module.ResolutionError, match="invalid.*constraint"): + module._load_distribution(config) + + @pytest.mark.parametrize( "editable, url", ((False, None), (True, "https://example.com/vllm")), @@ -699,3 +842,120 @@ def read_text(self, filename): return json.dumps({"url": url, "dir_info": {"editable": editable}}) assert module._locate_editable_distribution_root(FakeDistribution()) is None + + +def test_environment_selection_recognizes_ops_json(monkeypatch, tmp_path): + module = _load_resolver_module() + config_path = tmp_path / "ops.JSON" + monkeypatch.setenv("INFINI_OPS_OPS", str(config_path)) + + assert module._selection_from_environment(None, None) == (None, config_path) + + +def test_environment_selection_keeps_inline_allowlist(monkeypatch): + module = _load_resolver_module() + monkeypatch.setenv("INFINI_OPS_OPS", "add,gemm") + + assert module._selection_from_environment(None, None) == (["add,gemm"], None) + + +def test_explicit_selection_precedes_environment(monkeypatch, tmp_path): + module = _load_resolver_module() + env_config = tmp_path / "environment.json" + explicit_config = tmp_path / "explicit.json" + monkeypatch.setenv("INFINI_OPS_OPS", str(env_config)) + + assert module._selection_from_environment(["add"], None) == (["add"], None) + assert module._selection_from_environment(None, explicit_config) == ( + None, + explicit_config, + ) + + +def test_resolve_tvm_ffi_cuda_source_with_link_dependency(monkeypatch, tmp_path): + module = _load_resolver_module() + source_root = tmp_path / "linked" + platform = source_root / "tvm_ffi" / "nvidia" + op_dir = platform / "ops" / "sampling" + op_dir.mkdir(parents=True) + (platform / "sampling.yaml").write_text( + "python_distribution_package: flashinfer-jit-cache\n" + "library_glob: flashinfer_jit_cache/jit_cache/sampling/sampling.so\n" + ) + (platform / "tvm_ffi.yaml").write_text( + "python_distribution_package: apache-tvm-ffi\n" + "library_glob: tvm_ffi/lib/libtvm_ffi.so\n" + "include_glob: tvm_ffi/include\n" + ) + (op_dir / "flashinfer.yaml").write_text( + "library: sampling\n" + "link_libraries:\n" + " - tvm_ffi\n" + "required_symbols:\n" + " - __tvm_ffi_softmax\n" + ) + (op_dir / "flashinfer.h").write_text("// declaration\n") + (op_dir / "flashinfer.cu").write_text("// definition\n") + + libraries = { + "sampling": tmp_path / "sampling.so", + "tvm_ffi": tmp_path / "libtvm_ffi.so", + } + for library in libraries.values(): + library.touch() + include_dir = tmp_path / "include" + include_dir.mkdir() + monkeypatch.setattr( + module, + "_locate_distribution_library", + lambda config: libraries[config.name], + ) + monkeypatch.setattr( + module, + "_locate_distribution_include", + lambda config: include_dir if config.name == "tvm_ffi" else None, + ) + exported = {"__tvm_ffi_softmax"} + monkeypatch.setattr( + module, + "_inspect_dynamic_symbols", + lambda *args: (exported, exported), + ) + + output_dir = tmp_path / "generated" + payload = module.resolve_linked_ops( + ["nvidia"], + ["sampling"], + source_root=source_root, + output_dir=output_dir, + ) + + assert {entry["name"] for entry in payload["libraries"]} == { + "sampling", + "tvm_ffi", + } + assert payload["operators"] == [ + { + "device": "nvidia", + "transport": "tvm_ffi", + "implementation": "flashinfer", + "library": "sampling", + "link_libraries": ["tvm_ffi"], + "name": "sampling", + "required_symbols": ["__tvm_ffi_softmax"], + "source": str((op_dir / "flashinfer.cu").resolve()), + } + ] + manifest = (output_dir / "manifest.cmake").read_text() + tvm_sources = manifest.split("set(INFINI_OPS_LINKED_TVM_FFI_SOURCES", maxsplit=1)[ + 1 + ].split(")", maxsplit=1)[0] + torch_sources = manifest.split("set(INFINI_OPS_LINKED_TORCH_SOURCES", maxsplit=1)[ + 1 + ].split(")", maxsplit=1)[0] + include_dirs = manifest.split("set(INFINI_OPS_LINKED_INCLUDE_DIRS", maxsplit=1)[ + 1 + ].split(")", maxsplit=1)[0] + assert "flashinfer.cu" in tvm_sources + assert "flashinfer.cu" not in torch_sources + assert str(include_dir).replace("\\", "/") in include_dirs diff --git a/tests/test_top_k_top_p_sampling_from_logits.py b/tests/test_top_k_top_p_sampling_from_logits.py index d3c9b79af..ae44f7696 100644 --- a/tests/test_top_k_top_p_sampling_from_logits.py +++ b/tests/test_top_k_top_p_sampling_from_logits.py @@ -50,6 +50,173 @@ def test_top_k_top_p_sampling_from_logits( assert torch.all(torch.isin(first, allowed_tensor)) +def test_flashinfer_sampling_joint_host_indices(device, implementation_index): + if implementation_index != 16: + pytest.skip("FlashInfer linked-provider coverage") + + logits = torch.tensor( + ( + (9.0, 1.0, 0.0, -1.0), + (0.0, 8.0, 1.0, -1.0), + (-1.0, 0.0, 1.0, 7.0), + ), + dtype=torch.float32, + device=device, + ) + indices = torch.tensor((2, 0, 2, 1, 0, 1, 2), dtype=torch.int64) + batch_size = indices.numel() + top_k = torch.ones(batch_size, dtype=torch.int64) + top_p = torch.ones(batch_size, dtype=torch.float32) + out = torch.empty(batch_size, dtype=torch.int64, device=device) + + _top_k_top_p_sampling_from_logits( + logits, + top_k, + top_p, + 1234, + 9, + out, + implementation_index, + indices=indices, + filter_apply_order="joint", + ) + + expected = torch.tensor((3, 0, 3, 1, 0, 1, 3), dtype=torch.int64, device=device) + assert torch.equal(out, expected) + + +def test_flashinfer_sampling_top_k_first_cuda_indices(device, implementation_index): + if implementation_index != 16: + pytest.skip("FlashInfer linked-provider coverage") + + logits = torch.tensor( + ( + (9.0, 1.0, 0.0, -1.0), + (0.0, 8.0, 1.0, -1.0), + (-1.0, 0.0, 1.0, 7.0), + ), + dtype=torch.bfloat16, + device=device, + ) + indices = torch.tensor((2, 0, 2, 1, 0), dtype=torch.int32, device=device) + batch_size = indices.numel() + top_k = torch.ones(batch_size, dtype=torch.int32) + top_p = torch.ones(batch_size, dtype=torch.float64) + out = torch.empty(batch_size, dtype=torch.int32, device=device) + + _top_k_top_p_sampling_from_logits( + logits, + top_k, + top_p, + 1234, + 9, + out, + implementation_index, + indices=indices, + filter_apply_order="top_k_first", + ) + + expected = torch.tensor((3, 0, 3, 1, 0), dtype=torch.int32, device=device) + assert torch.equal(out, expected) + + +def test_flashinfer_sampling_offset(device, implementation_index): + if implementation_index != 16: + pytest.skip("FlashInfer linked-provider coverage") + + batch_size = 256 + logits = torch.zeros((batch_size, 4), dtype=torch.float32, device=device) + top_k = torch.full((batch_size,), 4, dtype=torch.int32) + top_p = torch.ones(batch_size, dtype=torch.float32) + first = torch.empty(batch_size, dtype=torch.int32, device=device) + repeated = torch.empty_like(first) + different_offset = torch.empty_like(first) + + _top_k_top_p_sampling_from_logits( + logits, + top_k, + top_p, + 1234, + 9, + first, + implementation_index, + filter_apply_order="joint", + ) + _top_k_top_p_sampling_from_logits( + logits, + top_k, + top_p, + 1234, + 9, + repeated, + implementation_index, + filter_apply_order="joint", + ) + _top_k_top_p_sampling_from_logits( + logits, + top_k, + top_p, + 1234, + 10, + different_offset, + implementation_index, + filter_apply_order="joint", + ) + + assert torch.equal(first, repeated) + assert not torch.equal(first, different_offset) + + +def test_flashinfer_sampling_uses_handle_stream(device, implementation_index): + if device != "cuda" or implementation_index != 16: + pytest.skip("FlashInfer linked-provider stream coverage") + + batch_size = 64 + logits = torch.zeros((batch_size, 4), dtype=torch.float32, device=device) + logits[:, 0] = 1.0 + top_k = torch.ones(batch_size, dtype=torch.int32) + top_p = torch.ones(batch_size, dtype=torch.float32) + out = torch.full((batch_size,), -1, dtype=torch.int32, device=device) + stream = torch.cuda.Stream() + + def call_sampling(): + infini.ops.top_k_top_p_sampling_from_logits( + logits, + top_k, + top_p, + None, + "joint", + True, + False, + 1234, + 9, + out, + stream=stream.cuda_stream, + implementation_index=implementation_index, + ) + + try: + call_sampling() + stream.synchronize() + out.fill_(-1) + torch.cuda.synchronize() + + with torch.cuda.stream(stream): + torch.cuda._sleep(50_000_000) + call_sampling() + + default_stream = torch.cuda.default_stream() + with torch.cuda.stream(default_stream): + snapshot = out.clone() + default_stream.synchronize() + assert torch.all(snapshot == -1) + + stream.synchronize() + assert torch.all(out == 0) + finally: + torch.cuda.synchronize() + + def _top_k_top_p_sampling_from_logits( logits, top_k, @@ -58,13 +225,16 @@ def _top_k_top_p_sampling_from_logits( offset, out, implementation_index, + *, + indices=None, + filter_apply_order="top_k_first", ): infini.ops.top_k_top_p_sampling_from_logits( logits, top_k, top_p, - None, - "top_k_first", + indices, + filter_apply_order, True, False, seed,