diff --git a/docs/linked-operators.md b/docs/linked-operators.md index 8d95c4449..c06d71bc8 100644 --- a/docs/linked-operators.md +++ b/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 `.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: diff --git a/scripts/resolve_linked_ops.py b/scripts/resolve_linked_ops.py index e898937bf..cb5645cab 100644 --- a/scripts/resolve_linked_ops.py +++ b/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,6 +575,11 @@ 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] @@ -485,6 +587,7 @@ def resolve_linked_ops( 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( diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 749cd9f82..579f7a3af 100644 --- a/src/CMakeLists.txt +++ b/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}") diff --git a/src/linked/torch/nvidia/ops/gptq_marlin_repack/vllm.cc b/src/linked/torch/nvidia/ops/gptq_marlin_repack/vllm.cc new file mode 100644 index 000000000..b7901f230 --- /dev/null +++ b/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 +#include +#include + +#include +#include + +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 diff --git a/src/linked/torch/nvidia/ops/gptq_marlin_repack/vllm.h b/src/linked/torch/nvidia/ops/gptq_marlin_repack/vllm.h new file mode 100644 index 000000000..5a92d836f --- /dev/null +++ b/src/linked/torch/nvidia/ops/gptq_marlin_repack/vllm.h @@ -0,0 +1,39 @@ +#ifndef INFINI_OPS_LINKED_TORCH_NVIDIA_OPS_GPTQ_MARLIN_REPACK_VLLM_H_ +#define INFINI_OPS_LINKED_TORCH_NVIDIA_OPS_GPTQ_MARLIN_REPACK_VLLM_H_ + +#include "linked/torch/nvidia/c10.h" +#include "linked/torch/ops/gptq_marlin_repack.h" + +namespace infini::ops::linked::torch::nvidia { + +struct VllmGptqMarlinRepack : C10 { + static at::Tensor 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); +}; + +} // namespace infini::ops::linked::torch::nvidia + +namespace infini::ops::linked::torch { + +extern template class TorchGptqMarlinRepack< + ::infini::ops::linked::torch::nvidia::VllmGptqMarlinRepack>; + +} // namespace infini::ops::linked::torch + +namespace infini::ops { + +template <> +class Operator + : public linked::torch::TorchGptqMarlinRepack< + linked::torch::nvidia::VllmGptqMarlinRepack> { + public: + using linked::torch::TorchGptqMarlinRepack< + linked::torch::nvidia::VllmGptqMarlinRepack>::TorchGptqMarlinRepack; + + using linked::torch::TorchGptqMarlinRepack< + linked::torch::nvidia::VllmGptqMarlinRepack>::operator(); +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_LINKED_TORCH_NVIDIA_OPS_GPTQ_MARLIN_REPACK_VLLM_H_ diff --git a/src/linked/torch/nvidia/ops/gptq_marlin_repack/vllm.yaml b/src/linked/torch/nvidia/ops/gptq_marlin_repack/vllm.yaml new file mode 100644 index 000000000..a7733aa0e --- /dev/null +++ b/src/linked/torch/nvidia/ops/gptq_marlin_repack/vllm.yaml @@ -0,0 +1,5 @@ +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 diff --git a/src/linked/torch/nvidia/vllm.yaml b/src/linked/torch/nvidia/vllm.yaml new file mode 100644 index 000000000..b3335f759 --- /dev/null +++ b/src/linked/torch/nvidia/vllm.yaml @@ -0,0 +1,2 @@ +python_distribution_package: vllm +library_glob: vllm/_C.*.so diff --git a/src/linked/torch/ops/gptq_marlin_repack.h b/src/linked/torch/ops/gptq_marlin_repack.h new file mode 100644 index 000000000..3d4c4a1f2 --- /dev/null +++ b/src/linked/torch/ops/gptq_marlin_repack.h @@ -0,0 +1,45 @@ +#ifndef INFINI_OPS_LINKED_TORCH_OPS_GPTQ_MARLIN_REPACK_H_ +#define INFINI_OPS_LINKED_TORCH_OPS_GPTQ_MARLIN_REPACK_H_ + +#include + +#include "base/gptq_marlin_repack.h" +#include "torch/tensor_.h" + +namespace infini::ops::linked::torch { + +template +class TorchGptqMarlinRepack : public ::infini::ops::GptqMarlinRepack { + public: + using ::infini::ops::GptqMarlinRepack::GptqMarlinRepack; + using ::infini::ops::GptqMarlinRepack::operator(); + + void operator()(const Tensor b_q_weight, const Tensor perm, + const int64_t size_k, const int64_t size_n, + const int64_t num_bits, const bool is_a_8bit, + Tensor out) const override { + ValidateCallMetadata(b_q_weight, perm, size_k, size_n, num_bits, is_a_8bit, + out); + + const typename Backend::StreamGuard stream_guard{ + Backend::GetStreamFromExternal(stream_, device_index_)}; + auto at_b_q_weight = ToAtenTensor( + const_cast(b_q_weight.data()), b_q_weight_metadata_.shape(), + b_q_weight_metadata_.strides(), b_q_weight_metadata_.dtype(), + device_index_); + auto at_perm = ToAtenTensor( + const_cast(perm.data()), perm_metadata_.shape(), + perm_metadata_.strides(), perm_metadata_.dtype(), device_index_); + auto at_out = ToAtenTensor( + out.data(), out_metadata_.shape(), out_metadata_.strides(), + out_metadata_.dtype(), device_index_); + + auto result = Backend::Call(std::move(at_b_q_weight), std::move(at_perm), + size_k, size_n, num_bits, is_a_8bit); + at_out.copy_(result); + } +}; + +} // namespace infini::ops::linked::torch + +#endif // INFINI_OPS_LINKED_TORCH_OPS_GPTQ_MARLIN_REPACK_H_ diff --git a/src/native/cuda/nvidia/ops/gptq_marlin_repack/kernel.cu b/src/native/cuda/nvidia/ops/gptq_marlin_repack/kernel.cu deleted file mode 100644 index 2213dbdcf..000000000 --- a/src/native/cuda/nvidia/ops/gptq_marlin_repack/kernel.cu +++ /dev/null @@ -1,128 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright contributors to the vLLM project -// Adapted from vLLM at commit ffc4f08c8ee130d4ea6347c1bf31ffd4f8af28ab: -// csrc/libtorch_stable/quantization/marlin/gptq_marlin_repack.cu - -#include - -#include -#include - -#include "native/cuda/nvidia/ops/gptq_marlin_repack/kernel.cuh" -#include "native/cuda/nvidia/ops/gptq_marlin_repack/kernel.h" - -namespace infini::ops { -namespace { - -class DeviceGuard { - public: - explicit DeviceGuard(int device_index) { - auto status = cudaGetDevice(&previous_device_); - assert(status == cudaSuccess && - "`GptqMarlinRepack` failed to query the current CUDA device"); - - if (previous_device_ != device_index) { - status = cudaSetDevice(device_index); - assert(status == cudaSuccess && - "`GptqMarlinRepack` failed to select the input CUDA device"); - restore_ = true; - } - } - - ~DeviceGuard() { - if (restore_) { - const auto status = cudaSetDevice(previous_device_); - assert(status == cudaSuccess && - "`GptqMarlinRepack` failed to restore the CUDA device"); - } - } - - private: - int previous_device_{0}; - - bool restore_{false}; -}; - -template -void Launch(const uint32_t* b_q_weight, const uint32_t* perm, uint32_t* out, - int size_k, int size_n, int blocks, int shared_memory_bytes, - cudaStream_t stream) { - const auto attribute_status = cudaFuncSetAttribute( - gptq_marlin_repack_detail::GptqMarlinRepackKernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, shared_memory_bytes); - assert(attribute_status == cudaSuccess && - "`GptqMarlinRepack` failed to configure dynamic shared memory"); - - gptq_marlin_repack_detail::GptqMarlinRepackKernel - <<>>(b_q_weight, perm, out, size_k, size_n); -} - -} // namespace - -void Operator::operator()( - const Tensor b_q_weight, const Tensor perm, const int64_t size_k, - const int64_t size_n, const int64_t num_bits, const bool is_a_8bit, - Tensor out) const { - ValidateCallMetadata(b_q_weight, perm, size_k, size_n, num_bits, is_a_8bit, - out); - - DeviceGuard device_guard{device_index_}; - int blocks = 0; - auto status = cudaDeviceGetAttribute(&blocks, cudaDevAttrMultiProcessorCount, - device_index_); - assert(status == cudaSuccess && blocks > 0 && - "`GptqMarlinRepack` failed to query CUDA multiprocessor count"); - - int shared_memory_bytes = 0; - status = cudaDeviceGetAttribute(&shared_memory_bytes, - cudaDevAttrMaxSharedMemoryPerBlockOptin, - device_index_); - assert(status == cudaSuccess && shared_memory_bytes > 0 && - "`GptqMarlinRepack` failed to query CUDA shared memory capacity"); - - const auto* b_q_weight_ptr = - reinterpret_cast(b_q_weight.data()); - const auto* perm_ptr = - has_perm_ ? reinterpret_cast(perm.data()) : nullptr; - auto* out_ptr = reinterpret_cast(out.data()); - const auto stream = static_cast(stream_ ? stream_ : 0); - const auto kernel_size_k = static_cast(size_k_); - const auto kernel_size_n = static_cast(size_n_); - - if (is_a_8bit_) { - if (num_bits_ == 4) { - Launch<4, false, true>(b_q_weight_ptr, perm_ptr, out_ptr, kernel_size_k, - kernel_size_n, blocks, shared_memory_bytes, - stream); - } else { - Launch<8, false, true>(b_q_weight_ptr, perm_ptr, out_ptr, kernel_size_k, - kernel_size_n, blocks, shared_memory_bytes, - stream); - } - } else if (has_perm_) { - if (num_bits_ == 4) { - Launch<4, true, false>(b_q_weight_ptr, perm_ptr, out_ptr, kernel_size_k, - kernel_size_n, blocks, shared_memory_bytes, - stream); - } else { - Launch<8, true, false>(b_q_weight_ptr, perm_ptr, out_ptr, kernel_size_k, - kernel_size_n, blocks, shared_memory_bytes, - stream); - } - } else if (num_bits_ == 4) { - Launch<4, false, false>(b_q_weight_ptr, perm_ptr, out_ptr, kernel_size_k, - kernel_size_n, blocks, shared_memory_bytes, stream); - } else { - Launch<8, false, false>(b_q_weight_ptr, perm_ptr, out_ptr, kernel_size_k, - kernel_size_n, blocks, shared_memory_bytes, stream); - } - - status = cudaGetLastError(); - assert(status == cudaSuccess && - "`GptqMarlinRepack` CUDA kernel launch failed"); -} - -} // namespace infini::ops diff --git a/src/native/cuda/nvidia/ops/gptq_marlin_repack/kernel.cuh b/src/native/cuda/nvidia/ops/gptq_marlin_repack/kernel.cuh deleted file mode 100644 index e9d9f6e23..000000000 --- a/src/native/cuda/nvidia/ops/gptq_marlin_repack/kernel.cuh +++ /dev/null @@ -1,285 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright contributors to the vLLM project -// Adapted from vLLM at commit ffc4f08c8ee130d4ea6347c1bf31ffd4f8af28ab: -// csrc/libtorch_stable/quantization/marlin/gptq_marlin_repack.cu -// csrc/libtorch_stable/quantization/marlin/marlin.cuh - -#ifndef INFINI_OPS_NVIDIA_GPTQ_MARLIN_REPACK_KERNEL_CUH_ -#define INFINI_OPS_NVIDIA_GPTQ_MARLIN_REPACK_KERNEL_CUH_ - -#include -#include - -#include - -namespace infini::ops::gptq_marlin_repack_detail { - -constexpr int kRepackStages = 8; -constexpr int kRepackThreads = 256; -constexpr int kTileSize = 16; -constexpr int kTileKSize = kTileSize; -constexpr int kTileNSize = kTileKSize * 4; - -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 - -__device__ inline void CpAsync4(void* shared_ptr, const void* global_ptr) { - reinterpret_cast(shared_ptr)[0] = - reinterpret_cast(global_ptr)[0]; -} - -__device__ inline void CpAsyncFence() {} - -template -__device__ inline void CpAsyncWait() {} - -#else - -__device__ inline void CpAsync4(void* shared_ptr, const void* global_ptr) { - constexpr int kBytes = 16; - const auto shared_address = - static_cast(__cvta_generic_to_shared(shared_ptr)); - asm volatile("cp.async.cg.shared.global [%0], [%1], %2;\n" - : - : "r"(shared_address), "l"(global_ptr), "n"(kBytes)); -} - -__device__ inline void CpAsyncFence() { - asm volatile("cp.async.commit_group;\n" ::); -} - -template -__device__ inline void CpAsyncWait() { - asm volatile("cp.async.wait_group %0;\n" : : "n"(kCount)); -} - -#endif - -template -__global__ void GptqMarlinRepackKernel( - const uint32_t* __restrict__ b_q_weight_ptr, - const uint32_t* __restrict__ perm_ptr, uint32_t* __restrict__ out_ptr, - int size_k, int size_n) { - constexpr int kPackFactor = 32 / kNumBits; - constexpr int kTargetTileNSize = kTileNSize / (kIsA8Bit ? 2 : 1); - constexpr int kTargetTileKSize = kTileKSize * (kIsA8Bit ? 2 : 1); - const int k_tiles = size_k / kTargetTileKSize; - const int n_tiles = size_n / kTargetTileNSize; - const int block_k_tiles = (k_tiles + gridDim.x - 1) / gridDim.x; - - const auto start_k_tile = blockIdx.x * block_k_tiles; - if (start_k_tile >= k_tiles) { - return; - } - - const int finish_k_tile = min(start_k_tile + block_k_tiles, k_tiles); - - const auto wait_for_stage = [&]() { - CpAsyncWait(); - __syncthreads(); - }; - - extern __shared__ int4 shared[]; - constexpr int kPermSize = kTargetTileKSize / 4; - int4* shared_perm_ptr = shared; - int4* shared_pipe_ptr = shared_perm_ptr; - if constexpr (kHasPerm) { - shared_pipe_ptr += kPermSize; - } - - constexpr int kTileInts = kTargetTileKSize / kPackFactor; - constexpr int kStageNThreads = kTargetTileNSize / 4; - constexpr int kStageKThreads = kHasPerm ? kTargetTileKSize : kTileInts; - constexpr int kStageSize = kStageKThreads * kStageNThreads; - - const auto load_perm_to_shared = [&](int k_tile_id) { - const int first_k_int4 = (k_tile_id * kTargetTileKSize) / 4; - const auto* perm_int4_ptr = reinterpret_cast(perm_ptr); - - if (threadIdx.x < kPermSize) { - shared_perm_ptr[threadIdx.x] = perm_int4_ptr[first_k_int4 + threadIdx.x]; - } - __syncthreads(); - }; - - const auto fetch_to_shared = [&](int pipe, int k_tile_id, int n_tile_id) { - if (n_tile_id >= n_tiles) { - CpAsyncFence(); - return; - } - - const int first_n = n_tile_id * kTargetTileNSize; - int4* shared_ptr = shared_pipe_ptr + kStageSize * pipe; - - if constexpr (kHasPerm) { - if (threadIdx.x < kStageSize) { - const auto k_id = threadIdx.x / kStageNThreads; - const auto n_id = threadIdx.x % kStageNThreads; - const auto* shared_perm_int_ptr = - reinterpret_cast(shared_perm_ptr); - const int src_k = shared_perm_int_ptr[k_id]; - const int src_k_packed = src_k / kPackFactor; - - CpAsync4( - &shared_ptr[k_id * kStageNThreads + n_id], - reinterpret_cast( - &b_q_weight_ptr[src_k_packed * size_n + first_n + n_id * 4])); - } - } else if (threadIdx.x < kStageSize) { - const auto k_id = threadIdx.x / kStageNThreads; - const auto n_id = threadIdx.x % kStageNThreads; - const int first_k = k_tile_id * kTargetTileKSize; - const int first_k_packed = first_k / kPackFactor; - - CpAsync4(&shared_ptr[k_id * kStageNThreads + n_id], - reinterpret_cast( - &b_q_weight_ptr[(first_k_packed + k_id) * size_n + first_n + - n_id * 4])); - } - - CpAsyncFence(); - }; - - const auto repack_tile = [&](int pipe, int k_tile_id, int n_tile_id) { - if (n_tile_id >= n_tiles) { - return; - } - - const auto warp_id = threadIdx.x / 32; - const auto thread_id = threadIdx.x % 32; - if (warp_id >= 4) { - return; - } - - const int tensor_core_column = thread_id / 4; - const int tensor_core_row = (thread_id % 4) * (kIsA8Bit ? 4 : 2); - constexpr int kTensorCoreOffsets[4] = {0, 1, 8, 9}; - const int current_n = - (warp_id / (kIsA8Bit ? 2 : 1)) * 16 + tensor_core_column; - - constexpr int kSharedStride = kTargetTileNSize; - constexpr uint32_t kMask = (1U << kNumBits) - 1; - int4* shared_stage_ptr = shared_pipe_ptr + kStageSize * pipe; - auto* shared_stage_int_ptr = reinterpret_cast(shared_stage_ptr); - auto* shared_perm_int_ptr = reinterpret_cast(shared_perm_ptr); - uint32_t values[8]; - - if constexpr (kHasPerm) { - static_assert(!kIsA8Bit); - for (int i = 0; i < 4; ++i) { - const int k_index = tensor_core_row + kTensorCoreOffsets[i]; - const uint32_t src_k = shared_perm_int_ptr[k_index]; - const uint32_t src_k_position = src_k % kPackFactor; - const uint32_t first_value = - shared_stage_int_ptr[k_index * kSharedStride + current_n]; - const uint32_t second_value = - shared_stage_int_ptr[k_index * kSharedStride + current_n + 8]; - - values[i] = (first_value >> (src_k_position * kNumBits)) & kMask; - values[4 + i] = (second_value >> (src_k_position * kNumBits)) & kMask; - } - } else { - uint32_t first_values[kTileInts]; - uint32_t second_values[kTileInts]; - -#pragma unroll - for (int i = 0; i < kTileInts; ++i) { - if constexpr (kIsA8Bit) { - first_values[i] = shared_stage_int_ptr[current_n + kSharedStride * i + - (warp_id % 2) * 8]; - } else { - first_values[i] = shared_stage_int_ptr[current_n + kSharedStride * i]; - second_values[i] = - shared_stage_int_ptr[current_n + 8 + kSharedStride * i]; - } - } - -#pragma unroll - for (int i = 0; i < 4; ++i) { - const int current_element = - tensor_core_row + (kIsA8Bit ? i : kTensorCoreOffsets[i]); - const int current_int = current_element / kPackFactor; - const int current_position = current_element % kPackFactor; - values[i] = - (first_values[current_int] >> (current_position * kNumBits)) & - kMask; - if constexpr (kIsA8Bit) { - values[4 + i] = (first_values[current_int + kTileInts / 2] >> - (current_position * kNumBits)) & - kMask; - } else { - values[4 + i] = - (second_values[current_int] >> (current_position * kNumBits)) & - kMask; - } - } - } - - constexpr int kTileElements = - kTargetTileKSize * kTargetTileNSize / kPackFactor; - const int out_offset = (k_tile_id * n_tiles + n_tile_id) * kTileElements; - - // Matches FasterTransformer's interleaved numeric conversion layout: - // https://github.com/NVIDIA/FasterTransformer/blob/main/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h - if constexpr (!kIsA8Bit && kNumBits == 4) { - constexpr int kPackIndices[8] = {0, 2, 4, 6, 1, 3, 5, 7}; - uint32_t result = 0; -#pragma unroll - for (int i = 0; i < 8; ++i) { - result |= values[kPackIndices[i]] << (i * 4); - } - out_ptr[out_offset + thread_id * 4 + warp_id] = result; - } else if constexpr (kIsA8Bit && kNumBits == 4) { - constexpr int kPackIndices[8] = {0, 4, 1, 5, 2, 6, 3, 7}; - uint32_t result = 0; -#pragma unroll - for (int i = 0; i < 8; ++i) { - result |= values[kPackIndices[i]] << (i * 4); - } - out_ptr[out_offset + thread_id * 4 + warp_id] = result; - } else { - constexpr int kPackIndices[4] = {0, 2, 1, 3}; - uint32_t first_result = 0; - uint32_t second_result = 0; -#pragma unroll - for (int i = 0; i < 4; ++i) { - const int index = kIsA8Bit ? i : kPackIndices[i]; - first_result |= values[index] << (i * 8); - second_result |= values[4 + index] << (i * 8); - } - out_ptr[out_offset + thread_id * 8 + warp_id * 2] = first_result; - out_ptr[out_offset + thread_id * 8 + warp_id * 2 + 1] = second_result; - } - }; - - const auto start_pipes = [&](int k_tile_id, int n_tile_id) { -#pragma unroll - for (int pipe = 0; pipe < kRepackStages - 1; ++pipe) { - fetch_to_shared(pipe, k_tile_id, n_tile_id + pipe); - } - wait_for_stage(); - }; - -#pragma unroll - for (int k_tile_id = start_k_tile; k_tile_id < finish_k_tile; ++k_tile_id) { - int n_tile_id = 0; - if constexpr (kHasPerm) { - load_perm_to_shared(k_tile_id); - } - start_pipes(k_tile_id, n_tile_id); - - while (n_tile_id < n_tiles) { -#pragma unroll - for (int pipe = 0; pipe < kRepackStages; ++pipe) { - fetch_to_shared((pipe + kRepackStages - 1) % kRepackStages, k_tile_id, - n_tile_id + pipe + kRepackStages - 1); - repack_tile(pipe, k_tile_id, n_tile_id + pipe); - wait_for_stage(); - } - n_tile_id += kRepackStages; - } - } -} - -} // namespace infini::ops::gptq_marlin_repack_detail - -#endif // INFINI_OPS_NVIDIA_GPTQ_MARLIN_REPACK_KERNEL_CUH_ diff --git a/src/native/cuda/nvidia/ops/gptq_marlin_repack/kernel.h b/src/native/cuda/nvidia/ops/gptq_marlin_repack/kernel.h deleted file mode 100644 index 764f5261e..000000000 --- a/src/native/cuda/nvidia/ops/gptq_marlin_repack/kernel.h +++ /dev/null @@ -1,22 +0,0 @@ -#ifndef INFINI_OPS_NVIDIA_GPTQ_MARLIN_REPACK_KERNEL_H_ -#define INFINI_OPS_NVIDIA_GPTQ_MARLIN_REPACK_KERNEL_H_ - -#include "base/gptq_marlin_repack.h" - -namespace infini::ops { - -template <> -class Operator - : public GptqMarlinRepack { - public: - using GptqMarlinRepack::GptqMarlinRepack; - - void operator()(const Tensor b_q_weight, const Tensor perm, - const int64_t size_k, const int64_t size_n, - const int64_t num_bits, const bool is_a_8bit, - Tensor out) const override; -}; - -} // namespace infini::ops - -#endif // INFINI_OPS_NVIDIA_GPTQ_MARLIN_REPACK_KERNEL_H_ diff --git a/tests/test_resolve_linked_ops.py b/tests/test_resolve_linked_ops.py index 3538ee8db..ac583b632 100644 --- a/tests/test_resolve_linked_ops.py +++ b/tests/test_resolve_linked_ops.py @@ -85,6 +85,7 @@ def test_resolve_collects_selected_implementation_and_library(monkeypatch, tmp_p } ] assert payload["libraries"][0]["path"] == str(library_path) + assert not payload["libraries"][0]["force_load"] assert json.loads((output_dir / "resolved.json").read_text()) == payload manifest = (output_dir / "manifest.cmake").read_text() @@ -95,6 +96,77 @@ def test_resolve_collects_selected_implementation_and_library(monkeypatch, tmp_p assert "ignored" not in manifest +def test_resolve_validates_dispatcher_contract_and_force_loads_library( + monkeypatch, tmp_path +): + module = _load_resolver_module() + source_root = tmp_path / "linked" + platform = source_root / "torch" / "nvidia" + op_dir = platform / "ops" / "gptq_marlin_repack" + op_dir.mkdir(parents=True) + (platform / "vllm.yaml").write_text( + "python_distribution_package: vllm\nlibrary_glob: vllm/_C*.so\n" + ) + 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" + ) + (op_dir / "vllm.yaml").write_text( + f"library: vllm\noperator_schema: {schema}\ndispatch_key: CUDA\n" + ) + (op_dir / "vllm.h").write_text("// declaration\n") + (op_dir / "vllm.cc").write_text("// definition\n") + library_path = tmp_path / "site-packages" / "vllm" / "_C.abi3.so" + library_path.parent.mkdir(parents=True) + library_path.touch() + + monkeypatch.setattr( + module, "_locate_distribution_library", lambda config: library_path + ) + contracts = [] + monkeypatch.setattr( + module, + "_verify_dispatcher_contracts", + lambda resolved: contracts.extend(resolved), + ) + monkeypatch.setattr( + module, + "_inspect_dynamic_symbols", + lambda *args: pytest.fail("Dispatcher bindings do not inspect symbols"), + ) + + output_dir = tmp_path / "generated" + payload = module.resolve_linked_ops( + ["nvidia"], + ["gptq_marlin_repack"], + source_root=source_root, + output_dir=output_dir, + ) + + assert len(contracts) == 1 + contract, resolved_path = contracts[0] + assert (contract.operator_schema, contract.dispatch_key) == (schema, "CUDA") + assert resolved_path == library_path + assert payload["libraries"][0]["force_load"] + assert payload["operators"] == [ + { + "device": "nvidia", + "transport": "torch", + "implementation": "vllm", + "library": "vllm", + "name": "gptq_marlin_repack", + "operator_schema": schema, + "dispatch_key": "CUDA", + "source": str((op_dir / "vllm.cc").resolve()), + } + ] + manifest = (output_dir / "manifest.cmake").read_text() + force_load_block = manifest.split( + "set(INFINI_OPS_LINKED_FORCE_LOAD_LIBRARIES", maxsplit=1 + )[1].split(")", maxsplit=1)[0] + assert str(library_path).replace("\\", "/") in force_load_block + + def test_resolve_supports_multiple_implementations_for_one_operator( monkeypatch, tmp_path ): @@ -154,6 +226,41 @@ def test_resolve_supports_multiple_implementations_for_one_operator( ] +@pytest.mark.parametrize( + ("binding", "message"), + ( + ( + "library: vllm\n" + "required_symbols:\n" + " - symbol()\n" + "operator_schema: _C::op() -> Tensor\n" + "dispatch_key: CUDA\n", + "exactly one of required_symbols or operator_schema", + ), + ( + "library: vllm\noperator_schema: _C::op() -> Tensor\n", + "operator_schema requires dispatch_key", + ), + ( + "library: vllm\nrequired_symbols:\n - symbol()\ndispatch_key: CUDA\n", + "dispatch_key requires operator_schema", + ), + ), +) +def test_resolve_requires_one_complete_binding_contract(tmp_path, binding, message): + module = _load_resolver_module() + source_root = tmp_path / "linked" + _, op_dir = _write_linked_config(source_root) + (op_dir / "vllm.yaml").write_text(binding) + + with pytest.raises(module.ResolutionError, match=message): + module.resolve_linked_ops( + ["metax"], + source_root=source_root, + output_dir=tmp_path / "generated", + ) + + @pytest.mark.parametrize( ("library_extra", "binding_extra", "unknown_key"), ( @@ -390,6 +497,73 @@ def fake_run(command, **kwargs): } +def test_dispatcher_contract_validation_uses_one_isolated_process( + monkeypatch, tmp_path +): + module = _load_resolver_module() + schema = "_C::op(Tensor input) -> Tensor" + config = module.BindingConfig( + device="nvidia", + name="op", + implementation="vllm", + path=tmp_path / "vllm.yaml", + transport="torch", + source=tmp_path / "vllm.cc", + library="vllm", + required_symbols=(), + operator_schema=schema, + dispatch_key="CUDA", + ) + other_schema = "other::op(Tensor input) -> Tensor" + other_config = module.BindingConfig( + device="nvidia", + name="other_op", + implementation="other", + path=tmp_path / "other.yaml", + transport="torch", + source=tmp_path / "other.cc", + library="other", + required_symbols=(), + operator_schema=other_schema, + dispatch_key="CUDA", + ) + library_path = tmp_path / "_C.so" + other_library_path = tmp_path / "other.so" + calls = [] + + def fake_run(command, **kwargs): + calls.append((command, kwargs)) + return subprocess.CompletedProcess(command, 0, stdout="", stderr="") + + monkeypatch.setattr(module.subprocess, "run", fake_run) + module._verify_dispatcher_contracts( + [(config, library_path), (other_config, other_library_path)] + ) + + assert len(calls) == 1 + assert calls[0][0][0:2] == [sys.executable, "-c"] + payload = json.loads(calls[0][0][-1]) + assert payload == [ + { + "binding_path": str(config.path), + "library_path": str(library_path), + "schema": schema, + "dispatch_key": "CUDA", + }, + { + "binding_path": str(other_config.path), + "library_path": str(other_library_path), + "schema": other_schema, + "dispatch_key": "CUDA", + }, + ] + assert calls[0][1] == { + "check": True, + "capture_output": True, + "text": True, + } + + def test_library_glob_matches_path_segments(): module = _load_resolver_module()