feat(linked): add FlashInfer sampling provider

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.
 

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,16 +15,25 @@ src/linked/<transport>/<device>/
   ops/<operator>/
     <implementation>.yaml
     <implementation>.h
-    <implementation>.cc
+    <implementation>.{cc,cu}
 ```
 
+CUDA providers may use `<implementation>.cu` instead of `.cc`.
+
 The platform library file contains DSO discovery information:
 
 ```yaml
 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 `<implementation>.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 `<implementation>.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
 

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]

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,13 +1647,42 @@ 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(
         "--ops",
         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]] = []

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,