refactor(linked): migrate gptq_marlin_repack provider

docs/linked-operators.md

@@ -1,8 +1,9 @@
 # Linked Operators
 
-The linked backend calls operator symbols exported by an installed third-party
-shared library. It applies when a platform package exposes a usable C++ ABI but
-does not provide source code or a stable C API.
+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.
 
 ## Source Layout
 
@@ -37,6 +38,20 @@ required_symbols:
   - silu_and_mul(at::Tensor&, at::Tensor&)
 ```
 
+A Dispatcher implementation instead declares its exact schema and required
+dispatch key:
+
+```yaml
+library: vllm
+operator_schema: >-
+  _C::gptq_marlin_repack(Tensor b_q_weight, Tensor perm, SymInt size_k,
+  SymInt size_n, int num_bits, bool is_a_8bit) -> Tensor
+dispatch_key: CUDA
+```
+
+Each implementation uses exactly one contract form. The resolver rejects a
+partial Dispatcher contract or a binding that mixes both forms.
+
 ## Adapter Boundary
 
 Keep ABI behavior in `<implementation>.cc`, not in YAML. Shared operator
@@ -51,10 +66,15 @@ its provider-specific `Call` ABI.
 ## Configuration
 
 At configure time, `scripts/resolve_linked_ops.py` locates the installed Python
-distribution and verifies every required symbol with both `nm` and `readelf`.
-Raw `readelf` symbols are demangled with GNU or LLVM `c++filt` before exact
-comparison. The resolver writes a CMake manifest and diagnostic JSON under
-`generated/linked/`. Generated files are not committed.
+distribution. It verifies C++ symbols with both `nm` and `readelf`; raw
+`readelf` symbols are demangled with GNU or LLVM `c++filt` before exact
+comparison. Dispatcher contracts are verified in an isolated Python process by
+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.
 
 Enable the backend independently of generated ATen implementations:
 

scripts/resolve_linked_ops.py

@@ -19,7 +19,12 @@
     "python_distribution_package",
     "library_glob",
 }
-_BINDING_KEYS = {"library", "required_symbols"}
+_BINDING_KEYS = {
+    "library",
+    "required_symbols",
+    "operator_schema",
+    "dispatch_key",
+}
 _SUPPORTED_TRANSPORTS = {"torch"}
 
 
@@ -72,9 +77,11 @@ class BindingConfig:
     source: pathlib.Path
     library: str
     required_symbols: tuple[str, ...]
+    operator_schema: str | None
+    dispatch_key: str | None
 
 
-def _load_yaml_mapping(path, expected_keys):
+def _load_yaml_mapping(path, expected_keys, required_keys=None):
     try:
         data = yaml.load(path.read_text(encoding="utf-8"), Loader=_StrictLoader)
     except (OSError, yaml.YAMLError) as error:
@@ -85,7 +92,8 @@ def _load_yaml_mapping(path, expected_keys):
 
     keys = set(data)
     unknown_keys = sorted(keys - expected_keys)
-    missing_keys = sorted(expected_keys - keys)
+    required_keys = expected_keys if required_keys is None else required_keys
+    missing_keys = sorted(required_keys - keys)
     if unknown_keys:
         raise ResolutionError(
             f"{path} contains unknown keys: {', '.join(unknown_keys)}"
@@ -155,21 +163,40 @@ def _load_bindings(platform_dir, device, transport, selected_ops):
         if selected_ops is not None and name not in selected_ops:
             continue
 
-        data = _load_yaml_mapping(path, _BINDING_KEYS)
-        symbols = data["required_symbols"]
-        if (
-            not isinstance(symbols, list)
-            or not symbols
-            or any(
-                not isinstance(symbol, str) or not symbol.strip() for symbol in symbols
-            )
-        ):
+        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")
+
+        if (symbols is None) == (operator_schema is None):
             raise ResolutionError(
-                f"{path}: required_symbols must be a non-empty list of strings"
+                f"{path} must define exactly one of required_symbols or operator_schema"
             )
-        symbols = tuple(symbol.strip() for symbol in symbols)
-        if len(symbols) != len(set(symbols)):
-            raise ResolutionError(f"{path}: required_symbols contains duplicates")
+
+        if symbols is not None:
+            if dispatch_key is not None:
+                raise ResolutionError(f"{path}: dispatch_key requires operator_schema")
+            if (
+                not isinstance(symbols, list)
+                or not symbols
+                or any(
+                    not isinstance(symbol, str) or not symbol.strip()
+                    for symbol in symbols
+                )
+            ):
+                raise ResolutionError(
+                    f"{path}: required_symbols must be a non-empty list of strings"
+                )
+            symbols = tuple(symbol.strip() for symbol in symbols)
+            if len(symbols) != len(set(symbols)):
+                raise ResolutionError(f"{path}: required_symbols contains duplicates")
+            operator_schema = None
+        else:
+            symbols = ()
+            operator_schema = _require_string(data, "operator_schema", path)
+            if dispatch_key is None:
+                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")
@@ -188,6 +215,8 @@ def _load_bindings(platform_dir, device, transport, selected_ops):
                 source=source.resolve(),
                 library=_require_string(data, "library", path),
                 required_symbols=symbols,
+                operator_schema=operator_schema,
+                dispatch_key=dispatch_key,
             )
         )
 
@@ -346,6 +375,64 @@ def _verify_required_symbols(config, library_path, nm_symbols, readelf_symbols):
             )
 
 
+def _verify_dispatcher_contracts(contracts):
+    payload = [
+        {
+            "binding_path": str(config.path),
+            "library_path": str(library_path),
+            "schema": config.operator_schema,
+            "dispatch_key": config.dispatch_key,
+        }
+        for config, library_path in contracts
+    ]
+    script = (
+        "import json\n"
+        "import sys\n"
+        "import torch\n"
+        "contracts = json.loads(sys.argv[1])\n"
+        "loaded = set()\n"
+        "for contract in contracts:\n"
+        "    library_path = contract['library_path']\n"
+        "    if library_path not in loaded:\n"
+        "        torch.ops.load_library(library_path)\n"
+        "        loaded.add(library_path)\n"
+        "for contract in contracts:\n"
+        "    expected = torch._C.parse_schema(contract['schema'])\n"
+        "    actual = torch._C._dispatch_find_schema_or_throw(\n"
+        "        expected.name, expected.overload_name\n"
+        "    ).schema()\n"
+        "    if str(actual) != str(expected):\n"
+        "        sys.exit(\n"
+        "            f\"{contract['binding_path']}: expected {expected}, \"\n"
+        "            f'found {actual}'\n"
+        "        )\n"
+        "    qualified_name = expected.name\n"
+        "    if expected.overload_name:\n"
+        "        qualified_name += f'.{expected.overload_name}'\n"
+        "    dispatch_key = contract['dispatch_key']\n"
+        "    if not torch._C._dispatch_has_kernel_for_dispatch_key(\n"
+        "        qualified_name, dispatch_key\n"
+        "    ):\n"
+        "        sys.exit(\n"
+        "            f\"{contract['binding_path']}: {qualified_name} has no \"\n"
+        "            f'{dispatch_key} kernel'\n"
+        "        )\n"
+    )
+    try:
+        subprocess.run(
+            [sys.executable, "-c", script, json.dumps(payload)],
+            check=True,
+            capture_output=True,
+            text=True,
+        )
+    except (OSError, subprocess.CalledProcessError) as error:
+        stderr = getattr(error, "stderr", "") or ""
+        detail = f": {stderr.strip()}" if stderr.strip() else ""
+        raise ResolutionError(
+            f"dispatcher contracts are not provided by the resolved libraries{detail}"
+        ) from error
+
+
 def _cmake_quote(value):
     return str(value).replace("\\", "/").replace(";", "\\;").replace('"', '\\"')
 
@@ -358,6 +445,9 @@ def _render_cmake_manifest(payload):
         "INFINI_OPS_LINKED_LIBRARIES": [
             library["path"] for library in payload["libraries"]
         ],
+        "INFINI_OPS_LINKED_FORCE_LOAD_LIBRARIES": [
+            library["path"] for library in payload["libraries"] if library["force_load"]
+        ],
         "INFINI_OPS_LINKED_RUNTIME_DIRS": [
             library["runtime_dir"] for library in payload["libraries"]
         ],
@@ -433,6 +523,7 @@ def resolve_linked_ops(
     )
     resolved_libraries = {}
     inspected_symbols = {}
+    dispatcher_contracts = []
     for binding in bindings:
         key = (binding.transport, binding.device, binding.library)
         library_config = library_configs.get(key)
@@ -445,14 +536,20 @@ def resolve_linked_ops(
         if key not in resolved_libraries:
             library_path = _locate_distribution_library(library_config)
             resolved_libraries[key] = library_path
-            inspected_symbols[key] = _inspect_dynamic_symbols(
-                library_path, nm, readelf, cxxfilt
-            )
 
         library_path = resolved_libraries[key]
-        nm_symbols, readelf_symbols = inspected_symbols[key]
-        _verify_required_symbols(binding, library_path, nm_symbols, readelf_symbols)
-
+        if binding.required_symbols:
+            if key not in inspected_symbols:
+                inspected_symbols[key] = _inspect_dynamic_symbols(
+                    library_path, nm, readelf, cxxfilt
+                )
+            nm_symbols, readelf_symbols = inspected_symbols[key]
+            _verify_required_symbols(binding, library_path, nm_symbols, readelf_symbols)
+        else:
+            dispatcher_contracts.append((binding, library_path))
+
+    if dispatcher_contracts:
+        _verify_dispatcher_contracts(dispatcher_contracts)
     required_symbols = {
         symbol for binding in bindings for symbol in binding.required_symbols
     }
@@ -478,13 +575,19 @@ def resolve_linked_ops(
                 f"{previous} and {library_path}"
             )
 
+    force_load_keys = {
+        (binding.transport, binding.device, binding.library)
+        for binding in bindings
+        if binding.operator_schema is not None
+    }
     libraries = []
     for key in sorted(resolved_libraries):
         config = library_configs[key]
         library_path = resolved_libraries[key]
         libraries.append(
             {
                 "device": config.device,
+                "force_load": key in force_load_keys,
                 "name": config.name,
                 "path": str(library_path),
                 "python_distribution_package": (config.python_distribution_package),
@@ -493,21 +596,27 @@ def resolve_linked_ops(
             }
         )
 
+    operators = []
+    for binding in bindings:
+        operator = {
+            "device": binding.device,
+            "transport": binding.transport,
+            "library": binding.library,
+            "implementation": binding.implementation,
+            "name": binding.name,
+            "source": str(binding.source),
+        }
+        if binding.required_symbols:
+            operator["required_symbols"] = list(binding.required_symbols)
+        else:
+            operator["operator_schema"] = binding.operator_schema
+            operator["dispatch_key"] = binding.dispatch_key
+        operators.append(operator)
+
     payload = {
         "devices": list(devices),
         "libraries": libraries,
-        "operators": [
-            {
-                "device": binding.device,
-                "transport": binding.transport,
-                "library": binding.library,
-                "implementation": binding.implementation,
-                "name": binding.name,
-                "required_symbols": list(binding.required_symbols),
-                "source": str(binding.source),
-            }
-            for binding in bindings
-        ],
+        "operators": operators,
     }
     _write_if_changed(output_dir / "manifest.cmake", _render_cmake_manifest(payload))
     _write_if_changed(

src/CMakeLists.txt

@@ -512,6 +512,7 @@ endif()
 
 set(INFINI_OPS_LINKED_SOURCES "")
 set(INFINI_OPS_LINKED_LIBRARIES "")
+set(INFINI_OPS_LINKED_FORCE_LOAD_LIBRARIES "")
 set(INFINI_OPS_LINKED_RUNTIME_DIRS "")
 set(INFINI_OPS_LINKED_TRANSPORTS "")
 set(_infini_ops_linked_uses_torch FALSE)
@@ -598,10 +599,20 @@ if(WITH_LINKED)
 
         set(_linked_library_target
             "infini_ops_linked_library_${_linked_library_index}")
+        list(FIND INFINI_OPS_LINKED_FORCE_LOAD_LIBRARIES
+            "${_linked_library}" _linked_force_load)
+        if(_linked_force_load EQUAL -1)
+            set(_linked_library_link_items "-l:${_linked_library_name}")
+        else()
+            set(_linked_library_link_items
+                "-Wl,--push-state,--no-as-needed"
+                "-l:${_linked_library_name}"
+                "-Wl,--pop-state")
+        endif()
         add_library("${_linked_library_target}" INTERFACE IMPORTED)
         set_target_properties("${_linked_library_target}" PROPERTIES
             INTERFACE_LINK_DIRECTORIES "${_linked_library_dir}"
-            INTERFACE_LINK_LIBRARIES "-l:${_linked_library_name}")
+            INTERFACE_LINK_LIBRARIES "${_linked_library_link_items}")
         # Python extension DSOs generally have no SONAME. Link the exact
         # resolved basename while keeping the install location configurable.
         target_link_libraries(infiniops PRIVATE "${_linked_library_target}")

src/linked/torch/nvidia/ops/gptq_marlin_repack/vllm.cc

@@ -0,0 +1,38 @@
+#include "linked/torch/nvidia/ops/gptq_marlin_repack/vllm.h"
+
+#include <ATen/core/dispatch/Dispatcher.h>
+#include <ATen/core/stack.h>
+#include <c10/core/SymInt.h>
+
+#include <cassert>
+#include <utility>
+
+namespace infini::ops::linked::torch::nvidia {
+
+at::Tensor VllmGptqMarlinRepack::Call(at::Tensor b_q_weight, at::Tensor perm,
+                                      int64_t size_k, int64_t size_n,
+                                      int64_t num_bits, bool is_a_8bit) {
+  static const auto op = c10::Dispatcher::singleton().findSchemaOrThrow(
+      "_C::gptq_marlin_repack", "");
+  c10::Stack stack;
+  stack.emplace_back(std::move(b_q_weight));
+  stack.emplace_back(std::move(perm));
+  stack.emplace_back(c10::SymInt{size_k});
+  stack.emplace_back(c10::SymInt{size_n});
+  stack.emplace_back(num_bits);
+  stack.emplace_back(is_a_8bit);
+  op.callBoxed(&stack);
+
+  assert(stack.size() == 1 &&
+         "`gptq_marlin_repack` returned an unexpected number of values");
+  return std::move(stack.front()).toTensor();
+}
+
+}  // namespace infini::ops::linked::torch::nvidia
+
+namespace infini::ops::linked::torch {
+
+template class TorchGptqMarlinRepack<
+    ::infini::ops::linked::torch::nvidia::VllmGptqMarlinRepack>;
+
+}  // namespace infini::ops::linked::torch