diff --git a/docs/linked-operators.md b/docs/linked-operators.md index 7e5afee9f..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 @@ -98,10 +111,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..4c20608d3 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,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( @@ -531,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"] ], @@ -540,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"] ], @@ -611,11 +720,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 +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: @@ -691,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, @@ -702,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: @@ -713,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: diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 36f946d8d..0e52d7811 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -564,11 +564,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) @@ -578,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 @@ -631,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}`.") @@ -755,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 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_resolve_linked_ops.py b/tests/test_resolve_linked_ops.py index 3444b008b..7fddc19cf 100644 --- a/tests/test_resolve_linked_ops.py +++ b/tests/test_resolve_linked_ops.py @@ -766,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")), @@ -811,3 +870,92 @@ 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", + } + 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,