From d3199f9d00f95aac5c8638687d586a6f5b780745 Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Thu, 13 Aug 2026 18:30:30 +0800 Subject: [PATCH] feat(linked): support TVM FFI providers --- docs/linked-operators.md | 47 +++++-- pyproject.toml | 2 +- scripts/resolve_linked_ops.py | 155 ++++++++++++++++++--- src/CMakeLists.txt | 18 ++- tests/test_resolve_linked_ops.py | 222 +++++++++++++++++++++++++++++++ 5 files changed, 414 insertions(+), 30 deletions(-) diff --git a/docs/linked-operators.md b/docs/linked-operators.md index 7e5afee9f..edf57f8e0 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 @@ -73,8 +86,10 @@ loading the DSO, comparing the registered schema exactly, and checking the requested dispatch key. The resolver writes a CMake manifest and diagnostic JSON under `generated/linked/`. Generated files are not committed. -DSOs that provide Dispatcher registrations are force-loaded only for their own -link item so that the linker cannot discard their static registration code. +DSOs that provide Dispatcher registrations and DSOs named in +`link_libraries` are retained with `--no-as-needed` only for their own link +items. This prevents the linker from discarding registration code or an +explicit dependency while preserving its state for unrelated libraries. Enable the backend independently of generated ATen implementations: @@ -86,6 +101,11 @@ cmake -S . -B build \ -DINFINI_OPS_OPS=silu_and_mul ``` +`WITH_TORCH=OFF` disables the generated ATen backend, but `WITH_LINKED=ON` +currently still requires an installed `torch` package. Linked configuration +shares its Python interpreter and C++ ABI setup with the existing `torch` +transport. + 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 @@ -98,10 +118,11 @@ 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/resolve_linked_ops.py b/scripts/resolve_linked_ops.py index 66d63d9a8..1541b1011 100644 --- a/scripts/resolve_linked_ops.py +++ b/scripts/resolve_linked_ops.py @@ -10,6 +10,8 @@ import sys import urllib.parse import urllib.request +from packaging.specifiers import InvalidSpecifier, SpecifierSet +from packaging.version import InvalidVersion, Version import yaml @@ -21,15 +23,18 @@ _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): @@ -69,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) @@ -83,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): @@ -147,7 +155,11 @@ def _load_libraries(platform_dir, device, transport, selected_libraries=None): for path in sorted(platform_dir.glob("*.yaml")): if selected_libraries is not None and path.stem not in selected_libraries: continue - data = _load_yaml_mapping(path, _LIBRARY_KEYS) + data = _load_yaml_mapping( + path, + _LIBRARY_KEYS, + {"python_distribution_package", "library_glob"}, + ) libraries[path.stem] = LibraryConfig( device=device, transport=transport, @@ -157,6 +169,16 @@ def _load_libraries(platform_dir, device, transport, selected_libraries=None): 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 @@ -213,14 +235,34 @@ def _load_bindings(platform_dir, device, transport, selected_ops, config): continue source = path.with_suffix(".cc") + cuda_source = path.with_suffix(".cu") if not source.is_file(): - raise ResolutionError(f"{path}: missing sibling {source.name}") + 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" @@ -263,6 +305,7 @@ def _load_bindings(platform_dir, device, transport, selected_ops, config): required_symbols=symbols, operator_schema=operator_schema, dispatch_key=dispatch_key, + link_libraries=link_libraries, ) ) @@ -315,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 @@ -326,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() @@ -365,6 +428,36 @@ 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( @@ -531,6 +624,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"] ], @@ -540,6 +643,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"] ], @@ -611,11 +719,16 @@ def resolve_linked_ops( 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, - {binding.library for binding in platform_bindings}, + selected_libraries, ) for name, library_config in libraries.items(): library_configs[(transport, device, name)] = library_config @@ -632,17 +745,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: @@ -687,10 +802,16 @@ def resolve_linked_ops( for binding in bindings if binding.operator_schema is not None } + force_load_keys.update( + (binding.transport, binding.device, library) + for binding in bindings + for library in binding.link_libraries + ) libraries = [] 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, @@ -702,6 +823,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: @@ -713,6 +836,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: diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index c39ad6bef..9f21e55d5 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -596,11 +596,15 @@ if(INFINI_OPS_OPS AND NOT INFINI_OPS_OPS MATCHES "\\.json$") 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) @@ -610,6 +614,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 @@ -663,6 +668,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}`.") @@ -787,9 +794,18 @@ 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}) +endif() + +if(INFINI_OPS_LINKED_INCLUDE_DIRS) + target_include_directories(infiniops PRIVATE + ${INFINI_OPS_LINKED_INCLUDE_DIRS}) +endif() if(WITH_CAMBRICON AND TORCH_SOURCES) execute_process( COMMAND "${_TORCH_PYTHON}" -c diff --git a/tests/test_resolve_linked_ops.py b/tests/test_resolve_linked_ops.py index 3444b008b..90dcee9c3 100644 --- a/tests/test_resolve_linked_ops.py +++ b/tests/test_resolve_linked_ops.py @@ -391,6 +391,44 @@ def test_resolve_requires_matching_implementation_sources(tmp_path, missing_suff ) +def test_resolve_rejects_both_cpp_and_cuda_sources(tmp_path): + module = _load_resolver_module() + source_root = tmp_path / "linked" + _, op_dir = _write_linked_config(source_root) + (op_dir / "vllm.cu").write_text("// duplicate definition\n") + + with pytest.raises( + module.ResolutionError, + match=r"both vllm\.cc and vllm\.cu are present", + ): + module.resolve_linked_ops( + ["metax"], + source_root=source_root, + output_dir=tmp_path / "generated", + ) + + +def test_resolve_rejects_unknown_link_library(monkeypatch, tmp_path): + module = _load_resolver_module() + source_root = tmp_path / "linked" + _write_linked_config( + source_root, + binding_extra="link_libraries:\n - missing\n", + ) + monkeypatch.setattr( + module, + "_locate_distribution_library", + lambda config: tmp_path / "_C.so", + ) + + with pytest.raises(module.ResolutionError, match="unknown library 'missing'"): + module.resolve_linked_ops( + ["metax"], + source_root=source_root, + output_dir=tmp_path / "generated", + ) + + def test_resolve_rejects_symbol_missing_from_either_tool(monkeypatch, tmp_path): module = _load_resolver_module() source_root = tmp_path / "linked" @@ -766,6 +804,94 @@ def read_text(self, filename): assert module._locate_distribution_library(config) == library.resolve() +def test_locate_distribution_include_matches_installed_directory(monkeypatch, tmp_path): + module = _load_resolver_module() + site_packages = tmp_path / "site-packages" + include_dir = site_packages / "tvm_ffi" / "include" + include_dir.mkdir(parents=True) + + class FakeDistribution: + def locate_file(self, entry): + return site_packages / entry + + def read_text(self, filename): + return None + + monkeypatch.setattr( + module.importlib.metadata, "distribution", lambda name: FakeDistribution() + ) + config = module.LibraryConfig( + device="nvidia", + name="tvm_ffi", + path=tmp_path / "tvm_ffi.yaml", + transport="tvm_ffi", + python_distribution_package="apache-tvm-ffi", + library_glob="tvm_ffi/lib/libtvm_ffi.so", + include_glob="tvm_ffi/include", + ) + + assert module._locate_distribution_include(config) == include_dir.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")), @@ -811,3 +937,99 @@ def test_explicit_selection_precedes_environment(monkeypatch, tmp_path): 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", + } + libraries_by_name = {entry["name"]: entry for entry in payload["libraries"]} + assert libraries_by_name["sampling"]["force_load"] is False + assert libraries_by_name["tvm_ffi"]["force_load"] is True + 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] + force_load_libraries = manifest.split( + "set(INFINI_OPS_LINKED_FORCE_LOAD_LIBRARIES", 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 + assert "libtvm_ffi.so" in force_load_libraries