From 6666aaa38eef0a429ba8e73965af41a260434019 Mon Sep 17 00:00:00 2001 From: whjthu Date: Tue, 13 Oct 2026 00:05:01 +0800 Subject: [PATCH 01/25] refactor: enable native dense inference --- CMakeLists.txt | 6 + cmake/NativeFlashAttention.cmake | 52 +++ ..._ops_nvidia_dense_sampling_torch_free.json | 21 + scripts/generate_public_operator_header.py | 98 +++- scripts/generate_wrappers.py | 55 ++- scripts/prepare_flash_attention_inference.py | 126 +++++ src/CMakeLists.txt | 21 + src/base/copy.h | 5 + src/base/embedding.h | 4 +- src/base/fill.h | 20 + src/base/flash_attn_varlen_func.h | 71 +++ src/base/flash_attn_with_kvcache.h | 142 ++++++ src/base/gemm.h | 118 +++++ src/base/reshape_and_cache_flash.h | 21 + src/config.h | 15 + src/dispatcher.h | 3 +- src/handle.h | 13 + src/native/cuda/nvidia/blas_utils.h | 8 + src/native/cuda/nvidia/ops/add/kernel.h | 6 + src/native/cuda/nvidia/ops/argmax/kernel.cuh | 124 +++++ src/native/cuda/nvidia/ops/argmax/kernel.h | 179 ++++++++ .../cuda/nvidia/ops/binary_view_capture.h | 128 ++++++ src/native/cuda/nvidia/ops/cast/kernel.h | 134 ++++++ src/native/cuda/nvidia/ops/copy/kernel.h | 5 + src/native/cuda/nvidia/ops/embedding/kernel.h | 1 + src/native/cuda/nvidia/ops/fill/kernel.h | 1 + .../ops/flash_attn_varlen_func/kernel.cc | 429 ++++++++++++++++++ .../ops/flash_attn_varlen_func/kernel.h | 64 +++ .../ops/flash_attn_varlen_func/launch.cu.in | 9 + .../infini_flash_cuda.h | 15 + .../ops/flash_attn_with_kvcache/kernel.cc | 320 +++++++++++++ .../ops/flash_attn_with_kvcache/kernel.h | 175 +++++++ .../ops/flash_attn_with_kvcache/launch.cu.in | 9 + src/native/cuda/nvidia/ops/flash_head_dim.h | 21 + src/native/cuda/nvidia/ops/flash_workspace.h | 104 +++++ .../nvidia/ops/fused_add_rms_norm/kernel.h | 1 + src/native/cuda/nvidia/ops/gemm/cublas.h | 20 +- src/native/cuda/nvidia/ops/gemm/cublaslt.h | 3 +- .../cuda/nvidia/ops/gemm/handle_owner.h | 111 +++++ src/native/cuda/nvidia/ops/mul/kernel.h | 5 + .../ops/reshape_and_cache_flash/kernel.h | 1 + src/native/cuda/nvidia/ops/rms_norm/kernel.h | 23 +- .../cuda/nvidia/ops/rotary_embedding/kernel.h | 1 + src/native/cuda/nvidia/ops/silu/kernel.h | 1 + .../cuda/nvidia/ops/silu_and_mul/kernel.h | 1 + src/native/cuda/nvidia/ops/softmax/kernel.h | 1 + src/native/cuda/nvidia/ops/tril/kernel.h | 20 + src/native/cuda/nvidia/ops/triu/kernel.h | 20 + src/native/cuda/ops/add/kernel.h | 94 ++-- src/native/cuda/ops/copy/kernel.h | 80 +++- src/native/cuda/ops/embedding/kernel.cuh | 8 +- src/native/cuda/ops/embedding/kernel.h | 106 +++-- src/native/cuda/ops/fill/kernel.h | 59 ++- src/native/cuda/ops/gemm/blas.h | 42 +- src/native/cuda/ops/mul/kernel.h | 89 ++-- src/native/cuda/ops/silu/kernel.h | 76 +++- src/native/cuda/ops/silu_and_mul/kernel.h | 80 +++- src/native/cuda/ops/softmax/kernel.h | 74 ++- src/native/cuda/ops/triangular/kernel.cuh | 73 +++ src/native/cuda/ops/triangular/kernel.h | 198 ++++++++ src/operator.h | 214 +++++++-- tests/test_generate_public_operator_header.py | 16 + tests/test_generate_wrappers.py | 14 + 63 files changed, 3695 insertions(+), 259 deletions(-) create mode 100644 cmake/NativeFlashAttention.cmake create mode 100644 scripts/configs/infiniops_ops_nvidia_dense_sampling_torch_free.json create mode 100644 scripts/prepare_flash_attention_inference.py create mode 100644 src/native/cuda/nvidia/ops/argmax/kernel.cuh create mode 100644 src/native/cuda/nvidia/ops/argmax/kernel.h create mode 100644 src/native/cuda/nvidia/ops/binary_view_capture.h create mode 100644 src/native/cuda/nvidia/ops/cast/kernel.h create mode 100644 src/native/cuda/nvidia/ops/flash_attn_varlen_func/kernel.cc create mode 100644 src/native/cuda/nvidia/ops/flash_attn_varlen_func/kernel.h create mode 100644 src/native/cuda/nvidia/ops/flash_attn_varlen_func/launch.cu.in create mode 100644 src/native/cuda/nvidia/ops/flash_attn_with_kvcache/infini_flash_cuda.h create mode 100644 src/native/cuda/nvidia/ops/flash_attn_with_kvcache/kernel.cc create mode 100644 src/native/cuda/nvidia/ops/flash_attn_with_kvcache/kernel.h create mode 100644 src/native/cuda/nvidia/ops/flash_attn_with_kvcache/launch.cu.in create mode 100644 src/native/cuda/nvidia/ops/flash_head_dim.h create mode 100644 src/native/cuda/nvidia/ops/flash_workspace.h create mode 100644 src/native/cuda/nvidia/ops/gemm/handle_owner.h create mode 100644 src/native/cuda/nvidia/ops/tril/kernel.h create mode 100644 src/native/cuda/nvidia/ops/triu/kernel.h create mode 100644 src/native/cuda/ops/triangular/kernel.cuh create mode 100644 src/native/cuda/ops/triangular/kernel.h diff --git a/CMakeLists.txt b/CMakeLists.txt index 0d8c3e5b3..39e17c3d3 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -26,6 +26,12 @@ option(WITH_ASCEND "Enable Ascend backend" OFF) option(WITH_TORCH "Enable PyTorch C++ backend" OFF) option(WITH_LINKED "Enable linked third-party operator backends" OFF) +option(BUILD_NVIDIA_FLASH_INFERENCE "Build pinned Torch-free NVIDIA inference Attention" OFF) +set(FLASH_ATTN_SOURCE_ARCHIVE "" CACHE FILEPATH "Exact flash_attn-2.8.3 source archive") +if(BUILD_NVIDIA_FLASH_INFERENCE AND NOT WITH_NVIDIA) + message(FATAL_ERROR "Native FlashAttention requires WITH_NVIDIA") +endif() + option(WITH_NINETOOTHED "Enable NineToothed-generated kernels" OFF) option(WITH_TRITON "Enable the Triton JIT backend" OFF) diff --git a/cmake/NativeFlashAttention.cmake b/cmake/NativeFlashAttention.cmake new file mode 100644 index 000000000..7cdef6024 --- /dev/null +++ b/cmake/NativeFlashAttention.cmake @@ -0,0 +1,52 @@ +# Build the pinned inference kernels separately. No Torch include/link flags or +# Python runtime enter this target; Python is used only to prepare pinned source. +function(infiniops_add_native_flash target) + find_package(Python3 REQUIRED COMPONENTS Interpreter) + find_package(CUDAToolkit REQUIRED) + if(NOT EXISTS "${FLASH_ATTN_SOURCE_ARCHIVE}") + message(FATAL_ERROR "Native FlashAttention requires the pinned source archive") + endif() + set(_provider "${CMAKE_CURRENT_FUNCTION_LIST_DIR}/../src/native/cuda/nvidia/ops/flash_attn_with_kvcache") + set(_varlen_provider "${CMAKE_CURRENT_FUNCTION_LIST_DIR}/../src/native/cuda/nvidia/ops/flash_attn_varlen_func") + set(_source "${CMAKE_CURRENT_FUNCTION_LIST_DIR}/..") + set(_prepared "${CMAKE_CURRENT_BINARY_DIR}/native-flash-inference") + execute_process(COMMAND "${Python3_EXECUTABLE}" + "${_source}/scripts/prepare_flash_attention_inference.py" + --archive "${FLASH_ATTN_SOURCE_ARCHIVE}" --output "${_prepared}" + RESULT_VARIABLE _status) + if(NOT _status EQUAL 0) + message(FATAL_ERROR "Pinned FlashAttention preparation failed") + endif() + set(_launches) + foreach(_flash_kind fp16 bf16) + if(_flash_kind STREQUAL fp16) + set(_flash_type "cutlass::half_t") + else() + set(_flash_type "cutlass::bfloat16_t") + endif() + foreach(_flash_causal false true) + foreach(_flash_hdim 32 64 96 128 192 256) + set(_launch "${CMAKE_CURRENT_BINARY_DIR}/flash_${_flash_hdim}_${_flash_kind}_${_flash_causal}.cu") + configure_file("${_provider}/launch.cu.in" "${_launch}" @ONLY) + list(APPEND _launches "${_launch}") + set(_varlen_launch "${CMAKE_CURRENT_BINARY_DIR}/flash_varlen_${_flash_hdim}_${_flash_kind}_${_flash_causal}.cu") + configure_file("${_varlen_provider}/launch.cu.in" "${_varlen_launch}" @ONLY) + list(APPEND _launches "${_varlen_launch}") + endforeach() + endforeach() + endforeach() + add_library(${target} OBJECT "${_provider}/kernel.cc" + "${_varlen_provider}/kernel.cc" ${_launches}) + set_target_properties(${target} PROPERTIES POSITION_INDEPENDENT_CODE ON CUDA_STANDARD 17) + target_compile_features(${target} PRIVATE cxx_std_17) + target_compile_definitions(${target} PRIVATE + FLASH_NAMESPACE=infini_flash FLASHATTENTION_DISABLE_DROPOUT) + target_include_directories(${target} PRIVATE + "${_source}/src" "${_source}/include" "${_prepared}" + "${_prepared}/cutlass" "${_provider}" "${_varlen_provider}" + ${INFINI_RT_INCLUDE_DIRS}) + target_compile_options(${target} PRIVATE + # Match the pinned upstream setup.py numerical compilation mode. + "$<$:--expt-relaxed-constexpr;--expt-extended-lambda;--use_fast_math>") + target_link_libraries(${target} PRIVATE CUDA::cudart) +endfunction() diff --git a/scripts/configs/infiniops_ops_nvidia_dense_sampling_torch_free.json b/scripts/configs/infiniops_ops_nvidia_dense_sampling_torch_free.json new file mode 100644 index 000000000..adce808f7 --- /dev/null +++ b/scripts/configs/infiniops_ops_nvidia_dense_sampling_torch_free.json @@ -0,0 +1,21 @@ +{ + "add": {"implementations": [0]}, + "argmax": {"implementations": [0]}, + "cast": {"implementations": [0]}, + "copy": {"implementations": [0]}, + "embedding": {"implementations": [0]}, + "fill": {"implementations": [0]}, + "flash_attn_varlen_func": {"implementations": [0]}, + "flash_attn_with_kvcache": {"implementations": [0]}, + "fused_add_rms_norm": {"implementations": [0]}, + "gemm": {"implementations": [0]}, + "mul": {"implementations": [0]}, + "reshape_and_cache_flash": {"implementations": [0]}, + "rms_norm": {"implementations": [0]}, + "rotary_embedding": {"implementations": [0]}, + "silu": {"implementations": [0]}, + "silu_and_mul": {"implementations": [0]}, + "softmax": {"implementations": [0]}, + "tril": {"implementations": [0]}, + "triu": {"implementations": [0]} +} diff --git a/scripts/generate_public_operator_header.py b/scripts/generate_public_operator_header.py index eb6a3add4..8c6d2edba 100644 --- a/scripts/generate_public_operator_header.py +++ b/scripts/generate_public_operator_header.py @@ -16,6 +16,90 @@ re.MULTILINE, ) +# Provider implementations are visible only while building InfiniOps. Their +# lookup and construction bodies must not be optimized in a downstream TU, +# even when that TU sees an `extern template` declaration. +_PROVIDER_SIGNATURES = ( + ( + "Operator::CanMake(Tensor)", + re.compile( + r"(?P^[ \t]*)template\s*\s*\n" + r"(?P=indent)static\s+bool\s+CanMake\s*\(\s*" + r"const\s+Config\s*&\s*config\s*,\s*const\s+Tensor\s+tensor\s*,\s*" + r"const\s+Args\s*&\s*\.\.\.\s*args\s*\)\s*", + re.MULTILINE, + ), + ), + ( + "Operator::Make(Tensor)", + re.compile( + r"(?P^[ \t]*)template\s*\s*\n" + r"(?P=indent)static\s+std::unique_ptr\s+Make\s*\(\s*" + r"const\s+Config\s*&\s*config\s*,\s*const\s+Tensor\s+tensor\s*,\s*" + r"Args\s*&&\s*\.\.\.\s*args\s*\)\s*", + re.MULTILINE, + ), + ), + ( + "Operator::Make(vector)", + re.compile( + r"(?P^[ \t]*)template\s*\s*\n" + r"(?P=indent)static\s+std::unique_ptr\s+Make\s*\(\s*" + r"const\s+Config\s*&\s*config\s*,\s*" + r"const\s+std::vector\s+tensors\s*,\s*" + r"Args\s*&&\s*\.\.\.\s*args\s*\)\s*", + re.MULTILINE, + ), + ), + ( + "Operator::active_implementation_indices", + re.compile( + r"(?P^[ \t]*)static\s+std::vector\s+" + r"active_implementation_indices\s*\(\s*Device::Type\s+dev_type\s*\)\s*", + re.MULTILINE, + ), + ), + ( + "Operator::DefaultConfig", + re.compile( + r"(?P^[ \t]*)static\s+Config\s+" + r"DefaultConfig\s*\(\s*Device::Type\s+dev_type\s*\)\s*", + re.MULTILINE, + ), + ), + ( + "Operator::DefaultImplementationIndex", + re.compile( + r"(?P^[ \t]*)static\s+std::size_t\s+" + r"DefaultImplementationIndex\s*\(\s*Device::Type\s+dev_type\s*\)\s*", + re.MULTILINE, + ), + ), + ( + "Operator::MakeWithDevice", + re.compile( + r"(?P^[ \t]*)template\s*\s*\n" + r"(?P=indent)static\s+std::unique_ptr\s+" + r"MakeWithDevice\s*\(\s*const\s+Config\s*&\s*config\s*,\s*" + r"Device::Type\s+dispatch_device_type\s*,\s*" + r"Args\s*&&\s*\.\.\.\s*args\s*\)\s*", + re.MULTILINE, + ), + ), + ( + "Operator::SelectWithDevice", + re.compile( + r"(?P^[ \t]*)template\s*\s*\n" + r"(?P=indent)static\s+bool\s+SelectWithDevice\s*\(\s*" + r"const\s+Config\s*&\s*config\s*,\s*" + r"Device::Type\s+dispatch_device_type\s*,\s*" + r"bool\s+strict\s*,\s*Selected\s*&&\s*selected\s*,\s*" + r"Args\s*&&\s*\.\.\.\s*args\s*\)\s*", + re.MULTILINE, + ), + ), +) + def _find_matching_brace(text: str, open_brace: int) -> int: depth = 0 @@ -111,19 +195,19 @@ def _find_matching_brace(text: str, open_brace: int) -> int: raise RuntimeError("Could not find the end of `Operator::Call`.") -def _hide_operator_call_definition(text: str) -> str: - match = _CALL_SIGNATURE_RE.search(text) +def _hide_definition(text: str, signature: re.Pattern[str], name: str) -> str: + match = signature.search(text) if match is None: - raise RuntimeError("Could not find the `Operator::Call` signature.") + raise RuntimeError(f"Could not find the `{name}` signature.") body_start = text.find("{", match.end()) if body_start == -1: - raise RuntimeError("Could not find the `Operator::Call` body.") + raise RuntimeError(f"Could not find the `{name}` body.") if text[match.end() : body_start].strip(): - raise RuntimeError("Unexpected tokens before the `Operator::Call` body.") + raise RuntimeError(f"Unexpected tokens before the `{name}` body.") body_end = _find_matching_brace(text, body_start) signature = text[match.start() : match.end()].rstrip() @@ -132,7 +216,9 @@ def _hide_operator_call_definition(text: str) -> str: def generate_public_operator_header(source: Path, output: Path) -> None: - public_text = _hide_operator_call_definition(source.read_text()) + public_text = _hide_definition(source.read_text(), _CALL_SIGNATURE_RE, "Operator::Call") + for name, signature in _PROVIDER_SIGNATURES: + public_text = _hide_definition(public_text, signature, name) public_text = public_text.replace('#include "host_range_profiler.h"\n', "") output.parent.mkdir(parents=True, exist_ok=True) output.write_text(public_text) diff --git a/scripts/generate_wrappers.py b/scripts/generate_wrappers.py index 4b98dfb9b..459abda47 100644 --- a/scripts/generate_wrappers.py +++ b/scripts/generate_wrappers.py @@ -969,12 +969,7 @@ def _overload_order_key(node): namespace infini::ops {{ std::size_t DefaultImplementationIndexFor{symbol_name}(Device::Type dev_type) {{ - auto indices = generated_dispatch::ActiveImplementationIndicesFor{symbol_name}(dev_type); - if (indices.empty()) {{ - throw py::value_error("No active implementation for {symbol_name} on device " + - std::string{{Device::StringFromType(dev_type)}}); - }} - return indices.front(); + return {op_type}::DefaultConfig(dev_type).implementation_index(); }} void Bind{symbol_name}(py::module& m) {{ @@ -1331,7 +1326,17 @@ def _append_optional_params(prefix, params): for node in operator.constructors: params = _generate_params(node) - args = _generate_arguments(node) + constructor_args = [ + arg.spelling for arg in node.get_arguments() if arg.spelling != "stream" + ] + # Operator::Make forwards its trailing arguments. Named by-value + # dispatch parameters are mutable lvalues, while the generated public + # instantiations use const T& for every trailing argument. Canonicalize + # this call so Python bindings request the exported specialization. + args = ", ".join( + [constructor_args[0]] + + [f"std::as_const({name})" for name in constructor_args[1:]] + ) make_params = _append_optional_params("const Config& config", params) if make_params in emitted_make_params: @@ -1423,6 +1428,8 @@ def _generate_generated_dispatch_source(impl_paths, definitions): return f"""#include "generated_dispatch.h" #include "host_range_profiler.h" +#include + // clang-format off {impl_includes} // clang-format on @@ -1615,6 +1622,13 @@ def _append_unique(declaration, definition): f"Operator<{op_type}>::active_implementation_indices(Device::Type);", ) + _append_unique( + f"extern template Config " + f"Operator<{op_type}>::DefaultConfig(Device::Type);", + f"template Config " + f"Operator<{op_type}>::DefaultConfig(Device::Type);", + ) + _append_unique( f"extern template std::size_t " f"Operator<{op_type}>::DefaultImplementationIndex(Device::Type);", @@ -1646,6 +1660,13 @@ def _append_unique(declaration, definition): make_params, ) make_instantiation = f"Operator<{op_type}>::Make<{make_template_arguments}>({make_function_params})" + can_make_template_arguments = ", ".join( + _normalized_type(arg) for arg in rest_args + ) + can_make_instantiation = ( + f"Operator<{op_type}>::CanMake<{can_make_template_arguments}>" + f"({make_function_params})" + ) operator_instantiation = ( f"Operator<{op_type}>::operator()<{template_arguments}>({params}) const" @@ -1659,6 +1680,11 @@ def _append_unique(declaration, definition): f"extern template std::unique_ptr> {make_instantiation};", f"template std::unique_ptr> {make_instantiation};", ) + if operator.name in ("flash_attn_with_kvcache", "reshape_and_cache_flash"): + _append_unique( + f"extern template bool {can_make_instantiation};", + f"template bool {can_make_instantiation};", + ) _append_unique( f"extern template void {operator_instantiation};", f"template void {operator_instantiation};", @@ -1909,6 +1935,7 @@ def _get_all_ops( with_ninetoothed=False, with_linked=False, with_triton=False, + with_native_flash=False, ): scan_dirs = set(devices) @@ -1939,6 +1966,17 @@ def _get_all_ops( impl_roots.append(_GENERATION_DIR) impl_headers_by_operator = _index_impl_headers(impl_roots, scan_dirs) + if not with_native_flash: + # These optional providers require their exact pinned kernel build. + # Never emit a specialization merely because a source header exists. + for name, headers in impl_headers_by_operator.items(): + impl_headers_by_operator[name] = [ + header for header in headers + if header.path.as_posix() not in { + "src/native/cuda/nvidia/ops/flash_attn_with_kvcache/kernel.h", + "src/native/cuda/nvidia/ops/flash_attn_varlen_func/kernel.h", + } + ] for base_dir in base_dirs: for file_path in base_dir.iterdir(): @@ -2134,6 +2172,8 @@ def _dispatch_gen_batch_size(): action="store_true", help="Include linked third-party backend implementations.", ) + parser.add_argument("--with-native-flash", action="store_true", + help="Include the pinned native inference Attention provider.") parser.add_argument( "--ops", nargs="+", @@ -2170,6 +2210,7 @@ def _dispatch_gen_batch_size(): with_ninetoothed=args.with_ninetoothed, with_linked=args.with_linked, with_triton=args.with_triton, + with_native_flash=args.with_native_flash, ) if config_path is not None: diff --git a/scripts/prepare_flash_attention_inference.py b/scripts/prepare_flash_attention_inference.py new file mode 100644 index 000000000..bbe5222e3 --- /dev/null +++ b/scripts/prepare_flash_attention_inference.py @@ -0,0 +1,126 @@ +"""Prepare pinned upstream forward kernels without a Torch runtime dependency. + +This is a build-time, inference-only adaptation, not another Attention algorithm. +Do not link these altered parameter layouts to the original Torch extension. +""" + +import argparse +import hashlib +import json +import pathlib +import re +import tarfile + + +ARCHIVE_SHA256 = "1e71dd64a9e0280e0447b8a0c2541bad4bf6ac65bdeaa2f90e51a9e57de0370d" +ROOT = "flash_attn-2.8.3/" +KERNEL_ROOT = ROOT + "csrc/flash_attn/src/" + + +def replace_once(text, before, after): + if text.count(before) != 1: + raise ValueError(f"pinned FlashAttention source contract changed: {before!r}") + return text.replace(before, after, 1) + + +def write_if_changed(path, content): + # CMake may prepare the same pinned closure again during an incremental + # build. Preserve timestamps so hundreds of CUDA kernels are not rebuilt + # merely because configure ran; changed content still invalidates them. + content = content.encode("utf-8") if isinstance(content, str) else content + if not path.exists() or path.read_bytes() != content: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(content) + + +def adapt(name, text): + if name == "flash.h": + text = replace_once(text, "#include // For at::Generator and at::PhiloxCudaState", "") + text = replace_once(text, " at::PhiloxCudaState philox_args;", " // Inference-only build: dropout and its Torch RNG state are absent.") + elif name == "flash_fwd_kernel.h": + text = replace_once(text, '#include "philox_unpack.cuh" // For at::cuda::philox::unpack', "") + text = replace_once(text, " auto seed_offset = at::cuda::philox::unpack(params.philox_args);", " static_assert(!Is_dropout, \"InfiniOps inference build does not support dropout\");\n const auto seed_offset = std::make_tuple(uint64_t{0}, uint64_t{0});") + elif name == "flash_fwd_launch_template.h": + text = replace_once(text, "#include // For C10_CUDA_CHECK and C10_CUDA_KERNEL_LAUNCH_CHECK", '#include "infini_flash_cuda.h"') + text = text.replace("C10_CUDA_KERNEL_LAUNCH_CHECK", "INFINI_FLASH_KERNEL_CHECK") + text = text.replace("C10_CUDA_CHECK", "INFINI_FLASH_CUDA_CHECK") + elif name == "hardware_info.h": + text = replace_once(text, "#include ", '#include \n#include "infini_flash_cuda.h"') + text = replace_once(text, " exit(1);", " INFINI_FLASH_CUDA_CHECK(status_);") + return text + + +def prepare(archive, output): + archive, output = pathlib.Path(archive), pathlib.Path(output) + digest = hashlib.sha256(archive.read_bytes()).hexdigest() + if digest != ARCHIVE_SHA256: + raise ValueError(f"expected pinned flash_attn-2.8.3 archive {ARCHIVE_SHA256}, got {digest}") + with tarfile.open(archive, "r:gz") as source: + members = {member.name: member for member in source.getmembers()} + + def read(name): + member = members.get(name) + if member is None or not member.isfile(): + raise ValueError(f"missing regular pinned source file: {name}") + return source.extractfile(member).read() + + # Only the transitive forward include closure is needed. No backward, + # ATen wrapper, Python extension or Torch allocator is compiled. + pending = ["flash_fwd_launch_template.h"] + headers = {} + while pending: + name = pending.pop() + if name in headers: + continue + if pathlib.PurePosixPath(name).name != name: + raise ValueError(f"unexpected upstream include: {name}") + text = adapt(name, read(KERNEL_ROOT + name).decode()) + if re.search(r"#\s*include\s*[<\"](?:ATen|c10|torch)/|\bat::|\bc10::", text): + raise ValueError(f"unremoved Torch dependency in {name}") + headers[name] = text + for include in re.findall(r'#\s*include\s*"([^"]+)"', text): + if KERNEL_ROOT + include in members: + pending.append(include) + + # Reuse the pinned host scheduling heuristic as well as the kernels; + # do not invent a second split-KV tuning policy in LM. + api = read(ROOT + "csrc/flash_attn/flash_api.cpp").decode() + begin = api.index("inline int num_splits_heuristic(") + end = api.index("\nstd::tuple set_params_splitkv", begin) + headers["num_splits.h"] = ( + "// Copyright (c) 2023, Tri Dao. See accompanying LICENSE.\n" + "#pragma once\n#include \n#include \n#include \n" + "namespace infini_flash {\n" + api[begin:end] + "\n}\n" + ) + + output.mkdir(parents=True, exist_ok=True) + for name, text in headers.items(): + write_if_changed(output / name, text) + for member in source.getmembers(): + prefix = ROOT + "csrc/cutlass/include/" + if member.name.startswith(prefix) and member.isfile(): + relative = pathlib.PurePosixPath(member.name[len(prefix):]) + if relative.is_absolute() or ".." in relative.parts: + raise ValueError("unsafe CUTLASS archive path") + path = output / "cutlass" / pathlib.Path(*relative.parts) + path.parent.mkdir(parents=True, exist_ok=True) + write_if_changed(path, read(member.name)) + for name in ("LICENSE", "AUTHORS"): + write_if_changed(output / name, read(ROOT + name)) + manifest = { + "upstream": "flash-attn", "version": "2.8.3", + "archive_sha256": digest, "namespace": "infini_flash", + "inference_only": True, "dropout": False, + "headers_sha256": {name: hashlib.sha256(text.encode()).hexdigest() + for name, text in sorted(headers.items())}, + } + write_if_changed(output / "source-manifest.json", json.dumps(manifest, indent=2) + "\n") + return manifest + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--archive", required=True) + parser.add_argument("--output", required=True) + args = parser.parse_args() + prepare(args.archive, args.output) diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index ce9f75770..f658d6402 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -95,6 +95,16 @@ if(WITH_CPU) find_package(OpenMP REQUIRED COMPONENTS CXX) target_link_libraries(infiniops PRIVATE OpenMP::OpenMP_CXX) + # Generated Call/dispatch translation units use CUDA when NVIDIA is + # enabled, but also instantiate CPU kernels. OpenMP::OpenMP_CXX only adds + # compile flags to CXX sources; forward them to NVCC's host compiler too. + if(CMAKE_CUDA_COMPILER_ID STREQUAL "NVIDIA") + separate_arguments(_cpu_openmp_flags NATIVE_COMMAND "${OpenMP_CXX_FLAGS}") + foreach(_flag IN LISTS _cpu_openmp_flags) + target_compile_options(infiniops PRIVATE + "$<$:-Xcompiler=${_flag}>") + endforeach() + endif() list(APPEND DEVICE_LIST "cpu") endif() @@ -148,6 +158,9 @@ if(WITH_NVIDIA) ) file(GLOB_RECURSE NVIDIA_SOURCES CONFIGURE_DEPENDS ${NVIDIA_PATTERNS}) + # Optional pinned kernels have their own include closure and CUDA TUs. + list(FILTER NVIDIA_SOURCES EXCLUDE REGEX "/ops/flash_attn_with_kvcache/") + list(FILTER NVIDIA_SOURCES EXCLUDE REGEX "/ops/flash_attn_varlen_func/") # The CUTLASS INT8 Tensor Core kernel needs SM75 device compilation. Scope # the extra code objects to this provider source so other CUDA operators @@ -159,6 +172,11 @@ if(WITH_NVIDIA) "--generate-code=arch=compute_75,code=sm_75;--generate-code=arch=compute_75,code=compute_75") enable_language(CUDA) + if(BUILD_NVIDIA_FLASH_INFERENCE) + include("${PROJECT_SOURCE_DIR}/cmake/NativeFlashAttention.cmake") + infiniops_add_native_flash(infiniops_native_flash) + target_link_libraries(infiniops PRIVATE infiniops_native_flash) + endif() target_compile_definitions(infiniops PUBLIC WITH_NVIDIA=1) target_sources(infiniops PRIVATE ${NVIDIA_SOURCES}) @@ -1220,6 +1238,9 @@ if(GENERATE_OPERATOR_CALL_INSTANTIATIONS OR GENERATE_PYTHON_BINDINGS) if(WITH_LINKED) list(APPEND GENERATOR_ARGS --with-linked) endif() + if(BUILD_NVIDIA_FLASH_INFERENCE) + list(APPEND GENERATOR_ARGS --with-native-flash) + endif() if(WITH_NINETOOTHED) list(APPEND GENERATOR_ARGS --with-ninetoothed) endif() diff --git a/src/base/copy.h b/src/base/copy.h index fe6b33efe..6b13eb65b 100644 --- a/src/base/copy.h +++ b/src/base/copy.h @@ -33,6 +33,11 @@ class Copy : public Operator { Tensor out) const = 0; protected: + static bool SameView(const Tensor& src, const Tensor& out) { + return src.data() == out.data() && src.shape() == out.shape() && + src.strides() == out.strides(); + } + static Tensor::Strides BroadcastStrides(const Tensor src, const Tensor out) { assert(src.ndim() <= out.ndim() && "`Copy` input rank must not exceed output rank"); diff --git a/src/base/embedding.h b/src/base/embedding.h index 8d10243bb..45a9188f2 100644 --- a/src/base/embedding.h +++ b/src/base/embedding.h @@ -46,7 +46,9 @@ class Embedding : public Operator { assert(out.size(-1) == embedding_dim_ && "`Embedding` output last dim must equal `weight` embedding dim"); assert((input_dtype_ == DataType::kInt32 || - input_dtype_ == DataType::kInt64) && + input_dtype_ == DataType::kInt64 || + (input.device().type() == Device::Type::kCpu && + (input_dtype_ == DataType::kUInt32 || input_dtype_ == DataType::kUInt64))) && "`Embedding` supports int32 and int64 indices only"); assert((weight_dtype_ == DataType::kFloat32 || weight_dtype_ == DataType::kFloat16 || diff --git a/src/base/fill.h b/src/base/fill.h index 0545513ef..6c2bee7a1 100644 --- a/src/base/fill.h +++ b/src/base/fill.h @@ -1,12 +1,31 @@ #ifndef INFINI_OPS_BASE_FILL_H_ #define INFINI_OPS_BASE_FILL_H_ + +#include "data_type.h" #include "operator.h" namespace infini::ops { class Fill : public Operator { public: + // Output metadata contract; require storage at Call/RT memset, not planning. + // This does not select a provider or qualify its scalar/capture domain. + static bool valid_output(const Tensor& out, bool require_storage = true) { + const auto dtype = static_cast(out.dtype()); + if (dtype < 0 || dtype > static_cast(DataType::kFloat64) || + out.shape().size() != out.strides().size()) return false; + for (std::size_t i = 0; i < out.ndim(); ++i) + if (out.size(i) > 1 && out.stride(i) == 0) return false; + return !require_storage || !out.numel() || out.data(); + } + + static bool supports(const Tensor& input, double, const Tensor& out) { + return input.shape().size() == input.strides().size() && + input.shape() == out.shape() && input.dtype() == out.dtype() && + input.device() == out.device() && valid_output(out, false); + } + Fill(const Tensor input, const double value, Tensor out) : input_shape_{input.shape()}, input_strides_{input.strides()}, @@ -57,6 +76,7 @@ class Fill : public Operator { DataType out_type_; int device_index_{0}; + }; } // namespace infini::ops diff --git a/src/base/flash_attn_varlen_func.h b/src/base/flash_attn_varlen_func.h index 35ebb5aad..c4d2da8fa 100644 --- a/src/base/flash_attn_varlen_func.h +++ b/src/base/flash_attn_varlen_func.h @@ -14,6 +14,77 @@ namespace infini::ops { // `flash_attn_varlen_func` public interface. class FlashAttnVarlenFunc : public Operator { public: + static constexpr bool kAutomaticDefault = true; + + // Common metadata, not native layout/capacity or capture qualification. + // Planning may use null data; invocation checks actual nonempty storage. + static bool supports( + const Tensor q, const Tensor k, const Tensor v, const Tensor cu_seqlens_q, + const Tensor cu_seqlens_k, const std::optional alibi_slopes, + const std::optional block_table, const int64_t max_seqlen_q, + const int64_t max_seqlen_k, const double dropout_p, + const std::optional softmax_scale, const bool causal, + const std::vector window_size, const double softcap, + const bool deterministic, const bool return_attn_probs, Tensor out, + std::optional softmax_lse, std::optional s_dmask) { + if (q.ndim() != 3 || k.ndim() != (block_table ? 4 : 3) || + v.ndim() != k.ndim() || k.shape() != v.shape() || out.shape() != q.shape() || + (q.dtype() != DataType::kFloat16 && q.dtype() != DataType::kBFloat16) || + k.dtype() != q.dtype() || v.dtype() != q.dtype() || out.dtype() != q.dtype() || + cu_seqlens_q.ndim() != 1 || cu_seqlens_k.ndim() != 1 || + cu_seqlens_q.shape() != cu_seqlens_k.shape() || cu_seqlens_q.size(0) < 2 || + cu_seqlens_q.dtype() != DataType::kInt32 || + cu_seqlens_k.dtype() != DataType::kInt32 || + (alibi_slopes && ((alibi_slopes->ndim() != 1 && alibi_slopes->ndim() != 2) || + alibi_slopes->dtype() != DataType::kFloat32)) || + (block_table && (block_table->ndim() != 2 || block_table->dtype() != DataType::kInt32)) || + softmax_lse.has_value() != s_dmask.has_value() || + return_attn_probs != softmax_lse.has_value() || + (softmax_lse && (softmax_lse->dtype() != DataType::kFloat32 || s_dmask->dtype() != q.dtype())) || + max_seqlen_q <= 0 || max_seqlen_k <= 0 || + window_size.size() != 2 || window_size[0] < -1 || window_size[1] < -1 || + dropout_p != 0.0 || softcap != 0.0 || deterministic) return false; + const auto metadata = [&](const Tensor& tensor) { + return tensor.device().type() == q.device().type() && + tensor.device().index() == q.device().index(); + }; + for (const auto& tensor : {q, k, v, cu_seqlens_q, cu_seqlens_k, out}) + if (!metadata(tensor)) return false; + for (const auto& tensor : {alibi_slopes, block_table, softmax_lse, s_dmask}) + if (tensor && !metadata(*tensor)) return false; + // Span validation above establishes stride cardinality before axis access. + if (q.size(1) == 0 || k.size(-2) == 0 || q.size(1) % k.size(-2) != 0 || + q.size(2) == 0 || q.size(2) > 256 || q.size(2) % 8 != 0 || + q.size(2) != k.size(-1) || q.stride(-1) != 1 || k.stride(-1) != 1 || + v.stride(-1) != 1 || out.stride(-1) != 1 || + !cu_seqlens_q.IsContiguous() || !cu_seqlens_k.IsContiguous()) return false; + if (softmax_lse && + (softmax_lse->shape() != Tensor::Shape{q.size(1), q.size(0)} || + s_dmask->shape() != Tensor::Shape{0})) return false; + if (block_table && + (block_table->size(0) != cu_seqlens_q.size(0) - 1 || + !block_table->IsContiguous() || k.size(1) % 256 != 0)) return false; + if (alibi_slopes && + (!alibi_slopes->IsContiguous() || + !((alibi_slopes->ndim() == 1 && alibi_slopes->size(0) == q.size(1)) || + (alibi_slopes->ndim() == 2 && + alibi_slopes->size(0) == cu_seqlens_q.size(0) - 1 && + alibi_slopes->size(1) == q.size(1))))) return false; + (void)softmax_scale; + (void)causal; + return true; + } + + static bool supports(const Tensor q, const Tensor k, const Tensor v, + const Tensor cu_seqlens_q, const Tensor cu_seqlens_k, + const int64_t max_seqlen_q, const int64_t max_seqlen_k, + Tensor out) { + return supports(q, k, v, cu_seqlens_q, cu_seqlens_k, std::nullopt, + std::nullopt, max_seqlen_q, max_seqlen_k, 0.0, std::nullopt, + false, {-1, -1}, 0.0, false, false, out, std::nullopt, + std::nullopt); + } + FlashAttnVarlenFunc(const Tensor q, const Tensor k, const Tensor v, const Tensor cu_seqlens_q, const Tensor cu_seqlens_k, const int64_t max_seqlen_q, const int64_t max_seqlen_k, diff --git a/src/base/flash_attn_with_kvcache.h b/src/base/flash_attn_with_kvcache.h index ecd2c14c9..711101422 100644 --- a/src/base/flash_attn_with_kvcache.h +++ b/src/base/flash_attn_with_kvcache.h @@ -4,16 +4,145 @@ #include #include #include +#include +#include #include #include "operator.h" namespace infini::ops { +class FlashAttnWithKvcacheConfig : public Config { + public: + explicit FlashAttnWithKvcacheConfig(int64_t logical_kv_upper_bound = 0) + : logical_kv_upper_bound_{logical_kv_upper_bound} { + if (logical_kv_upper_bound < 0) { + throw std::invalid_argument("negative logical KV upper bound"); + } + } + + std::unique_ptr Clone() const override { + return std::make_unique(*this); + } + + bool requires_config_aware_provider() const override { + return logical_kv_upper_bound_ != 0; + } + + int64_t logical_kv_upper_bound() const { return logical_kv_upper_bound_; } + + static int64_t BoundFrom(const Config& config) { + const auto* typed = dynamic_cast(&config); + return typed ? typed->logical_kv_upper_bound() : 0; + } + + private: + int64_t logical_kv_upper_bound_; +}; + // Inference attention with an optional in-place KV-cache update, aligned with // Dao-AILab FlashAttention's `flash_attn_with_kvcache` public interface. class FlashAttnWithKvcache : public Operator { public: + static constexpr bool kAutomaticDefault = true; + + // Common metadata contract only; each provider retains its own qualification. + // Null storage is legal for planning, but never for a nonempty invocation. + template + static bool supports( + const Tensor q, Tensor k_cache, Tensor v_cache, + const std::optional k, const std::optional v, + const std::optional rotary_cos, + const std::optional rotary_sin, const Lengths cache_seqlens, + const std::optional cache_batch_idx, + const std::optional cache_leftpad, + const std::optional block_table, + const std::optional alibi_slopes, + const std::optional softmax_scale, const bool causal, + const std::vector window_size, const double softcap, + const bool rotary_interleaved, const int64_t num_splits, + const bool return_softmax_lse, Tensor out, + std::optional softmax_lse) { + const auto lengths = [&] { + if constexpr (std::is_convertible_v) + return std::optional{}; + else + return std::optional{cache_seqlens}; + }(); + if constexpr (std::is_convertible_v) + if (static_cast(cache_seqlens) < 0) return false; + const auto metadata = [&](const Tensor& tensor) { + return tensor.device().type() == q.device().type() && + tensor.device().index() == q.device().index(); + }; + for (const auto& tensor : {q, k_cache, v_cache, out}) + if (!metadata(tensor)) return false; + for (const auto& tensor : {k, v, rotary_cos, rotary_sin, lengths, + cache_batch_idx, cache_leftpad, block_table, + alibi_slopes, softmax_lse}) + if (tensor && !metadata(*tensor)) return false; + if (q.ndim() != 4 || k_cache.ndim() != 4 || v_cache.ndim() != 4 || + out.shape() != q.shape() || k_cache.shape() != v_cache.shape() || + (q.dtype() != DataType::kFloat16 && q.dtype() != DataType::kBFloat16) || + k_cache.dtype() != q.dtype() || v_cache.dtype() != q.dtype() || + out.dtype() != q.dtype() || q.size(0) == 0 || q.size(1) == 0 || + q.size(2) == 0 || k_cache.size(0) == 0 || k_cache.size(1) == 0 || + k_cache.size(2) == 0 || q.size(2) % k_cache.size(2) != 0 || + q.size(3) == 0 || q.size(3) > 256 || q.size(3) != k_cache.size(3) || + q.stride(-1) != 1 || k_cache.stride(-1) != 1 || + v_cache.stride(-1) != 1 || out.stride(-1) != 1 || + k.has_value() != v.has_value() || + rotary_cos.has_value() != rotary_sin.has_value() || + return_softmax_lse != softmax_lse.has_value() || + window_size.size() != 2 || window_size[0] < -1 || window_size[1] < -1 || + !(softcap >= 0.0) || num_splits < 0) return false; + if (k && (k->ndim() != 4 || k->shape() != v->shape() || + k->size(0) != q.size(0) || k->size(2) != k_cache.size(2) || + k->size(3) != q.size(3) || k->dtype() != q.dtype() || + v->dtype() != q.dtype() || k->stride(-1) != 1 || + v->stride(-1) != 1)) return false; + if (rotary_cos && + (!k || rotary_cos->ndim() != 2 || + rotary_cos->shape() != rotary_sin->shape() || + rotary_cos->dtype() != q.dtype() || rotary_sin->dtype() != q.dtype() || + rotary_cos->size(1) == 0 || rotary_cos->size(1) > q.size(3) / 2 || + rotary_cos->size(1) % 8 != 0)) return false; + const auto index_vector = [&](const std::optional& tensor) { + return !tensor || (tensor->ndim() == 1 && tensor->size(0) == q.size(0) && + tensor->dtype() == DataType::kInt32 && + tensor->IsContiguous()); + }; + if (!index_vector(lengths) || !index_vector(cache_batch_idx) || + !index_vector(cache_leftpad)) return false; + if (block_table) { + if (block_table->ndim() != 2 || block_table->size(0) != q.size(0) || + block_table->dtype() != DataType::kInt32 || + !block_table->IsContiguous() || k_cache.size(1) % 256 != 0) return false; + } else if (!cache_batch_idx && k_cache.size(0) < q.size(0)) return false; + if (alibi_slopes && + (alibi_slopes->dtype() != DataType::kFloat32 || + !alibi_slopes->IsContiguous() || + !((alibi_slopes->ndim() == 1 && alibi_slopes->size(0) == q.size(2)) || + (alibi_slopes->ndim() == 2 && alibi_slopes->size(0) == q.size(0) && + alibi_slopes->size(1) == q.size(2))))) return false; + if (softmax_lse && + (softmax_lse->dtype() != DataType::kFloat32 || + softmax_lse->shape() != Tensor::Shape{q.size(0), q.size(2), q.size(1)})) + return false; + (void)softmax_scale; + (void)causal; + (void)rotary_interleaved; + return true; + } + + static bool supports(const Tensor q, Tensor k_cache, Tensor v_cache, Tensor out) { + return supports(q, k_cache, v_cache, std::nullopt, std::nullopt, + std::nullopt, std::nullopt, std::optional{}, + std::nullopt, std::nullopt, std::nullopt, std::nullopt, + std::nullopt, false, {-1, -1}, 0.0, true, 0, false, out, + std::nullopt); + } + FlashAttnWithKvcache(const Tensor q, Tensor k_cache, Tensor v_cache, Tensor out) : FlashAttnWithKvcache{q, @@ -467,6 +596,19 @@ class FlashAttnWithKvcache : public Operator { int device_index_{0}; }; +template <> +struct CacheKeyBuilder { + template + detail::CacheKey operator()(const Config& config, const Args&... args) const { + const auto bound = FlashAttnWithKvcacheConfig::BoundFrom(config); + if (bound == 0) { + return detail::CacheKey::Build(config.implementation_index(), args...); + } + return detail::CacheKey::Build(config.implementation_index(), bound, + args...); + } +}; + } // namespace infini::ops #endif // INFINI_OPS_BASE_FLASH_ATTN_WITH_KVCACHE_H_ diff --git a/src/base/gemm.h b/src/base/gemm.h index fe6df2655..58fa25151 100644 --- a/src/base/gemm.h +++ b/src/base/gemm.h @@ -4,13 +4,102 @@ #include #include #include +#include +#include +#include #include "operator.h" namespace infini::ops { +class GemmConfig : public Config { + public: + enum class Fp32Precision { DefaultFast, FullFp32 }; + + explicit GemmConfig(Fp32Precision precision = Fp32Precision::DefaultFast) + : fp32_precision_{precision} { + if (!IsValid(precision)) { + throw std::invalid_argument("invalid GEMM FP32 precision mode"); + } + } + + std::unique_ptr Clone() const override { + return std::make_unique(*this); + } + + Fp32Precision fp32_precision() const { return fp32_precision_; } + + static Fp32Precision PrecisionFrom(const Config& config) { + const auto* gemm_config = dynamic_cast(&config); + const auto precision = gemm_config ? gemm_config->fp32_precision() + : Fp32Precision::DefaultFast; + if (!IsValid(precision)) { + throw std::invalid_argument("invalid GEMM FP32 precision mode"); + } + return precision; + } + + static bool SupportsFp32Precision(Device::Type device, + std::size_t implementation_index, + Fp32Precision precision) { + if (!IsValid(precision)) return false; + return precision == Fp32Precision::DefaultFast || + (device == Device::Type::kNvidia && + (implementation_index == 0 || implementation_index == 1)); + } + + static void RequireSupported(const Config& config, Device::Type device) { + if (!SupportsFp32Precision(device, config.implementation_index(), + PrecisionFrom(config))) { + throw std::invalid_argument( + "FullFp32 GEMM requires NVIDIA implementation 0 or 1"); + } + } + + static bool SupportsMetadata(const Config& config, const Tensor& a, + const Tensor& b, const Tensor& y) { + const auto precision = PrecisionFrom(config); + if (precision == Fp32Precision::DefaultFast) return true; + return SupportsFp32Precision(a.device().type(), + config.implementation_index(), precision) && + a.dtype() == DataType::kFloat32 && + b.dtype() == DataType::kFloat32 && + y.dtype() == DataType::kFloat32; + } + + static void RequireSupported(const Config& config, const Tensor& a, + const Tensor& b, const Tensor& y) { + RequireSupported(config, a.device().type()); + if (!SupportsMetadata(config, a, b, y)) { + throw std::invalid_argument( + "FullFp32 GEMM requires float32 A, B, and output"); + } + } + + private: + static bool IsValid(Fp32Precision precision) { + return precision == Fp32Precision::DefaultFast || + precision == Fp32Precision::FullFp32; + } + + Fp32Precision fp32_precision_; +}; + class Gemm : public Operator { public: + using Operator::Make; + + template + static std::unique_ptr> Make(const Config& config, + const Tensor a, const Tensor b, + Tail&&... tail) { + static_assert(sizeof...(Tail) == 1 || sizeof...(Tail) == 6, + "GEMM requires short or full argument signature"); + const auto& y = std::get(std::forward_as_tuple(tail...)); + GemmConfig::RequireSupported(config, a, b, y); + return Operator::Make(config, a, b, std::forward(tail)...); + } + Gemm(const Tensor a, const Tensor b, const std::optional c, std::optional alpha, std::optional beta, std::optional trans_a, std::optional trans_b, Tensor y) @@ -66,6 +155,20 @@ class Gemm : public Operator { } protected: + GemmConfig::Fp32Precision fp32_precision(const Tensor& a, const Tensor& b, + const Tensor& y) const { + const auto precision = config_ptr_ + ? GemmConfig::PrecisionFrom(*config_ptr_) + : GemmConfig::Fp32Precision::DefaultFast; + if (precision == GemmConfig::Fp32Precision::FullFp32 && + (a.dtype() != DataType::kFloat32 || b.dtype() != DataType::kFloat32 || + y.dtype() != DataType::kFloat32)) { + throw std::invalid_argument( + "FullFp32 GEMM requires float32 A, B, and output"); + } + return precision; + } + static float EffectiveBeta(const std::optional& c, std::optional beta) { static_cast(beta); @@ -114,6 +217,21 @@ class Gemm : public Operator { Tensor::Stride batch_stride_y_{0}; }; +template <> +struct CacheKeyBuilder { + template + detail::CacheKey operator()(const Config& config, const Tensor a, + const Tensor b, const Tail&... tail) const { + static_assert(sizeof...(Tail) == 1 || sizeof...(Tail) == 6, + "GEMM requires short or full argument signature"); + const auto& y = std::get(std::forward_as_tuple(tail...)); + GemmConfig::RequireSupported(config, a, b, y); + return detail::CacheKey::Build( + config.implementation_index(), + static_cast(GemmConfig::PrecisionFrom(config)), a, b, tail...); + } +}; + } // namespace infini::ops #endif diff --git a/src/base/reshape_and_cache_flash.h b/src/base/reshape_and_cache_flash.h index a11c997b6..53762e98e 100644 --- a/src/base/reshape_and_cache_flash.h +++ b/src/base/reshape_and_cache_flash.h @@ -14,6 +14,27 @@ namespace infini::ops { // Aligned with vLLM `_custom_ops.reshape_and_cache_flash`. class ReshapeAndCacheFlash : public Operator { public: + static bool supports(const Tensor& key, const Tensor& value, + const Tensor& slots, const Tensor& k_scale, + const Tensor& v_scale, const std::string& mode, + const Tensor& key_cache, const Tensor& value_cache) { + if (key.ndim() != 3 || value.ndim() != 3 || slots.ndim() != 1 || + key_cache.ndim() != 4 || value_cache.ndim() != 4) return false; + return key.shape() == value.shape() && key.dtype() == value.dtype() && + (key.dtype() == DataType::kFloat16 || key.dtype() == DataType::kBFloat16 || + key.dtype() == DataType::kFloat32) && slots.dtype() == DataType::kInt64 && + slots.stride(0) == 1 && slots.size(0) <= key.size(0) && key.size(1) > 0 && + key.size(2) > 0 && key_cache.shape() == value_cache.shape() && + key_cache.size(1) > 0 && key_cache.size(2) == key.size(1) && + key_cache.size(3) == key.size(2) && key.stride(2) == 1 && + value.stride(2) == 1 && key_cache.stride(3) == 1 && value_cache.stride(3) == 1 && + mode == "auto" && key_cache.dtype() == key.dtype() && value_cache.dtype() == key.dtype() && + !key_cache.HasBroadcastDim() && !value_cache.HasBroadcastDim() && + k_scale.shape() == v_scale.shape() && + (k_scale.numel() == 1 || k_scale.numel() == key.size(1)) && + k_scale.dtype() == DataType::kFloat32 && v_scale.dtype() == DataType::kFloat32; + } + ReshapeAndCacheFlash(const Tensor key, const Tensor value, const Tensor slot_mapping, const Tensor k_scale, const Tensor v_scale, const std::string kv_cache_dtype, diff --git a/src/config.h b/src/config.h index ac371e13e..40e788d55 100644 --- a/src/config.h +++ b/src/config.h @@ -2,6 +2,7 @@ #define INFINI_OPS_CONFIG_H_ #include +#include #include #include @@ -17,6 +18,10 @@ class Config { return std::make_unique(*this); } + // A typed configuration may require a provider that consumes its metadata + // during construction. Legacy configurations and providers are unchanged. + virtual bool requires_config_aware_provider() const { return false; } + std::size_t implementation_index() const { return implementation_index_.value_or(0); } @@ -29,6 +34,16 @@ class Config { return !implementation_index_.has_value(); } + // Opt-in metadata-based selection. Explicit indices retain their meaning. + static constexpr std::size_t kAutomaticImplementation = + std::numeric_limits::max(); + + static Config Automatic() { + Config config; + config.set_implementation_index(kAutomaticImplementation); + return config; + } + private: std::optional implementation_index_; }; diff --git a/src/dispatcher.h b/src/dispatcher.h index 1d7cbdffc..dc816fb0f 100644 --- a/src/dispatcher.h +++ b/src/dispatcher.h @@ -289,7 +289,8 @@ template auto DispatchFuncListAliasImpl(ValueType value, Functor&& func, std::string_view context_str, List, Args&&... args) { - return DispatchFunc>(items)...>( + return DispatchFunc, + static_cast>(items)...>( value, std::forward(func), context_str, std::forward(args)...); } diff --git a/src/handle.h b/src/handle.h index 27211f0ec..68f32c9d6 100644 --- a/src/handle.h +++ b/src/handle.h @@ -10,12 +10,25 @@ namespace infini::ops { class Handle { public: + // The shared pointer owns the storage and points at its device-visible + // address. Its deleter must defer reuse until work already enqueued on this + // Handle's stream has finished (or use a stream-ordered allocator). + struct WorkspaceLease { + std::shared_ptr data; + std::size_t size_in_bytes{0}; + }; + virtual ~Handle() = default; virtual std::unique_ptr Clone() const { return std::make_unique(*this); } + // Called by Operator::Call only after the selected/cached instance reports + // its required size. Make and direct instance invocation do not acquire a + // transient lease: graph owners must retain their own stable workspace. + virtual WorkspaceLease acquire_workspace(std::size_t) const { return {}; } + void* stream() const { return stream_; } void* workspace() const { return workspace_; } diff --git a/src/native/cuda/nvidia/blas_utils.h b/src/native/cuda/nvidia/blas_utils.h index 45644ae38..ab6f3a65e 100644 --- a/src/native/cuda/nvidia/blas_utils.h +++ b/src/native/cuda/nvidia/blas_utils.h @@ -6,6 +6,7 @@ // clang-format on #include "data_type.h" +#include "base/gemm.h" #include "native/cuda/blas_utils.h" namespace infini::ops { @@ -23,6 +24,13 @@ struct BlasUtils { return CUBLAS_COMPUTE_32F; return CUBLAS_COMPUTE_32F_FAST_TF32; } + + static auto GetComputeType(DataType dtype, + GemmConfig::Fp32Precision precision) { + if (precision == GemmConfig::Fp32Precision::FullFp32) + return CUBLAS_COMPUTE_32F; + return GetComputeType(dtype); + } }; } // namespace infini::ops diff --git a/src/native/cuda/nvidia/ops/add/kernel.h b/src/native/cuda/nvidia/ops/add/kernel.h index 373f5a775..c1ac2f141 100644 --- a/src/native/cuda/nvidia/ops/add/kernel.h +++ b/src/native/cuda/nvidia/ops/add/kernel.h @@ -14,6 +14,12 @@ class Operator : public CudaAdd> { public: using CudaAdd>::CudaAdd; + bool capture_safe() const override { + return out_type_ == DataType::kFloat16 || + out_type_ == DataType::kBFloat16 || + out_type_ == DataType::kFloat32 || + out_type_ == DataType::kInt32; + } }; } // namespace infini::ops diff --git a/src/native/cuda/nvidia/ops/argmax/kernel.cuh b/src/native/cuda/nvidia/ops/argmax/kernel.cuh new file mode 100644 index 000000000..df61797d2 --- /dev/null +++ b/src/native/cuda/nvidia/ops/argmax/kernel.cuh @@ -0,0 +1,124 @@ +#ifndef INFINI_OPS_NVIDIA_ARGMAX_KERNEL_CUH_ +#define INFINI_OPS_NVIDIA_ARGMAX_KERNEL_CUH_ + +#include +#include +#include +#include + +#include +#include + +namespace infini::ops::nvidia_argmax_detail { + +constexpr int kThreads = 256; +constexpr int kChunkWidth = 2048; + +template +__device__ float ToFloat(T value) { return static_cast(value); } + +template <> +__device__ inline float ToFloat(half value) { return __half2float(value); } + +template <> +__device__ inline float ToFloat<__nv_bfloat16>(__nv_bfloat16 value) { + return __bfloat162float(value); +} + +struct Candidate { + float value; + int index; +}; + +static_assert(sizeof(Candidate) == 8); + +// The frozen PyTorch Argmax contract prefers NaN over finite values and the +// first index among equal maxima (including multiple NaNs). +__device__ inline Candidate Better(Candidate a, Candidate b) { + if (b.index == INT_MAX) return a; + if (a.index == INT_MAX) return b; + const bool a_nan = isnan(a.value); + const bool b_nan = isnan(b.value); + if (a_nan != b_nan) return b_nan ? b : a; + if (a_nan || a.value == b.value) return b.index < a.index ? b : a; + return b.value > a.value ? b : a; +} + +template +__global__ void RowsKernel(const T* input, std::int64_t* output, int width) { + const int row = blockIdx.x; + const T* values = input + static_cast(row) * width; + Candidate best{-CUDART_INF_F, INT_MAX}; + for (std::int64_t column = threadIdx.x; column < width; + column += blockDim.x) { + best = Better(best, {ToFloat(values[column]), static_cast(column)}); + } + __shared__ Candidate partial[kThreads]; + partial[threadIdx.x] = best; + __syncthreads(); + for (int offset = kThreads / 2; offset; offset /= 2) { + if (threadIdx.x < offset) { + partial[threadIdx.x] = + Better(partial[threadIdx.x], partial[threadIdx.x + offset]); + } + __syncthreads(); + } + if (threadIdx.x == 0) output[row] = partial[0].index; +} + +template +__global__ void PartialsKernel(const T* input, Candidate* partials, + int rows, int width, int chunks) { + const std::size_t tasks = static_cast(rows) * chunks; + __shared__ Candidate reduced[kThreads]; + for (std::size_t task = blockIdx.x; task < tasks; task += gridDim.x) { + const int row = static_cast(task / chunks); + const int chunk = static_cast(task % chunks); + const int begin = chunk * kChunkWidth; + const int end = begin + min(kChunkWidth, width - begin); + const T* values = input + static_cast(row) * width; + Candidate best{-CUDART_INF_F, INT_MAX}; + for (std::int64_t column = std::int64_t{begin} + threadIdx.x; column < end; + column += blockDim.x) { + best = Better(best, {ToFloat(values[column]), static_cast(column)}); + } + reduced[threadIdx.x] = best; + __syncthreads(); + for (int offset = kThreads / 2; offset; offset /= 2) { + if (threadIdx.x < offset) { + reduced[threadIdx.x] = + Better(reduced[threadIdx.x], reduced[threadIdx.x + offset]); + } + __syncthreads(); + } + if (threadIdx.x == 0) partials[task] = reduced[0]; + __syncthreads(); + } +} + +static __global__ void FinishKernel(const Candidate* partials, + std::int64_t* output, int rows, int chunks) { + __shared__ Candidate reduced[kThreads]; + for (std::int64_t row = blockIdx.x; row < rows; row += gridDim.x) { + const Candidate* values = partials + static_cast(row) * chunks; + Candidate best{-CUDART_INF_F, INT_MAX}; + for (int chunk = threadIdx.x; chunk < chunks; chunk += blockDim.x) { + best = Better(best, values[chunk]); + } + reduced[threadIdx.x] = best; + __syncthreads(); + for (int offset = kThreads / 2; offset; offset /= 2) { + if (threadIdx.x < offset) { + reduced[threadIdx.x] = + Better(reduced[threadIdx.x], reduced[threadIdx.x + offset]); + } + __syncthreads(); + } + if (threadIdx.x == 0) output[row] = reduced[0].index; + __syncthreads(); + } +} + +} // namespace infini::ops::nvidia_argmax_detail + +#endif // INFINI_OPS_NVIDIA_ARGMAX_KERNEL_CUH_ diff --git a/src/native/cuda/nvidia/ops/argmax/kernel.h b/src/native/cuda/nvidia/ops/argmax/kernel.h new file mode 100644 index 000000000..1d972a2e2 --- /dev/null +++ b/src/native/cuda/nvidia/ops/argmax/kernel.h @@ -0,0 +1,179 @@ +#ifndef INFINI_OPS_NVIDIA_ARGMAX_KERNEL_H_ +#define INFINI_OPS_NVIDIA_ARGMAX_KERNEL_H_ + +#include + +#include +#include +#include +#include +#include +#include + +#include "base/argmax.h" +#include "native/cuda/nvidia/ops/argmax/kernel.cuh" +#include "native/cuda/nvidia/runtime_.h" + +namespace infini::ops { + +template <> +class Operator : public Argmax { + public: + Operator(const Tensor input, const std::optional dim, + const bool keepdim, Tensor out) + : Argmax(input, dim, keepdim, out) { + Validate(input, dim, keepdim, out); + const auto rows = input.ndim() == 2 ? input.shape()[0] : 1; + const auto width = input.shape().back(); + if (width > nvidia_argmax_detail::kChunkWidth) { + const auto chunks = (width - 1) / nvidia_argmax_detail::kChunkWidth + 1; + if (rows > std::numeric_limits::max() / chunks / + sizeof(nvidia_argmax_detail::Candidate)) { + throw std::invalid_argument("NVIDIA Argmax workspace size overflows"); + } + workspace_bytes_ = rows * chunks * sizeof(nvidia_argmax_detail::Candidate); + } + } + + bool capture_safe() const override { return true; } + + std::size_t workspace_size_in_bytes() const override { return workspace_bytes_; } + + bool requires_workspace() const override { return workspace_bytes_ != 0; } + + void operator()(const Tensor input, const std::optional dim, + const bool keepdim, Tensor out) const override { + if (input.shape() != input_shape_ || input.strides() != input_strides_ || + input.dtype() != input_type_ || out.shape() != out_shape_ || + out.strides() != out_strides_ || out.dtype() != out_type_ || + dim != dim_ || keepdim != keepdim_ || + out.device().index() != device_index_) { + throw std::invalid_argument("NVIDIA Argmax call metadata changed after Make"); + } + Validate(input, dim, keepdim, out); + const auto rows = static_cast(input.ndim() == 2 ? input.shape()[0] : 1); + const auto width = static_cast(input.shape().back()); + const auto stream = static_cast(stream_ ? stream_ : 0); + if (workspace_bytes_) CheckWorkspace(input, out); + switch (input.dtype()) { + case DataType::kFloat16: + Launch(static_cast(input.data()), out, rows, width, stream, + workspace_); + break; + case DataType::kBFloat16: + Launch(static_cast(input.data()), out, rows, + width, stream, workspace_); + break; + case DataType::kFloat32: + Launch(static_cast(input.data()), out, rows, width, stream, + workspace_); + break; + default: + throw std::invalid_argument("NVIDIA Argmax unsupported dtype"); + } + } + + private: + void CheckWorkspace(const Tensor input, const Tensor out) const { + if (!workspace_ || workspace_size_in_bytes_ < workspace_bytes_ || + reinterpret_cast(workspace_) % + alignof(nvidia_argmax_detail::Candidate)) { + throw std::invalid_argument("NVIDIA Argmax workspace is missing or too small"); + } + const auto begin = reinterpret_cast(workspace_); + if (workspace_bytes_ > std::numeric_limits::max() - begin) { + throw std::invalid_argument("NVIDIA Argmax workspace span overflows"); + } + const auto end = begin + workspace_bytes_; + const auto overlaps = [begin, end](const Tensor tensor) { + const auto tensor_begin = reinterpret_cast(tensor.data()); + const auto tensor_end = tensor_begin + tensor.numel() * tensor.element_size(); + return begin < tensor_end && tensor_begin < end; + }; + if (overlaps(input) || overlaps(out)) { + throw std::invalid_argument("NVIDIA Argmax workspace overlaps a tensor"); + } + } + + static void Validate(const Tensor input, const std::optional dim, + bool keepdim, const Tensor out) { + if (input.device().type() != Device::Type::kNvidia || + input.device() != out.device() || + (input.ndim() != 1 && input.ndim() != 2) || keepdim || + (input.ndim() == 1 ? dim.has_value() : dim != 1) || + !input.IsContiguous() || !out.IsContiguous() || + (input.dtype() != DataType::kFloat16 && + input.dtype() != DataType::kBFloat16 && + input.dtype() != DataType::kFloat32) || + out.dtype() != DataType::kInt64) { + throw std::invalid_argument("unsupported NVIDIA Argmax metadata"); + } + const auto rows = input.ndim() == 2 ? input.shape()[0] : 1; + const auto width = input.shape().back(); + if (!rows || !width || rows > static_cast(INT_MAX) || + width > static_cast(INT_MAX) || + (input.ndim() == 1 ? out.ndim() != 0 + : out.ndim() != 1 || out.shape()[0] != rows) || + !input.data() || !out.data()) { + throw std::invalid_argument("NVIDIA Argmax has empty or invalid storage"); + } + if (rows > std::numeric_limits::max() / width || + rows * width > std::numeric_limits::max() / + input.element_size() || + rows > std::numeric_limits::max() / sizeof(std::int64_t)) { + throw std::invalid_argument("NVIDIA Argmax storage size overflows"); + } + const auto input_bytes = input.numel() * input.element_size(); + const auto output_bytes = out.numel() * out.element_size(); + const auto in_begin = reinterpret_cast(input.data()); + const auto out_begin = reinterpret_cast(out.data()); + if (input_bytes > std::numeric_limits::max() - in_begin || + output_bytes > std::numeric_limits::max() - out_begin || + (in_begin < out_begin + output_bytes && + out_begin < in_begin + input_bytes)) { + throw std::invalid_argument("NVIDIA Argmax input/output storage overlaps"); + } + int current_device = -1; + const auto status = cudaGetDevice(¤t_device); + if (status != cudaSuccess) throw std::runtime_error(cudaGetErrorString(status)); + if (current_device != input.device().index()) { + throw std::invalid_argument("NVIDIA Argmax tensor device is not current device"); + } + } + + template + static void Launch(const T* input, Tensor out, unsigned rows, int width, + cudaStream_t stream, void* workspace) { + if (width <= nvidia_argmax_detail::kChunkWidth) { + nvidia_argmax_detail::RowsKernel + <<>>( + input, static_cast(out.data()), width); + } else { + const int chunks = (width - 1) / nvidia_argmax_detail::kChunkWidth + 1; + const auto tasks = static_cast(rows) * chunks; + auto* partials = static_cast(workspace); + const auto partial_grid = static_cast( + std::min(tasks, 65535)); + const auto final_grid = std::min(rows, 65535); + nvidia_argmax_detail::PartialsKernel + <<>>( + input, partials, static_cast(rows), width, chunks); + const auto partial_status = cudaPeekAtLastError(); + if (partial_status != cudaSuccess) { + throw std::runtime_error(cudaGetErrorString(partial_status)); + } + nvidia_argmax_detail::FinishKernel + <<>>( + partials, static_cast(out.data()), + static_cast(rows), chunks); + } + const auto status = cudaPeekAtLastError(); + if (status != cudaSuccess) throw std::runtime_error(cudaGetErrorString(status)); + } + + std::size_t workspace_bytes_{0}; +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_NVIDIA_ARGMAX_KERNEL_H_ diff --git a/src/native/cuda/nvidia/ops/binary_view_capture.h b/src/native/cuda/nvidia/ops/binary_view_capture.h new file mode 100644 index 000000000..4e40741d3 --- /dev/null +++ b/src/native/cuda/nvidia/ops/binary_view_capture.h @@ -0,0 +1,128 @@ +#ifndef INFINI_OPS_NVIDIA_BINARY_VIEW_CAPTURE_H_ +#define INFINI_OPS_NVIDIA_BINARY_VIEW_CAPTURE_H_ + +#include +#include +#include +#include +#include +#include + +#include "tensor.h" + +namespace infini::ops::nvidia_binary_view { + +inline bool SpanBytes(const Tensor tensor, std::size_t* bytes) { + if (tensor.shape().size() != tensor.strides().size() || !tensor.element_size()) + return false; + for (auto dim : tensor.shape()) if (!dim) { *bytes = 0; return true; } + std::size_t last = 0; + for (std::size_t i = 0; i < tensor.ndim(); ++i) { + if (tensor.strides()[i] < 0) return false; + const auto stride = static_cast(tensor.strides()[i]); + const auto extent = tensor.shape()[i] - 1; + if (extent && stride > (std::numeric_limits::max() - last) / extent) + return false; + last += extent * stride; + } + if (last == std::numeric_limits::max() || + last + 1 > std::numeric_limits::max() / tensor.element_size()) + return false; + *bytes = (last + 1) * tensor.element_size(); + return !*bytes || tensor.data(); +} + +inline bool NonOverlappingOutput(const Tensor out) { + if (out.shape().size() != out.strides().size()) return false; + if (out.HasBroadcastDim()) return false; + for (auto dim : out.shape()) if (!dim) return true; + std::vector> dimensions; + for (std::size_t i = 0; i < out.ndim(); ++i) { + if (out.strides()[i] < 0) return false; + if (out.shape()[i] > 1) + dimensions.emplace_back(static_cast(out.strides()[i]), + out.shape()[i]); + } + std::sort(dimensions.begin(), dimensions.end()); + std::size_t span = 1; + for (const auto& [stride, extent] : dimensions) { + if (stride < span || + extent - 1 > (std::numeric_limits::max() - span) / stride) + return false; + span += (extent - 1) * stride; + } + return true; +} + +inline bool Overlaps(const Tensor a, std::size_t a_bytes, + const Tensor b, std::size_t b_bytes) { + if (!a_bytes || !b_bytes) return false; + const auto lhs = reinterpret_cast(a.data()); + const auto rhs = reinterpret_cast(b.data()); + return lhs <= rhs ? rhs - lhs < a_bytes : lhs - rhs < b_bytes; +} + +inline bool PointersSupported(const Tensor input, const Tensor other, Tensor out) { + if (input.device() != other.device() || input.device() != out.device()) + return false; + int current_device = -1; + if (cudaGetDevice(¤t_device) != cudaSuccess || + current_device != input.device().index()) return false; + std::size_t input_bytes, other_bytes, out_bytes; + if (!SpanBytes(input, &input_bytes) || !SpanBytes(other, &other_bytes) || + !SpanBytes(out, &out_bytes)) return false; + const auto safe_alias = [&](const Tensor source, std::size_t bytes) { + return !Overlaps(source, bytes, out, out_bytes) || + (source.data() == out.data() && source.shape() == out.shape() && + source.strides() == out.strides()); + }; + return safe_alias(input, input_bytes) && safe_alias(other, other_bytes); +} + +inline bool Supported(const Tensor input, const Tensor other, Tensor out) { + if (input.shape().size() != input.strides().size() || + other.shape().size() != other.strides().size() || + out.shape().size() != out.strides().size() || + input.device().type() != Device::Type::kNvidia || + input.dtype() != other.dtype() || input.dtype() != out.dtype() || + input.ndim() > out.ndim() || other.ndim() > out.ndim() || + !NonOverlappingOutput(out)) return false; + for (std::size_t i = 0; i < out.ndim(); ++i) { + const auto a = i < out.ndim() - input.ndim() ? 1 : + input.size(i + input.ndim() - out.ndim()); + const auto b = i < out.ndim() - other.ndim() ? 1 : + other.size(i + other.ndim() - out.ndim()); + // Current Add/Mul base constructors do not implement the valid 0/1 + // empty broadcast. Keep this provider fail-closed until that is fixed. + if ((a == 0 && b == 1) || (a == 1 && b == 0) || + (a != b && a != 1 && b != 1) || out.size(i) != std::max(a, b)) + return false; + } + return PointersSupported(input, other, out); +} + +struct Metadata { + Tensor::Shape input_shape, other_shape, out_shape; + Tensor::Strides input_strides, other_strides, out_strides; + DataType dtype; + Device device; + + Metadata(const Tensor input, const Tensor other, const Tensor out) + : input_shape(input.shape()), other_shape(other.shape()), + out_shape(out.shape()), input_strides(input.strides()), + other_strides(other.strides()), out_strides(out.strides()), + dtype(out.dtype()), device(out.device()) {} + + bool Matches(const Tensor input, const Tensor other, Tensor out) const { + return input.shape() == input_shape && other.shape() == other_shape && + out.shape() == out_shape && input.strides() == input_strides && + other.strides() == other_strides && out.strides() == out_strides && + input.dtype() == dtype && other.dtype() == dtype && + out.dtype() == dtype && out.device() == device && + PointersSupported(input, other, out); + } +}; + +} // namespace infini::ops::nvidia_binary_view + +#endif diff --git a/src/native/cuda/nvidia/ops/cast/kernel.h b/src/native/cuda/nvidia/ops/cast/kernel.h new file mode 100644 index 000000000..aa947487b --- /dev/null +++ b/src/native/cuda/nvidia/ops/cast/kernel.h @@ -0,0 +1,134 @@ +#ifndef INFINI_OPS_NVIDIA_CAST_KERNEL_H_ +#define INFINI_OPS_NVIDIA_CAST_KERNEL_H_ + +#include +#include +#include +#include +#include +#include + +#include "base/cast.h" +#include "native/cuda/nvidia/caster.cuh" + +namespace infini::ops { +namespace nvidia_cast_detail { +template +__global__ void Convert(const InT* input, OutT* out, std::size_t first, + std::size_t count) { + const auto local = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (local < count) + out[first + local] = Caster::template Cast(input[first + local]); +} +} // namespace nvidia_cast_detail + +template <> +class Operator : public Cast { + public: + static bool supports(const Tensor input, Tensor out) { + if (input.device().type() != Device::Type::kNvidia || + input.device() != out.device() || input.shape() != out.shape() || + !SupportedDtype(input.dtype()) || !SupportedDtype(out.dtype())) return false; + std::size_t input_bytes = 0, out_bytes = 0; + if (!ContiguousBytes(input, &input_bytes) || + !ContiguousBytes(out, &out_bytes) || + !ValidAddress(input, input_bytes) || !ValidAddress(out, out_bytes)) return false; + int current = -1; + if (cudaGetDevice(¤t) != cudaSuccess || current != input.device().index()) + return false; + const auto ip = reinterpret_cast(input.data()); + const auto op = reinterpret_cast(out.data()); + const bool overlap = input_bytes && out_bytes && + ip < op + out_bytes && op < ip + input_bytes; + return !overlap || SameView(input, out); + } + + Operator(const Tensor input, Tensor out) + : Cast{Checked(input, out), out}, device_{input.device()} {} + + bool capture_safe() const override { return false; } + + void operator()(const Tensor input, Tensor out) const override { + if (input.dtype() != input_dtype_ || out.dtype() != out_dtype_ || + input.shape() != input_shape_ || out.shape() != out_shape_ || + input.strides() != input_strides_ || out.strides() != out_strides_ || + input.device() != device_ || !supports(input, out)) + throw std::invalid_argument("NVIDIA Cast: call metadata or alias changed after Make"); + if (output_size_ == 0 || SameView(input, out)) return; + const auto stream = static_cast(stream_); + DispatchFunc, + List>( + {static_cast(input_dtype_), static_cast(out_dtype_)}, + [&](auto list_tag) { + using InT = TypeMapType(list_tag)>; + using OutT = TypeMapType(list_tag)>; + constexpr std::size_t cap = 65535u * 256u; + for (std::size_t first = 0; first < output_size_;) { + const auto count = std::min(cap, output_size_ - first); + const auto blocks = static_cast((count + 255) / 256); + nvidia_cast_detail::Convert<<>>( + static_cast(input.data()), static_cast(out.data()), + first, count); + const auto status = cudaGetLastError(); + if (status != cudaSuccess) + throw std::runtime_error(std::string("NVIDIA Cast: launch failed: ") + + cudaGetErrorString(status)); + first += count; + } + }, "NVIDIA Cast"); + } + + private: + static bool SupportedDtype(DataType type) { + return type == DataType::kFloat32 || type == DataType::kFloat16 || + type == DataType::kBFloat16; + } + + static std::size_t Width(DataType type) { + return type == DataType::kFloat32 ? 4u : 2u; + } + + static bool ContiguousBytes(const Tensor t, std::size_t* bytes) { + if (t.shape().size() != t.strides().size()) return false; + bool empty = false; + for (std::size_t axis = 0; axis < t.ndim(); ++axis) { + if (t.stride(axis) < 0) return false; + empty = empty || t.size(axis) == 0; + } + if (empty) { *bytes = 0; return true; } + const auto limit = static_cast(std::numeric_limits::max()) / + Width(t.dtype()); + std::size_t count = 1; + for (std::size_t axis = t.ndim(); axis > 0; --axis) { + const auto size = t.size(axis - 1); + if (size > 1 && t.stride(axis - 1) != static_cast(count)) + return false; + if (size > limit / count) return false; + count *= size; + } + *bytes = count * Width(t.dtype()); + return true; + } + + static bool ValidAddress(const Tensor t, std::size_t bytes) { + if (bytes == 0) return true; + const auto start = reinterpret_cast(t.data()); + return start && start % Width(t.dtype()) == 0 && + start <= std::numeric_limits::max() - bytes; + } + + static bool SameView(const Tensor input, const Tensor out) { + return input.dtype() == out.dtype() && input.data() == out.data() && + input.shape() == out.shape() && input.strides() == out.strides(); + } + + static Tensor Checked(const Tensor input, const Tensor out) { + if (!supports(input, out)) + throw std::invalid_argument("NVIDIA Cast: unsupported dtype/device/layout or alias"); + return input; + } + + Device device_; +}; +} // namespace infini::ops +#endif diff --git a/src/native/cuda/nvidia/ops/copy/kernel.h b/src/native/cuda/nvidia/ops/copy/kernel.h index ebc3f441f..a7fe91750 100644 --- a/src/native/cuda/nvidia/ops/copy/kernel.h +++ b/src/native/cuda/nvidia/ops/copy/kernel.h @@ -14,6 +14,11 @@ class Operator : public CudaCopy> { public: using CudaCopy>::CudaCopy; + bool capture_safe() const override { + return output_size_ != 0 && ndim_ != 0 && + (out_type_ == DataType::kFloat32 || out_type_ == DataType::kFloat16 || + out_type_ == DataType::kBFloat16 || out_type_ == DataType::kInt32); + } }; } // namespace infini::ops diff --git a/src/native/cuda/nvidia/ops/embedding/kernel.h b/src/native/cuda/nvidia/ops/embedding/kernel.h index f05f155ae..54bc0f3db 100644 --- a/src/native/cuda/nvidia/ops/embedding/kernel.h +++ b/src/native/cuda/nvidia/ops/embedding/kernel.h @@ -14,6 +14,7 @@ class Operator : public CudaEmbedding> { public: using CudaEmbedding>::CudaEmbedding; + bool capture_safe() const override { return !max_norm_.has_value(); } }; } // namespace infini::ops diff --git a/src/native/cuda/nvidia/ops/fill/kernel.h b/src/native/cuda/nvidia/ops/fill/kernel.h index 76fcfc40e..8e2636fc2 100644 --- a/src/native/cuda/nvidia/ops/fill/kernel.h +++ b/src/native/cuda/nvidia/ops/fill/kernel.h @@ -14,6 +14,7 @@ class Operator : public CudaFill> { public: using CudaFill>::CudaFill; + bool capture_safe() const override { return true; } }; } // namespace infini::ops diff --git a/src/native/cuda/nvidia/ops/flash_attn_varlen_func/kernel.cc b/src/native/cuda/nvidia/ops/flash_attn_varlen_func/kernel.cc new file mode 100644 index 000000000..b74c546c4 --- /dev/null +++ b/src/native/cuda/nvidia/ops/flash_attn_varlen_func/kernel.cc @@ -0,0 +1,429 @@ +#include "kernel.h" +#include "native/cuda/nvidia/ops/flash_attn_with_kvcache/infini_flash_cuda.h" +#include "flash.h" +#include "native/cuda/nvidia/ops/flash_workspace.h" +#include "native/cuda/nvidia/ops/flash_head_dim.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace infini::ops::native_flash { +void launch_32_fp16_false(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_32_fp16_true(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_32_bf16_false(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_32_bf16_true(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_64_fp16_false(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_64_fp16_true(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_64_bf16_false(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_64_bf16_true(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_96_fp16_false(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_96_fp16_true(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_96_bf16_false(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_96_bf16_true(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_128_fp16_false(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_128_fp16_true(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_128_bf16_false(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_128_bf16_true(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_192_fp16_false(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_192_fp16_true(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_192_bf16_false(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_192_bf16_true(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_256_fp16_false(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_256_fp16_true(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_256_bf16_false(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_256_bf16_true(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_varlen_32_fp16_false(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_varlen_32_fp16_true(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_varlen_32_bf16_false(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_varlen_32_bf16_true(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_varlen_64_fp16_false(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_varlen_64_fp16_true(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_varlen_64_bf16_false(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_varlen_64_bf16_true(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_varlen_96_fp16_false(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_varlen_96_fp16_true(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_varlen_96_bf16_false(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_varlen_96_bf16_true(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_varlen_128_fp16_false(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_varlen_128_fp16_true(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_varlen_128_bf16_false(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_varlen_128_bf16_true(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_varlen_192_fp16_false(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_varlen_192_fp16_true(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_varlen_192_bf16_false(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_varlen_192_bf16_true(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_varlen_256_fp16_false(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_varlen_256_fp16_true(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_varlen_256_bf16_false(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_varlen_256_bf16_true(infini_flash::Flash_fwd_params&, cudaStream_t); +} // namespace infini::ops::native_flash + +namespace infini::ops { +namespace { + +struct DeviceScope { + int previous; + explicit DeviceScope(int device) { + INFINI_FLASH_CUDA_CHECK(cudaGetDevice(&previous)); + INFINI_FLASH_CUDA_CHECK(cudaSetDevice(device)); + } + ~DeviceScope() { cudaSetDevice(previous); } +}; + +bool plain_layout(Tensor t, int dimensions) { + if (t.ndim() != dimensions || t.stride(-1) != 1) return false; + std::array axes{}; + int count = 0; + for (int axis = 0; axis < dimensions; ++axis) { + if (t.size(axis) <= 0 || + t.size(axis) > std::numeric_limits::max() - 127 || + t.stride(axis) < 0) return false; + if (t.size(axis) != 1) axes[count++] = axis; + } + std::sort(axes.begin(), axes.begin() + count, + [&](int a, int b) { return t.stride(a) < t.stride(b); }); + std::size_t extent = 1; + for (int index = 0; index < count; ++index) { + const auto axis = axes[index]; + const auto size = static_cast(t.size(axis)); + const auto stride = static_cast(t.stride(axis)); + if (stride < extent || (size - 1) > + (std::numeric_limits::max() - extent) / stride) return false; + extent += (size - 1) * stride; + } + return true; +} + +bool vector_aligned(Tensor t) { + if (t.data() && reinterpret_cast(t.data()) % 16) return false; + for (int axis = 0; axis < t.ndim() - 1; ++axis) + if (t.size(axis) > 1 && t.stride(axis) % 8) return false; + return true; +} + +bool int_aligned(Tensor t) { + return !t.data() || reinterpret_cast(t.data()) % alignof(int) == 0; +} + +std::size_t checked_product(std::initializer_list sizes) { + std::size_t result = 1; + for (auto size : sizes) { + if (size && result > std::numeric_limits::max() / size) + throw std::overflow_error("native varlen Attention scratch overflow"); + result *= size; + } + return result; +} + +} // namespace + +struct Operator::Impl { + infini_flash::Flash_fwd_params params{}; + std::size_t workspace_bytes = 0; + std::size_t lse_offset = 0; + std::optional scale; + bool requested_causal; + std::array requested_window; + bool paged; + int64_t max_q, max_k; +}; + +#define VARLEN_ARGS \ + const Tensor q, const Tensor k, const Tensor v, const Tensor cu_seqlens_q, \ + const Tensor cu_seqlens_k, const std::optional alibi_slopes, \ + const std::optional block_table, const int64_t max_seqlen_q, \ + const int64_t max_seqlen_k, const double dropout_p, \ + const std::optional softmax_scale, const bool causal, \ + const std::vector window_size, const double softcap, \ + const bool deterministic, const bool return_attn_probs, Tensor out, \ + std::optional softmax_lse, std::optional s_dmask +#define VARLEN_VALUES \ + q, k, v, cu_seqlens_q, cu_seqlens_k, alibi_slopes, block_table, max_seqlen_q,\ + max_seqlen_k, dropout_p, softmax_scale, causal, window_size, softcap, \ + deterministic, return_attn_probs, out, softmax_lse, s_dmask + +namespace { +bool supports_metadata(VARLEN_ARGS) { + if (!FlashAttnVarlenFunc::supports(VARLEN_VALUES)) return false; + const bool paged = block_table.has_value(); + const bool valid_window = + ((window_size[0] == -1 && window_size[1] == -1) || + (window_size[0] >= 0 && + window_size[0] <= std::numeric_limits::max() - 127 && + window_size[1] == 0)); + if (!plain_layout(q, 3) || !plain_layout(out, 3) || + !plain_layout(k, paged ? 4 : 3) || + !plain_layout(v, paged ? 4 : 3) || + !native_flash::head_dim_bin(q.size(2)) || + q.device().type() != Device::Type::kNvidia || + !vector_aligned(q) || !vector_aligned(k) || + !vector_aligned(v) || !vector_aligned(out) || + !plain_layout(cu_seqlens_q, 1) || !plain_layout(cu_seqlens_k, 1) || + !int_aligned(cu_seqlens_q) || !int_aligned(cu_seqlens_k) || + max_seqlen_q > std::numeric_limits::max() - 127 || + max_seqlen_k > std::numeric_limits::max() - 127 || + max_seqlen_q > q.size(0) || + alibi_slopes || return_attn_probs || + !valid_window) return false; + if (paged) { + if (!plain_layout(*block_table, 2) || + !int_aligned(*block_table) || + block_table->size(1) > + (std::numeric_limits::max() - 127) / k.size(1) || + max_seqlen_k > block_table->size(1) * k.size(1)) return false; + } else if (max_seqlen_k > k.size(0)) { + return false; + } + (void)softmax_scale; + (void)causal; + return true; +} +} // namespace + +bool Operator::supports( + VARLEN_ARGS) { + if (!supports_metadata(VARLEN_VALUES)) return false; + // Qualification must not pick this provider on pre-Ampere NVIDIA devices. + // This is a device-attribute query only; no provider or scratch is created. + int major = 0; + return cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, + q.device().index()) == cudaSuccess && major >= 8; +} + +void Operator::prepare( + VARLEN_ARGS) { + (void)alibi_slopes; (void)dropout_p; (void)softcap; + (void)deterministic; (void)return_attn_probs; (void)softmax_lse; (void)s_dmask; + DeviceScope device(device_index_); + int major; + INFINI_FLASH_CUDA_CHECK(cudaDeviceGetAttribute( + &major, cudaDevAttrComputeCapabilityMajor, device_index_)); + if (major < 8) throw std::invalid_argument("native varlen Attention requires SM80+"); + auto state = std::make_shared(); + state->scale = softmax_scale; + state->requested_causal = causal; + state->requested_window = {window_size[0], window_size[1]}; + state->paged = block_table.has_value(); + state->max_q = max_seqlen_q; + state->max_k = max_seqlen_k; + auto& p = state->params; + // Follow pinned flash_api.cpp set_params_fprop and mha_varlen_fwd. The + // upstream single-token GQA reshape is an optional optimization, omitted. + p.b = cu_seqlens_q.size(0) - 1; + p.h = q.size(1); p.h_k = k.size(k.ndim() - 2); + p.h_h_k_ratio = p.h / p.h_k; + p.d = q.size(2); p.d_rounded = native_flash::head_dim_bin(p.d); + p.seqlen_q = max_seqlen_q; p.seqlen_k = max_seqlen_k; + p.seqlen_q_rounded = ((p.seqlen_q + 127) / 128) * 128; + p.seqlen_k_rounded = ((p.seqlen_k + 127) / 128) * 128; + p.total_q = q.size(0); + p.q_row_stride = q.stride(0); p.q_head_stride = q.stride(1); + p.k_row_stride = k.stride(k.ndim() - 3); + p.k_head_stride = k.stride(k.ndim() - 2); + p.v_row_stride = v.stride(v.ndim() - 3); + p.v_head_stride = v.stride(v.ndim() - 2); + p.o_row_stride = out.stride(0); p.o_head_stride = out.stride(1); + p.k_batch_stride = state->paged ? k.stride(0) : 0; + p.v_batch_stride = state->paged ? v.stride(0) : 0; + p.block_table_batch_stride = state->paged ? block_table->stride(0) : 0; + p.page_block_size = state->paged ? k.size(1) : 1; + p.is_bf16 = q.dtype() == DataType::kBFloat16; + const float scale = static_cast(softmax_scale.value_or( + 1.0 / std::sqrt(static_cast(p.d)))); + if (!std::isfinite(scale)) throw std::invalid_argument("nonfinite varlen Attention scale"); + p.scale_softmax = scale; + p.scale_softmax_log2 = scale * M_LOG2E; + p.p_dropout = p.rp_dropout = 1.f; + p.p_dropout_in_uint8_t = 255; + p.scale_softmax_rp_dropout = scale; + const bool finite_window = window_size[0] >= 0; + // FA's causal specialization ignores the left bound. Finite causal windows + // must use its existing local specialization, including single-token Q. + p.is_causal = !finite_window && causal && max_seqlen_q != 1; + p.window_size_left = finite_window ? static_cast(window_size[0]) + : (p.is_causal ? max_seqlen_k : -1); + p.window_size_right = finite_window ? 0 : (p.is_causal ? 0 : -1); + p.is_seqlens_k_cumulative = true; + p.unpadded_lse = true; + p.num_splits = 0; // Paged KV still uses upstream split-KV kernel, one split. + native_flash::WorkspaceLayout layout; + state->lse_offset = layout.append(checked_product({ + static_cast(p.h), static_cast(p.total_q), + sizeof(float)})); + state->workspace_bytes = layout.required(); + impl_ = std::move(state); +} + +std::size_t Operator:: +workspace_size_in_bytes() const { return impl_->workspace_bytes; } + +Operator::~Operator() { + impl_.reset(); +} + +void Operator::operator()( + VARLEN_ARGS) const { + if (!supports_metadata(VARLEN_VALUES) || q.shape() != q_shape_ || + q.strides() != q_strides_ || k.shape() != k_shape_ || + k.strides() != k_strides_ || v.shape() != v_shape_ || + v.strides() != v_strides_ || out.shape() != out_shape_ || + out.strides() != out_strides_ || + cu_seqlens_q.shape() != cu_seqlens_q_shape_ || + cu_seqlens_q.strides() != cu_seqlens_q_strides_ || + cu_seqlens_k.shape() != cu_seqlens_k_shape_ || + cu_seqlens_k.strides() != cu_seqlens_k_strides_ || + (block_table ? Tensor::Shape{block_table->shape()} : Tensor::Shape{}) != block_table_shape_ || + (block_table ? Tensor::Strides{block_table->strides()} : Tensor::Strides{}) != block_table_strides_ || + q.device().index() != device_index_ || q.dtype() != q_dtype_ || + softmax_scale != impl_->scale || causal != impl_->requested_causal || + window_size[0] != impl_->requested_window[0] || + window_size[1] != impl_->requested_window[1] || + max_seqlen_q != impl_->max_q || max_seqlen_k != impl_->max_k) + throw std::invalid_argument("native varlen Attention prepared metadata changed"); + DeviceScope device(device_index_); + auto stream = static_cast(stream_); + native_flash::WorkspaceLease scratch(workspace_, workspace_size_in_bytes_, + impl_->workspace_bytes, stream, + device_index_); + auto p = impl_->params; + p.softmax_lse_ptr = static_cast(scratch.data()) + impl_->lse_offset; + p.q_ptr = const_cast(q.data()); + p.k_ptr = const_cast(k.data()); + p.v_ptr = const_cast(v.data()); + p.o_ptr = out.data(); + // The pinned FA parameter ABI uses mutable pointers, but inference reads + // cumulative lengths and the page table without modifying them. + p.cu_seqlens_q = const_cast(static_cast(cu_seqlens_q.data())); + p.cu_seqlens_k = const_cast(static_cast(cu_seqlens_k.data())); + p.block_table = block_table + ? const_cast(static_cast(block_table->data())) : nullptr; + // Pinned flash_api.cpp forces split-KV for paged KV; non-split kernels do + // not read block_table. This one-split case needs no accumulation scratch. + if (impl_->paged) { + switch (p.d_rounded) { + case 32: + if (p.is_bf16) { + if (p.is_causal) native_flash::launch_32_bf16_true(p, stream); + else native_flash::launch_32_bf16_false(p, stream); + } else { + if (p.is_causal) native_flash::launch_32_fp16_true(p, stream); + else native_flash::launch_32_fp16_false(p, stream); + } + break; + case 64: + if (p.is_bf16) { + if (p.is_causal) native_flash::launch_64_bf16_true(p, stream); + else native_flash::launch_64_bf16_false(p, stream); + } else { + if (p.is_causal) native_flash::launch_64_fp16_true(p, stream); + else native_flash::launch_64_fp16_false(p, stream); + } + break; + case 96: + if (p.is_bf16) { + if (p.is_causal) native_flash::launch_96_bf16_true(p, stream); + else native_flash::launch_96_bf16_false(p, stream); + } else { + if (p.is_causal) native_flash::launch_96_fp16_true(p, stream); + else native_flash::launch_96_fp16_false(p, stream); + } + break; + case 128: + if (p.is_bf16) { + if (p.is_causal) native_flash::launch_128_bf16_true(p, stream); + else native_flash::launch_128_bf16_false(p, stream); + } else { + if (p.is_causal) native_flash::launch_128_fp16_true(p, stream); + else native_flash::launch_128_fp16_false(p, stream); + } + break; + case 192: + if (p.is_bf16) { + if (p.is_causal) native_flash::launch_192_bf16_true(p, stream); + else native_flash::launch_192_bf16_false(p, stream); + } else { + if (p.is_causal) native_flash::launch_192_fp16_true(p, stream); + else native_flash::launch_192_fp16_false(p, stream); + } + break; + case 256: + if (p.is_bf16) { + if (p.is_causal) native_flash::launch_256_bf16_true(p, stream); + else native_flash::launch_256_bf16_false(p, stream); + } else { + if (p.is_causal) native_flash::launch_256_fp16_true(p, stream); + else native_flash::launch_256_fp16_false(p, stream); + } + break; + default: throw std::logic_error("unqualified native Flash head dimension"); + } + } else { + switch (p.d_rounded) { + case 32: + if (p.is_bf16) { + if (p.is_causal) native_flash::launch_varlen_32_bf16_true(p, stream); + else native_flash::launch_varlen_32_bf16_false(p, stream); + } else { + if (p.is_causal) native_flash::launch_varlen_32_fp16_true(p, stream); + else native_flash::launch_varlen_32_fp16_false(p, stream); + } + break; + case 64: + if (p.is_bf16) { + if (p.is_causal) native_flash::launch_varlen_64_bf16_true(p, stream); + else native_flash::launch_varlen_64_bf16_false(p, stream); + } else { + if (p.is_causal) native_flash::launch_varlen_64_fp16_true(p, stream); + else native_flash::launch_varlen_64_fp16_false(p, stream); + } + break; + case 96: + if (p.is_bf16) { + if (p.is_causal) native_flash::launch_varlen_96_bf16_true(p, stream); + else native_flash::launch_varlen_96_bf16_false(p, stream); + } else { + if (p.is_causal) native_flash::launch_varlen_96_fp16_true(p, stream); + else native_flash::launch_varlen_96_fp16_false(p, stream); + } + break; + case 128: + if (p.is_bf16) { + if (p.is_causal) native_flash::launch_varlen_128_bf16_true(p, stream); + else native_flash::launch_varlen_128_bf16_false(p, stream); + } else { + if (p.is_causal) native_flash::launch_varlen_128_fp16_true(p, stream); + else native_flash::launch_varlen_128_fp16_false(p, stream); + } + break; + case 192: + if (p.is_bf16) { + if (p.is_causal) native_flash::launch_varlen_192_bf16_true(p, stream); + else native_flash::launch_varlen_192_bf16_false(p, stream); + } else { + if (p.is_causal) native_flash::launch_varlen_192_fp16_true(p, stream); + else native_flash::launch_varlen_192_fp16_false(p, stream); + } + break; + case 256: + if (p.is_bf16) { + if (p.is_causal) native_flash::launch_varlen_256_bf16_true(p, stream); + else native_flash::launch_varlen_256_bf16_false(p, stream); + } else { + if (p.is_causal) native_flash::launch_varlen_256_fp16_true(p, stream); + else native_flash::launch_varlen_256_fp16_false(p, stream); + } + break; + default: throw std::logic_error("unqualified native Flash head dimension"); + } + } + scratch.finish(); +} + +#undef VARLEN_ARGS +#undef VARLEN_VALUES +} // namespace infini::ops diff --git a/src/native/cuda/nvidia/ops/flash_attn_varlen_func/kernel.h b/src/native/cuda/nvidia/ops/flash_attn_varlen_func/kernel.h new file mode 100644 index 000000000..4b91cf311 --- /dev/null +++ b/src/native/cuda/nvidia/ops/flash_attn_varlen_func/kernel.h @@ -0,0 +1,64 @@ +#pragma once + +#include + +#include "base/flash_attn_varlen_func.h" + +#define INFINI_NATIVE_VARLEN_ARGS \ + const Tensor q, const Tensor k, const Tensor v, const Tensor cu_seqlens_q, \ + const Tensor cu_seqlens_k, const std::optional alibi_slopes, \ + const std::optional block_table, const int64_t max_seqlen_q, \ + const int64_t max_seqlen_k, const double dropout_p, \ + const std::optional softmax_scale, const bool causal, \ + const std::vector window_size, const double softcap, \ + const bool deterministic, const bool return_attn_probs, Tensor out, \ + std::optional softmax_lse, std::optional s_dmask +#define INFINI_NATIVE_VARLEN_VALUES \ + q, k, v, cu_seqlens_q, cu_seqlens_k, alibi_slopes, block_table, max_seqlen_q, \ + max_seqlen_k, dropout_p, softmax_scale, causal, window_size, softcap, \ + deterministic, return_attn_probs, out, softmax_lse, s_dmask + +namespace infini::ops { + +// Inference-only qualified head dimensions are checked by supports(). +// Automatic selection leaves other configurations to an existing provider; +// construction failures propagate. +template <> +class Operator + : public FlashAttnVarlenFunc { + public: + static bool supports(INFINI_NATIVE_VARLEN_ARGS); + + Operator(INFINI_NATIVE_VARLEN_ARGS) + : FlashAttnVarlenFunc(checked_q(INFINI_NATIVE_VARLEN_VALUES), k, v, + cu_seqlens_q, cu_seqlens_k, alibi_slopes, + block_table, max_seqlen_q, max_seqlen_k, + dropout_p, softmax_scale, causal, window_size, + softcap, deterministic, return_attn_probs, out, + softmax_lse, s_dmask) { + prepare(INFINI_NATIVE_VARLEN_VALUES); + } + ~Operator() override; + using FlashAttnVarlenFunc::operator(); + + void operator()(INFINI_NATIVE_VARLEN_ARGS) const override; + + // Graph callers must supply stable external workspace for the graph lifetime. + bool capture_safe() const override { return true; } + std::size_t workspace_size_in_bytes() const override; + + private: + static Tensor checked_q(INFINI_NATIVE_VARLEN_ARGS) { + if (!supports(INFINI_NATIVE_VARLEN_VALUES)) + throw std::invalid_argument("unsupported native varlen Attention configuration"); + return q; + } + void prepare(INFINI_NATIVE_VARLEN_ARGS); + struct Impl; + std::shared_ptr impl_; +}; + +} // namespace infini::ops + +#undef INFINI_NATIVE_VARLEN_ARGS +#undef INFINI_NATIVE_VARLEN_VALUES diff --git a/src/native/cuda/nvidia/ops/flash_attn_varlen_func/launch.cu.in b/src/native/cuda/nvidia/ops/flash_attn_varlen_func/launch.cu.in new file mode 100644 index 000000000..926ecc22d --- /dev/null +++ b/src/native/cuda/nvidia/ops/flash_attn_varlen_func/launch.cu.in @@ -0,0 +1,9 @@ +#include "flash_fwd_launch_template.h" + +namespace infini::ops::native_flash { +void launch_varlen_@_flash_hdim@_@_flash_kind@_@_flash_causal@( + infini_flash::Flash_fwd_params& params, cudaStream_t stream) { + infini_flash::run_mha_fwd_hdim@_flash_hdim@<@_flash_type@, @_flash_causal@>( + params, stream); +} +} // namespace infini::ops::native_flash diff --git a/src/native/cuda/nvidia/ops/flash_attn_with_kvcache/infini_flash_cuda.h b/src/native/cuda/nvidia/ops/flash_attn_with_kvcache/infini_flash_cuda.h new file mode 100644 index 000000000..789c807b4 --- /dev/null +++ b/src/native/cuda/nvidia/ops/flash_attn_with_kvcache/infini_flash_cuda.h @@ -0,0 +1,15 @@ +#pragma once + +#include +#include +#include + +inline void infini_flash_check(cudaError_t error, const char* operation) { + if (error != cudaSuccess) { + throw std::runtime_error(std::string(operation) + ": " + + cudaGetErrorString(error)); + } +} + +#define INFINI_FLASH_CUDA_CHECK(call) infini_flash_check((call), #call) +#define INFINI_FLASH_KERNEL_CHECK() INFINI_FLASH_CUDA_CHECK(cudaGetLastError()) diff --git a/src/native/cuda/nvidia/ops/flash_attn_with_kvcache/kernel.cc b/src/native/cuda/nvidia/ops/flash_attn_with_kvcache/kernel.cc new file mode 100644 index 000000000..b9a056536 --- /dev/null +++ b/src/native/cuda/nvidia/ops/flash_attn_with_kvcache/kernel.cc @@ -0,0 +1,320 @@ +#include "kernel.h" +#include "infini_flash_cuda.h" +#include "flash.h" +#include "num_splits.h" +#include "native/cuda/nvidia/ops/flash_workspace.h" +#include "native/cuda/nvidia/ops/flash_head_dim.h" + +#include +#include +#include +#include +#include +#include + +namespace infini::ops::native_flash { +void launch_32_fp16_false(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_32_fp16_true(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_32_bf16_false(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_32_bf16_true(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_64_fp16_false(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_64_fp16_true(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_64_bf16_false(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_64_bf16_true(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_96_fp16_false(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_96_fp16_true(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_96_bf16_false(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_96_bf16_true(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_128_fp16_false(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_128_fp16_true(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_128_bf16_false(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_128_bf16_true(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_192_fp16_false(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_192_fp16_true(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_192_bf16_false(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_192_bf16_true(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_256_fp16_false(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_256_fp16_true(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_256_bf16_false(infini_flash::Flash_fwd_params&, cudaStream_t); +void launch_256_bf16_true(infini_flash::Flash_fwd_params&, cudaStream_t); +} // namespace infini::ops::native_flash + +namespace infini::ops { +namespace { +struct DeviceScope { + int previous; + explicit DeviceScope(int device) { + INFINI_FLASH_CUDA_CHECK(cudaGetDevice(&previous)); + INFINI_FLASH_CUDA_CHECK(cudaSetDevice(device)); + } + ~DeviceScope() { cudaSetDevice(previous); } +}; + +// Candidate ordinary MHA-only extension. Keep true d/strides, use existing 96 +// specialization's Is_even_K=false masking. No physical padding or new algorithm. +int ordinary_head_dim_bin(int dimension) { + return dimension == 72 || dimension == 80 ? 96 : native_flash::head_dim_bin(dimension); +} + +bool plain_layout(Tensor t, int dimensions) { + if (dimensions < 1 || dimensions > 4 || t.ndim() != dimensions || + t.stride(-1) != 1) return false; + std::array axes{}; + int count = 0; + for (int axis = 0; axis < dimensions; ++axis) { + if (t.size(axis) <= 0 || t.size(axis) > std::numeric_limits::max() - 127 || + t.stride(axis) < 0) + return false; + // A singleton axis is never indexed away from zero. LM unsqueeze views + // legitimately use stride zero here; this is not a broadcasted data axis. + if (t.size(axis) == 1) continue; + axes[count++] = axis; + } + // Logical dimension order need not match storage order. Static KV is a + // head-major view; the upstream kernel already receives its row/head strides. + std::sort(axes.begin(), axes.begin() + count, + [&](int a, int b) { return t.stride(a) < t.stride(b); }); + std::size_t extent = 1; + for (int index = 0; index < count; ++index) { + const int axis = axes[index]; + const auto size = static_cast(t.size(axis)); + const auto stride = static_cast(t.stride(axis)); + if (stride < extent || (size - 1) > + (std::numeric_limits::max() - extent) / stride) return false; + extent += (size - 1) * stride; + } + return true; +} + +std::size_t product(std::initializer_list sizes) { + std::size_t result = 1; + for (auto size : sizes) { + if (size && result > std::numeric_limits::max() / size) + throw std::overflow_error("native Attention scratch size overflow"); + result *= size; + } + return result; +} + +} // namespace + +struct Operator::Impl { + infini_flash::Flash_fwd_params params{}; + std::size_t workspace_bytes = 0; + std::size_t lse_offset = 0, lse_accum_offset = 0, out_accum_offset = 0; + std::optional scale; + int64_t requested_splits; + bool requested_causal; + std::array requested_window; + bool lengths_present; +}; + +bool Operator::supported_layout( + const Tensor q, Tensor k, Tensor v, std::optional lengths, + std::optional blocks, Tensor out) { + if (!plain_layout(q, 4) || !plain_layout(k, 4) || !plain_layout(v, 4) || + !plain_layout(out, 4) || !ordinary_head_dim_bin(q.size(3)) || + q.device().type() != Device::Type::kNvidia) + return false; + // New dimensions are restricted to ordinary dense MHA, not decode/GQA, + // length-masked or paged cache domains that require separate qualification. + if ((q.size(3) == 72 || q.size(3) == 80) && + (lengths || blocks || q.size(0) != k.size(0) || q.size(2) != k.size(2))) + return false; + // Upstream vectorized loads/stores require 16-byte-aligned rows and heads. + // Metadata-only descriptor views may have null data until invocation. + for (auto tensor : {q, k, v, out}) { + if (tensor.data() && reinterpret_cast(tensor.data()) % 16) + return false; + for (int axis = 0; axis < 3; ++axis) + if (tensor.size(axis) > 1 && tensor.stride(axis) % 8 != 0) return false; + } + if ((lengths && !plain_layout(*lengths, 1)) || + (blocks && (!plain_layout(*blocks, 2) || + blocks->size(1) > + (std::numeric_limits::max() - 127) / k.size(1)))) + return false; + return true; +} + +void Operator::prepare( + const Tensor q, Tensor cache, std::optional lengths, + std::optional blocks, std::optional scale, + bool causal, const std::vector& window_size, + int64_t splits, Tensor out, int64_t logical_bound) { + DeviceScope device(device_index_); + int major, sms; + INFINI_FLASH_CUDA_CHECK(cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, device_index_)); + INFINI_FLASH_CUDA_CHECK(cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, device_index_)); + if (major < 8) throw std::invalid_argument("native Attention requires SM80 or newer"); + auto state = std::make_shared(); + state->scale = scale; + state->requested_splits = splits; + state->requested_causal = causal; + state->requested_window = {window_size[0], window_size[1]}; + state->lengths_present = lengths.has_value(); + auto& p = state->params; + p.b = q.size(0); p.h = q.size(2); p.h_k = cache.size(2); + p.d = q.size(3); p.d_rounded = ordinary_head_dim_bin(p.d); + p.seqlen_q = q.size(1); + p.seqlen_k = logical_bound > 0 ? logical_bound : + (blocks ? blocks->size(1) * cache.size(1) : cache.size(1)); + p.q_batch_stride = q.stride(0); p.q_row_stride = q.stride(1); p.q_head_stride = q.stride(2); + p.o_batch_stride = out.stride(0); p.o_row_stride = out.stride(1); p.o_head_stride = out.stride(2); + // Preserve upstream's no-copy decode GQA transpose, including strided Q. + // It reinterprets heads as query positions and is valid only globally. + const bool finite_window = window_size[0] >= 0; + if (!finite_window && p.seqlen_q == 1 && p.h > p.h_k) { + const int groups = p.h / p.h_k; + p.seqlen_q = groups; p.h = p.h_k; + p.q_row_stride = q.stride(2); p.q_head_stride = q.stride(2) * groups; + p.o_row_stride = out.stride(2); p.o_head_stride = out.stride(2) * groups; + } + p.h_h_k_ratio = p.h / p.h_k; + p.k_batch_stride = cache.stride(0); p.k_row_stride = cache.stride(1); p.k_head_stride = cache.stride(2); + p.v_batch_stride = v_cache_strides_[0]; p.v_row_stride = v_cache_strides_[1]; p.v_head_stride = v_cache_strides_[2]; + p.seqlen_q_rounded = ((p.seqlen_q + 127) / 128) * 128; + p.seqlen_k_rounded = ((p.seqlen_k + 127) / 128) * 128; + // Finite right-zero masking is causal via the local kernel, not the global + // causal specialization. Keep the real single-token query geometry. + p.is_causal = !finite_window && causal && q.size(1) != 1; + p.window_size_left = finite_window ? static_cast(window_size[0]) + : (p.is_causal ? p.seqlen_k : -1); + p.window_size_right = finite_window ? 0 : (p.is_causal ? 0 : -1); + p.is_bf16 = q.dtype() == DataType::kBFloat16; + p.scale_softmax = scale.value_or(1.0 / std::sqrt(static_cast(p.d))); + if (!std::isfinite(p.scale_softmax)) throw std::invalid_argument("nonfinite Attention scale"); + // Preserve upstream's double-precision M_LOG2E multiplication before the + // final float assignment rather than introducing another rounded constant. + p.scale_softmax_log2 = p.scale_softmax * M_LOG2E; + p.p_dropout = p.rp_dropout = 1.f; p.p_dropout_in_uint8_t = 255; + p.scale_softmax_rp_dropout = p.scale_softmax; + p.is_seqlens_k_cumulative = !lengths; + p.page_block_size = blocks ? cache.size(1) : 1; + p.block_table_batch_stride = blocks ? blocks->stride(0) : 0; + const auto work_blocks = product({static_cast(p.b), + static_cast(p.h), + static_cast((p.seqlen_q + 63) / 64)}); + if (work_blocks > std::numeric_limits::max()) + throw std::overflow_error("native Attention scheduling size overflow"); + const int block_n = native_flash::split_block_n(p.d_rounded); + p.num_splits = splits > 0 ? splits : infini_flash::num_splits_heuristic( + static_cast(work_blocks), sms * 2, + (p.seqlen_k + block_n - 1) / block_n, 128); + auto bytes = product({static_cast(p.b), + static_cast(p.h), + static_cast(p.seqlen_q), sizeof(float)}); + native_flash::WorkspaceLayout layout; + state->lse_offset = layout.append(bytes); + if (p.num_splits > 1) { + const auto split_bytes = product({bytes, static_cast(p.num_splits)}); + state->lse_accum_offset = layout.append(split_bytes); + state->out_accum_offset = layout.append(product({ + split_bytes, static_cast(p.d_rounded)})); + } + state->workspace_bytes = layout.required(); + impl_ = std::move(state); +} + +std::size_t Operator:: +workspace_size_in_bytes() const { return impl_->workspace_bytes; } + +Operator::~Operator() { + impl_.reset(); +} + +void Operator::execute( + const Tensor q, Tensor k, Tensor v, std::optional lengths, + std::optional blocks, std::optional scale, + bool causal, const std::vector& window_size, + int64_t splits, Tensor out) const { + if (q.shape() != q_shape_ || q.strides() != q_strides_ || + k.shape() != k_cache_shape_ || k.strides() != k_cache_strides_ || + v.shape() != v_cache_shape_ || v.strides() != v_cache_strides_ || + out.shape() != out_shape_ || out.strides() != out_strides_ || + (blocks ? Tensor::Shape{blocks->shape()} : Tensor::Shape{}) != block_table_shape_ || + (blocks ? Tensor::Strides{blocks->strides()} : Tensor::Strides{}) != block_table_strides_ || + scale != impl_->scale || causal != impl_->requested_causal || + window_size.size() != 2 || + window_size[0] != impl_->requested_window[0] || + window_size[1] != impl_->requested_window[1] || + splits != impl_->requested_splits || lengths.has_value() != impl_->lengths_present || + q.device().index() != device_index_ || q.dtype() != q_dtype_) + throw std::invalid_argument("native Attention prepared metadata changed"); + DeviceScope device(device_index_); + auto stream = static_cast(stream_); + native_flash::WorkspaceLease scratch(workspace_, workspace_size_in_bytes_, + impl_->workspace_bytes, stream, + device_index_); + auto p = impl_->params; // Dynamic addresses, never cached input values. + auto* workspace = static_cast(scratch.data()); + p.softmax_lse_ptr = workspace + impl_->lse_offset; + if (p.num_splits > 1) { + p.softmax_lseaccum_ptr = workspace + impl_->lse_accum_offset; + p.oaccum_ptr = workspace + impl_->out_accum_offset; + } + // Upstream parameter struct uses void* even for the read-only query. + p.q_ptr = const_cast(q.data()); + p.k_ptr = k.data(); p.v_ptr = v.data(); p.o_ptr = out.data(); + p.cu_seqlens_k = lengths ? static_cast(lengths->data()) : nullptr; + p.block_table = blocks ? static_cast(blocks->data()) : nullptr; + switch (p.d_rounded) { + case 32: + if (p.is_bf16) { + if (p.is_causal) native_flash::launch_32_bf16_true(p, stream); + else native_flash::launch_32_bf16_false(p, stream); + } else { + if (p.is_causal) native_flash::launch_32_fp16_true(p, stream); + else native_flash::launch_32_fp16_false(p, stream); + } + break; + case 64: + if (p.is_bf16) { + if (p.is_causal) native_flash::launch_64_bf16_true(p, stream); + else native_flash::launch_64_bf16_false(p, stream); + } else { + if (p.is_causal) native_flash::launch_64_fp16_true(p, stream); + else native_flash::launch_64_fp16_false(p, stream); + } + break; + case 96: + if (p.is_bf16) { + if (p.is_causal) native_flash::launch_96_bf16_true(p, stream); + else native_flash::launch_96_bf16_false(p, stream); + } else { + if (p.is_causal) native_flash::launch_96_fp16_true(p, stream); + else native_flash::launch_96_fp16_false(p, stream); + } + break; + case 128: + if (p.is_bf16) { + if (p.is_causal) native_flash::launch_128_bf16_true(p, stream); + else native_flash::launch_128_bf16_false(p, stream); + } else { + if (p.is_causal) native_flash::launch_128_fp16_true(p, stream); + else native_flash::launch_128_fp16_false(p, stream); + } + break; + case 192: + if (p.is_bf16) { + if (p.is_causal) native_flash::launch_192_bf16_true(p, stream); + else native_flash::launch_192_bf16_false(p, stream); + } else { + if (p.is_causal) native_flash::launch_192_fp16_true(p, stream); + else native_flash::launch_192_fp16_false(p, stream); + } + break; + case 256: + if (p.is_bf16) { + if (p.is_causal) native_flash::launch_256_bf16_true(p, stream); + else native_flash::launch_256_bf16_false(p, stream); + } else { + if (p.is_causal) native_flash::launch_256_fp16_true(p, stream); + else native_flash::launch_256_fp16_false(p, stream); + } + break; + default: throw std::logic_error("unqualified native Flash head dimension"); + } + scratch.finish(); +} +} // namespace infini::ops diff --git a/src/native/cuda/nvidia/ops/flash_attn_with_kvcache/kernel.h b/src/native/cuda/nvidia/ops/flash_attn_with_kvcache/kernel.h new file mode 100644 index 000000000..974bf7c24 --- /dev/null +++ b/src/native/cuda/nvidia/ops/flash_attn_with_kvcache/kernel.h @@ -0,0 +1,175 @@ +#pragma once + +#include +#include +#include +#include "base/flash_attn_with_kvcache.h" + +// These macros only spell the existing public overloads, not a second API. +#define INFINI_NATIVE_KVCACHE_PARAMS(Lengths) \ + const Tensor q, Tensor k_cache, Tensor v_cache, \ + const std::optional k, const std::optional v, \ + const std::optional rotary_cos, \ + const std::optional rotary_sin, const Lengths cache_seqlens, \ + const std::optional cache_batch_idx, \ + const std::optional cache_leftpad, \ + const std::optional block_table, \ + const std::optional alibi_slopes, \ + const std::optional softmax_scale, const bool causal, \ + const std::vector window_size, const double softcap, \ + const bool rotary_interleaved, const int64_t num_splits, \ + const bool return_softmax_lse, Tensor out, \ + std::optional softmax_lse +#define INFINI_NATIVE_KVCACHE_VALUES \ + q, k_cache, v_cache, k, v, rotary_cos, rotary_sin, cache_seqlens, \ + cache_batch_idx, cache_leftpad, block_table, alibi_slopes, \ + softmax_scale, causal, window_size, softcap, rotary_interleaved, \ + num_splits, return_softmax_lse, out, softmax_lse + +namespace infini::ops { + +// Native inference provider: qualified heads, read-only cache, global or finite +// left-window/right-zero attention, no ALiBi/softcap. +// Other existing configurations remain with their previous providers. The +// explicit build option is not enabled in the accepted LM candidate yet. +template <> +class Operator + : public FlashAttnWithKvcache { + public: + template + static bool supports(INFINI_NATIVE_KVCACHE_PARAMS(Lengths)) { + if constexpr (!std::is_same_v>) { + return false; // Scalar metadata needs a separate prepared device buffer. + } else { + return FlashAttnWithKvcache::supports(INFINI_NATIVE_KVCACHE_VALUES) && + !k && !v && !rotary_cos && !rotary_sin && !cache_batch_idx && + !cache_leftpad && !alibi_slopes && + ((window_size[0] == -1 && window_size[1] == -1) || + (window_size[0] >= 0 && + window_size[0] <= std::numeric_limits::max() - 127 && + window_size[1] == 0)) && + softcap == 0.0 && num_splits <= 128 && !return_softmax_lse && + supported_layout(q, k_cache, v_cache, cache_seqlens, block_table, out) && + ((q.size(3) != 72 && q.size(3) != 80) || + (!causal && window_size[0] == -1)); + } + } + + static bool supports(const Tensor q, Tensor k_cache, Tensor v_cache, Tensor out) { + return FlashAttnWithKvcache::supports(q, k_cache, v_cache, out) && + supported_layout(q, k_cache, v_cache, {}, {}, out); + } + + template + static bool supports(const Config& config, + INFINI_NATIVE_KVCACHE_PARAMS(Lengths)) { + if constexpr (!std::is_same_v>) { + return false; + } else { + if (!supports(INFINI_NATIVE_KVCACHE_VALUES)) return false; + const auto bound = FlashAttnWithKvcacheConfig::BoundFrom(config); + if (bound == 0) return true; + const auto capacity = block_table + ? block_table->size(1) * k_cache.size(1) : k_cache.size(1); + return bound <= std::numeric_limits::max() - 127 && + static_cast(bound) <= capacity && + (cache_seqlens.has_value() || + static_cast(bound) == capacity); + } + } + + static bool supports(const Config& config, const Tensor q, Tensor k_cache, + Tensor v_cache, Tensor out) { + const auto bound = FlashAttnWithKvcacheConfig::BoundFrom(config); + return supports(q, k_cache, v_cache, out) && + (bound == 0 || + (bound <= std::numeric_limits::max() - 127 && + static_cast(bound) == k_cache.size(1))); + } + + template + Operator(INFINI_NATIVE_KVCACHE_PARAMS(Lengths)) + : Operator(Config{}, INFINI_NATIVE_KVCACHE_VALUES) {} + + template + Operator(const Config& config, INFINI_NATIVE_KVCACHE_PARAMS(Lengths)) + : FlashAttnWithKvcache{ + checked_q(config, INFINI_NATIVE_KVCACHE_VALUES), k_cache, v_cache, k, v, + rotary_cos, rotary_sin, cache_seqlens, cache_batch_idx, cache_leftpad, + block_table, alibi_slopes, softmax_scale, causal, window_size, + softcap, rotary_interleaved, num_splits, return_softmax_lse, out, + softmax_lse} { + if constexpr (std::is_same_v>) { + prepare(q, k_cache, cache_seqlens, block_table, softmax_scale, causal, + window_size, num_splits, out, + FlashAttnWithKvcacheConfig::BoundFrom(config)); + } + } + + Operator(const Tensor q, Tensor k_cache, Tensor v_cache, Tensor out) + : Operator(q, k_cache, v_cache, std::nullopt, std::nullopt, std::nullopt, + std::nullopt, std::optional{}, std::nullopt, + std::nullopt, std::nullopt, std::nullopt, std::nullopt, false, + {-1, -1}, 0.0, true, 0, false, out, std::nullopt) {} + Operator(const Config& config, const Tensor q, Tensor k_cache, + Tensor v_cache, Tensor out) + : Operator(config, q, k_cache, v_cache, std::nullopt, std::nullopt, + std::nullopt, std::nullopt, std::optional{}, + std::nullopt, std::nullopt, std::nullopt, std::nullopt, + std::nullopt, false, {-1, -1}, 0.0, true, 0, false, out, + std::nullopt) {} + ~Operator() override; + using FlashAttnWithKvcache::operator(); + + void operator()(INFINI_NATIVE_KVCACHE_PARAMS(int64_t)) const override { + throw std::invalid_argument("native KV Attention requires device length metadata"); + } + void operator()(INFINI_NATIVE_KVCACHE_PARAMS(std::optional)) const override { + checked_q(INFINI_NATIVE_KVCACHE_VALUES); + execute(q, k_cache, v_cache, cache_seqlens, block_table, softmax_scale, + causal, window_size, num_splits, out); + } + + // Graph callers must supply stable external workspace for the graph lifetime. + // The qualified native domain is enforced by supports(), not the GPU type. + bool capture_safe() const override { + // New ordinary head dimensions are SOURCE_ONLY, never capture-qualified. + return q_shape_[3] != 72 && q_shape_[3] != 80; + } + std::size_t workspace_size_in_bytes() const override; + + private: + template + static Tensor checked_q(INFINI_NATIVE_KVCACHE_PARAMS(Lengths)) { + if (!supports(INFINI_NATIVE_KVCACHE_VALUES)) { + throw std::invalid_argument("unsupported native KV Attention configuration"); + } + return q; + } + template + static Tensor checked_q(const Config& config, + INFINI_NATIVE_KVCACHE_PARAMS(Lengths)) { + if (!supports(config, INFINI_NATIVE_KVCACHE_VALUES)) { + throw std::invalid_argument("unsupported native KV Attention logical bound"); + } + return q; + } + static bool supported_layout(const Tensor q, Tensor k_cache, Tensor v_cache, + std::optional lengths, + std::optional blocks, Tensor out); + void prepare(const Tensor q, Tensor cache, std::optional lengths, + std::optional blocks, std::optional scale, + bool causal, const std::vector& window_size, + int64_t splits, Tensor out, int64_t logical_bound); + void execute(const Tensor q, Tensor k_cache, Tensor v_cache, + std::optional lengths, std::optional blocks, + std::optional scale, bool causal, + const std::vector& window_size, int64_t splits, Tensor out) const; + struct Impl; + std::shared_ptr impl_; // Complete type is confined to the provider TU. +}; + +} // namespace infini::ops + +#undef INFINI_NATIVE_KVCACHE_PARAMS +#undef INFINI_NATIVE_KVCACHE_VALUES diff --git a/src/native/cuda/nvidia/ops/flash_attn_with_kvcache/launch.cu.in b/src/native/cuda/nvidia/ops/flash_attn_with_kvcache/launch.cu.in new file mode 100644 index 000000000..7abcedfbf --- /dev/null +++ b/src/native/cuda/nvidia/ops/flash_attn_with_kvcache/launch.cu.in @@ -0,0 +1,9 @@ +#include "flash_fwd_launch_template.h" + +namespace infini::ops::native_flash { +void launch_@_flash_hdim@_@_flash_kind@_@_flash_causal@(infini_flash::Flash_fwd_params& p, + cudaStream_t stream) { + infini_flash::run_mha_fwd_splitkv_dispatch<@_flash_type@, @_flash_hdim@, + @_flash_causal@>(p, stream); +} +} // namespace infini::ops::native_flash diff --git a/src/native/cuda/nvidia/ops/flash_head_dim.h b/src/native/cuda/nvidia/ops/flash_head_dim.h new file mode 100644 index 000000000..7c59de3b0 --- /dev/null +++ b/src/native/cuda/nvidia/ops/flash_head_dim.h @@ -0,0 +1,21 @@ +#pragma once + +namespace infini::ops::native_flash { + +// This intentionally names only the head dimensions qualified by this native +// provider. The pinned FA 2.8.3 kernels round 160 and 224 to the next 64. +inline int head_dim_bin(int dimension) { + switch (dimension) { + case 32: case 64: case 96: case 128: + case 192: case 256: return dimension; + case 160: return 192; + case 224: return 256; + default: return 0; + } +} + +inline int split_block_n(int bin) { + return bin <= 64 ? 256 : (bin <= 128 ? 128 : 64); +} + +} // namespace infini::ops::native_flash diff --git a/src/native/cuda/nvidia/ops/flash_workspace.h b/src/native/cuda/nvidia/ops/flash_workspace.h new file mode 100644 index 000000000..5b5599193 --- /dev/null +++ b/src/native/cuda/nvidia/ops/flash_workspace.h @@ -0,0 +1,104 @@ +#pragma once + +#include + +#include +#include +#include +#include +#include +#include + +namespace infini::ops::native_flash { + +constexpr std::size_t kWorkspaceAlignment = 256; + +inline std::size_t checked_add(std::size_t a, std::size_t b) { + if (b > std::numeric_limits::max() - a) + throw std::overflow_error("native FlashAttention workspace size overflow"); + return a + b; +} + +inline std::size_t align_workspace(std::size_t bytes) { + return checked_add(bytes, kWorkspaceAlignment - 1) & + ~(kWorkspaceAlignment - 1); +} + +struct WorkspaceLayout { + std::size_t size = 0; + std::size_t append(std::size_t bytes) { + const auto offset = align_workspace(size); + size = checked_add(offset, bytes); + return offset; + } + std::size_t required() const { return align_workspace(size); } +}; + +// A stream-ordered eager lease. Graph callers supply stable external storage. +// This object never owns workspace for a prepared operator's lifetime. +class WorkspaceLease { + public: + WorkspaceLease(void* external, std::size_t available, std::size_t required, + cudaStream_t stream, int device) + : data_{external}, stream_{stream} { + if (external || available) { + if (!external || available < required) + throw std::invalid_argument("native FlashAttention workspace is missing or too small"); + if (reinterpret_cast(external) % kWorkspaceAlignment) + throw std::invalid_argument("native FlashAttention workspace is not 256-byte aligned"); + return; + } + cudaStreamCaptureStatus capture{cudaStreamCaptureStatusNone}; + check(cudaStreamIsCapturing(stream, &capture), "cudaStreamIsCapturing"); + if (capture != cudaStreamCaptureStatusNone) + throw std::invalid_argument( + "native FlashAttention graph capture requires stable external workspace"); + const auto status = cudaMallocAsync(&data_, required, stream); + if (status != cudaSuccess) { + int supported = 0; + const auto attr = cudaDeviceGetAttribute( + &supported, cudaDevAttrMemoryPoolsSupported, device); + if (attr == cudaSuccess && !supported) + throw std::runtime_error( + "native FlashAttention eager workspace requires CUDA async memory pools"); + check(status, "cudaMallocAsync"); + } + owned_ = true; + } + + WorkspaceLease(const WorkspaceLease&) = delete; + WorkspaceLease& operator=(const WorkspaceLease&) = delete; + + ~WorkspaceLease() noexcept { + if (!owned_) return; + const auto status = cudaFreeAsync(data_, stream_); + if (status == cudaSuccess) return; + // Failure path only: wait for any submitted work before conventional free. + const auto synchronized = cudaStreamSynchronize(stream_); + const auto freed = synchronized == cudaSuccess ? cudaFree(data_) : synchronized; + std::fprintf(stderr, + "native FlashAttention async release failed: %s; cleanup: %s\n", + cudaGetErrorString(status), cudaGetErrorString(freed)); + } + + void* data() const { return data_; } + + void finish() { + if (!owned_) return; + check(cudaFreeAsync(data_, stream_), "cudaFreeAsync"); + owned_ = false; + data_ = nullptr; + } + + private: + static void check(cudaError_t status, const char* operation) { + if (status != cudaSuccess) + throw std::runtime_error(std::string("native FlashAttention ") + + operation + ": " + cudaGetErrorString(status)); + } + void* data_ = nullptr; + cudaStream_t stream_ = nullptr; + bool owned_ = false; +}; + +} // namespace infini::ops::native_flash diff --git a/src/native/cuda/nvidia/ops/fused_add_rms_norm/kernel.h b/src/native/cuda/nvidia/ops/fused_add_rms_norm/kernel.h index c4664ea8b..fe2e7037a 100644 --- a/src/native/cuda/nvidia/ops/fused_add_rms_norm/kernel.h +++ b/src/native/cuda/nvidia/ops/fused_add_rms_norm/kernel.h @@ -15,6 +15,7 @@ class Operator public: using CudaFusedAddRmsNorm< Runtime>::CudaFusedAddRmsNorm; + bool capture_safe() const override { return true; } }; } // namespace infini::ops diff --git a/src/native/cuda/nvidia/ops/gemm/cublas.h b/src/native/cuda/nvidia/ops/gemm/cublas.h index eefe0b7af..aafd9b7c8 100644 --- a/src/native/cuda/nvidia/ops/gemm/cublas.h +++ b/src/native/cuda/nvidia/ops/gemm/cublas.h @@ -2,6 +2,7 @@ #define INFINI_OPS_NVIDIA_GEMM_CUBLAS_H_ #include "native/cuda/nvidia/blas.h" +#include "native/cuda/nvidia/ops/gemm/handle_owner.h" #include "native/cuda/ops/gemm/blas.h" namespace infini::ops { @@ -10,7 +11,24 @@ template <> class Operator : public BlasGemm> { public: - using BlasGemm>::BlasGemm; + Operator(const Tensor a, const Tensor b, const std::optional c, + std::optional alpha, std::optional beta, + std::optional trans_a, std::optional trans_b, Tensor y) + : BlasGemm>{ + a, b, c, alpha, beta, trans_a, trans_b, y}, + handle_{NvidiaGemmHandle::Acquire(y.device().index())} {} + + Operator(const Tensor a, const Tensor b, Tensor y) + : Operator{a, b, std::nullopt, std::nullopt, std::nullopt, + std::nullopt, std::nullopt, y} {} + + bool capture_safe() const override { return true; } + + protected: + cublasHandle_t GetHandle() const override { return handle_->get(); } + + private: + std::shared_ptr handle_; }; } // namespace infini::ops diff --git a/src/native/cuda/nvidia/ops/gemm/cublaslt.h b/src/native/cuda/nvidia/ops/gemm/cublaslt.h index 774fba1bd..f6de4e90e 100644 --- a/src/native/cuda/nvidia/ops/gemm/cublaslt.h +++ b/src/native/cuda/nvidia/ops/gemm/cublaslt.h @@ -80,7 +80,8 @@ class Operator : public Gemm { cublasLtMatmulDesc_t op_desc{}; auto status = cublasLtMatmulDescCreate( - &op_desc, BlasUtils::GetComputeType(y.dtype()), + &op_desc, BlasUtils::GetComputeType( + y.dtype(), fp32_precision(a, b, y)), CUDA_R_32F); assert(status == CUBLAS_STATUS_SUCCESS && "failed to create cuBLASLt matmul descriptor"); diff --git a/src/native/cuda/nvidia/ops/gemm/handle_owner.h b/src/native/cuda/nvidia/ops/gemm/handle_owner.h new file mode 100644 index 000000000..46bdbb98c --- /dev/null +++ b/src/native/cuda/nvidia/ops/gemm/handle_owner.h @@ -0,0 +1,111 @@ +#ifndef INFINI_OPS_NVIDIA_GEMM_HANDLE_OWNER_H_ +#define INFINI_OPS_NVIDIA_GEMM_HANDLE_OWNER_H_ + +#include +#include + +#include +#include +#include +#include +#include +#include + +namespace infini::ops { + +// Prepared GEMMs (including the existing Call cache) own leases. The registry +// only shares handles within a worker/context; it never keeps a handle alive. +// Keep prepared instances alive until their graphs and in-flight work finish. +// The caller must not reset/destroy the CUDA context while such instances live. +class NvidiaGemmHandle { + public: + static std::shared_ptr Acquire(int device) { + const auto context = CurrentContext(); + CUdevice current_device{}; + const auto status = cuCtxGetDevice(¤t_device); + if (status != CUDA_SUCCESS || current_device != device) { + throw std::invalid_argument("GEMM tensor device differs from the active CUDA context"); + } + static thread_local std::map> owners; + for (auto it = owners.begin(); it != owners.end();) { + if (it->second.expired()) it = owners.erase(it); + else ++it; + } + if (auto it = owners.find(context); it != owners.end()) { + if (auto owner = it->second.lock()) return owner; + } + auto owner = std::shared_ptr(new NvidiaGemmHandle(context)); + owners[context] = owner; + return owner; + } + + NvidiaGemmHandle(const NvidiaGemmHandle&) = delete; + NvidiaGemmHandle& operator=(const NvidiaGemmHandle&) = delete; + + ~NvidiaGemmHandle() noexcept { + // Destruction may follow a device switch or occur on a cleanup thread. + // Do not destroy/recreate a handle merely because another device is active. + const auto push = cuCtxPushCurrent(context_); + if (push != CUDA_SUCCESS) { + ReportCleanupError("selecting owning context", static_cast(push)); + return; + } + const auto status = cublasDestroy(handle_); + if (status != CUBLAS_STATUS_SUCCESS) { + ReportCleanupError("destroying handle", static_cast(status)); + } + CUcontext popped{}; + const auto pop = cuCtxPopCurrent(&popped); + if (pop != CUDA_SUCCESS) { + ReportCleanupError("restoring context", static_cast(pop)); + } + } + + cublasHandle_t get() const { + if (worker_ != std::this_thread::get_id()) { + throw std::invalid_argument("GEMM instance used by a different worker thread"); + } + if (CurrentContext() != context_) { + throw std::invalid_argument("GEMM instance used with a different CUDA context"); + } + return handle_; + } + + private: + explicit NvidiaGemmHandle(CUcontext context) + : context_(context), worker_(std::this_thread::get_id()) { + const auto status = cublasCreate(&handle_); + if (status != CUBLAS_STATUS_SUCCESS || handle_ == nullptr) { + if (handle_ != nullptr) { + const auto cleanup = cublasDestroy(handle_); + if (cleanup != CUBLAS_STATUS_SUCCESS) { + ReportCleanupError("cleaning failed creation", static_cast(cleanup)); + } + } + throw std::runtime_error("creating GEMM BLAS handle failed: " + + std::to_string(static_cast(status))); + } + } + + static CUcontext CurrentContext() { + CUcontext context{}; + const auto status = cuCtxGetCurrent(&context); + if (status != CUDA_SUCCESS || context == nullptr) { + throw std::runtime_error("GEMM requires an active CUDA context: " + + std::to_string(static_cast(status))); + } + return context; + } + + static void ReportCleanupError(const char* operation, int status) noexcept { + std::fprintf(stderr, "InfiniOps GEMM cleanup: %s failed (%d)\n", operation, status); + } + + CUcontext context_{}; + std::thread::id worker_; + cublasHandle_t handle_{}; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/nvidia/ops/mul/kernel.h b/src/native/cuda/nvidia/ops/mul/kernel.h index 547dd3270..c936644b1 100644 --- a/src/native/cuda/nvidia/ops/mul/kernel.h +++ b/src/native/cuda/nvidia/ops/mul/kernel.h @@ -14,6 +14,11 @@ class Operator : public CudaMul> { public: using CudaMul>::CudaMul; + bool capture_safe() const override { + return out_type_ == DataType::kFloat16 || + out_type_ == DataType::kBFloat16 || + out_type_ == DataType::kFloat32; + } }; } // namespace infini::ops diff --git a/src/native/cuda/nvidia/ops/reshape_and_cache_flash/kernel.h b/src/native/cuda/nvidia/ops/reshape_and_cache_flash/kernel.h index ffd08ed0f..70b504dc3 100644 --- a/src/native/cuda/nvidia/ops/reshape_and_cache_flash/kernel.h +++ b/src/native/cuda/nvidia/ops/reshape_and_cache_flash/kernel.h @@ -15,6 +15,7 @@ class Operator public: using CudaReshapeAndCacheFlash< Runtime>::CudaReshapeAndCacheFlash; + bool capture_safe() const override { return true; } }; } // namespace infini::ops diff --git a/src/native/cuda/nvidia/ops/rms_norm/kernel.h b/src/native/cuda/nvidia/ops/rms_norm/kernel.h index 74cfab17a..a3df6e4f8 100644 --- a/src/native/cuda/nvidia/ops/rms_norm/kernel.h +++ b/src/native/cuda/nvidia/ops/rms_norm/kernel.h @@ -1,7 +1,7 @@ #ifndef INFINI_OPS_NVIDIA_RMS_NORM_KERNEL_H_ #define INFINI_OPS_NVIDIA_RMS_NORM_KERNEL_H_ -#include +#include #include "native/cuda/nvidia/caster.cuh" #include "native/cuda/nvidia/runtime_.h" @@ -13,7 +13,26 @@ template <> class Operator : public CudaRmsNorm> { public: - using CudaRmsNorm>::CudaRmsNorm; + using Base = CudaRmsNorm>; + + Operator(const Tensor input, const Tensor weight, float eps, Tensor out) + : Base(input, weight, eps, out), + qualified_(input.numel() != 0 && std::isfinite(eps) && eps >= 0 && + input.dtype() == weight.dtype()) {} + + Operator(const Tensor input, const Tensor weight, Tensor out) + : Operator(input, weight, 1e-6f, out) {} + + bool capture_safe() const override { return qualified_; } + + void operator()(const Tensor input, const Tensor weight, float eps, + Tensor out) const override { + if (out.numel() == 0) return; + Base::operator()(input, weight, eps, out); + } + + private: + bool qualified_; }; } // namespace infini::ops diff --git a/src/native/cuda/nvidia/ops/rotary_embedding/kernel.h b/src/native/cuda/nvidia/ops/rotary_embedding/kernel.h index 2ee381c8a..1955484d5 100644 --- a/src/native/cuda/nvidia/ops/rotary_embedding/kernel.h +++ b/src/native/cuda/nvidia/ops/rotary_embedding/kernel.h @@ -13,6 +13,7 @@ class Operator public: using CudaRotaryEmbedding< Runtime>::CudaRotaryEmbedding; + bool capture_safe() const override { return true; } }; } // namespace infini::ops diff --git a/src/native/cuda/nvidia/ops/silu/kernel.h b/src/native/cuda/nvidia/ops/silu/kernel.h index 7f5001f71..52687560f 100644 --- a/src/native/cuda/nvidia/ops/silu/kernel.h +++ b/src/native/cuda/nvidia/ops/silu/kernel.h @@ -14,6 +14,7 @@ class Operator : public CudaSilu> { public: using CudaSilu>::CudaSilu; + bool capture_safe() const override { return true; } }; } // namespace infini::ops diff --git a/src/native/cuda/nvidia/ops/silu_and_mul/kernel.h b/src/native/cuda/nvidia/ops/silu_and_mul/kernel.h index 85d311717..8ca6a3740 100644 --- a/src/native/cuda/nvidia/ops/silu_and_mul/kernel.h +++ b/src/native/cuda/nvidia/ops/silu_and_mul/kernel.h @@ -14,6 +14,7 @@ class Operator : public CudaSiluAndMul> { public: using CudaSiluAndMul>::CudaSiluAndMul; + bool capture_safe() const override { return true; } }; } // namespace infini::ops diff --git a/src/native/cuda/nvidia/ops/softmax/kernel.h b/src/native/cuda/nvidia/ops/softmax/kernel.h index 5c4278fa6..307ed338c 100644 --- a/src/native/cuda/nvidia/ops/softmax/kernel.h +++ b/src/native/cuda/nvidia/ops/softmax/kernel.h @@ -14,6 +14,7 @@ class Operator : public CudaSoftmax> { public: using CudaSoftmax>::CudaSoftmax; + bool capture_safe() const override { return true; } }; } // namespace infini::ops diff --git a/src/native/cuda/nvidia/ops/tril/kernel.h b/src/native/cuda/nvidia/ops/tril/kernel.h new file mode 100644 index 000000000..3f1bd6da3 --- /dev/null +++ b/src/native/cuda/nvidia/ops/tril/kernel.h @@ -0,0 +1,20 @@ +#ifndef INFINI_OPS_NVIDIA_TRIL_KERNEL_H_ +#define INFINI_OPS_NVIDIA_TRIL_KERNEL_H_ + +#include "native/cuda/nvidia/runtime_.h" +#include "native/cuda/ops/triangular/kernel.h" + +namespace infini::ops { + +template <> +class Operator + : public triangular_detail::CudaTriangular< + Runtime, Tril, false> { + public: + using triangular_detail::CudaTriangular< + Runtime, Tril, false>::CudaTriangular; +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_NVIDIA_TRIL_KERNEL_H_ diff --git a/src/native/cuda/nvidia/ops/triu/kernel.h b/src/native/cuda/nvidia/ops/triu/kernel.h new file mode 100644 index 000000000..b3cfa9240 --- /dev/null +++ b/src/native/cuda/nvidia/ops/triu/kernel.h @@ -0,0 +1,20 @@ +#ifndef INFINI_OPS_NVIDIA_TRIU_KERNEL_H_ +#define INFINI_OPS_NVIDIA_TRIU_KERNEL_H_ + +#include "native/cuda/nvidia/runtime_.h" +#include "native/cuda/ops/triangular/kernel.h" + +namespace infini::ops { + +template <> +class Operator + : public triangular_detail::CudaTriangular< + Runtime, Triu, true> { + public: + using triangular_detail::CudaTriangular< + Runtime, Triu, true>::CudaTriangular; +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_NVIDIA_TRIU_KERNEL_H_ diff --git a/src/native/cuda/ops/add/kernel.h b/src/native/cuda/ops/add/kernel.h index 911d706a2..956b1a113 100644 --- a/src/native/cuda/ops/add/kernel.h +++ b/src/native/cuda/ops/add/kernel.h @@ -4,6 +4,9 @@ #include #include #include +#include +#include +#include #include #include "base/add.h" @@ -21,45 +24,73 @@ class CudaAdd : public Add { CudaAdd(const Tensor input, const Tensor other, const double alpha, Tensor out) : Add{input, other, alpha, out} { - size_t shape_size = ndim_ * sizeof(*d_input_shape_); - size_t strides_size = ndim_ * sizeof(*d_input_strides_); + const size_t shape_size = ndim_ * sizeof(*d_input_shape_); + const size_t strides_size = ndim_ * sizeof(*d_input_strides_); const size_t metadata_size = 3 * (shape_size + strides_size); + if (Backend::GetDevice(&owner_device_) != Backend::kSuccess) + throw std::runtime_error("CudaAdd cannot query metadata owner device"); + if (!metadata_size) return; std::vector metadata(metadata_size); - Backend::Malloc((void**)&d_metadata_, metadata_size); - - size_t offset = 0; - d_input_shape_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, input_shape_.data(), shape_size); - offset += shape_size; - - d_other_shape_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, other_shape_.data(), shape_size); - offset += shape_size; - - d_out_shape_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, out_shape_.data(), shape_size); - offset += shape_size; - - d_input_strides_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, input_strides_.data(), strides_size); - offset += strides_size; - - d_other_strides_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, other_strides_.data(), strides_size); - offset += strides_size; - - d_out_strides_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, out_strides_.data(), strides_size); - - Backend::Memcpy(d_metadata_, metadata.data(), metadata_size, - Backend::kMemcpyHostToDevice); + void* allocation = nullptr; + const auto allocation_status = Backend::Malloc(&allocation, metadata_size); + if (allocation_status != Backend::kSuccess || !allocation) { + if (allocation && Backend::Free(allocation) != Backend::kSuccess) + std::terminate(); + throw std::bad_alloc(); + } + d_metadata_ = static_cast(allocation); + + try { + size_t offset = 0; + d_input_shape_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, input_shape_.data(), shape_size); + offset += shape_size; + + d_other_shape_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, other_shape_.data(), shape_size); + offset += shape_size; + + d_out_shape_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, out_shape_.data(), shape_size); + offset += shape_size; + + d_input_strides_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, input_strides_.data(), strides_size); + offset += strides_size; + + d_other_strides_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, other_strides_.data(), strides_size); + offset += strides_size; + + d_out_strides_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, out_strides_.data(), strides_size); + + if (Backend::Memcpy(d_metadata_, metadata.data(), metadata_size, + Backend::kMemcpyHostToDevice) != Backend::kSuccess) + throw std::runtime_error("CudaAdd metadata copy failed"); + } catch (...) { + if (Backend::Free(d_metadata_) != Backend::kSuccess) std::terminate(); + d_metadata_ = nullptr; + throw; + } } CudaAdd(const Tensor input, const Tensor other, Tensor out) : CudaAdd{input, other, 1.0, out} {} - ~CudaAdd() { Backend::Free(d_metadata_); } + ~CudaAdd() { + if (!d_metadata_) return; + int previous = -1; + if (Backend::GetDevice(&previous) != Backend::kSuccess) std::terminate(); + const bool switched = previous != owner_device_; + if (switched && Backend::SetDevice(owner_device_) != Backend::kSuccess) + std::terminate(); + const auto freed = Backend::Free(d_metadata_); + if (switched && Backend::SetDevice(previous) != Backend::kSuccess) + std::terminate(); + if (freed != Backend::kSuccess) std::terminate(); + } void operator()(const Tensor input, const Tensor other, const double alpha, Tensor out) const override { @@ -99,6 +130,7 @@ class CudaAdd : public Add { private: std::byte* d_metadata_{nullptr}; + int owner_device_{-1}; Tensor::Size* d_input_shape_{nullptr}; diff --git a/src/native/cuda/ops/copy/kernel.h b/src/native/cuda/ops/copy/kernel.h index 5caaac264..8a41b34ae 100644 --- a/src/native/cuda/ops/copy/kernel.h +++ b/src/native/cuda/ops/copy/kernel.h @@ -4,6 +4,9 @@ #include #include #include +#include +#include +#include #include #include "base/copy.h" @@ -21,40 +24,68 @@ class CudaCopy : public Copy { public: CudaCopy(const Tensor input, const bool non_blocking, Tensor out) : Copy{input, non_blocking, out} { + if (Backend::GetDevice(&owner_device_) != Backend::kSuccess) + throw std::runtime_error("CudaCopy cannot query metadata owner device"); + if (!output_size_ || (is_input_contiguous_ && is_out_contiguous_)) return; size_t shape_size = ndim_ * sizeof(*d_input_shape_); size_t strides_size = ndim_ * sizeof(*d_input_strides_); const size_t metadata_size = 2 * (shape_size + strides_size); std::vector metadata(metadata_size); - Backend::Malloc((void**)&d_metadata_, metadata_size); - - size_t offset = 0; - d_input_shape_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, input_shape_.data(), shape_size); - offset += shape_size; - - d_out_shape_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, out_shape_.data(), shape_size); - offset += shape_size; - - d_input_strides_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, input_strides_.data(), strides_size); - offset += strides_size; - - d_out_strides_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, out_strides_.data(), strides_size); - - Backend::Memcpy(d_metadata_, metadata.data(), metadata_size, - Backend::kMemcpyHostToDevice); + void* allocation = nullptr; + const auto allocation_status = Backend::Malloc(&allocation, metadata_size); + if (allocation_status != Backend::kSuccess || !allocation) { + if (allocation && Backend::Free(allocation) != Backend::kSuccess) + std::terminate(); + throw std::bad_alloc(); + } + d_metadata_ = static_cast(allocation); + + try { + size_t offset = 0; + d_input_shape_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, input_shape_.data(), shape_size); + offset += shape_size; + + d_out_shape_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, out_shape_.data(), shape_size); + offset += shape_size; + + d_input_strides_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, input_strides_.data(), strides_size); + offset += strides_size; + + d_out_strides_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, out_strides_.data(), strides_size); + + const bool copied = + Backend::Memcpy(d_metadata_, metadata.data(), metadata_size, + Backend::kMemcpyHostToDevice) == Backend::kSuccess; + if (!copied) + throw std::runtime_error("CudaCopy metadata copy failed"); + } catch (...) { + if (Backend::Free(d_metadata_) != Backend::kSuccess) std::terminate(); + d_metadata_ = nullptr; + throw; + } } - ~CudaCopy() { Backend::Free(d_metadata_); } + ~CudaCopy() { + if (!d_metadata_) return; + int previous = -1; + if (Backend::GetDevice(&previous) != Backend::kSuccess) std::terminate(); + const bool switched = previous != owner_device_; + if (switched && Backend::SetDevice(owner_device_) != Backend::kSuccess) + std::terminate(); + const auto freed = Backend::Free(d_metadata_); + if (switched && Backend::SetDevice(previous) != Backend::kSuccess) + std::terminate(); + if (freed != Backend::kSuccess) std::terminate(); + } void operator()(const Tensor input, const bool /*non_blocking*/, Tensor out) const override { - if (output_size_ == 0) { - return; - } + if (output_size_ == 0 || SameView(input, out)) return; auto cuda_stream = static_cast(stream_ ? stream_ : 0); @@ -80,6 +111,7 @@ class CudaCopy : public Copy { private: std::byte* d_metadata_{nullptr}; + int owner_device_{-1}; Tensor::Size* d_input_shape_{nullptr}; diff --git a/src/native/cuda/ops/embedding/kernel.cuh b/src/native/cuda/ops/embedding/kernel.cuh index d80b66564..142537d91 100644 --- a/src/native/cuda/ops/embedding/kernel.cuh +++ b/src/native/cuda/ops/embedding/kernel.cuh @@ -138,13 +138,17 @@ __forceinline__ __device__ void CopyRowContiguous(T* __restrict__ dst, bool aligned_16 = IsAligned(src, 16) && IsAligned(dst, 16); if (aligned_16 && embedding_dim >= 4 && embedding_dim % 4 == 0) { CopyVectorizedFloat4(dst, src, embedding_dim); - } else if (embedding_dim >= 2 && embedding_dim % 2 == 0) { + } else if (embedding_dim >= 2 && embedding_dim % 2 == 0 && + IsAligned(src, alignof(float2)) && + IsAligned(dst, alignof(float2))) { CopyVectorizedFloat2(dst, src, embedding_dim); } else { CopyScalar(dst, src, embedding_dim); } } else if constexpr (IsFP16 || IsBFloat16) { - if (embedding_dim >= 2 && embedding_dim % 2 == 0) { + if (embedding_dim >= 2 && embedding_dim % 2 == 0 && + IsAligned(src, alignof(uint32_t)) && + IsAligned(dst, alignof(uint32_t))) { CopyVectorized16(dst, src, embedding_dim); } else { CopyScalar(dst, src, embedding_dim); diff --git a/src/native/cuda/ops/embedding/kernel.h b/src/native/cuda/ops/embedding/kernel.h index d5d8f7e28..95c081d60 100644 --- a/src/native/cuda/ops/embedding/kernel.h +++ b/src/native/cuda/ops/embedding/kernel.h @@ -4,6 +4,9 @@ #include #include #include +#include +#include +#include #include #include "base/embedding.h" @@ -31,42 +34,58 @@ class CudaEmbedding : public Embedding { is_out_contiguous_{out.IsContiguous()}, weight_row_stride_{weight.stride(0)}, weight_col_stride_{weight.stride(1)} { - size_t input_shape_size = input_ndim_ * sizeof(*d_input_shape_); - size_t input_strides_size = input_ndim_ * sizeof(*d_input_strides_); - size_t out_shape_size = out_ndim_ * sizeof(*d_out_shape_); - size_t out_strides_size = out_ndim_ * sizeof(*d_out_strides_); + if (Backend::GetDevice(&owner_device_) != Backend::kSuccess) + throw std::runtime_error("CudaEmbedding cannot query metadata owner device"); + const size_t input_shape_size = input_ndim_ * sizeof(*d_input_shape_); + const size_t input_strides_size = input_ndim_ * sizeof(*d_input_strides_); + const size_t out_shape_size = out_ndim_ * sizeof(*d_out_shape_); + const size_t out_strides_size = out_ndim_ * sizeof(*d_out_strides_); const size_t metadata_size = input_shape_size + input_strides_size + out_shape_size + out_strides_size; - - std::vector metadata(metadata_size); - - Backend::Malloc((void**)&d_metadata_, metadata_size); - - size_t offset = 0; - d_input_shape_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, input_shape_.data(), - input_shape_size); - offset += input_shape_size; - - d_input_strides_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, input_strides_.data(), - input_strides_size); - offset += input_strides_size; - - d_out_shape_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, out_shape_.data(), out_shape_size); - offset += out_shape_size; - - d_out_strides_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, out_strides_.data(), - out_strides_size); - - Backend::Memcpy(d_metadata_, metadata.data(), metadata_size, - Backend::kMemcpyHostToDevice); - - if (max_norm.has_value() && vocab_size_ > 0) { - Backend::Malloc(reinterpret_cast(&d_visited_), - vocab_size_ * sizeof(*d_visited_)); + try { + if (num_indices_ != 0 && !(is_input_contiguous_ && is_out_contiguous_)) { + std::vector metadata(metadata_size); + void* allocation = nullptr; + const auto status = Backend::Malloc(&allocation, metadata_size); + if (status != Backend::kSuccess || !allocation) { + if (allocation && Backend::Free(allocation) != Backend::kSuccess) + std::terminate(); + throw std::bad_alloc(); + } + d_metadata_ = static_cast(allocation); + size_t offset = 0; + d_input_shape_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, input_shape_.data(), input_shape_size); + offset += input_shape_size; + d_input_strides_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, input_strides_.data(), input_strides_size); + offset += input_strides_size; + d_out_shape_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, out_shape_.data(), out_shape_size); + offset += out_shape_size; + d_out_strides_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, out_strides_.data(), out_strides_size); + if (Backend::Memcpy(d_metadata_, metadata.data(), metadata_size, + Backend::kMemcpyHostToDevice) != Backend::kSuccess) + throw std::runtime_error("CudaEmbedding metadata copy failed"); + } + if (max_norm.has_value() && vocab_size_ > 0) { + void* allocation = nullptr; + const auto status = Backend::Malloc(&allocation, + vocab_size_ * sizeof(*d_visited_)); + if (status != Backend::kSuccess || !allocation) { + if (allocation && Backend::Free(allocation) != Backend::kSuccess) + std::terminate(); + throw std::bad_alloc(); + } + d_visited_ = static_cast(allocation); + } + } catch (...) { + if (d_visited_ && Backend::Free(d_visited_) != Backend::kSuccess) + std::terminate(); + if (d_metadata_ && Backend::Free(d_metadata_) != Backend::kSuccess) + std::terminate(); + throw; } } @@ -83,12 +102,18 @@ class CudaEmbedding : public Embedding { : CudaEmbedding(input, weight, padding_idx, std::nullopt, 2.0, scale_grad_by_freq, sparse, out) {} - ~CudaEmbedding() { - Backend::Free(d_metadata_); - - if (d_visited_) { - Backend::Free(d_visited_); - } + ~CudaEmbedding() override { + if (!d_metadata_ && !d_visited_) return; + int previous = -1; + if (Backend::GetDevice(&previous) != Backend::kSuccess) std::terminate(); + const bool switched = previous != owner_device_; + if (switched && Backend::SetDevice(owner_device_) != Backend::kSuccess) + std::terminate(); + const auto visited_freed = !d_visited_ || Backend::Free(d_visited_) == Backend::kSuccess; + const auto metadata_freed = !d_metadata_ || Backend::Free(d_metadata_) == Backend::kSuccess; + if (switched && Backend::SetDevice(previous) != Backend::kSuccess) + std::terminate(); + if (!visited_freed || !metadata_freed) std::terminate(); } void operator()(const Tensor input, const Tensor weight, @@ -171,6 +196,7 @@ class CudaEmbedding : public Embedding { int* d_visited_{nullptr}; std::byte* d_metadata_{nullptr}; + int owner_device_{-1}; Tensor::Size* d_input_shape_{nullptr}; diff --git a/src/native/cuda/ops/fill/kernel.h b/src/native/cuda/ops/fill/kernel.h index 2dbd4e3ed..1bea676a0 100644 --- a/src/native/cuda/ops/fill/kernel.h +++ b/src/native/cuda/ops/fill/kernel.h @@ -5,6 +5,9 @@ #include #include #include +#include +#include +#include #include #include "base/fill.h" @@ -33,27 +36,55 @@ class CudaFill : public Fill { "`CudaFill` requires `input` and `out` on the same device."); assert(!out.HasBroadcastDim() && "`CudaFill` output must not have broadcast dimensions."); + if (Backend::GetDevice(&owner_device_) != Backend::kSuccess) + throw std::runtime_error("CudaFill cannot query metadata owner device"); + if (output_size_ == 0 || is_out_contiguous_) return; size_t shape_size = ndim_ * sizeof(*d_out_shape_); size_t strides_size = ndim_ * sizeof(*d_out_strides_); const size_t metadata_size = shape_size + strides_size; std::vector metadata(metadata_size); - Backend::Malloc((void**)&d_metadata_, metadata_size); - - size_t offset = 0; - d_out_shape_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, out_shape_.data(), shape_size); - offset += shape_size; - - d_out_strides_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, out_strides_.data(), strides_size); - - Backend::Memcpy(d_metadata_, metadata.data(), metadata_size, - Backend::kMemcpyHostToDevice); + void* allocation = nullptr; + const auto allocation_status = Backend::Malloc(&allocation, metadata_size); + if (allocation_status != Backend::kSuccess || !allocation) { + if (allocation && Backend::Free(allocation) != Backend::kSuccess) + std::terminate(); + throw std::bad_alloc(); + } + d_metadata_ = static_cast(allocation); + + try { + size_t offset = 0; + d_out_shape_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, out_shape_.data(), shape_size); + offset += shape_size; + + d_out_strides_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, out_strides_.data(), strides_size); + + if (Backend::Memcpy(d_metadata_, metadata.data(), metadata_size, + Backend::kMemcpyHostToDevice) != Backend::kSuccess) + throw std::runtime_error("CudaFill metadata copy failed"); + } catch (...) { + if (Backend::Free(d_metadata_) != Backend::kSuccess) std::terminate(); + d_metadata_ = nullptr; + throw; + } } - ~CudaFill() override { Backend::Free(d_metadata_); } + ~CudaFill() override { + if (!d_metadata_) return; + int previous = -1; + if (Backend::GetDevice(&previous) != Backend::kSuccess) std::terminate(); + const bool switched = previous != owner_device_; + if (switched && Backend::SetDevice(owner_device_) != Backend::kSuccess) + std::terminate(); + const auto freed = Backend::Free(d_metadata_); + if (switched && Backend::SetDevice(previous) != Backend::kSuccess) + std::terminate(); + if (freed != Backend::kSuccess) std::terminate(); + } void operator()(const Tensor input, const double value, Tensor out) const override { @@ -101,6 +132,8 @@ class CudaFill : public Fill { std::byte* d_metadata_{nullptr}; + int owner_device_{-1}; + Tensor::Size* d_out_shape_{nullptr}; Tensor::Stride* d_out_strides_{nullptr}; diff --git a/src/native/cuda/ops/gemm/blas.h b/src/native/cuda/ops/gemm/blas.h index ec28a0e1e..3abb42649 100644 --- a/src/native/cuda/ops/gemm/blas.h +++ b/src/native/cuda/ops/gemm/blas.h @@ -2,6 +2,8 @@ #define INFINI_OPS_CUDA_GEMM_BLAS_H_ #include +#include +#include #include "base/gemm.h" #include "native/cuda/blas_utils.h" @@ -37,8 +39,10 @@ class BlasGemm : public Gemm { std::optional alpha, std::optional beta, std::optional trans_a, std::optional trans_b, Tensor y) const override { - Backend::BlasSetStream(GetHandle(), - static_cast(stream_)); + const auto handle = GetHandle(); + CheckStatus(Backend::BlasSetStream( + handle, static_cast(stream_)), + "setting the GEMM stream"); const auto& alpha_value{alpha.value_or(alpha_)}; const auto beta_value{EffectiveBeta(c, beta)}; @@ -49,9 +53,14 @@ class BlasGemm : public Gemm { auto op_b{GetOpB(trans_a_value, trans_b_value)}; const void* alpha_ptr{GetAlphaPtr(alpha_value, y.dtype())}; const void* beta_ptr{GetBetaPtr(beta_value, y.dtype())}; + auto compute_type = BlasUtils::GetComputeType(y.dtype()); + if constexpr (Backend::kDeviceType == Device::Type::kNvidia) { + compute_type = BlasUtils::GetComputeType( + y.dtype(), fp32_precision(a, b, y)); + } - Backend::BlasGemmStridedBatchedEx( - GetHandle(), op_a, op_b, swap_a_and_b_ ? n_ : m_, + CheckStatus(Backend::BlasGemmStridedBatchedEx( + handle, op_a, op_b, swap_a_and_b_ ? n_ : m_, swap_a_and_b_ ? m_ : n_, k_, alpha_ptr, swap_a_and_b_ ? b.data() : a.data(), BlasUtils::GetDataType(swap_a_and_b_ ? b.dtype() @@ -65,8 +74,8 @@ class BlasGemm : public Gemm { swap_a_and_b_ ? batch_stride_a_ : batch_stride_b_, beta_ptr, y.data(), BlasUtils::GetDataType(y.dtype()), ldy_, batch_stride_y_, batch_count_, - BlasUtils::GetComputeType(y.dtype()), - Backend::BLAS_GEMM_DEFAULT); + compute_type, + Backend::BLAS_GEMM_DEFAULT), "submitting GEMM"); } protected: @@ -97,15 +106,30 @@ class BlasGemm : public Gemm { : Backend::BLAS_OP_N; } - static typename Backend::BlasHandle& GetHandle() { + protected: + virtual typename Backend::BlasHandle GetHandle() const { thread_local typename Backend::BlasHandle handle = []() { - typename Backend::BlasHandle h; - Backend::BlasCreate(&h); + typename Backend::BlasHandle h{}; + const auto status = Backend::BlasCreate(&h); + if (status != 0 || h == nullptr) { + if (h != nullptr) Backend::BlasDestroy(h); + throw std::runtime_error("creating GEMM BLAS handle failed: " + + std::to_string(static_cast(status))); + } return h; }(); return handle; } + private: + template + static void CheckStatus(Status status, const char* operation) { + if (status != 0) { + throw std::runtime_error(std::string(operation) + " failed: " + + std::to_string(static_cast(status))); + } + } + bool a_is_col_major_{false}; bool b_is_col_major_{false}; diff --git a/src/native/cuda/ops/mul/kernel.h b/src/native/cuda/ops/mul/kernel.h index 216cf8df6..4c41b6f5d 100644 --- a/src/native/cuda/ops/mul/kernel.h +++ b/src/native/cuda/ops/mul/kernel.h @@ -4,6 +4,9 @@ #include #include #include +#include +#include +#include #include #include "base/mul.h" @@ -22,46 +25,69 @@ class CudaMul : public Mul { size_t shape_size = ndim_ * sizeof(*d_input_shape_); size_t strides_size = ndim_ * sizeof(*d_input_strides_); const size_t metadata_size = 3 * (shape_size + strides_size); + if (Backend::GetDevice(&owner_device_) != Backend::kSuccess) + throw std::runtime_error("CudaMul cannot query metadata owner device"); if (metadata_size == 0) { return; } std::vector metadata(metadata_size); - Backend::Malloc((void**)&d_metadata_, metadata_size); - - size_t offset = 0; - d_input_shape_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, input_shape_.data(), shape_size); - offset += shape_size; - - d_other_shape_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, other_shape_.data(), shape_size); - offset += shape_size; - - d_out_shape_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, out_shape_.data(), shape_size); - offset += shape_size; - - d_input_strides_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, input_strides_.data(), strides_size); - offset += strides_size; - - d_other_strides_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, other_strides_.data(), strides_size); - offset += strides_size; - - d_out_strides_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, out_strides_.data(), strides_size); - - Backend::Memcpy(d_metadata_, metadata.data(), metadata_size, - Backend::kMemcpyHostToDevice); + void* allocation = nullptr; + const auto allocation_status = Backend::Malloc(&allocation, metadata_size); + if (allocation_status != Backend::kSuccess || !allocation) { + if (allocation && Backend::Free(allocation) != Backend::kSuccess) + std::terminate(); + throw std::bad_alloc(); + } + d_metadata_ = static_cast(allocation); + + try { + size_t offset = 0; + d_input_shape_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, input_shape_.data(), shape_size); + offset += shape_size; + + d_other_shape_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, other_shape_.data(), shape_size); + offset += shape_size; + + d_out_shape_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, out_shape_.data(), shape_size); + offset += shape_size; + + d_input_strides_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, input_strides_.data(), strides_size); + offset += strides_size; + + d_other_strides_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, other_strides_.data(), strides_size); + offset += strides_size; + + d_out_strides_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, out_strides_.data(), strides_size); + + if (Backend::Memcpy(d_metadata_, metadata.data(), metadata_size, + Backend::kMemcpyHostToDevice) != Backend::kSuccess) + throw std::runtime_error("CudaMul metadata copy failed"); + } catch (...) { + if (Backend::Free(d_metadata_) != Backend::kSuccess) std::terminate(); + d_metadata_ = nullptr; + throw; + } } ~CudaMul() { - if (d_metadata_ != nullptr) { - Backend::Free(d_metadata_); - } + if (!d_metadata_) return; + int previous = -1; + if (Backend::GetDevice(&previous) != Backend::kSuccess) std::terminate(); + const bool switched = previous != owner_device_; + if (switched && Backend::SetDevice(owner_device_) != Backend::kSuccess) + std::terminate(); + const auto freed = Backend::Free(d_metadata_); + if (switched && Backend::SetDevice(previous) != Backend::kSuccess) + std::terminate(); + if (freed != Backend::kSuccess) std::terminate(); } void operator()(const Tensor input, const Tensor other, @@ -99,6 +125,7 @@ class CudaMul : public Mul { private: std::byte* d_metadata_{nullptr}; + int owner_device_{-1}; Tensor::Size* d_input_shape_{nullptr}; diff --git a/src/native/cuda/ops/silu/kernel.h b/src/native/cuda/ops/silu/kernel.h index f883de566..95de3459d 100644 --- a/src/native/cuda/ops/silu/kernel.h +++ b/src/native/cuda/ops/silu/kernel.h @@ -4,6 +4,9 @@ #include #include #include +#include +#include +#include #include #include "base/silu.h" @@ -20,34 +23,64 @@ template class CudaSilu : public Silu { public: CudaSilu(const Tensor input, Tensor out) : Silu{input, out} { + if (Backend::GetDevice(&owner_device_) != Backend::kSuccess) + throw std::runtime_error("CudaSilu cannot query metadata owner device"); + if (output_size_ == 0 || (is_input_contiguous_ && is_out_contiguous_)) + return; + size_t shape_size = ndim_ * sizeof(*d_input_shape_); size_t strides_size = ndim_ * sizeof(*d_input_strides_); const size_t metadata_size = 2 * (shape_size + strides_size); std::vector metadata(metadata_size); - Backend::Malloc((void**)&d_metadata_, metadata_size); - - size_t offset = 0; - d_input_shape_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, input_shape_.data(), shape_size); - offset += shape_size; - - d_out_shape_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, out_shape_.data(), shape_size); - offset += shape_size; - - d_input_strides_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, input_strides_.data(), strides_size); - offset += strides_size; - - d_out_strides_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, out_strides_.data(), strides_size); - - Backend::Memcpy(d_metadata_, metadata.data(), metadata_size, - Backend::kMemcpyHostToDevice); + void* allocation = nullptr; + const auto allocation_status = Backend::Malloc(&allocation, metadata_size); + if (allocation_status != Backend::kSuccess || !allocation) { + if (allocation && Backend::Free(allocation) != Backend::kSuccess) + std::terminate(); + throw std::bad_alloc(); + } + d_metadata_ = static_cast(allocation); + + try { + size_t offset = 0; + d_input_shape_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, input_shape_.data(), shape_size); + offset += shape_size; + + d_out_shape_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, out_shape_.data(), shape_size); + offset += shape_size; + + d_input_strides_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, input_strides_.data(), strides_size); + offset += strides_size; + + d_out_strides_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, out_strides_.data(), strides_size); + + if (Backend::Memcpy(d_metadata_, metadata.data(), metadata_size, + Backend::kMemcpyHostToDevice) != Backend::kSuccess) + throw std::runtime_error("CudaSilu metadata copy failed"); + } catch (...) { + if (Backend::Free(d_metadata_) != Backend::kSuccess) std::terminate(); + d_metadata_ = nullptr; + throw; + } } - ~CudaSilu() { Backend::Free(d_metadata_); } + ~CudaSilu() { + if (!d_metadata_) return; + int previous = -1; + if (Backend::GetDevice(&previous) != Backend::kSuccess) std::terminate(); + const bool switched = previous != owner_device_; + if (switched && Backend::SetDevice(owner_device_) != Backend::kSuccess) + std::terminate(); + const auto freed = Backend::Free(d_metadata_); + if (switched && Backend::SetDevice(previous) != Backend::kSuccess) + std::terminate(); + if (freed != Backend::kSuccess) std::terminate(); + } void operator()(const Tensor input, Tensor out) const override { if (output_size_ == 0) { @@ -81,6 +114,7 @@ class CudaSilu : public Silu { private: std::byte* d_metadata_{nullptr}; + int owner_device_{-1}; Tensor::Size* d_input_shape_{nullptr}; diff --git a/src/native/cuda/ops/silu_and_mul/kernel.h b/src/native/cuda/ops/silu_and_mul/kernel.h index 0e9b5e620..aa3dbb9c3 100644 --- a/src/native/cuda/ops/silu_and_mul/kernel.h +++ b/src/native/cuda/ops/silu_and_mul/kernel.h @@ -4,6 +4,9 @@ #include #include #include +#include +#include +#include #include #include "base/silu_and_mul.h" @@ -17,37 +20,71 @@ template class CudaSiluAndMul : public SiluAndMul { public: CudaSiluAndMul(const Tensor input, Tensor out) : SiluAndMul{input, out} { + if (Backend::GetDevice(&owner_device_) != Backend::kSuccess) + throw std::runtime_error("CudaSiluAndMul cannot query metadata owner device"); + if (output_size_ == 0 || (is_input_contiguous_ && is_out_contiguous_)) + return; + size_t shape_size = ndim_ * sizeof(*d_input_shape_); size_t strides_size = ndim_ * sizeof(*d_input_strides_); const size_t metadata_size = 2 * (shape_size + strides_size); std::vector metadata(metadata_size); - Backend::Malloc((void**)&d_metadata_, metadata_size); - - size_t offset = 0; - d_input_shape_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, input_shape_.data(), shape_size); - offset += shape_size; - - d_out_shape_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, out_shape_.data(), shape_size); - offset += shape_size; - - d_input_strides_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, input_strides_.data(), strides_size); - offset += strides_size; - - d_out_strides_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, out_strides_.data(), strides_size); - - Backend::Memcpy(d_metadata_, metadata.data(), metadata_size, - Backend::kMemcpyHostToDevice); + void* allocation = nullptr; + const auto allocation_status = Backend::Malloc(&allocation, metadata_size); + if (allocation_status != Backend::kSuccess || !allocation) { + if (allocation && Backend::Free(allocation) != Backend::kSuccess) + std::terminate(); + throw std::bad_alloc(); + } + d_metadata_ = static_cast(allocation); + + try { + size_t offset = 0; + d_input_shape_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, input_shape_.data(), shape_size); + offset += shape_size; + + d_out_shape_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, out_shape_.data(), shape_size); + offset += shape_size; + + d_input_strides_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, input_strides_.data(), strides_size); + offset += strides_size; + + d_out_strides_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, out_strides_.data(), strides_size); + + if (Backend::Memcpy(d_metadata_, metadata.data(), metadata_size, + Backend::kMemcpyHostToDevice) != Backend::kSuccess) + throw std::runtime_error("CudaSiluAndMul metadata copy failed"); + } catch (...) { + if (Backend::Free(d_metadata_) != Backend::kSuccess) std::terminate(); + d_metadata_ = nullptr; + throw; + } } - ~CudaSiluAndMul() { Backend::Free(d_metadata_); } + ~CudaSiluAndMul() { + if (!d_metadata_) return; + int previous = -1; + if (Backend::GetDevice(&previous) != Backend::kSuccess) std::terminate(); + const bool switched = previous != owner_device_; + if (switched && Backend::SetDevice(owner_device_) != Backend::kSuccess) + std::terminate(); + const auto freed = Backend::Free(d_metadata_); + if (switched && Backend::SetDevice(previous) != Backend::kSuccess) + std::terminate(); + if (freed != Backend::kSuccess) std::terminate(); + } void operator()(const Tensor input, Tensor out) const override { + if (output_size_ == 0) { + return; + } + int block_size = RuntimeUtils::GetOptimalBlockSize(); DispatchFunc( {static_cast(out_type_), block_size}, @@ -74,6 +111,7 @@ class CudaSiluAndMul : public SiluAndMul { private: std::byte* d_metadata_{nullptr}; + int owner_device_{-1}; Tensor::Size* d_input_shape_{nullptr}; diff --git a/src/native/cuda/ops/softmax/kernel.h b/src/native/cuda/ops/softmax/kernel.h index b2807ca5b..131c91ca6 100644 --- a/src/native/cuda/ops/softmax/kernel.h +++ b/src/native/cuda/ops/softmax/kernel.h @@ -5,6 +5,9 @@ #include #include #include +#include +#include +#include #include #include "base/softmax.h" @@ -36,6 +39,12 @@ class CudaSoftmax : public Softmax { assert(!out.HasBroadcastDim() && "`Softmax` output must not have broadcasted dimensions"); + if (Backend::GetDevice(&owner_device_) != Backend::kSuccess) + throw std::runtime_error("CudaSoftmax cannot query metadata owner device"); + block_size_ = std::min( + RuntimeUtils::GetOptimalBlockSize(), 1024); + if (row_count_ == 0 || dim_size_ == 0) return; + const size_t shape_size = ndim_ * sizeof(*d_shape_); const size_t strides_size = ndim_ * sizeof(*d_input_strides_); const size_t metadata_size = shape_size + 2 * strides_size; @@ -44,28 +53,49 @@ class CudaSoftmax : public Softmax { } std::vector metadata(metadata_size); - Backend::Malloc(reinterpret_cast(&d_metadata_), metadata_size); - - size_t offset = 0; - d_shape_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, out_shape_.data(), shape_size); - offset += shape_size; - - d_input_strides_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, input_strides_.data(), strides_size); - offset += strides_size; - - d_out_strides_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, out_strides_.data(), strides_size); - - Backend::Memcpy(d_metadata_, metadata.data(), metadata_size, - Backend::kMemcpyHostToDevice); + void* allocation = nullptr; + const auto status = Backend::Malloc(&allocation, metadata_size); + if (status != Backend::kSuccess || !allocation) { + if (allocation && Backend::Free(allocation) != Backend::kSuccess) + std::terminate(); + throw std::bad_alloc(); + } + d_metadata_ = static_cast(allocation); + + try { + size_t offset = 0; + d_shape_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, out_shape_.data(), shape_size); + offset += shape_size; + + d_input_strides_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, input_strides_.data(), strides_size); + offset += strides_size; + + d_out_strides_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, out_strides_.data(), strides_size); + + if (Backend::Memcpy(d_metadata_, metadata.data(), metadata_size, + Backend::kMemcpyHostToDevice) != Backend::kSuccess) + throw std::runtime_error("CudaSoftmax metadata copy failed"); + } catch (...) { + if (Backend::Free(d_metadata_) != Backend::kSuccess) std::terminate(); + d_metadata_ = nullptr; + throw; + } } ~CudaSoftmax() { - if (d_metadata_ != nullptr) { - Backend::Free(d_metadata_); - } + if (!d_metadata_) return; + int previous = -1; + if (Backend::GetDevice(&previous) != Backend::kSuccess) std::terminate(); + const bool switched = previous != owner_device_; + if (switched && Backend::SetDevice(owner_device_) != Backend::kSuccess) + std::terminate(); + const auto freed = Backend::Free(d_metadata_); + if (switched && Backend::SetDevice(previous) != Backend::kSuccess) + std::terminate(); + if (freed != Backend::kSuccess) std::terminate(); } void operator()(const Tensor input, const int64_t, @@ -76,11 +106,9 @@ class CudaSoftmax : public Softmax { auto cuda_stream = static_cast(stream_ ? stream_ : 0); - int block_size = std::min( - RuntimeUtils::GetOptimalBlockSize(), 1024); DispatchFunc>( - {static_cast(out_type_), block_size}, + {static_cast(out_type_), block_size_}, [&](auto list_tag) { using T = TypeMapType(list_tag)>; constexpr int kBlockSize = ListGet<1>(list_tag); @@ -102,6 +130,8 @@ class CudaSoftmax : public Softmax { private: std::byte* d_metadata_{nullptr}; + int owner_device_{-1}; + int block_size_{0}; Tensor::Size* d_shape_{nullptr}; diff --git a/src/native/cuda/ops/triangular/kernel.cuh b/src/native/cuda/ops/triangular/kernel.cuh new file mode 100644 index 000000000..14960e65a --- /dev/null +++ b/src/native/cuda/ops/triangular/kernel.cuh @@ -0,0 +1,73 @@ +#ifndef INFINI_OPS_CUDA_TRIANGULAR_KERNEL_CUH_ +#define INFINI_OPS_CUDA_TRIANGULAR_KERNEL_CUH_ + +#include + +#include +#include + +namespace infini::ops::triangular_detail { + +constexpr int kMaxRank = 8; + +struct Metadata { + int rank{}; + std::size_t shape[kMaxRank]{}; + std::int64_t input_strides[kMaxRank]{}; + std::int64_t output_strides[kMaxRank]{}; + std::size_t rows{}; + std::size_t columns{}; + std::size_t elements{}; + std::int64_t diagonal{}; +}; + +template +__global__ void TriangularKernel(const T* input, T* output, Metadata meta) { + const std::size_t step = static_cast(gridDim.x) * blockDim.x; + for (std::size_t index = static_cast(blockIdx.x) * blockDim.x + + threadIdx.x; + index < meta.elements; index += step) { + const std::size_t column = index % meta.columns; + const std::size_t row = index / meta.columns % meta.rows; + std::int64_t input_offset = 0; + std::int64_t output_offset = 0; + if (meta.rank <= 3) { + const std::size_t batch = meta.rank == 3 + ? index / (meta.columns * meta.rows) + : 0; + input_offset = static_cast(row) * + meta.input_strides[meta.rank - 2] + + static_cast(column) * + meta.input_strides[meta.rank - 1]; + output_offset = static_cast(row) * + meta.output_strides[meta.rank - 2] + + static_cast(column) * + meta.output_strides[meta.rank - 1]; + if (meta.rank == 3) { + input_offset += static_cast(batch) * + meta.input_strides[0]; + output_offset += static_cast(batch) * + meta.output_strides[0]; + } + } else { + std::size_t remaining = index; + for (int dimension = meta.rank - 1; dimension >= 0; --dimension) { + const std::size_t coordinate = remaining % meta.shape[dimension]; + remaining /= meta.shape[dimension]; + input_offset += static_cast(coordinate) * + meta.input_strides[dimension]; + output_offset += static_cast(coordinate) * + meta.output_strides[dimension]; + } + } + const std::int64_t relative = static_cast(column) - + static_cast(row); + const bool keep = kUpper ? relative >= meta.diagonal + : relative <= meta.diagonal; + output[output_offset] = keep ? input[input_offset] : T{}; + } +} + +} // namespace infini::ops::triangular_detail + +#endif // INFINI_OPS_CUDA_TRIANGULAR_KERNEL_CUH_ diff --git a/src/native/cuda/ops/triangular/kernel.h b/src/native/cuda/ops/triangular/kernel.h new file mode 100644 index 000000000..0905b8d82 --- /dev/null +++ b/src/native/cuda/ops/triangular/kernel.h @@ -0,0 +1,198 @@ +#ifndef INFINI_OPS_CUDA_TRIANGULAR_KERNEL_H_ +#define INFINI_OPS_CUDA_TRIANGULAR_KERNEL_H_ + +#include + +#include +#include +#include +#include +#include +#include +#include + +#include "base/tril.h" +#include "base/triu.h" +#include "native/cuda/ops/triangular/kernel.cuh" + +namespace infini::ops::triangular_detail { + +inline std::size_t CheckedElementCount(const Tensor& tensor) { + std::size_t count = 1; + for (const auto extent : tensor.shape()) { + if (extent > static_cast( + std::numeric_limits::max())) { + throw std::invalid_argument("triangular shape exceeds signed indexing range"); + } + if (extent && count > std::numeric_limits::max() / extent) { + throw std::invalid_argument("triangular element count overflows"); + } + count *= extent; + } + return count; +} + +inline std::size_t MaxOffset(const Tensor& tensor) { + std::size_t offset = 0; + for (std::size_t dim = 0; dim < tensor.ndim(); ++dim) { + const auto stride = tensor.strides()[dim]; + if (stride < 0) { + throw std::invalid_argument("triangular negative strides are unsupported"); + } + const std::size_t extent = tensor.shape()[dim]; + const std::size_t positive_stride = static_cast(stride); + if (extent > 1 && positive_stride > + (static_cast( + std::numeric_limits::max()) - offset) / + (extent - 1)) { + throw std::invalid_argument("triangular stride span overflows"); + } + offset += extent ? (extent - 1) * positive_stride : 0; + } + return offset; +} + +inline bool SameView(const Tensor& input, const Tensor& output) { + return input.data() == output.data() && + input.shape() == output.shape() && + input.strides() == output.strides() && + input.dtype() == output.dtype() && + input.device() == output.device(); +} + +inline std::pair ByteSpan( + const Tensor& tensor, std::size_t elements) { + if (!elements) return {0, 0}; + const auto maximum = MaxOffset(tensor); + const auto bytes_per_element = tensor.element_size(); + if (maximum >= std::numeric_limits::max() / + bytes_per_element) { + throw std::invalid_argument("triangular byte span overflows"); + } + const auto bytes = (maximum + 1) * bytes_per_element; + const auto begin = reinterpret_cast(tensor.data()); + if (bytes > std::numeric_limits::max() - begin) { + throw std::invalid_argument("triangular address span overflows"); + } + return {begin, begin + bytes}; +} + +inline void ValidateOutputNonOverlapping(const Tensor& output) { + std::array, kMaxRank> dimensions{}; + std::size_t count = 0; + for (std::size_t dim = 0; dim < output.ndim(); ++dim) { + if (output.shape()[dim] > 1) { + dimensions[count++] = {static_cast(output.strides()[dim]), + output.shape()[dim]}; + } + } + std::sort(dimensions.begin(), dimensions.begin() + count); + std::size_t required = 1; + for (std::size_t dim = 0; dim < count; ++dim) { + const auto [stride, extent] = dimensions[dim]; + if (stride < required || + stride > std::numeric_limits::max() / extent) { + throw std::invalid_argument("triangular output has overlapping strides"); + } + required = stride * extent; + } +} + +template +Metadata Validate(const Tensor& input, std::int64_t diagonal, + const Tensor& output) { + static_assert(std::is_same_v || std::is_same_v); + if (input.ndim() < 2 || input.ndim() > kMaxRank || + input.shape() != output.shape() || input.dtype() != output.dtype() || + input.device() != output.device()) { + throw std::invalid_argument( + "triangular requires matching rank-2 through rank-8 tensors"); + } + const auto element_size = input.element_size(); + if (element_size != 1 && element_size != 2 && element_size != 4 && + element_size != 8) { + throw std::invalid_argument("triangular unsupported element size"); + } + const auto elements = CheckedElementCount(input); + if (elements && (!input.data() || !output.data())) { + throw std::invalid_argument("triangular nonempty tensors require storage"); + } + MaxOffset(input); + MaxOffset(output); + ValidateOutputNonOverlapping(output); + const auto input_span = ByteSpan(input, elements); + const auto output_span = ByteSpan(output, elements); + if (input_span.first < output_span.second && + output_span.first < input_span.second && !SameView(input, output)) { + throw std::invalid_argument( + "triangular input and output partially overlap"); + } + Metadata meta{}; + meta.rank = static_cast(input.ndim()); + meta.rows = input.shape()[input.ndim() - 2]; + meta.columns = input.shape()[input.ndim() - 1]; + meta.elements = elements; + meta.diagonal = diagonal; + for (int dim = 0; dim < meta.rank; ++dim) { + meta.shape[dim] = input.shape()[dim]; + meta.input_strides[dim] = input.strides()[dim]; + meta.output_strides[dim] = output.strides()[dim]; + } + return meta; +} + +template +void Launch(const Tensor& input, Tensor output, Metadata meta, + cudaStream_t stream) { + constexpr unsigned kThreads = 256; + const std::size_t blocks = 1 + (meta.elements - 1) / kThreads; + const unsigned grid = static_cast( + std::min(blocks, 65535)); + TriangularKernel<<>>( + static_cast(input.data()), static_cast(output.data()), + meta); + const auto status = cudaPeekAtLastError(); + if (status != cudaSuccess) { + throw std::runtime_error(cudaGetErrorString(status)); + } +} + +template +class CudaTriangular : public Key { + public: + CudaTriangular(const Tensor input, std::int64_t diagonal, Tensor output) + : Key(input, diagonal, output), + metadata_{Validate(input, diagonal, output)} {} + + bool capture_safe() const override { return true; } + + void operator()(const Tensor input, std::int64_t /*diagonal*/, + Tensor output) const override { + if (metadata_.elements == 0) return; + const auto stream = static_cast( + this->stream_ ? this->stream_ : 0); + switch (input.element_size()) { + case 1: + Launch(input, output, metadata_, stream); + break; + case 2: + Launch(input, output, metadata_, stream); + break; + case 4: + Launch(input, output, metadata_, stream); + break; + case 8: + Launch(input, output, metadata_, stream); + break; + default: + throw std::invalid_argument("triangular unsupported element size"); + } + } + + private: + Metadata metadata_{}; +}; + +} // namespace infini::ops::triangular_detail + +#endif // INFINI_OPS_CUDA_TRIANGULAR_KERNEL_H_ diff --git a/src/operator.h b/src/operator.h index 6e77ac8b8..85803b793 100644 --- a/src/operator.h +++ b/src/operator.h @@ -11,6 +11,7 @@ #include #include #include +#include #include #include #include @@ -28,6 +29,41 @@ namespace infini::ops::detail { +// A key with providers covering different domains may opt into metadata-based +// defaults. Explicit implementation indices and other keys remain unchanged. +template +struct AutomaticDefault : std::false_type {}; + +template +struct AutomaticDefault> + : std::bool_constant {}; + +// Providers may reject unsupported metadata before construction or execution. +// Legacy providers retain their existing domain; this is not capture eligibility. +template +auto ProviderSupports(int, const Args&... args) + -> decltype(Provider::supports(args...), bool{}) { + return Provider::supports(args...); +} + +template +bool ProviderSupports(long, const Args&...) { + return true; +} + +template +class HasConfigAwareSupports { + template + static auto Test(int) + -> decltype(P::supports(std::declval(), + std::declval()...), std::true_type{}); + template + static std::false_type Test(long); + + public: + static constexpr bool value = decltype(Test(0))::value; +}; + struct CacheKey { std::size_t hash; @@ -254,6 +290,14 @@ class OperatorBase { virtual std::size_t workspace_size_in_bytes() const { return 0; } + // Existing providers may treat workspace as optional. Only providers that + // opt in here make an absent/undersized Call workspace an error. + virtual bool requires_workspace() const { return false; } + + // Capture is opt-in for each prepared implementation. In particular, an + // operator's device type alone does not establish resource lifetime safety. + virtual bool capture_safe() const { return false; } + void set_handle(const Handle& handle) { handle_ptr_ = handle.Clone(); } void set_config(const Config& config) { config_ptr_ = config.Clone(); } @@ -343,6 +387,18 @@ class Operator : public OperatorBase { std::as_const(args)...); } + // Metadata eligibility only. This never constructs a provider and does not + // predict allocation success or capture safety. No eligible provider returns + // false; an inactive explicit implementation index throws invalid_argument. + // Providers without supports() retain their legacy, interface-only eligibility. + template + static bool CanMake(const Config& config, const Tensor tensor, + const Args&... args) { + return SelectWithDevice(config, tensor.device().type(), false, + [](auto, auto, auto, auto&...) {}, tensor, + args...); + } + template static void Call(const Handle& handle, const Config& config, const Args&... args) { @@ -411,7 +467,36 @@ class Operator : public OperatorBase { [[maybe_unused]] HostRangeScope host_range_operator_invoke{ HostRangeLayer::kOperatorInvoke}; - return (*op)(handle, args...); + const auto required_workspace = op->workspace_size_in_bytes(); + if (required_workspace == 0) { + return (*op)(handle, args...); + } + // An explicit workspace takes precedence. Preserve legacy optional + // workspace behavior unless this provider declares it mandatory. + if (handle.workspace() != nullptr || handle.workspace_size_in_bytes() != 0) { + if (op->requires_workspace() && + (handle.workspace() == nullptr || + handle.workspace_size_in_bytes() < required_workspace)) { + throw std::invalid_argument("operator workspace is missing or too small"); + } + return (*op)(handle, args...); + } + auto lease = handle.acquire_workspace(required_workspace); + if (!lease.data) { + if (op->requires_workspace()) { + throw std::invalid_argument("operator workspace lease is missing"); + } + return (*op)(handle, args...); + } + if (lease.size_in_bytes < required_workspace) { + throw std::invalid_argument("operator workspace lease is missing or too small"); + } + auto invocation_handle = handle.Clone(); + invocation_handle->set_workspace(lease.data.get()); + invocation_handle->set_workspace_size_in_bytes(lease.size_in_bytes); + // Keep the lease alive through submission. Its owner implements the + // stream-aware deferred release needed after an asynchronous kernel. + return (*op)(*invocation_handle, args...); } template @@ -427,6 +512,16 @@ class Operator : public OperatorBase { return CallReturning(tensor, args...); } + static Config DefaultConfig(Device::Type dev_type) { + const auto default_index = DefaultImplementationIndex(dev_type); + if constexpr (detail::AutomaticDefault::value) { + return Config::Automatic(); + } + Config config; + config.set_implementation_index(default_index); + return config; + } + static std::vector active_implementation_indices( Device::Type dev_type) { if (!detail::ListContains(dev_type, ActiveDevices{})) { @@ -492,13 +587,6 @@ class Operator : public OperatorBase { return default_index; } - static Config DefaultConfig(Device::Type dev_type) { - Config config; - config.set_implementation_index(DefaultImplementationIndex(dev_type)); - - return config; - } - static Config ImplicitConfig(Device::Type dev_type) { if (TuningManager::Instance().IsEnabled()) return Config{}; return DefaultConfig(dev_type); @@ -516,41 +604,115 @@ class Operator : public OperatorBase { static std::unique_ptr MakeWithDevice( const Config& config, Device::Type dispatch_device_type, Args&&... args) { std::unique_ptr op_ptr; + SelectWithDevice( + config, dispatch_device_type, true, + [&](auto device_tag, auto implementation_tag, auto config_aware_tag, + auto&... selected_args) { + constexpr Device::Type kDev = decltype(device_tag)::value; + constexpr std::size_t kImplementationIndex = + decltype(implementation_tag)::value; + using Provider = Operator; + if constexpr (decltype(config_aware_tag)::value) { + op_ptr = std::make_unique(config, selected_args...); + } else { + op_ptr = std::make_unique(selected_args...); + } + op_ptr->set_config(config); + op_ptr->config_ptr_->set_implementation_index(kImplementationIndex); + }, + std::forward(args)...); + + if (!op_ptr) { + throw std::invalid_argument("no active operator implementation supports these parameters"); + } + + return op_ptr; + } + + template + static bool SelectWithDevice(const Config& config, + Device::Type dispatch_device_type, + bool strict, Selected&& selected, + Args&&... args) { + bool matched = false; auto cache_args = std::forward_as_tuple(args...); + if (!detail::ListContains(dispatch_device_type, ActiveDevices{})) { + if (strict) { + throw std::invalid_argument("operator device is not enabled"); + } + return false; + } + DispatchFunc>( dispatch_device_type, [&](auto device_tag) { constexpr Device::Type kDev = decltype(device_tag)::value; - detail::DispatchImplementation( - config.implementation_index(), - [&](auto implementation_tag) { + using Implementations = typename ActiveImplementations::type; + const bool automatic = config.implementation_index() == + Config::kAutomaticImplementation; + const auto construct = [&](auto implementation_tag) { constexpr std::size_t kImplementationIndex = decltype(implementation_tag)::value; - if constexpr (std::is_constructible_v< - Operator, - Args...>) { + using Provider = Operator; + constexpr bool kConfigAware = + std::is_constructible_v && + detail::HasConfigAwareSupports::value; + if constexpr (kConfigAware || + std::is_constructible_v) { + bool supported = false; + if constexpr (kConfigAware) { + supported = std::apply( + [&](const auto&... cached_args) { + return Provider::supports(config, cached_args...); + }, cache_args); + } else if (!config.requires_config_aware_provider()) { + supported = std::apply( + [](const auto&... cached_args) { + return detail::ProviderSupports(0, cached_args...); + }, cache_args); + } + if (!supported) { + if (automatic) return; + if (strict) { + throw std::invalid_argument( + "requested operator implementation does not support these parameters"); + } + return; + } std::apply( [&](auto&... cached_args) { - op_ptr = std::make_unique< - Operator>( - cached_args...); + matched = true; + selected(device_tag, implementation_tag, + std::bool_constant{}, + cached_args...); }, cache_args); } else { - assert(false && - "operator is not implemented for this device and " - "implementation index"); + if (!automatic && strict) { + throw std::invalid_argument("operator implementation has no matching interface"); + } } - }, - "Operator::Make(implementation_index)", - typename ActiveImplementations::type{}); + }; + if (automatic) { + // Selection has no kernel side effects. A constructor failure is + // propagated, never treated as permission to try another provider. + for (const auto index : detail::ListToVector(Implementations{})) { + detail::DispatchImplementation(index, construct, + "Operator::Make(automatic)", Implementations{}); + if (matched) break; + } + } else { + if (!detail::ListContains(config.implementation_index(), Implementations{})) { + throw std::invalid_argument("requested operator implementation is not active"); + } + detail::DispatchImplementation(config.implementation_index(), construct, + "Operator::Make(implementation_index)", Implementations{}); + } }, "Operator::Make"); - op_ptr->set_config(config); - - return op_ptr; + return matched; } static inline std::atomic cache_generation_{0}; diff --git a/tests/test_generate_public_operator_header.py b/tests/test_generate_public_operator_header.py index e7e2656bb..a9a56954e 100644 --- a/tests/test_generate_public_operator_header.py +++ b/tests/test_generate_public_operator_header.py @@ -1,5 +1,6 @@ import importlib.util import pathlib +import re def _load_generator_module(): @@ -29,3 +30,18 @@ def test_public_operator_header_omits_host_range_profiler(tmp_path): assert ( "static void Call(const Handle& handle, const Config& config," in public_header ) + for name in ( + "Make", + "CanMake", + "active_implementation_indices", + "DefaultConfig", + "DefaultImplementationIndex", + "MakeWithDevice", + "SelectWithDevice", + ): + assert re.search(rf"\b{name}\s*\([^{{;]*\)\s*;", public_header) + assert "DispatchFunc>(\n" not in public_header + assert "Config::Automatic()" not in public_header + public_section = public_header.split("class Operator : public OperatorBase {", 1)[1] + public_section = public_section.split(" protected:", 1)[0] + assert "static Config DefaultConfig(Device::Type dev_type);" in public_section diff --git a/tests/test_generate_wrappers.py b/tests/test_generate_wrappers.py index aa9b25366..9ff7148a4 100644 --- a/tests/test_generate_wrappers.py +++ b/tests/test_generate_wrappers.py @@ -106,6 +106,14 @@ def test_operator_call_instantiations_externalize_default_implementation_lookup( assert f"extern template {signature}" in declarations assert f"template {signature}" in definitions + default_config_signature = ( + "Config Operator<::infini::ops::Abs>::DefaultConfig(Device::Type);" + ) + assert declarations.count(f"extern template {default_config_signature}") == 1 + assert definitions.count(f"template {default_config_signature}") == 1 + assert not any("::CanMake<" in entry for entry in declarations) + assert not any("::CanMake<" in entry for entry in definitions) + def test_operator_call_instantiations_externalize_active_implementation_query(): module = _load_generator_module() @@ -371,6 +379,12 @@ class Mul { assert '#include "tuning.h"' in text assert "std::size_t DefaultImplementationIndexForMul" in text + assert ( + "return ::infini::ops::Mul::DefaultConfig(dev_type).implementation_index();" + ) in text + assert "detail::AutomaticDefault<::infini::ops::Mul>::value" not in text + assert "return Config::kAutomaticImplementation;" not in text + assert "return indices.front();" not in text assert ( "config.set_implementation_index(" "DefaultImplementationIndexForMul(DeviceFromPybind11Handle(input).type()))" From 1055dca1ef0ff5e130ceea6cff8fa61dfd9bec6a Mon Sep 17 00:00:00 2001 From: whjthu Date: Tue, 13 Oct 2026 00:05:01 +0800 Subject: [PATCH 02/25] refactor: add general dense and vision operators --- CMakeLists.txt | 11 +- ..._ops_nvidia_dense_sampling_torch_free.json | 18 +- src/CMakeLists.txt | 12 +- src/base/cat.h | 14 ++ src/base/layer_norm.h | 41 ++++ src/base/mrope.h | 28 +++ src/base/mul_scalar.h | 47 +++++ src/base/pad_constant.h | 39 ++++ src/base/quick_gelu.h | 29 +++ src/base/row_l2_normalize.h | 45 +++++ src/base/separate_table_rope.h | 68 +++++++ src/native/cpu/ops/embedding/kernel.h | 100 ++++++++++ src/native/cpu/ops/fill/fill.h | 77 ++++++++ src/native/cpu/ops/mul_scalar/mul_scalar.h | 140 +++++++++++++ .../ops/row_l2_normalize/row_l2_normalize.h | 161 +++++++++++++++ .../cpu/ops/separate_table_rope/kernel.h | 52 +++++ src/native/cpu/ops/separate_table_rope/math.h | 66 +++++++ src/native/cuda/nvidia/ops/cat/kernel.cu | 34 ++++ src/native/cuda/nvidia/ops/cat/kernel.h | 105 ++++++++++ .../cuda/nvidia/ops/convolution/kernel.h | 2 + src/native/cuda/nvidia/ops/gelu/kernel.h | 1 + .../cuda/nvidia/ops/layer_norm/kernel.cu | 89 +++++++++ .../cuda/nvidia/ops/layer_norm/kernel.h | 154 +++++++++++++++ src/native/cuda/nvidia/ops/mrope/kernel.h | 10 + .../cuda/nvidia/ops/mul_scalar/kernel.cu | 97 +++++++++ .../cuda/nvidia/ops/mul_scalar/kernel.h | 145 ++++++++++++++ .../cuda/nvidia/ops/pad_constant/kernel.cu | 71 +++++++ .../cuda/nvidia/ops/pad_constant/kernel.h | 156 +++++++++++++++ .../cuda/nvidia/ops/quick_gelu/kernel.cu | 82 ++++++++ .../cuda/nvidia/ops/quick_gelu/kernel.h | 108 ++++++++++ .../nvidia/ops/separate_table_rope/kernel.h | 24 +++ src/native/cuda/nvidia/ops/sigmoid/kernel.h | 1 + src/native/cuda/nvidia/ops/sum/kernel.cu | 41 ++++ src/native/cuda/nvidia/ops/sum/kernel.h | 157 +++++++++++++++ src/native/cuda/nvidia/ops/take/kernel.cu | 60 ++++++ src/native/cuda/nvidia/ops/take/kernel.h | 146 ++++++++++++++ .../nvidia/ops/upsample_bilinear2d/kernel.cu | 40 ++++ .../nvidia/ops/upsample_bilinear2d/kernel.h | 184 ++++++++++++++++++ src/native/cuda/ops/convolution/kernel.h | 18 +- src/native/cuda/ops/gelu/kernel.h | 80 +++++--- src/native/cuda/ops/mrope/kernel.cuh | 42 ++++ src/native/cuda/ops/mrope/kernel.h | 26 +++ .../cuda/ops/separate_table_rope/kernel.h | 68 +++++++ .../cuda/ops/separate_table_rope/layout.h | 42 ++++ src/native/cuda/ops/sigmoid/kernel.h | 76 +++++--- src/native/cuda/ops/sum/kernel.cuh | 57 ++++++ .../cuda/ops/upsample_bilinear2d/kernel.cuh | 86 ++++++++ 47 files changed, 3093 insertions(+), 57 deletions(-) create mode 100644 src/base/layer_norm.h create mode 100644 src/base/mrope.h create mode 100644 src/base/mul_scalar.h create mode 100644 src/base/pad_constant.h create mode 100644 src/base/quick_gelu.h create mode 100644 src/base/row_l2_normalize.h create mode 100644 src/base/separate_table_rope.h create mode 100644 src/native/cpu/ops/embedding/kernel.h create mode 100644 src/native/cpu/ops/fill/fill.h create mode 100644 src/native/cpu/ops/mul_scalar/mul_scalar.h create mode 100644 src/native/cpu/ops/row_l2_normalize/row_l2_normalize.h create mode 100644 src/native/cpu/ops/separate_table_rope/kernel.h create mode 100644 src/native/cpu/ops/separate_table_rope/math.h create mode 100644 src/native/cuda/nvidia/ops/cat/kernel.cu create mode 100644 src/native/cuda/nvidia/ops/cat/kernel.h create mode 100644 src/native/cuda/nvidia/ops/layer_norm/kernel.cu create mode 100644 src/native/cuda/nvidia/ops/layer_norm/kernel.h create mode 100644 src/native/cuda/nvidia/ops/mrope/kernel.h create mode 100644 src/native/cuda/nvidia/ops/mul_scalar/kernel.cu create mode 100644 src/native/cuda/nvidia/ops/mul_scalar/kernel.h create mode 100644 src/native/cuda/nvidia/ops/pad_constant/kernel.cu create mode 100644 src/native/cuda/nvidia/ops/pad_constant/kernel.h create mode 100644 src/native/cuda/nvidia/ops/quick_gelu/kernel.cu create mode 100644 src/native/cuda/nvidia/ops/quick_gelu/kernel.h create mode 100644 src/native/cuda/nvidia/ops/separate_table_rope/kernel.h create mode 100644 src/native/cuda/nvidia/ops/sum/kernel.cu create mode 100644 src/native/cuda/nvidia/ops/sum/kernel.h create mode 100644 src/native/cuda/nvidia/ops/take/kernel.cu create mode 100644 src/native/cuda/nvidia/ops/take/kernel.h create mode 100644 src/native/cuda/nvidia/ops/upsample_bilinear2d/kernel.cu create mode 100644 src/native/cuda/nvidia/ops/upsample_bilinear2d/kernel.h create mode 100644 src/native/cuda/ops/mrope/kernel.cuh create mode 100644 src/native/cuda/ops/mrope/kernel.h create mode 100644 src/native/cuda/ops/separate_table_rope/kernel.h create mode 100644 src/native/cuda/ops/separate_table_rope/layout.h create mode 100644 src/native/cuda/ops/sum/kernel.cuh create mode 100644 src/native/cuda/ops/upsample_bilinear2d/kernel.cuh diff --git a/CMakeLists.txt b/CMakeLists.txt index 39e17c3d3..cb0b0e54b 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -245,7 +245,10 @@ if(AUTO_DETECT_BACKENDS) endif() endif() -if(WITH_TORCH OR WITH_LINKED) + # A linked backend is not necessarily a Torch backend: TVM-FFI transports + # consume DLPack views and their own runtime. Discover Torch lazily only when + # native Torch wrappers or a selected linked Torch transport need it. +macro(infiniops_find_torch) if(WITH_MARS AND WITH_TORCH AND WITH_LINKED) find_package(Python COMPONENTS Interpreter Development.Embed REQUIRED) else() @@ -273,7 +276,7 @@ if(WITH_TORCH OR WITH_LINKED) if(NOT _torch_result EQUAL 0) message(FATAL_ERROR - "`WITH_TORCH=ON` or `WITH_LINKED=ON` requires an installed `torch` package.") + "The selected Torch backend/transport requires an installed `torch` package.") endif() execute_process( @@ -366,6 +369,10 @@ if(WITH_TORCH OR WITH_LINKED) endif() message(STATUS "Found PyTorch: ${TORCH_INCLUDE_DIRS}") +endmacro() + +if(WITH_TORCH) + infiniops_find_torch() endif() # Only one CUDA-like GPU backend can be enabled at a time. diff --git a/scripts/configs/infiniops_ops_nvidia_dense_sampling_torch_free.json b/scripts/configs/infiniops_ops_nvidia_dense_sampling_torch_free.json index adce808f7..0d442f542 100644 --- a/scripts/configs/infiniops_ops_nvidia_dense_sampling_torch_free.json +++ b/scripts/configs/infiniops_ops_nvidia_dense_sampling_torch_free.json @@ -2,20 +2,36 @@ "add": {"implementations": [0]}, "argmax": {"implementations": [0]}, "cast": {"implementations": [0]}, + "cat": {"implementations": [0]}, + "convolution": {"implementations": [0]}, "copy": {"implementations": [0]}, "embedding": {"implementations": [0]}, "fill": {"implementations": [0]}, "flash_attn_varlen_func": {"implementations": [0]}, "flash_attn_with_kvcache": {"implementations": [0]}, "fused_add_rms_norm": {"implementations": [0]}, + "gelu": {"implementations": [0]}, "gemm": {"implementations": [0]}, + "layer_norm": {"implementations": [0]}, + "mrope": {"implementations": [0]}, "mul": {"implementations": [0]}, + "mul_scalar": {"implementations": [0]}, + "pad_constant": {"implementations": [0]}, + "quick_gelu": {"implementations": [0]}, + "relu": {"implementations": [0]}, "reshape_and_cache_flash": {"implementations": [0]}, "rms_norm": {"implementations": [0]}, "rotary_embedding": {"implementations": [0]}, + "row_l2_normalize": {"implementations": [0]}, + "separate_table_rope": {"implementations": [0]}, + "sigmoid": {"implementations": [0]}, "silu": {"implementations": [0]}, "silu_and_mul": {"implementations": [0]}, "softmax": {"implementations": [0]}, + "sum": {"implementations": [0]}, + "take": {"implementations": [0]}, + "top_k_top_p_sampling_from_logits": {"implementations": [16]}, "tril": {"implementations": [0]}, - "triu": {"implementations": [0]} + "triu": {"implementations": [0]}, + "upsample_bilinear2d": {"implementations": [0]} } diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index f658d6402..dff1c8928 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -688,6 +688,12 @@ set(_infini_ops_linked_uses_torch FALSE) set(_infini_ops_linked_uses_tvm_ffi FALSE) if(WITH_LINKED) + find_package(Python COMPONENTS Interpreter REQUIRED) + if(_TORCH_PYTHON) + set(_linked_python "${_TORCH_PYTHON}") + else() + set(_linked_python "${Python_EXECUTABLE}") + endif() if(NOT DEVICE_LIST) message(FATAL_ERROR "`WITH_LINKED=ON` requires an enabled device backend.") @@ -734,7 +740,7 @@ if(WITH_LINKED) endif() execute_process( - COMMAND "${_TORCH_PYTHON}" ${_linked_resolver_args} + COMMAND "${_linked_python}" ${_linked_resolver_args} WORKING_DIRECTORY "${PROJECT_SOURCE_DIR}" RESULT_VARIABLE _linked_resolution_result ) @@ -757,6 +763,10 @@ if(WITH_LINKED) endif() endforeach() + if(_infini_ops_linked_uses_torch AND NOT WITH_TORCH) + infiniops_find_torch() + endif() + list(REMOVE_DUPLICATES INFINI_OPS_LINKED_LIBRARIES) set(_linked_library_basenames) set(_linked_library_index 0) diff --git a/src/base/cat.h b/src/base/cat.h index b7f2c9f4c..1785ef547 100644 --- a/src/base/cat.h +++ b/src/base/cat.h @@ -9,6 +9,20 @@ namespace infini::ops { class Cat : public Operator { public: + static bool supports(const std::vector& tensors, int64_t dim, Tensor out) { + if (tensors.empty() || dim < -static_cast(out.ndim()) || + dim >= static_cast(out.ndim())) return false; + dim = dim < 0 ? dim + static_cast(out.ndim()) : dim; + Tensor::Size total = 0; + for (const auto& tensor : tensors) { + if (tensor.ndim() != out.ndim() || tensor.dtype() != out.dtype()) return false; + for (Tensor::Size axis = 0; axis < out.ndim(); ++axis) + if (axis != static_cast(dim) && tensor.size(axis) != out.size(axis)) return false; + total += tensor.size(dim); + } + return total == out.size(dim); + } + Cat(const std::vector tensors, const int64_t dim, Tensor out) : out_shape_{out.shape()}, out_strides_{out.strides()}, diff --git a/src/base/layer_norm.h b/src/base/layer_norm.h new file mode 100644 index 000000000..86d2bf034 --- /dev/null +++ b/src/base/layer_norm.h @@ -0,0 +1,41 @@ +#ifndef INFINI_OPS_BASE_LAYER_NORM_H_ +#define INFINI_OPS_BASE_LAYER_NORM_H_ + +#include + +#include "operator.h" + +namespace infini::ops { + +// Inference LayerNorm over the last dimension. The legacy Core training +// intermediates are deliberately not part of this operator's outputs. +class LayerNorm : public Operator { + public: + static constexpr bool kAutomaticDefault = true; + + LayerNorm(const Tensor input, const Tensor weight, + const std::optional bias, float eps, Tensor out) + : input_shape_{input.shape()}, input_strides_{input.strides()}, + weight_shape_{weight.shape()}, weight_strides_{weight.strides()}, + bias_shape_{bias ? Tensor::Shape{bias->shape()} : Tensor::Shape{}}, + bias_strides_{bias ? Tensor::Strides{bias->strides()} : Tensor::Strides{}}, + out_shape_{out.shape()}, out_strides_{out.strides()}, + dtype_{input.dtype()}, device_{input.device()}, has_bias_{bias.has_value()}, + eps_{eps} {} + + virtual void operator()(const Tensor input, const Tensor weight, + const std::optional bias, float eps, + Tensor out) const = 0; + + protected: + Tensor::Shape input_shape_, weight_shape_, bias_shape_, out_shape_; + Tensor::Strides input_strides_, weight_strides_, bias_strides_, out_strides_; + DataType dtype_; + Device device_; + bool has_bias_; + float eps_; +}; + +} // namespace infini::ops + +#endif diff --git a/src/base/mrope.h b/src/base/mrope.h new file mode 100644 index 000000000..026f3c5b6 --- /dev/null +++ b/src/base/mrope.h @@ -0,0 +1,28 @@ +#pragma once +#include +#include +#include "operator.h" +namespace infini::ops { +class Mrope : public Operator { + public: + Mrope(const Tensor q, const Tensor k, const Tensor cos, const Tensor sin, const Tensor p, + int head, int rot, int st, int sh, int sw, bool interleaved, Tensor qo, Tensor ko) { + if (head <= 0 || rot <= 0 || rot > head || rot % 2 || st <= 0 || sh <= 0 || sw <= 0 || int64_t(st)+sh+sw != rot/2) + throw std::invalid_argument("Mrope invalid sections"); + for (const auto &t : {k, cos, sin, p, qo, ko}) + if (t.device() != q.device()) throw std::invalid_argument("Mrope devices differ"); + for (const auto &t : {q, k, qo, ko}) + if (t.ndim() != 2 || t.stride(1) != 1 || t.size(0) != q.size(0) || t.dtype() != q.dtype() || t.size(1)%head) + throw std::invalid_argument("Mrope invalid q/k metadata"); + if (q.shape() != qo.shape() || k.shape() != ko.shape() + || (q.dtype() != DataType::kFloat16 && q.dtype() != DataType::kBFloat16 && q.dtype() != DataType::kFloat32 && q.dtype() != DataType::kFloat64) + || cos.ndim() != 2 || sin.ndim() != 2 || cos.shape() != sin.shape() || cos.size(1) != rot/2 + || cos.stride(1) != 1 || sin.stride(1) != 1 || cos.dtype() != q.dtype() || sin.dtype() != q.dtype() + || (p.dtype() != DataType::kInt32 && p.dtype() != DataType::kInt64) + || !((p.ndim() == 1 && p.size(0) == q.size(0)) || (p.ndim() == 2 && p.size(0) == 3 && p.size(1) == q.size(0)))) + throw std::invalid_argument("Mrope invalid tables/positions/outputs"); + } + virtual void operator()(const Tensor q, const Tensor k, const Tensor cos, const Tensor sin, const Tensor p, + int head, int rot, int st, int sh, int sw, bool interleaved, Tensor qo, Tensor ko) const = 0; +}; +} diff --git a/src/base/mul_scalar.h b/src/base/mul_scalar.h new file mode 100644 index 000000000..f63715a13 --- /dev/null +++ b/src/base/mul_scalar.h @@ -0,0 +1,47 @@ +#ifndef INFINI_OPS_BASE_MUL_SCALAR_H_ +#define INFINI_OPS_BASE_MUL_SCALAR_H_ + +#include +#include + +#include "operator.h" + +namespace infini::ops { + +class MulScalar : public Operator { + public: + static constexpr bool kAutomaticDefault = true; + + MulScalar(const Tensor input, const double alpha, Tensor out) + : input_shape_(input.shape()), input_strides_(input.strides()), + out_shape_(out.shape()), out_strides_(out.strides()), + dtype_(input.dtype()), device_(input.device()), alpha_(alpha) {} + + virtual void operator()(const Tensor input, const double alpha, + Tensor out) const = 0; + + protected: + Tensor::Shape input_shape_, out_shape_; + Tensor::Strides input_strides_, out_strides_; + DataType dtype_; + Device device_; + double alpha_; +}; + +// The scalar is an execution argument, not a device tensor. Hash its exact +// bits so a cached instance cannot substitute +0 for -0 or another NaN payload. +template <> +struct CacheKeyBuilder { + detail::CacheKey operator()(const Config& config, const Tensor input, + const double alpha, Tensor out) const { + std::uint64_t bits; + static_assert(sizeof(bits) == sizeof(alpha)); + std::memcpy(&bits, &alpha, sizeof(bits)); + return detail::CacheKey::Build(config.implementation_index(), input, + bits, out); + } +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_BASE_MUL_SCALAR_H_ diff --git a/src/base/pad_constant.h b/src/base/pad_constant.h new file mode 100644 index 000000000..969df3888 --- /dev/null +++ b/src/base/pad_constant.h @@ -0,0 +1,39 @@ +#ifndef INFINI_OPS_BASE_PAD_CONSTANT_H_ +#define INFINI_OPS_BASE_PAD_CONSTANT_H_ + +#include +#include + +#include "operator.h" + +namespace infini::ops { + +// Constant padding in Torch's trailing-dimension pair order. This is an +// inference primitive, not the old generic reflect/replicate/circular Pad. +class PadConstant : public Operator { + public: + static constexpr bool kAutomaticDefault = true; + + PadConstant(const Tensor input, const std::vector padding, + const double value, Tensor out) + : input_shape_(input.shape()), input_strides_(input.strides()), + out_shape_(out.shape()), out_strides_(out.strides()), + dtype_(input.dtype()), device_(input.device()), + padding_(padding), value_(value) {} + + virtual void operator()(const Tensor input, + const std::vector padding, + const double value, Tensor out) const = 0; + + protected: + Tensor::Shape input_shape_, out_shape_; + Tensor::Strides input_strides_, out_strides_; + DataType dtype_; + Device device_; + std::vector padding_; + double value_; +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_BASE_PAD_CONSTANT_H_ diff --git a/src/base/quick_gelu.h b/src/base/quick_gelu.h new file mode 100644 index 000000000..9f2fe895e --- /dev/null +++ b/src/base/quick_gelu.h @@ -0,0 +1,29 @@ +#ifndef INFINI_OPS_BASE_QUICK_GELU_H_ +#define INFINI_OPS_BASE_QUICK_GELU_H_ + +#include "operator.h" + +namespace infini::ops { + +// Inference y = x * sigmoid(1.702f * x). Native providers specify the +// dtype-specific intermediate rounding; this is not Gelu's none/tanh mode. +class QuickGelu : public Operator { + public: + static constexpr bool kAutomaticDefault = true; + + QuickGelu(const Tensor input, Tensor out) + : input_shape_{input.shape()}, input_strides_{input.strides()}, + out_shape_{out.shape()}, out_strides_{out.strides()}, + dtype_{input.dtype()}, device_{input.device()} {} + + virtual void operator()(const Tensor input, Tensor out) const = 0; + + protected: + Tensor::Shape input_shape_, out_shape_; + Tensor::Strides input_strides_, out_strides_; + DataType dtype_; + Device device_; +}; + +} // namespace infini::ops +#endif diff --git a/src/base/row_l2_normalize.h b/src/base/row_l2_normalize.h new file mode 100644 index 000000000..678a5118b --- /dev/null +++ b/src/base/row_l2_normalize.h @@ -0,0 +1,45 @@ +#ifndef INFINI_OPS_BASE_ROW_L2_NORMALIZE_H_ +#define INFINI_OPS_BASE_ROW_L2_NORMALIZE_H_ + +#include +#include + +#include "operator.h" + +namespace infini::ops { + +class RowL2Normalize : public Operator { + public: + static constexpr bool kAutomaticDefault = true; + + RowL2Normalize(const Tensor input, const double eps, Tensor out) + : input_shape_(input.shape()), input_strides_(input.strides()), + out_shape_(out.shape()), out_strides_(out.strides()), + dtype_(input.dtype()), device_(input.device()), eps_(eps) {} + + virtual void operator()(const Tensor input, const double eps, + Tensor out) const = 0; + + protected: + Tensor::Shape input_shape_, out_shape_; + Tensor::Strides input_strides_, out_strides_; + DataType dtype_; + Device device_; + double eps_; +}; + +template <> +struct CacheKeyBuilder { + detail::CacheKey operator()(const Config& config, const Tensor input, + const double eps, Tensor out) const { + std::uint64_t bits; + static_assert(sizeof(bits) == sizeof(eps)); + std::memcpy(&bits, &eps, sizeof(bits)); + return detail::CacheKey::Build(config.implementation_index(), input, + bits, out); + } +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_BASE_ROW_L2_NORMALIZE_H_ diff --git a/src/base/separate_table_rope.h b/src/base/separate_table_rope.h new file mode 100644 index 000000000..4f6a57c12 --- /dev/null +++ b/src/base/separate_table_rope.h @@ -0,0 +1,68 @@ +#ifndef INFINI_OPS_BASE_SEPARATE_TABLE_ROPE_H_ +#define INFINI_OPS_BASE_SEPARATE_TABLE_ROPE_H_ +#include +#include +#include +#include +#include +#include "data_type.h" +#include "operator.h" + +namespace infini::ops { +// Separate tables, full head rotation, distinct output: NOT packed RotaryEmbedding. +// bool is_neox=false is GPT-J, true is GPT-NeoX. No implicit algorithm fallback. +class SeparateTableRope : public Operator { + public: + static constexpr bool kAutomaticDefault = true; + static bool supports(const Tensor input, const Tensor positions, + const Tensor sin_table, const Tensor cos_table, + bool /*is_neox*/, const Tensor out) { + const auto rank = input.ndim(); + // Preserve the installed Rs4 CPU domain. NVIDIA additionally retains LM's + // independently typed tables and padded per-batch position rows. + const bool nvidia = input.device().type() == Device::Type::kNvidia; + if ((rank != 3 && rank != 4) || out.shape() != input.shape() || + out.dtype() != input.dtype() || cos_table.dtype() != sin_table.dtype() || + (!nvidia && sin_table.dtype() != input.dtype())) return false; + if (input.dtype() != DataType::kFloat16 && input.dtype() != DataType::kBFloat16 && + input.dtype() != DataType::kFloat32 && input.dtype() != DataType::kFloat64) return false; + if (nvidia && (input.dtype() == DataType::kFloat64 || + (sin_table.dtype() != DataType::kFloat16 && sin_table.dtype() != DataType::kBFloat16 && + sin_table.dtype() != DataType::kFloat32))) return false; + if (positions.dtype() != DataType::kInt8 && positions.dtype() != DataType::kInt16 && + positions.dtype() != DataType::kInt32 && positions.dtype() != DataType::kInt64 && + positions.dtype() != DataType::kUInt8 && positions.dtype() != DataType::kUInt16 && + positions.dtype() != DataType::kUInt32 && positions.dtype() != DataType::kUInt64) return false; + if (nvidia && positions.dtype() != DataType::kInt32 && positions.dtype() != DataType::kInt64) return false; + if (input.device() != out.device() || input.device() != positions.device() || + input.device() != sin_table.device() || input.device() != cos_table.device()) return false; + const auto batch = rank == 4 ? input.size(0) : 1; + const auto seq = input.size(rank - 3); + if (positions.ndim() == 1) { + if (positions.size(0) != seq || positions.stride(0) != 1) return false; + } else if (positions.ndim() == 2) { + if (positions.size(0) != batch || positions.size(1) != seq || + positions.stride(1) != 1 || (!nvidia && positions.stride(0) != static_cast(seq))) return false; + } else return false; + if (sin_table.ndim() != 2 || cos_table.shape() != sin_table.shape() || + sin_table.size(1) > std::numeric_limits::max() / 2 || + input.size(rank - 1) != 2 * sin_table.size(1) || + sin_table.stride(1) != 1 || cos_table.stride(1) != 1 || + sin_table.stride(0) != static_cast(sin_table.size(1)) || + cos_table.stride(0) != sin_table.stride(0) || + input.stride(rank - 1) != 1 || out.stride(rank - 1) != 1) return false; + if (nvidia && sin_table.size(0) < seq) return false; + return true; + } + SeparateTableRope(const Tensor input, const Tensor positions, + const Tensor sin_table, const Tensor cos_table, + bool is_neox, Tensor out) { + if (!supports(input, positions, sin_table, cos_table, is_neox, out)) + throw std::invalid_argument("SeparateTableRope: unsupported metadata/alias"); + } + virtual void operator()(const Tensor input, const Tensor positions, + const Tensor sin_table, const Tensor cos_table, + bool is_neox, Tensor out) const = 0; +}; +} // namespace infini::ops +#endif diff --git a/src/native/cpu/ops/embedding/kernel.h b/src/native/cpu/ops/embedding/kernel.h new file mode 100644 index 000000000..fbe8a7641 --- /dev/null +++ b/src/native/cpu/ops/embedding/kernel.h @@ -0,0 +1,100 @@ +#ifndef INFINI_OPS_CPU_EMBEDDING_KERNEL_H_ +#define INFINI_OPS_CPU_EMBEDDING_KERNEL_H_ + +#include +#include +#include +#include +#include "base/embedding.h" + +namespace infini::ops { + +template <> +class Operator : public Embedding { + public: + Operator(const Tensor input, const Tensor weight, + const std::optional padding_idx, + const std::optional max_norm, const double norm_type, + const bool scale_grad_by_freq, const bool sparse, Tensor out) + : Embedding(input, weight, padding_idx, max_norm, norm_type, + scale_grad_by_freq, sparse, out) { + Validate(input, weight, out, max_norm); + } + Operator(const Tensor input, const Tensor weight, Tensor out) + : Operator(input, weight, std::nullopt, std::nullopt, 2.0, false, false, out) {} + Operator(const Tensor input, const Tensor weight, const int64_t padding_idx, + const bool scale_grad_by_freq, const bool sparse, Tensor out) + : Operator(input, weight, padding_idx, std::nullopt, 2.0, + scale_grad_by_freq, sparse, out) {} + + using Embedding::operator(); + + void operator()(const Tensor input, const Tensor weight, + const std::optional padding_idx, + const std::optional max_norm, const double norm_type, + const bool scale_grad_by_freq, const bool sparse, + Tensor out) const override { + if (input.shape() != input_shape_ || input.dtype() != input_dtype_ || + weight.shape() != weight_shape_ || weight.dtype() != weight_dtype_ || + out.shape() != out_shape_ || out.dtype() != out_dtype_) + throw std::invalid_argument("CPU Embedding call metadata differs from Make"); + Validate(input, weight, out, max_norm); + if (out.numel() == 0) return; + std::vector rows(num_indices_); + for (std::size_t i = 0; i < num_indices_; ++i) { + const auto index = ReadIndex(input, i); + if (index >= vocab_size_) + throw std::out_of_range("CPU Embedding index exceeds weight rows"); + rows[i] = static_cast(index); + } + const auto row_bytes = embedding_dim_ * weight.element_size(); + const auto* src = static_cast(weight.data()); + auto* dst = static_cast(out.data()); + for (std::size_t i = 0; i < num_indices_; ++i) + std::memcpy(dst + i * row_bytes, src + rows[i] * row_bytes, row_bytes); + (void)padding_idx; + (void)norm_type; + (void)scale_grad_by_freq; + (void)sparse; + } + + private: + static void Validate(const Tensor input, const Tensor weight, const Tensor out, + const std::optional max_norm) { + if (input.device().type() != Device::Type::kCpu || input.device() != weight.device() || + input.device() != out.device() || max_norm || !input.IsContiguous() || + !weight.IsContiguous() || !out.IsContiguous() || + (out.numel() && (!input.data() || !weight.data() || !out.data()))) + throw std::invalid_argument("CPU Embedding requires contiguous same-device plain lookup"); + const auto overlaps = [](const Tensor a, const Tensor b) { + const auto left = reinterpret_cast(a.data()); + const auto right = reinterpret_cast(b.data()); + const auto a_bytes = a.numel() * a.element_size(); + const auto b_bytes = b.numel() * b.element_size(); + return a_bytes && b_bytes && + (left <= right ? right - left < a_bytes : left - right < b_bytes); + }; + if (overlaps(input, out) || overlaps(weight, out)) + throw std::invalid_argument("CPU Embedding output overlaps an input"); + } + template + static uint64_t LoadIndex(const Tensor input, std::size_t offset) { + T value; + std::memcpy(&value, static_cast(input.data()) + offset * sizeof(T), sizeof(T)); + if constexpr (std::is_signed_v) + if (value < 0) throw std::out_of_range("CPU Embedding index is negative"); + return static_cast(value); + } + static uint64_t ReadIndex(const Tensor input, std::size_t offset) { + switch (input.dtype()) { + case DataType::kInt32: return LoadIndex(input, offset); + case DataType::kInt64: return LoadIndex(input, offset); + case DataType::kUInt32: return LoadIndex(input, offset); + case DataType::kUInt64: return LoadIndex(input, offset); + default: throw std::invalid_argument("CPU Embedding indices must be 32/64-bit integers"); + } + } +}; + +} // namespace infini::ops +#endif diff --git a/src/native/cpu/ops/fill/fill.h b/src/native/cpu/ops/fill/fill.h new file mode 100644 index 000000000..ed2b6ab27 --- /dev/null +++ b/src/native/cpu/ops/fill/fill.h @@ -0,0 +1,77 @@ +#ifndef INFINI_OPS_CPU_FILL_FILL_H_ +#define INFINI_OPS_CPU_FILL_FILL_H_ + +#include "base/fill.h" +#include "native/cpu/caster_.h" +#include +#include + +namespace infini::ops { + +// Native provider for the former LM zeros/ones CPU domain. Scalar values +// outside this domain and Tensor-valued Fill remain explicitly unsupported. +template <> +class Operator : public Fill, + Caster { + public: + static bool supports(const Tensor input, const double value, Tensor out) { + return input.device().type() == Device::Type::kCpu && + (value == 0.0 || value == 1.0) && Fill::supports(input, value, out) && + valid_output(out); + } + + static bool supports(const Tensor, const Tensor, Tensor) { return false; } + + Operator(const Tensor input, const double value, Tensor out) + : Fill{Checked(input, value, out), value, out} {} + + Operator(const Tensor input, const Tensor value, Tensor out) + : Fill{input, value, out} { + throw std::invalid_argument("CPU Fill does not supply a Tensor-valued provider"); + } + + // The former CPU fill boundary qualified host replay and reported true here. + bool capture_safe() const override { return true; } + + void operator()(const Tensor input, const double value, Tensor out) const override { + if (input.shape() != input_shape_ || input.strides() != input_strides_ || + input.dtype() != input_type_ || out.shape() != out_shape_ || + out.strides() != out_strides_ || out.dtype() != out_type_ || + out.device().index() != device_index_ || value != value_ || + !supports(input, value, out)) + throw std::invalid_argument("CPU Fill metadata/value changed after Make"); + if (out.numel() == 0) return; + DispatchFunc(out.dtype(), [&](auto tag) { + using T = typename decltype(tag)::type; + const T scalar = Cast(value); + auto *data = static_cast(out.data()); + // Contiguous fastpath is conditional on the actual output layout. + if (out.IsContiguous()) { + for (size_t index = 0; index < out.numel(); ++index) + std::memcpy(data + index * sizeof(T), &scalar, sizeof(T)); + return; + } + for (size_t linear = 0; linear < out.numel(); ++linear) { + size_t remaining = linear, offset = 0; + for (size_t dim = out.ndim(); dim-- > 0;) { + offset += (remaining % out.size(dim)) * static_cast(out.stride(dim)); + remaining /= out.size(dim); + } + std::memcpy(data + offset * sizeof(T), &scalar, sizeof(T)); + } + }, "CPU Fill zeros/ones"); + } + + void operator()(const Tensor, const Tensor, Tensor) const override { + throw std::invalid_argument("CPU Fill does not supply a Tensor-valued provider"); + } + + private: + static Tensor Checked(const Tensor input, double value, Tensor out) { + if (!supports(input, value, out)) + throw std::invalid_argument("CPU Fill only supplies zero/one scalar fill with valid metadata"); + return input; + } +}; +} // namespace infini::ops +#endif diff --git a/src/native/cpu/ops/mul_scalar/mul_scalar.h b/src/native/cpu/ops/mul_scalar/mul_scalar.h new file mode 100644 index 000000000..2cf8b2af9 --- /dev/null +++ b/src/native/cpu/ops/mul_scalar/mul_scalar.h @@ -0,0 +1,140 @@ +#ifndef INFINI_OPS_CPU_MUL_SCALAR_MUL_SCALAR_H_ +#define INFINI_OPS_CPU_MUL_SCALAR_MUL_SCALAR_H_ + +#include +#include +#include +#include +#include +#include + +#include "base/mul_scalar.h" +#include "native/cpu/caster_.h" + +namespace infini::ops { + +template <> +class Operator : public MulScalar, + Caster { + public: + static bool supports(const Tensor input, const double, Tensor out) { + if (input.device().type() != Device::Type::kCpu || + input.device() != out.device() || input.dtype() != out.dtype() || + input.shape() != out.shape() || !SupportedDtype(input.dtype())) + return false; + std::size_t input_bytes = 0, out_bytes = 0; + if (!ContiguousBytes(input, &input_bytes) || + !ContiguousBytes(out, &out_bytes) || + !ValidAddress(input, input_bytes) || !ValidAddress(out, out_bytes)) + return false; + const auto ip = reinterpret_cast(input.data()); + const auto op = reinterpret_cast(out.data()); + const bool overlap = input_bytes && out_bytes && + ip < op + out_bytes && op < ip + input_bytes; + return !overlap || (ip == op && input.strides() == out.strides()); + } + + Operator(const Tensor input, const double alpha, Tensor out) + : MulScalar{Checked(input, alpha, out), alpha, out} { + std::size_t bytes = 0; + ContiguousBytes(input, &bytes); + count_ = bytes / Width(dtype_); + } + + bool capture_safe() const override { return false; } + + void operator()(const Tensor input, const double alpha, + Tensor out) const override { + if (input.shape() != input_shape_ || out.shape() != out_shape_ || + input.strides() != input_strides_ || out.strides() != out_strides_ || + input.dtype() != dtype_ || out.dtype() != dtype_ || + input.device() != device_ || out.device() != device_ || + std::memcmp(&alpha, &alpha_, sizeof(alpha))) + throw std::invalid_argument( + "CPU MulScalar: call metadata or alpha changed after Make"); + if (!supports(input, alpha, out)) + throw std::invalid_argument( + "CPU MulScalar: call pointer, layout or alias changed after Make"); + if (count_ == 0) return; + // CPU scalar arithmetic retains float opmath precision even for half/BF16. + const float multiplier = FloatAlpha(alpha_); + DispatchFunc>( + dtype_, [&](auto tag) { + using T = typename decltype(tag)::type; + const auto* src = static_cast(input.data()); + auto* dst = static_cast(out.data()); +#pragma omp parallel for + for (std::size_t i = 0; i < count_; ++i) { + const float product = Cast(src[i]) * multiplier; + if constexpr (IsBFloat16) { + // Match CPU scalar BF16 conversion, not the shared caster's + // sign/payload-preserving NaN conversion. + dst[i] = std::isnan(product) ? T{std::uint16_t{0x7FC0}} + : Cast(product); + } else { + dst[i] = Cast(product); + } + } + }, "CPU MulScalar"); + } + + private: + static bool SupportedDtype(DataType type) { + return type == DataType::kFloat32 || type == DataType::kFloat16 || + type == DataType::kBFloat16; + } + + static std::size_t Width(DataType type) { + return type == DataType::kFloat32 ? 4u : 2u; + } + + static bool ContiguousBytes(const Tensor t, std::size_t* bytes) { + if (t.shape().size() != t.strides().size()) return false; + const auto limit = static_cast( + std::numeric_limits::max()) / Width(t.dtype()); + bool empty = false; + for (std::size_t axis = 0; axis < t.ndim(); ++axis) { + if (t.stride(axis) < 0 || t.size(axis) > limit) return false; + empty = empty || t.size(axis) == 0; + } + if (empty) { *bytes = 0; return true; } + std::size_t count = 1; + for (std::size_t axis = t.ndim(); axis > 0; --axis) { + const auto size = t.size(axis - 1); + if ((size > 1 && t.stride(axis - 1) != + static_cast(count)) || + size > limit / count) + return false; + count *= size; + } + *bytes = count * Width(t.dtype()); + return true; + } + + static bool ValidAddress(const Tensor t, std::size_t bytes) { + if (bytes == 0) return true; + const auto start = reinterpret_cast(t.data()); + return start && start % Width(t.dtype()) == 0 && + start <= std::numeric_limits::max() - bytes; + } + + static float FloatAlpha(double alpha) { + if (std::isfinite(alpha) && + std::fabs(alpha) > std::numeric_limits::max()) + return std::copysign(std::numeric_limits::infinity(), alpha); + return static_cast(alpha); + } + + static Tensor Checked(const Tensor input, const double alpha, Tensor out) { + if (!supports(input, alpha, out)) + throw std::invalid_argument( + "CPU MulScalar: unsupported dtype/device/layout or alias"); + return input; + } + + std::size_t count_ = 0; +}; + +} // namespace infini::ops +#endif // INFINI_OPS_CPU_MUL_SCALAR_MUL_SCALAR_H_ diff --git a/src/native/cpu/ops/row_l2_normalize/row_l2_normalize.h b/src/native/cpu/ops/row_l2_normalize/row_l2_normalize.h new file mode 100644 index 000000000..8bfb24eb0 --- /dev/null +++ b/src/native/cpu/ops/row_l2_normalize/row_l2_normalize.h @@ -0,0 +1,161 @@ +#ifndef INFINI_OPS_CPU_ROW_L2_NORMALIZE_H_ +#define INFINI_OPS_CPU_ROW_L2_NORMALIZE_H_ + +#include +#include +#include +#include +#include +#include + +#include "base/row_l2_normalize.h" +#include "native/cpu/caster_.h" + +namespace infini::ops { + +template <> +class Operator + : public RowL2Normalize, Caster { + public: + static bool supports(const Tensor input, const double eps, Tensor out) { + if (input.device().type() != Device::Type::kCpu || + input.device() != out.device() || input.dtype() != out.dtype() || + input.shape() != out.shape() || input.ndim() != 2 || + !SupportedDtype(input.dtype()) || !std::isfinite(eps) || eps < 0) + return false; + std::size_t input_bytes = 0, out_bytes = 0; + if (!ContiguousBytes(input, &input_bytes) || + !ContiguousBytes(out, &out_bytes) || + !ValidAddress(input, input_bytes) || !ValidAddress(out, out_bytes)) + return false; + const auto ip = reinterpret_cast(input.data()); + const auto op = reinterpret_cast(out.data()); + const bool overlap = input_bytes && out_bytes && + ip < op + out_bytes && op < ip + input_bytes; + return !overlap || (ip == op && input.strides() == out.strides()); + } + + Operator(const Tensor input, const double eps, Tensor out) + : RowL2Normalize{Checked(input, eps, out), eps, out} {} + + bool capture_safe() const override { return false; } + + void operator()(const Tensor input, const double eps, + Tensor out) const override { + if (input.shape() != input_shape_ || out.shape() != out_shape_ || + input.strides() != input_strides_ || out.strides() != out_strides_ || + input.dtype() != dtype_ || out.dtype() != dtype_ || + input.device() != device_ || out.device() != device_ || + std::memcmp(&eps, &eps_, sizeof(eps))) + throw std::invalid_argument( + "CPU RowL2Normalize: call metadata or eps changed after Make"); + if (!supports(input, eps, out)) + throw std::invalid_argument( + "CPU RowL2Normalize: call pointer, layout or alias changed after Make"); + if (input.size(0) == 0 || input.size(1) == 0) return; + const float epsilon = eps_ > std::numeric_limits::max() + ? std::numeric_limits::infinity() : static_cast(eps_); + DispatchFunc>( + dtype_, [&](auto tag) { + using T = typename decltype(tag)::type; + const auto* src = static_cast(input.data()); + auto* dst = static_cast(out.data()); + const auto rows = input.size(0), cols = input.size(1); + const float eps_t = Cast(Round(epsilon)); +#pragma omp parallel for + for (std::size_t row = 0; row < rows; ++row) { + const auto offset = row * cols; + const float norm = Cast(Round( + std::sqrt(SumSquares(src + offset, cols)))); + // Preserve NaNs and the norm's T rounding before division. + const float denominator = std::isnan(norm) ? norm + : (norm < eps_t ? eps_t : norm); + for (std::size_t col = 0; col < cols; ++col) { + const float value = Cast(src[offset + col]); + dst[offset + col] = Round(value / denominator); + } + } + }, "CPU RowL2Normalize"); + } + + private: + template + static T Round(float value) { + if constexpr (IsBFloat16) { + // CPU BF16 arithmetic canonicalizes NaNs to a positive quiet NaN. + return std::isnan(value) ? T{std::uint16_t{0x7FC0}} : Cast(value); + } else { + return Cast(value); + } + } + + template + static float SumSquares(const T* src, std::size_t count) { + if (count > 32) { + const auto half = count / 2; + const float left = SumSquares(src, half); + const float right = SumSquares(src + half, count - half); + return left + right; + } + // Bounded leaf storage, never a float copy of a row or checkpoint head. + float partial[32] = {}; + for (std::size_t i = 0; i < count; ++i) { + const float value = Cast(src[i]); + partial[i] = value * value; + } + for (std::size_t step = 1; step < 32; step *= 2) + for (std::size_t i = 0; i < 32; i += 2 * step) + partial[i] = partial[i] + partial[i + step]; + return partial[0]; + } + + static bool SupportedDtype(DataType type) { + return type == DataType::kFloat32 || type == DataType::kFloat16 || + type == DataType::kBFloat16; + } + + static std::size_t Width(DataType type) { + return type == DataType::kFloat32 ? 4u : 2u; + } + + static bool ContiguousBytes(const Tensor t, std::size_t* bytes) { + if (t.shape().size() != t.strides().size()) return false; + const auto limit = static_cast( + std::numeric_limits::max()) / Width(t.dtype()); + bool empty = false; + for (std::size_t axis = 0; axis < t.ndim(); ++axis) { + if (t.stride(axis) < 0 || t.size(axis) > limit) return false; + empty = empty || t.size(axis) == 0; + } + if (empty) { *bytes = 0; return true; } + std::size_t count = 1; + for (std::size_t axis = t.ndim(); axis > 0; --axis) { + const auto size = t.size(axis - 1); + if ((size > 1 && t.stride(axis - 1) != + static_cast(count)) || + size > limit / count) + return false; + count *= size; + } + *bytes = count * Width(t.dtype()); + return true; + } + + static bool ValidAddress(const Tensor t, std::size_t bytes) { + if (bytes == 0) return true; + const auto start = reinterpret_cast(t.data()); + return start && start % Width(t.dtype()) == 0 && + start <= std::numeric_limits::max() - bytes; + } + + static Tensor Checked(const Tensor input, const double eps, Tensor out) { + if (!supports(input, eps, out)) + throw std::invalid_argument( + "CPU RowL2Normalize: unsupported dtype/device/layout, eps or alias"); + return input; + } +}; + +} // namespace infini::ops +#endif // INFINI_OPS_CPU_ROW_L2_NORMALIZE_H_ diff --git a/src/native/cpu/ops/separate_table_rope/kernel.h b/src/native/cpu/ops/separate_table_rope/kernel.h new file mode 100644 index 000000000..9ece98e93 --- /dev/null +++ b/src/native/cpu/ops/separate_table_rope/kernel.h @@ -0,0 +1,52 @@ +#ifndef INFINI_OPS_CPU_SEPARATE_TABLE_ROPE_KERNEL_H_ +#define INFINI_OPS_CPU_SEPARATE_TABLE_ROPE_KERNEL_H_ +#include "base/separate_table_rope.h" +#include "native/cpu/caster_.h" +#include "math.h" +namespace infini::ops { +template<> +class Operator : public SeparateTableRope { + public: + static bool supports(const Tensor x, const Tensor p, const Tensor s, + const Tensor c, bool neox, const Tensor y) { + return x.device().type() == Device::Type::kCpu && + SeparateTableRope::supports(x, p, s, c, neox, y); + } + Operator(const Tensor x, const Tensor p, const Tensor s, const Tensor c, bool neox, Tensor y) + : SeparateTableRope(x, p, s, c, neox, y) { + if (!supports(x, p, s, c, neox, y)) throw std::invalid_argument("SeparateTableRope: CPU only"); + } + void operator()(const Tensor x, const Tensor p, const Tensor s, const Tensor c, + bool neox, Tensor y) const override { + if (x.numel() == 0) return; + DispatchFunc(x.dtype(), [&](auto data_tag) { + using Data = typename decltype(data_tag)::type; + DispatchFunc(p.dtype(), [&](auto index_tag) { + using Index = typename decltype(index_tag)::type; + using Compute = std::conditional_t || IsBFloat16, float, Data>; + const auto rank = x.ndim(); + const separate_rope_cpu::Layout m{rank == 4 ? x.size(0) : 1, x.size(rank - 3), + x.size(rank - 2), x.size(rank - 1) / 2, s.size(0), + rank == 4 ? static_cast(x.stride(0)) : 0, + static_cast(x.stride(rank - 3)), static_cast(x.stride(rank - 2)), + rank == 4 ? static_cast(y.stride(0)) : 0, + static_cast(y.stride(rank - 3)), static_cast(y.stride(rank - 2)), p.ndim() == 2, neox}; + auto load = [](Data v) -> Compute { + if constexpr (IsFP16 || IsBFloat16) return v.ToFloat(); + else return v; + }; + auto store = [](Compute v) -> Data { + if constexpr (IsFP16) return Data{separate_rope_cpu::legacy_f16(v)}; + else if constexpr (IsBFloat16) return Data{separate_rope_cpu::legacy_bf16(v)}; + else return v; + }; + separate_rope_cpu::compute(m, static_cast(y.data()), + static_cast(x.data()), static_cast(p.data()), + static_cast(s.data()), static_cast(c.data()), load, store); + }, "SeparateTableRope CPU positions"); + }, "SeparateTableRope CPU data"); + } + // workspace=0 and capture_safe=false inherited: CPU host path only. +}; +} // namespace infini::ops +#endif diff --git a/src/native/cpu/ops/separate_table_rope/math.h b/src/native/cpu/ops/separate_table_rope/math.h new file mode 100644 index 000000000..83cc94c1c --- /dev/null +++ b/src/native/cpu/ops/separate_table_rope/math.h @@ -0,0 +1,66 @@ +#ifndef INFINI_OPS_CPU_SEPARATE_TABLE_ROPE_MATH_H_ +#define INFINI_OPS_CPU_SEPARATE_TABLE_ROPE_MATH_H_ +#include +#include +#include +#include +#include +namespace infini::ops::separate_rope_cpu { +struct Layout { + std::size_t batch, seq, heads, half, table_rows; + std::size_t xb, xs, xh, yb, ys, yh; + bool per_batch, neox; +}; +// Preserve Core0.2.9 CPU output conversion, not RT's newer half RNE. +inline std::uint16_t legacy_f16(float value) { + std::uint32_t bits; std::memcpy(&bits, &value, 4); + const auto sign = static_cast((bits >> 16) & 0x8000); + const int exponent = static_cast((bits >> 23) & 255) - 127; + auto mantissa = bits & 0x7fffff; + if (exponent >= 16) + return static_cast(sign | ((exponent == 128 && mantissa) ? 0x7e00 : 0x7c00)); + if (exponent >= -14) + return static_cast(sign | ((exponent + 15) << 10) | (mantissa >> 13)); + if (exponent >= -24) { + mantissa |= 0x800000; mantissa >>= (-14 - exponent); + return static_cast(sign | (mantissa >> 13)); + } + return sign; +} +inline std::uint16_t legacy_bf16(float value) { + std::uint32_t bits; std::memcpy(&bits, &value, 4); + return static_cast((bits + 0x7fff + ((bits >> 16) & 1)) >> 16); +} +template +void compute(const Layout& m, Data* y, const Data* x, const Index* positions, + const Data* sin_table, const Data* cos_table, Load load, Store store) { + if (m.batch == 0 || m.seq == 0 || m.heads == 0 || m.half == 0) return; + // Validate all position values before any store; old negative/OOB values were + // unchecked memory access, not supported semantics. No partial mutation. + for (std::size_t i = 0; i < (m.per_batch ? m.batch * m.seq : m.seq); ++i) { + if constexpr (std::is_signed_v) + if (positions[i] < 0) throw std::invalid_argument("SeparateTableRope: negative position"); + if (static_cast(positions[i]) >= m.table_rows) + throw std::invalid_argument("SeparateTableRope: position outside tables"); + } +#ifdef _OPENMP +#pragma omp parallel for +#endif + for (std::size_t b = 0; b < m.batch; ++b) + for (std::size_t h = 0; h < m.heads; ++h) + for (std::size_t t = 0; t < m.seq; ++t) { + const auto xo = b * m.xb + t * m.xs + h * m.xh; + const auto yo = b * m.yb + t * m.ys + h * m.yh; + const auto table = static_cast(positions[(m.per_batch ? b * m.seq : 0) + t]) * m.half; + for (std::size_t i = 0; i < m.half; ++i) { + const auto a = m.neox ? i : 2 * i, z = m.neox ? i + m.half : 2 * i + 1; + // Both pair members and table entries read before either inplace write. + const Compute x0 = load(x[xo + a]), x1 = load(x[xo + z]); + const Compute s = load(sin_table[table + i]), c = load(cos_table[table + i]); + y[yo + a] = store(x0 * c - x1 * s); + y[yo + z] = store(x0 * s + x1 * c); + } + } +} +} // namespace infini::ops::separate_rope_cpu +#endif diff --git a/src/native/cuda/nvidia/ops/cat/kernel.cu b/src/native/cuda/nvidia/ops/cat/kernel.cu new file mode 100644 index 000000000..3c0405fbc --- /dev/null +++ b/src/native/cuda/nvidia/ops/cat/kernel.cu @@ -0,0 +1,34 @@ +#include "native/cuda/nvidia/ops/cat/kernel.h" + +#include +#include +#include + +namespace infini::ops::nvidia_cat_detail { +namespace { +__global__ void CatAxis0Bytes(const std::uint8_t* first, + const std::uint8_t* second, std::uint8_t* out, + std::size_t first_bytes, + std::size_t total_bytes) { + auto offset = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const auto step = static_cast(gridDim.x) * blockDim.x; + for (; offset < total_bytes; offset += step) + out[offset] = offset < first_bytes ? first[offset] + : second[offset - first_bytes]; +} +} // namespace + +void LaunchAxis0Bytes(const void* first, const void* second, void* out, + std::size_t first_bytes, std::size_t total_bytes, + cudaStream_t stream) { + constexpr unsigned kThreads = 256; + const auto blocks = static_cast( + std::min((total_bytes - 1) / kThreads + 1, 65535)); + CatAxis0Bytes<<>>( + static_cast(first), + static_cast(second), + static_cast(out), first_bytes, total_bytes); + const auto status = cudaPeekAtLastError(); + if (status != cudaSuccess) throw std::runtime_error(cudaGetErrorString(status)); +} +} // namespace infini::ops::nvidia_cat_detail diff --git a/src/native/cuda/nvidia/ops/cat/kernel.h b/src/native/cuda/nvidia/ops/cat/kernel.h new file mode 100644 index 000000000..6af6d4ddf --- /dev/null +++ b/src/native/cuda/nvidia/ops/cat/kernel.h @@ -0,0 +1,105 @@ +#ifndef INFINI_OPS_NVIDIA_CAT_KERNEL_H_ +#define INFINI_OPS_NVIDIA_CAT_KERNEL_H_ + +#include + +#include +#include +#include +#include +#include + +#include "base/cat.h" +#include "base/copy.h" +#include "native/cuda/nvidia/ops/copy/kernel.h" + +namespace infini::ops::nvidia_cat_detail { +void LaunchAxis0Bytes(const void* first, const void* second, void* out, + std::size_t first_bytes, std::size_t total_bytes, + cudaStream_t stream); +} // namespace infini::ops::nvidia_cat_detail + +namespace infini::ops { + +template <> +class Operator : public Cat { + public: + static bool supports(const std::vector& tensors, int64_t dim, + Tensor out) { + if (out.device().type() != Device::Type::kNvidia || + !Cat::supports(tensors, dim, out)) return false; + const auto normalized = dim < 0 ? dim + static_cast(out.ndim()) : dim; + if (FastAxis0(tensors, normalized, out)) return true; + return true; + } + + Operator(const std::vector tensors, int64_t dim, Tensor out) + : Cat(tensors, dim, out) { + if (!supports(tensors, dim, out)) + throw std::invalid_argument("NVIDIA Cat metadata or Copy provider unsupported"); + fast_axis0_ = FastAxis0(tensors, dim_, out); + if (fast_axis0_) { + first_bytes_ = tensors[0].numel() * out.element_size(); + total_bytes_ = out.numel() * out.element_size(); + capture_safe_ = true; + return; + } + const auto slices = Slices(tensors, dim_, out); + copies_.reserve(tensors.size()); + for (size_t i = 0; i < tensors.size(); ++i) + copies_.push_back(Copy::Make(Config::Automatic(), tensors[i], false, slices[i])); + capture_safe_ = !copies_.empty(); + for (const auto& copy : copies_) capture_safe_ &= copy->capture_safe(); + } + + bool capture_safe() const override { return capture_safe_; } + + void operator()(const std::vector tensors, int64_t dim, + Tensor out) const override { + if (fast_axis0_) { + nvidia_cat_detail::LaunchAxis0Bytes( + tensors[0].data(), tensors[1].data(), out.data(), first_bytes_, + total_bytes_, static_cast(stream_ ? stream_ : 0)); + return; + } + const auto slices = Slices(tensors, dim_, out); + for (size_t i = 0; i < tensors.size(); ++i) + (*copies_[i])(*handle_ptr_, tensors[i], false, slices[i]); + } + + private: + static bool FastAxis0(const std::vector& tensors, int64_t dim, + Tensor out) { + return dim == 0 && tensors.size() == 2 && out.numel() != 0 && + tensors[0].numel() != 0 && tensors[1].numel() != 0 && + out.IsContiguous() && tensors[0].IsContiguous() && + tensors[1].IsContiguous(); + } + + static std::vector Slices(const std::vector& tensors, + int64_t dim, Tensor out) { + std::vector slices; + slices.reserve(tensors.size()); + size_t offset = 0; + for (const auto& tensor : tensors) { + void* data = out.data(); + if (tensor.numel()) + data = static_cast(data) + + offset * static_cast(out.stride(dim)) * out.element_size(); + slices.emplace_back(data, tensor.shape(), out.dtype(), out.device(), + out.strides()); + offset += tensor.size(dim); + } + return slices; + } + + std::vector>> copies_; + bool capture_safe_{false}; + bool fast_axis0_{false}; + size_t first_bytes_{0}; + size_t total_bytes_{0}; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/nvidia/ops/convolution/kernel.h b/src/native/cuda/nvidia/ops/convolution/kernel.h index b411554b3..39cccf833 100644 --- a/src/native/cuda/nvidia/ops/convolution/kernel.h +++ b/src/native/cuda/nvidia/ops/convolution/kernel.h @@ -14,6 +14,8 @@ class Operator : public CudaConv, Convolution> { public: using CudaConv, Convolution>::CudaConv; + + bool capture_safe() const override { return true; } }; } // namespace infini::ops diff --git a/src/native/cuda/nvidia/ops/gelu/kernel.h b/src/native/cuda/nvidia/ops/gelu/kernel.h index fd8340bd4..18aea8657 100644 --- a/src/native/cuda/nvidia/ops/gelu/kernel.h +++ b/src/native/cuda/nvidia/ops/gelu/kernel.h @@ -14,6 +14,7 @@ class Operator : public CudaGelu> { public: using CudaGelu>::CudaGelu; + bool capture_safe() const override { return true; } }; } // namespace infini::ops diff --git a/src/native/cuda/nvidia/ops/layer_norm/kernel.cu b/src/native/cuda/nvidia/ops/layer_norm/kernel.cu new file mode 100644 index 000000000..d6ce6f10e --- /dev/null +++ b/src/native/cuda/nvidia/ops/layer_norm/kernel.cu @@ -0,0 +1,89 @@ +#include "native/cuda/nvidia/ops/layer_norm/kernel.h" + +#include +#include +#include +#include + +#include + +namespace infini::ops::nvidia_layer_norm_detail { +namespace { +constexpr unsigned kThreads = 256; + +template +__global__ void LayerNormKernel(const T* input, const T* weight, const T* bias, + T* out, Layout layout, float eps) { + std::size_t row = blockIdx.x; + std::size_t input_offset = 0; + std::size_t out_offset = 0; + for (unsigned i = layout.rank - 1; i > 0; --i) { + const auto axis = i - 1; + const auto coordinate = row % layout.shape[axis]; + row /= layout.shape[axis]; + input_offset += coordinate * layout.input_strides[axis]; + out_offset += coordinate * layout.out_strides[axis]; + } + + using Reduce = cub::BlockReduce; + __shared__ typename Reduce::TempStorage reduce_storage; + __shared__ float mean; + __shared__ float inv_std; + float partial = 0.0f; + for (std::size_t d = threadIdx.x; d < layout.dim; d += kThreads) + partial += static_cast(input[input_offset + + d * layout.input_strides[layout.rank - 1]]); + const float sum = Reduce(reduce_storage).Sum(partial); + if (threadIdx.x == 0) mean = sum / static_cast(layout.dim); + __syncthreads(); + + partial = 0.0f; + for (std::size_t d = threadIdx.x; d < layout.dim; d += kThreads) { + const float centered = static_cast(input[input_offset + + d * layout.input_strides[layout.rank - 1]]) - mean; + partial += centered * centered; + } + const float variance_sum = Reduce(reduce_storage).Sum(partial); + if (threadIdx.x == 0) + inv_std = rsqrtf(variance_sum / static_cast(layout.dim) + eps); + __syncthreads(); + + for (std::size_t d = threadIdx.x; d < layout.dim; d += kThreads) { + const float x = static_cast(input[input_offset + + d * layout.input_strides[layout.rank - 1]]); + const float scale = static_cast(weight[d * layout.weight_stride]); + const float shift = bias ? static_cast(bias[d * layout.bias_stride]) : 0.0f; + out[out_offset + d * layout.out_strides[layout.rank - 1]] = + static_cast((x - mean) * inv_std * scale + shift); + } +} + +template +void Submit(const void* input, const void* weight, const void* bias, void* out, + Layout layout, std::size_t rows, float eps, cudaStream_t stream) { + LayerNormKernel<<(rows), kThreads, 0, stream>>>( + static_cast(input), static_cast(weight), + static_cast(bias), static_cast(out), layout, eps); +} +} // namespace + +void Launch(const void* input, const void* weight, const void* bias, void* out, + DataType dtype, Layout layout, std::size_t rows, + float eps, cudaStream_t stream) { + switch (dtype) { + case DataType::kFloat16: + Submit<__half>(input, weight, bias, out, layout, rows, eps, stream); + break; + case DataType::kBFloat16: + Submit<__nv_bfloat16>(input, weight, bias, out, layout, rows, eps, stream); + break; + case DataType::kFloat32: + Submit(input, weight, bias, out, layout, rows, eps, stream); + break; + default: + throw std::invalid_argument("NVIDIA LayerNorm unsupported dtype"); + } + const auto status = cudaPeekAtLastError(); + if (status != cudaSuccess) throw std::runtime_error(cudaGetErrorString(status)); +} +} // namespace infini::ops::nvidia_layer_norm_detail diff --git a/src/native/cuda/nvidia/ops/layer_norm/kernel.h b/src/native/cuda/nvidia/ops/layer_norm/kernel.h new file mode 100644 index 000000000..ba146e9b3 --- /dev/null +++ b/src/native/cuda/nvidia/ops/layer_norm/kernel.h @@ -0,0 +1,154 @@ +#ifndef INFINI_OPS_NVIDIA_LAYER_NORM_KERNEL_H_ +#define INFINI_OPS_NVIDIA_LAYER_NORM_KERNEL_H_ + +#include + +#include +#include +#include +#include +#include +#include + +#include "base/layer_norm.h" +#include "native/cuda/nvidia/ops/binary_view_capture.h" +#include "native/cuda/nvidia/runtime_.h" + +namespace infini::ops { +namespace nvidia_layer_norm_detail { +struct Layout { + std::size_t shape[8]; + std::size_t input_strides[8]; + std::size_t out_strides[8]; + std::size_t weight_stride; + std::size_t bias_stride; + std::size_t dim; + unsigned rank; +}; + +void Launch(const void* input, const void* weight, const void* bias, void* out, + DataType dtype, Layout layout, std::size_t rows, + float eps, cudaStream_t stream); +} // namespace nvidia_layer_norm_detail + +template <> +class Operator : public LayerNorm { + public: + static bool supports(const Tensor input, const Tensor weight, + const std::optional bias, float, Tensor out) { + using namespace nvidia_binary_view; + const auto rank = input.shape().size(); + if (input.device().type() != Device::Type::kNvidia || + input.device() != weight.device() || input.device() != out.device() || + (bias && input.device() != bias->device()) || + rank < 1 || rank > 8 || rank != input.strides().size() || + out.shape() != input.shape() || out.strides().size() != rank || + weight.shape().size() != 1 || weight.strides().size() != 1 || + (bias && (bias->shape().size() != 1 || bias->strides().size() != 1)) || + input.shape()[rank - 1] == 0 || + weight.shape()[0] != input.shape()[rank - 1] || + (bias && bias->shape()[0] != input.shape()[rank - 1]) || + input.dtype() != weight.dtype() || input.dtype() != out.dtype() || + (bias && bias->dtype() != input.dtype()) || + (input.dtype() != DataType::kFloat16 && + input.dtype() != DataType::kBFloat16 && + input.dtype() != DataType::kFloat32) || + (input.shape()[rank - 1] > 1 && + (input.strides()[rank - 1] != 1 || out.strides()[rank - 1] != 1)) || + !NonOverlappingOutput(out)) return false; + + int current = -1; + if (cudaGetDevice(¤t) != cudaSuccess || + current != input.device().index()) return false; + if (input.dtype() == DataType::kBFloat16) { + cudaDeviceProp prop{}; + if (cudaGetDeviceProperties(&prop, current) != cudaSuccess || + prop.major < 8) return false; + } + std::size_t rows = 1; + for (std::size_t i = 0; i + 1 < rank; ++i) { + if (input.shape()[i] && + rows > static_cast(std::numeric_limits::max()) / + input.shape()[i]) + return false; + rows *= input.shape()[i]; + } + if (rows > static_cast(std::numeric_limits::max()) || + input.shape()[rank - 1] > + static_cast(std::numeric_limits::max()) || + weight.strides()[0] < 0 || + (bias && bias->strides()[0] < 0)) return false; + std::size_t input_bytes, weight_bytes, out_bytes, bias_bytes = 0; + if (!SpanBytes(input, &input_bytes) || !SpanBytes(weight, &weight_bytes) || + !SpanBytes(out, &out_bytes) || + (bias && !SpanBytes(*bias, &bias_bytes))) return false; + const auto width = input.element_size(); + if ((input_bytes && reinterpret_cast(input.data()) % width) || + (weight_bytes && reinterpret_cast(weight.data()) % width) || + (out_bytes && reinterpret_cast(out.data()) % width) || + (bias_bytes && reinterpret_cast(bias->data()) % width)) + return false; + const bool exact_inplace = input.data() == out.data() && + input.strides() == out.strides() && NonOverlappingOutput(input); + return !Overlaps(out, out_bytes, weight, weight_bytes) && + (!bias || !Overlaps(out, out_bytes, *bias, bias_bytes)) && + (!Overlaps(out, out_bytes, input, input_bytes) || exact_inplace); + } + + Operator(const Tensor input, const Tensor weight, + const std::optional bias, float eps, Tensor out) + : LayerNorm(checked(input, weight, bias, eps, out), weight, bias, eps, out) {} + + // Keep graph qualification closed until the independent CUDA/RT contract + // has run on the target architecture. Direct capture is tested separately. + bool capture_safe() const override { return false; } + + void operator()(const Tensor input, const Tensor weight, + const std::optional bias, float eps, + Tensor out) const override { + if (input.shape() != input_shape_ || input.strides() != input_strides_ || + weight.shape() != weight_shape_ || weight.strides() != weight_strides_ || + out.shape() != out_shape_ || out.strides() != out_strides_ || + input.dtype() != dtype_ || weight.dtype() != dtype_ || + out.dtype() != dtype_ || input.device() != device_ || + bias.has_value() != has_bias_ || + (bias && (bias->shape() != bias_shape_ || + bias->strides() != bias_strides_ || bias->dtype() != dtype_)) || + bits(eps) != bits(eps_) || !supports(input, weight, bias, eps, out)) + throw std::invalid_argument("NVIDIA LayerNorm call metadata or alias changed after Make"); + + std::size_t rows = 1; + nvidia_layer_norm_detail::Layout layout{}; + layout.rank = static_cast(input.ndim()); + layout.dim = input.size(-1); + layout.weight_stride = static_cast(weight.stride(0)); + layout.bias_stride = bias ? static_cast(bias->stride(0)) : 0; + for (std::size_t i = 0; i < input.ndim(); ++i) { + layout.shape[i] = input.shape()[i]; + layout.input_strides[i] = static_cast(input.strides()[i]); + layout.out_strides[i] = static_cast(out.strides()[i]); + if (i + 1 < input.ndim()) rows *= input.shape()[i]; + } + if (!rows) return; + nvidia_layer_norm_detail::Launch( + input.data(), weight.data(), bias ? bias->data() : nullptr, + out.data(), dtype_, layout, rows, eps, + static_cast(stream_ ? stream_ : 0)); + } + + private: + static std::uint32_t bits(float value) { + std::uint32_t result; + std::memcpy(&result, &value, sizeof(result)); + return result; + } + static Tensor checked(const Tensor input, const Tensor weight, + const std::optional bias, float eps, Tensor out) { + if (!supports(input, weight, bias, eps, out)) + throw std::invalid_argument("unsupported NVIDIA LayerNorm metadata, device or alias"); + return input; + } +}; +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/nvidia/ops/mrope/kernel.h b/src/native/cuda/nvidia/ops/mrope/kernel.h new file mode 100644 index 000000000..f9a8f0a2f --- /dev/null +++ b/src/native/cuda/nvidia/ops/mrope/kernel.h @@ -0,0 +1,10 @@ +#pragma once +#include "native/cuda/nvidia/caster.cuh" +#include "native/cuda/nvidia/runtime_.h" +#include "native/cuda/ops/mrope/kernel.h" +namespace infini::ops { +template<> class Operator : public CudaMrope> { + public: + using CudaMrope>::CudaMrope; +}; +} diff --git a/src/native/cuda/nvidia/ops/mul_scalar/kernel.cu b/src/native/cuda/nvidia/ops/mul_scalar/kernel.cu new file mode 100644 index 000000000..2d2d28604 --- /dev/null +++ b/src/native/cuda/nvidia/ops/mul_scalar/kernel.cu @@ -0,0 +1,97 @@ +#include "native/cuda/nvidia/ops/mul_scalar/kernel.h" + +#include +#include + +#include +#include + +namespace infini::ops::nvidia_mul_scalar_detail { +namespace { +template +__device__ __forceinline__ T Multiply(T value, T alpha) { + return value * alpha; +} + +template <> +__device__ __forceinline__ __half Multiply(__half value, __half alpha) { + return __hmul(value, alpha); +} + +template <> +__device__ __forceinline__ __nv_bfloat16 Multiply( + __nv_bfloat16 value, __nv_bfloat16 alpha) { + return __hmul(value, alpha); +} + +template +__global__ void MulScalarContiguousKernel(const T* input, T* out, T alpha, + std::size_t count) { + const auto start = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const auto step = static_cast(gridDim.x) * blockDim.x; + for (std::size_t linear = start; linear < count; linear += step) { + out[linear] = Multiply(input[linear], alpha); + } +} + +template +__global__ void MulScalarStridedKernel(const T* input, T* out, Layout layout, + T alpha, std::size_t count) { + const auto start = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const auto step = static_cast(gridDim.x) * blockDim.x; + for (std::size_t linear = start; linear < count; linear += step) { + std::size_t remaining = linear; + std::int64_t input_offset = 0, out_offset = 0; + for (std::int64_t d = layout.ndim; d-- > 0;) { + const auto coord = static_cast( + remaining % static_cast(layout.shape[d])); + remaining /= static_cast(layout.shape[d]); + input_offset += coord * layout.input_strides[d]; + out_offset += coord * layout.out_strides[d]; + } + out[out_offset] = Multiply(input[input_offset], alpha); + } +} + +template +void Submit(const void* input, void* out, const Layout* layout, T alpha, + std::size_t count, cudaStream_t stream) { + constexpr unsigned kThreads = 256; + const auto blocks = static_cast( + std::min((count - 1) / kThreads + 1, 65535)); + if (!layout) + MulScalarContiguousKernel<<>>( + static_cast(input), static_cast(out), alpha, + count); + else + MulScalarStridedKernel<<>>( + static_cast(input), static_cast(out), *layout, alpha, + count); +} +} // namespace + +void Launch(const void* input, void* out, DataType dtype, double alpha, + const Layout* layout, std::size_t count, cudaStream_t stream) { + if (!count) return; + switch (dtype) { + case DataType::kFloat16: + Submit<__half>(input, out, layout, + __float2half_rn(static_cast(alpha)), count, stream); + break; + case DataType::kBFloat16: + Submit<__nv_bfloat16>( + input, out, layout, + __float2bfloat16_rn(static_cast(alpha)), count, stream); + break; + case DataType::kFloat32: + Submit(input, out, layout, static_cast(alpha), count, + stream); + break; + case DataType::kFloat64: + Submit(input, out, layout, alpha, count, stream); + break; + default: + throw std::invalid_argument("NVIDIA MulScalar unsupported dtype"); + } +} +} // namespace infini::ops::nvidia_mul_scalar_detail diff --git a/src/native/cuda/nvidia/ops/mul_scalar/kernel.h b/src/native/cuda/nvidia/ops/mul_scalar/kernel.h new file mode 100644 index 000000000..1388c5a28 --- /dev/null +++ b/src/native/cuda/nvidia/ops/mul_scalar/kernel.h @@ -0,0 +1,145 @@ +#ifndef INFINI_OPS_NVIDIA_MUL_SCALAR_KERNEL_H_ +#define INFINI_OPS_NVIDIA_MUL_SCALAR_KERNEL_H_ + +#include + +#include +#include +#include +#include + +#include "base/mul_scalar.h" +#include "native/cuda/nvidia/ops/binary_view_capture.h" +#include "native/cuda/nvidia/runtime_.h" + +namespace infini::ops::nvidia_mul_scalar_detail { +struct Layout { + std::int64_t ndim = 0; + std::int64_t shape[16]{}; + std::int64_t input_strides[16]{}; + std::int64_t out_strides[16]{}; + bool contiguous = false; +}; + +void Launch(const void* input, void* out, DataType dtype, double alpha, + const Layout* layout, std::size_t count, cudaStream_t stream); +} // namespace infini::ops::nvidia_mul_scalar_detail + +namespace infini::ops { +template <> +class Operator : public MulScalar { + public: + static bool supports(const Tensor input, const double, Tensor out) { + using namespace nvidia_binary_view; + const auto ndim = input.shape().size(); + if (input.device().type() != Device::Type::kNvidia || + input.device() != out.device() || input.dtype() != out.dtype() || + (input.dtype() != DataType::kFloat16 && + input.dtype() != DataType::kBFloat16 && + input.dtype() != DataType::kFloat32 && + input.dtype() != DataType::kFloat64) || + ndim > 16 || out.shape() != input.shape() || + input.strides().size() != ndim || out.strides().size() != ndim || + !NonOverlappingOutput(out)) + return false; + int current_device = -1; + if (cudaGetDevice(¤t_device) != cudaSuccess || + current_device != input.device().index()) + return false; + if (input.dtype() == DataType::kBFloat16) { + static thread_local int cached_device = -1; + static thread_local bool cached_bf16 = false; + if (cached_device != current_device) { + int major = 0; + if (cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, + current_device) != cudaSuccess) + return false; + cached_bf16 = major >= 8; + cached_device = current_device; + } + if (!cached_bf16) return false; + } + + std::size_t count = 1; + constexpr auto max_offset = static_cast( + std::numeric_limits::max()); + for (std::size_t d = 0; d < ndim; ++d) { + const auto extent = input.shape()[d]; + if (extent > max_offset || input.strides()[d] < 0 || + out.strides()[d] < 0 || + static_cast(input.strides()[d]) > max_offset || + static_cast(out.strides()[d]) > max_offset || + (extent && count > max_offset / extent)) + return false; + count *= extent; + } + std::size_t input_bytes = 0, out_bytes = 0; + if (!SpanBytes(input, &input_bytes) || !SpanBytes(out, &out_bytes) || + input_bytes > max_offset || out_bytes > max_offset) + return false; + if (count && (reinterpret_cast(input.data()) % + input.element_size() || + reinterpret_cast(out.data()) % + out.element_size())) + return false; + const bool exact_inplace = input.data() == out.data() && + input.shape() == out.shape() && input.strides() == out.strides(); + return !Overlaps(input, input_bytes, out, out_bytes) || exact_inplace; + } + + Operator(const Tensor input, const double alpha, Tensor out) + : MulScalar(Checked(input, alpha, out), alpha, out) { + nvidia_binary_view::SpanBytes(input, &input_span_bytes_); + nvidia_binary_view::SpanBytes(out, &out_span_bytes_); + count_ = out.numel(); + layout_.ndim = static_cast(input.ndim()); + layout_.contiguous = input.IsContiguous() && out.IsContiguous(); + for (std::size_t d = 0; d < input.ndim(); ++d) { + layout_.shape[d] = static_cast(input.shape()[d]); + layout_.input_strides[d] = static_cast(input.strides()[d]); + layout_.out_strides[d] = static_cast(out.strides()[d]); + } + } + + bool capture_safe() const override { return true; } + + void operator()(const Tensor input, const double alpha, + Tensor out) const override { + if (input.shape() != input_shape_ || out.shape() != out_shape_ || + input.strides() != input_strides_ || out.strides() != out_strides_ || + input.dtype() != dtype_ || out.dtype() != dtype_ || + input.device() != device_ || out.device() != device_ || + std::memcmp(&alpha, &alpha_, sizeof(alpha))) + throw std::invalid_argument( + "NVIDIA MulScalar metadata or alpha changed after Make"); + if (count_ == 0) return; + int current_device = -1; + if (cudaGetDevice(¤t_device) != cudaSuccess || + current_device != device_.index() || !input.data() || !out.data() || + reinterpret_cast(input.data()) % input.element_size() || + reinterpret_cast(out.data()) % out.element_size() || + (nvidia_binary_view::Overlaps( + input, input_span_bytes_, out, out_span_bytes_) && + !(input.data() == out.data() && + input_strides_ == out_strides_))) + throw std::invalid_argument( + "NVIDIA MulScalar current device, pointer or alias changed after Make"); + nvidia_mul_scalar_detail::Launch( + input.data(), out.data(), dtype_, alpha_, + layout_.contiguous ? nullptr : &layout_, count_, + static_cast(stream_ ? stream_ : 0)); + } + + private: + static Tensor Checked(const Tensor input, const double alpha, Tensor out) { + if (!supports(input, alpha, out)) + throw std::invalid_argument( + "unsupported NVIDIA MulScalar metadata, device or alias"); + return input; + } + nvidia_mul_scalar_detail::Layout layout_{}; + std::size_t input_span_bytes_ = 0, out_span_bytes_ = 0, count_ = 0; +}; +} // namespace infini::ops + +#endif // INFINI_OPS_NVIDIA_MUL_SCALAR_KERNEL_H_ diff --git a/src/native/cuda/nvidia/ops/pad_constant/kernel.cu b/src/native/cuda/nvidia/ops/pad_constant/kernel.cu new file mode 100644 index 000000000..3d7f52ffa --- /dev/null +++ b/src/native/cuda/nvidia/ops/pad_constant/kernel.cu @@ -0,0 +1,71 @@ +#include "native/cuda/nvidia/ops/pad_constant/kernel.h" + +#include +#include + +#include +#include + +namespace infini::ops::nvidia_pad_constant_detail { +namespace { +template +__global__ void PadConstantKernel(const T* input, T* out, Layout layout, + T fill, std::size_t count) { + const auto start = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const auto step = static_cast(gridDim.x) * blockDim.x; + for (std::size_t linear = start; linear < count; linear += step) { + std::size_t remaining = linear; + std::int64_t input_offset = 0, output_offset = 0; + bool inside = true; + for (std::int64_t d = layout.ndim; d-- > 0;) { + const auto coord = static_cast( + remaining % static_cast(layout.output_shape[d])); + remaining /= static_cast(layout.output_shape[d]); + output_offset += coord * layout.output_strides[d]; + const auto source = coord - layout.left_padding[d]; + if (source < 0 || source >= layout.input_shape[d]) inside = false; + else input_offset += source * layout.input_strides[d]; + } + out[output_offset] = inside ? input[input_offset] : fill; + } +} + +template +void Submit(const void* input, void* out, Layout layout, T value, + std::size_t count, cudaStream_t stream) { + constexpr unsigned kThreads = 256; + const auto blocks = static_cast( + std::min((count - 1) / kThreads + 1, 65535)); + PadConstantKernel<<>>( + static_cast(input), static_cast(out), layout, value, count); +} +} // namespace + +void Launch(const void* input, void* out, DataType dtype, double value, + Layout layout, std::size_t output_count, cudaStream_t stream) { + if (!output_count) return; + switch (dtype) { + case DataType::kFloat16: + Submit<__half>(input, out, layout, + __float2half_rn(static_cast(value)), + output_count, stream); + break; + case DataType::kBFloat16: + Submit<__nv_bfloat16>(input, out, layout, + __float2bfloat16_rn(static_cast(value)), + output_count, stream); + break; + case DataType::kFloat32: + Submit(input, out, layout, static_cast(value), + output_count, stream); + break; + case DataType::kFloat64: + Submit(input, out, layout, value, output_count, stream); + break; + default: + throw std::invalid_argument("NVIDIA PadConstant unsupported dtype"); + } + const auto status = cudaPeekAtLastError(); + if (status != cudaSuccess) throw std::runtime_error(cudaGetErrorString(status)); +} +} // namespace infini::ops::nvidia_pad_constant_detail diff --git a/src/native/cuda/nvidia/ops/pad_constant/kernel.h b/src/native/cuda/nvidia/ops/pad_constant/kernel.h new file mode 100644 index 000000000..58c234038 --- /dev/null +++ b/src/native/cuda/nvidia/ops/pad_constant/kernel.h @@ -0,0 +1,156 @@ +#ifndef INFINI_OPS_NVIDIA_PAD_CONSTANT_KERNEL_H_ +#define INFINI_OPS_NVIDIA_PAD_CONSTANT_KERNEL_H_ + +#include + +#include +#include +#include +#include +#include + +#include "base/pad_constant.h" +#include "native/cuda/nvidia/ops/binary_view_capture.h" +#include "native/cuda/nvidia/runtime_.h" + +namespace infini::ops::nvidia_pad_constant_detail { +struct Layout { + std::int64_t ndim = 0; + std::int64_t input_shape[16]{}; + std::int64_t output_shape[16]{}; + std::int64_t input_strides[16]{}; + std::int64_t output_strides[16]{}; + std::int64_t left_padding[16]{}; +}; + +void Launch(const void* input, void* out, DataType dtype, double value, + Layout layout, std::size_t output_count, cudaStream_t stream); +} // namespace infini::ops::nvidia_pad_constant_detail + +namespace infini::ops { +template <> +class Operator + : public PadConstant { + public: + static bool supports(const Tensor input, + const std::vector padding, + const double, Tensor out) { + using namespace nvidia_binary_view; + const auto ndim = input.shape().size(); + if (input.device().type() != Device::Type::kNvidia || + input.device() != out.device() || input.dtype() != out.dtype() || + (input.dtype() != DataType::kFloat16 && + input.dtype() != DataType::kBFloat16 && + input.dtype() != DataType::kFloat32 && + input.dtype() != DataType::kFloat64) || + ndim == 0 || ndim > 16 || out.shape().size() != ndim || + input.strides().size() != ndim || out.strides().size() != ndim || + padding.empty() || padding.size() % 2 || padding.size() > 2 * ndim || + !NonOverlappingOutput(out)) + return false; + int current_device = -1; + if (cudaGetDevice(¤t_device) != cudaSuccess || + current_device != input.device().index()) return false; + if (input.dtype() == DataType::kBFloat16) { + // Device architecture is immutable, while supports() is on the Call hot + // path. Keep the metadata-only qualification local to each host thread. + static thread_local int cached_device = -1; + static thread_local bool cached_bf16 = false; + if (cached_device != current_device) { + int major = 0; + if (cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, + current_device) != cudaSuccess) + return false; + cached_bf16 = major >= 8; + cached_device = current_device; + } + if (!cached_bf16) return false; + } + + const auto max_i64 = static_cast( + std::numeric_limits::max()); + std::size_t input_count = 1, output_count = 1; + bool any_padding = false; + for (std::size_t d = 0; d < ndim; ++d) { + const auto in_size = input.shape()[d]; + const auto out_size = out.shape()[d]; + const auto in_stride = input.strides()[d]; + const auto out_stride = out.strides()[d]; + const auto pair = ndim - 1 - d; + const auto left = pair < padding.size() / 2 ? padding[2 * pair] : 0; + const auto right = pair < padding.size() / 2 ? padding[2 * pair + 1] : 0; + if (left < 0 || right < 0 || in_size > max_i64 || out_size > max_i64 || + in_stride < 0 || out_stride < 0 || + static_cast(in_stride) > max_i64 || + static_cast(out_stride) > max_i64 || + static_cast(left) > max_i64 - in_size || + static_cast(right) > + max_i64 - in_size - static_cast(left) || + out_size != in_size + static_cast(left) + + static_cast(right) || + (in_size && input_count > std::numeric_limits::max() / in_size) || + (out_size && output_count > std::numeric_limits::max() / out_size)) + return false; + input_count *= in_size; + output_count *= out_size; + any_padding |= left != 0 || right != 0; + } + std::size_t input_bytes = 0, output_bytes = 0; + if (!SpanBytes(input, &input_bytes) || !SpanBytes(out, &output_bytes)) + return false; + const bool exact_inplace = !any_padding && input.data() == out.data() && + input.shape() == out.shape() && input.strides() == out.strides(); + return !Overlaps(input, input_bytes, out, output_bytes) || exact_inplace; + } + + Operator(const Tensor input, const std::vector padding, + const double value, Tensor out) + : PadConstant(Checked(input, padding, value, out), padding, value, out) { + BuildLayout(input, padding, out); + } + + bool capture_safe() const override { return true; } + + void operator()(const Tensor input, + const std::vector padding, + const double value, Tensor out) const override { + if (input.shape() != input_shape_ || input.strides() != input_strides_ || + out.shape() != out_shape_ || out.strides() != out_strides_ || + input.dtype() != dtype_ || out.dtype() != dtype_ || + input.device() != device_ || out.device() != device_ || + padding != padding_ || std::memcmp(&value, &value_, sizeof(value)) || + !supports(input, padding, value, out)) + throw std::invalid_argument("NVIDIA PadConstant metadata or alias changed after Make"); + if (out.numel() == 0 || (input.data() == out.data() && + input.shape() == out.shape() && input.strides() == out.strides())) + return; + nvidia_pad_constant_detail::Launch( + input.data(), out.data(), dtype_, value_, layout_, out.numel(), + static_cast(stream_ ? stream_ : 0)); + } + + private: + static Tensor Checked(const Tensor input, + const std::vector& padding, + double value, Tensor out) { + if (!supports(input, padding, value, out)) + throw std::invalid_argument("unsupported NVIDIA PadConstant metadata, device or alias"); + return input; + } + void BuildLayout(const Tensor input, + const std::vector& padding, Tensor out) { + layout_.ndim = static_cast(input.ndim()); + for (std::size_t d = 0; d < input.ndim(); ++d) { + const auto pair = input.ndim() - 1 - d; + layout_.input_shape[d] = static_cast(input.shape()[d]); + layout_.output_shape[d] = static_cast(out.shape()[d]); + layout_.input_strides[d] = static_cast(input.strides()[d]); + layout_.output_strides[d] = static_cast(out.strides()[d]); + layout_.left_padding[d] = pair < padding.size() / 2 ? padding[2 * pair] : 0; + } + } + nvidia_pad_constant_detail::Layout layout_{}; +}; +} // namespace infini::ops + +#endif // INFINI_OPS_NVIDIA_PAD_CONSTANT_KERNEL_H_ diff --git a/src/native/cuda/nvidia/ops/quick_gelu/kernel.cu b/src/native/cuda/nvidia/ops/quick_gelu/kernel.cu new file mode 100644 index 000000000..b97cc21ae --- /dev/null +++ b/src/native/cuda/nvidia/ops/quick_gelu/kernel.cu @@ -0,0 +1,82 @@ +#include "native/cuda/nvidia/ops/quick_gelu/kernel.h" + +#include +#include +#include + +#include +#include +#include +#include + +namespace infini::ops::nvidia_quick_gelu_detail { +namespace { +// The three branches preserve Core 0.2.9 quickgelu/cuda/kernel.cuh semantics. +// In particular FP16 must not be promoted to the BF16 FP32-intermediate path. +template +__device__ __forceinline__ T Apply(T x) { + constexpr float alpha = 1.702f; + if constexpr (std::is_same_v) { + const __half ax = __hmul(__float2half(alpha), x); + const __half denominator = __hadd(__float2half(1.0f), hexp(__hneg(ax))); + const __half sigmoid = hrcp(denominator); + return __hmul(x, sigmoid); + } else if constexpr (std::is_same_v) { + const float xf = __bfloat162float(x); + const float ax = alpha * xf; + const float sigmoid = 1.0f / (1.0f + __expf(-ax)); + return __float2bfloat16(xf * sigmoid); + } else { + const float ax = alpha * x; + float sigmoid; + if (ax >= 0.0f) { + const float z = expf(-ax); + sigmoid = 1.0f / (1.0f + z); + } else { + const float z = expf(ax); + sigmoid = z / (1.0f + z); + } + return x * sigmoid; + } +} + +template +__global__ void QuickGeluKernel(const T* input, T* out, Layout layout, + std::size_t count) { + const std::size_t step = static_cast(blockDim.x) * gridDim.x; + for (std::size_t index = static_cast(blockIdx.x) * blockDim.x + + threadIdx.x; index < count; index += step) { + std::size_t remainder = index, input_offset = 0, out_offset = 0; + for (unsigned axis = layout.rank; axis > 0; --axis) { + const auto coordinate = remainder % layout.shape[axis - 1]; + remainder /= layout.shape[axis - 1]; + input_offset += coordinate * layout.input_strides[axis - 1]; + out_offset += coordinate * layout.out_strides[axis - 1]; + } + // No restrict qualifiers: exact non-overlapping in-place is supported. + out[out_offset] = Apply(input[input_offset]); + } +} + +template +void Submit(const void* input, void* out, Layout layout, std::size_t count, + cudaStream_t stream) { + constexpr unsigned threads = 256; + const auto blocks = std::min(1 + (count - 1) / threads, 65535); + QuickGeluKernel<<(blocks), threads, 0, stream>>>( + static_cast(input), static_cast(out), layout, count); +} +} // namespace + +void Launch(const void* input, void* out, DataType dtype, Layout layout, + std::size_t count, cudaStream_t stream) { + switch (dtype) { + case DataType::kFloat16: Submit<__half>(input, out, layout, count, stream); break; + case DataType::kBFloat16: Submit<__nv_bfloat16>(input, out, layout, count, stream); break; + case DataType::kFloat32: Submit(input, out, layout, count, stream); break; + default: throw std::invalid_argument("NVIDIA QuickGelu unsupported dtype"); + } + const auto status = cudaPeekAtLastError(); + if (status != cudaSuccess) throw std::runtime_error(cudaGetErrorString(status)); +} +} // namespace infini::ops::nvidia_quick_gelu_detail diff --git a/src/native/cuda/nvidia/ops/quick_gelu/kernel.h b/src/native/cuda/nvidia/ops/quick_gelu/kernel.h new file mode 100644 index 000000000..b9ec42edb --- /dev/null +++ b/src/native/cuda/nvidia/ops/quick_gelu/kernel.h @@ -0,0 +1,108 @@ +#ifndef INFINI_OPS_NVIDIA_QUICK_GELU_KERNEL_H_ +#define INFINI_OPS_NVIDIA_QUICK_GELU_KERNEL_H_ + +#include + +#include +#include +#include +#include + +#include "base/quick_gelu.h" +#include "native/cuda/nvidia/ops/binary_view_capture.h" +#include "native/cuda/nvidia/runtime_.h" + +namespace infini::ops { +namespace nvidia_quick_gelu_detail { +struct Layout { + std::size_t shape[8]; + std::size_t input_strides[8]; + std::size_t out_strides[8]; + unsigned rank; +}; +void Launch(const void* input, void* out, DataType dtype, Layout layout, + std::size_t count, cudaStream_t stream); +} // namespace nvidia_quick_gelu_detail + +template <> +class Operator : public QuickGelu { + public: + static bool supports(const Tensor input, Tensor out) { + using namespace nvidia_binary_view; + const auto rank = input.shape().size(); + if (input.device().type() != Device::Type::kNvidia || + input.device() != out.device() || rank < 1 || rank > 8 || + input.strides().size() != rank || out.strides().size() != rank || + input.shape() != out.shape() || input.dtype() != out.dtype() || + (input.dtype() != DataType::kFloat16 && + input.dtype() != DataType::kBFloat16 && + input.dtype() != DataType::kFloat32) || !NonOverlappingOutput(out)) + return false; + std::size_t count = 1; + for (std::size_t axis = 0; axis < rank; ++axis) { + if (input.strides()[axis] < 0 || out.strides()[axis] < 0 || + (input.shape()[axis] && count > + std::numeric_limits::max() / input.shape()[axis])) + return false; + count *= input.shape()[axis]; + } + std::size_t input_bytes = 0, out_bytes = 0; + if (!SpanBytes(input, &input_bytes) || !SpanBytes(out, &out_bytes) || + input_bytes > static_cast(std::numeric_limits::max()) || + out_bytes > static_cast(std::numeric_limits::max())) return false; + const auto width = input.element_size(); + if ((input_bytes && reinterpret_cast(input.data()) % width) || + (out_bytes && reinterpret_cast(out.data()) % width)) + return false; + const bool exact_inplace = input.data() == out.data() && + input.strides() == out.strides() && NonOverlappingOutput(input); + if (Overlaps(input, input_bytes, out, out_bytes) && !exact_inplace) + return false; + int current = -1; + if (cudaGetDevice(¤t) != cudaSuccess || + current != input.device().index()) return false; + if (input.dtype() != DataType::kFloat32) { + cudaDeviceProp prop{}; + if (cudaGetDeviceProperties(&prop, current) != cudaSuccess) return false; + const int sm = 10 * prop.major + prop.minor; + if (sm < (input.dtype() == DataType::kBFloat16 ? 80 : 53)) return false; + } + return true; + } + + Operator(const Tensor input, Tensor out) + : QuickGelu(checked(input, out), out) {} + + // QG1 has not yet passed standalone or installed-provider qualification. + bool capture_safe() const override { return false; } + + void operator()(const Tensor input, Tensor out) const override { + if (input.shape() != input_shape_ || input.strides() != input_strides_ || + out.shape() != out_shape_ || out.strides() != out_strides_ || + input.dtype() != dtype_ || out.dtype() != dtype_ || + input.device() != device_ || out.device() != device_ || + !supports(input, out)) + throw std::invalid_argument("NVIDIA QuickGelu call metadata or alias changed after Make"); + std::size_t count = 1; + nvidia_quick_gelu_detail::Layout layout{}; + layout.rank = static_cast(input.ndim()); + for (std::size_t axis = 0; axis < input.ndim(); ++axis) { + layout.shape[axis] = input.shape()[axis]; + layout.input_strides[axis] = static_cast(input.strides()[axis]); + layout.out_strides[axis] = static_cast(out.strides()[axis]); + count *= input.shape()[axis]; + } + if (!count) return; + nvidia_quick_gelu_detail::Launch(input.data(), out.data(), dtype_, layout, + count, static_cast(stream_ ? stream_ : 0)); + } + + private: + static Tensor checked(const Tensor input, Tensor out) { + if (!supports(input, out)) + throw std::invalid_argument("unsupported NVIDIA QuickGelu metadata, device or alias"); + return input; + } +}; +} // namespace infini::ops +#endif diff --git a/src/native/cuda/nvidia/ops/separate_table_rope/kernel.h b/src/native/cuda/nvidia/ops/separate_table_rope/kernel.h new file mode 100644 index 000000000..d82f3d38d --- /dev/null +++ b/src/native/cuda/nvidia/ops/separate_table_rope/kernel.h @@ -0,0 +1,24 @@ +#ifndef INFINI_OPS_NVIDIA_SEPARATE_TABLE_ROPE_KERNEL_H_ +#define INFINI_OPS_NVIDIA_SEPARATE_TABLE_ROPE_KERNEL_H_ +#include "native/cuda/nvidia/caster.cuh" +#include "native/cuda/nvidia/runtime_.h" +#include "native/cuda/ops/separate_table_rope/kernel.h" +namespace infini::ops { +template <> +class Operator + : public CudaSeparateTableRope> { + public: + using CudaSeparateTableRope>::CudaSeparateTableRope; + void operator()(const Tensor x, const Tensor p, const Tensor s, + const Tensor c, bool neox, Tensor y) const override { + CudaSeparateTableRope>::operator()(x,p,s,c,neox,y); + if (x.numel() == 0) return; + // Same native NVIDIA launch-error API used by the installed Argmax source. + // Peek reports immediate launch failure; it does not synchronize or claim + // that device data access / position bounds are valid. + const auto status = cudaPeekAtLastError(); + if (status != cudaSuccess) throw std::runtime_error(cudaGetErrorString(status)); + } +}; +} // namespace infini::ops +#endif diff --git a/src/native/cuda/nvidia/ops/sigmoid/kernel.h b/src/native/cuda/nvidia/ops/sigmoid/kernel.h index 233a8e1d6..29a8c04c8 100644 --- a/src/native/cuda/nvidia/ops/sigmoid/kernel.h +++ b/src/native/cuda/nvidia/ops/sigmoid/kernel.h @@ -14,6 +14,7 @@ class Operator : public CudaSigmoid> { public: using CudaSigmoid>::CudaSigmoid; + bool capture_safe() const override { return true; } }; } // namespace infini::ops diff --git a/src/native/cuda/nvidia/ops/sum/kernel.cu b/src/native/cuda/nvidia/ops/sum/kernel.cu new file mode 100644 index 000000000..7fe5da6ca --- /dev/null +++ b/src/native/cuda/nvidia/ops/sum/kernel.cu @@ -0,0 +1,41 @@ +#include "native/cuda/nvidia/ops/sum/kernel.h" +#include "native/cuda/ops/sum/kernel.cuh" + +#include + +#include +#include + +namespace infini::ops::nvidia_sum_detail { +namespace { +template +void Submit(const void* input, void* out, Layout layout, + std::size_t output_count, cudaStream_t stream) { + constexpr unsigned kThreads = 256; + const auto blocks = static_cast( + std::min((output_count - 1) / kThreads + 1, 65535)); + SumKernel<<>>( + static_cast(input), static_cast(out), layout, output_count); +} +} // namespace + +void Launch(const void* input, void* out, DataType dtype, Layout layout, + std::size_t output_count, cudaStream_t stream) { + if (!output_count) return; + switch (dtype) { + case DataType::kFloat16: + Submit<__half>(input, out, layout, output_count, stream); + break; + case DataType::kBFloat16: + Submit<__nv_bfloat16>(input, out, layout, output_count, stream); + break; + case DataType::kFloat32: + Submit(input, out, layout, output_count, stream); + break; + default: + throw std::invalid_argument("NVIDIA Sum unsupported dtype"); + } + const auto status = cudaPeekAtLastError(); + if (status != cudaSuccess) throw std::runtime_error(cudaGetErrorString(status)); +} +} // namespace infini::ops::nvidia_sum_detail diff --git a/src/native/cuda/nvidia/ops/sum/kernel.h b/src/native/cuda/nvidia/ops/sum/kernel.h new file mode 100644 index 000000000..3d63bff1d --- /dev/null +++ b/src/native/cuda/nvidia/ops/sum/kernel.h @@ -0,0 +1,157 @@ +#ifndef INFINI_OPS_NVIDIA_SUM_KERNEL_H_ +#define INFINI_OPS_NVIDIA_SUM_KERNEL_H_ + +#include + +#include +#include +#include +#include +#include +#include + +#include "base/sum.h" +#include "native/cuda/nvidia/runtime_.h" + +namespace infini::ops { +namespace nvidia_sum_detail { +struct Layout { + std::int64_t shape[4]{}; + std::int64_t stride[4]{}; + std::int32_t rank{}; + std::int32_t axis{}; + std::int64_t reduce_length{}; +}; +void Launch(const void* input, void* out, DataType dtype, Layout layout, + std::size_t output_count, cudaStream_t stream); +} // namespace nvidia_sum_detail + +template <> +class Operator : public Sum { + public: + static bool supports(const Tensor input, + const std::optional> dim, + const bool keepdim, + const std::optional dtype, Tensor out) { + if (input.device().type() != Device::Type::kNvidia || + input.device() != out.device() || input.dtype() != out.dtype() || + !dim || dim->size() != 1 || keepdim || dtype || + input.shape().size() < 2 || input.shape().size() > 4 || + out.shape().size() + 1 != input.shape().size()) return false; + if (input.dtype() != DataType::kFloat16 && + input.dtype() != DataType::kBFloat16 && + input.dtype() != DataType::kFloat32) return false; + int current_device = -1; + if (cudaGetDevice(¤t_device) != cudaSuccess || + current_device != input.device().index()) return false; + if (input.dtype() == DataType::kBFloat16) { + cudaDeviceProp properties{}; + if (cudaGetDeviceProperties(&properties, current_device) != cudaSuccess || + properties.major < 8) return false; + } + const auto rank = input.shape().size(); + if ((*dim)[0] < 0 || static_cast((*dim)[0]) >= rank) return false; + const auto axis = static_cast((*dim)[0]); + if (input.shape()[axis] == 0 || input.shape()[axis] > 64 || + input.strides().size() != rank || + out.strides().size() != out.shape().size()) return false; + std::size_t output_count = 1; + std::size_t output_dimension = 0; + std::size_t last_element = 0; + for (std::size_t d = 0; d < rank; ++d) { + const auto extent = input.shape()[d]; + const auto stride = input.strides()[d]; + if (!extent || stride < 0 || (extent > 1 && stride == 0) || + extent > static_cast(std::numeric_limits::max()) || + static_cast(stride) > + static_cast(std::numeric_limits::max())) return false; + if (extent - 1 > + (std::numeric_limits::max() - last_element) / + static_cast(stride ? stride : 1)) return false; + last_element += (extent - 1) * static_cast(stride); + if (d != axis) { + if (out.shape()[output_dimension++] != extent || + output_count > std::numeric_limits::max() / extent) return false; + output_count *= extent; + } + } + const auto width = input.element_size(); + if (!width || last_element > static_cast( + std::numeric_limits::max()) || + last_element == std::numeric_limits::max() || + last_element + 1 > std::numeric_limits::max() / width || + output_count > std::numeric_limits::max() || + output_count > std::numeric_limits::max() / width || + !CheckedContiguous(out) || !input.data() || !out.data() || + reinterpret_cast(input.data()) % width || + reinterpret_cast(out.data()) % width) return false; + return !Overlaps(input.data(), (last_element + 1) * width, + out.data(), output_count * width); + } + + Operator(const Tensor input, + const std::optional> dim, + const bool keepdim, const std::optional dtype, Tensor out) + : Sum(input, dim, keepdim, dtype, out) { + Validate(input, dim, keepdim, dtype, out); + } + + bool capture_safe() const override { return true; } + + void operator()(const Tensor input, + const std::optional> dim, + const bool keepdim, const std::optional dtype, + Tensor out) const override { + if (input.shape() != input_shape_ || input.strides() != input_strides_ || + input.dtype() != input_type_ || out.shape() != out_shape_ || + out.strides() != out_strides_ || out.dtype() != out_type_ || + dim != dim_ || keepdim != keepdim_ || dtype != dtype_ || + out.device().index() != device_index_) { + throw std::invalid_argument("NVIDIA Sum metadata changed after Make"); + } + Validate(input, dim, keepdim, dtype, out); + nvidia_sum_detail::Layout layout; + layout.rank = static_cast(input.shape().size()); + layout.axis = static_cast((*dim)[0]); + layout.reduce_length = static_cast(input.shape()[layout.axis]); + for (std::int32_t d = 0; d < layout.rank; ++d) { + layout.shape[d] = static_cast(input.shape()[d]); + layout.stride[d] = static_cast(input.strides()[d]); + } + nvidia_sum_detail::Launch(input.data(), out.data(), input.dtype(), layout, + out.numel(), static_cast(stream_ ? stream_ : 0)); + } + + private: + static bool CheckedContiguous(const Tensor tensor) { + std::size_t expected = 1; + for (std::size_t i = tensor.shape().size(); i > 0; --i) { + if (expected > static_cast( + std::numeric_limits::max()) || + tensor.strides()[i - 1] != static_cast(expected)) return false; + if (tensor.shape()[i - 1] && + expected > std::numeric_limits::max() / tensor.shape()[i - 1]) return false; + expected *= tensor.shape()[i - 1]; + } + return true; + } + static bool Overlaps(const void* a, std::size_t a_bytes, + const void* b, std::size_t b_bytes) { + const auto first = reinterpret_cast(a); + const auto second = reinterpret_cast(b); + const auto max = std::numeric_limits::max(); + return a_bytes > max - first || b_bytes > max - second || + (first < second + b_bytes && second < first + a_bytes); + } + static void Validate(const Tensor input, + const std::optional> dim, + const bool keepdim, const std::optional dtype, + Tensor out) { + if (!supports(input, dim, keepdim, dtype, out)) { + throw std::invalid_argument("unsupported NVIDIA Sum metadata, device or alias"); + } + } +}; +} // namespace infini::ops + +#endif // INFINI_OPS_NVIDIA_SUM_KERNEL_H_ diff --git a/src/native/cuda/nvidia/ops/take/kernel.cu b/src/native/cuda/nvidia/ops/take/kernel.cu new file mode 100644 index 000000000..6f39543e0 --- /dev/null +++ b/src/native/cuda/nvidia/ops/take/kernel.cu @@ -0,0 +1,60 @@ +#include "native/cuda/nvidia/ops/take/kernel.h" + +#include + +#include +#include +#include + +namespace infini::ops::nvidia_take_detail { +namespace { + +template +__global__ void TakeKernel(const Bits* input, const std::int64_t* index, + Bits* out, std::int64_t input_count, + std::size_t output_count) { + const auto step = static_cast(gridDim.x) * blockDim.x; + for (auto pos = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + pos < output_count; pos += step) { + auto source = index[pos]; + if (source < -input_count || source >= input_count) { + // Unlike a silent skip, an unconditional trap remains active in release + // builds. Callers must observe stream/graph completion; the CUDA context + // is poisoned after this error, so negative tests need a child process. + asm volatile("trap;"); + return; + } + if (source < 0) source += input_count; + out[pos] = input[source]; + } +} + +template +void Submit(const void* input, const std::int64_t* index, void* out, + std::size_t input_count, std::size_t output_count, + cudaStream_t stream) { + constexpr unsigned kThreads = 256; + const auto blocks = static_cast( + std::min((output_count - 1) / kThreads + 1, 65535)); + TakeKernel<<>>( + static_cast(input), index, static_cast(out), + static_cast(input_count), output_count); +} + +} // namespace + +void Launch(const void* input, const std::int64_t* index, void* out, + std::size_t input_count, std::size_t output_count, + std::size_t element_size, cudaStream_t stream) { + switch (element_size) { + case 1: Submit(input, index, out, input_count, output_count, stream); break; + case 2: Submit(input, index, out, input_count, output_count, stream); break; + case 4: Submit(input, index, out, input_count, output_count, stream); break; + case 8: Submit(input, index, out, input_count, output_count, stream); break; + default: throw std::invalid_argument("NVIDIA Take unsupported element size"); + } + const auto status = cudaPeekAtLastError(); + if (status != cudaSuccess) throw std::runtime_error(cudaGetErrorString(status)); +} + +} // namespace infini::ops::nvidia_take_detail diff --git a/src/native/cuda/nvidia/ops/take/kernel.h b/src/native/cuda/nvidia/ops/take/kernel.h new file mode 100644 index 000000000..b14a25747 --- /dev/null +++ b/src/native/cuda/nvidia/ops/take/kernel.h @@ -0,0 +1,146 @@ +#ifndef INFINI_OPS_NVIDIA_TAKE_KERNEL_H_ +#define INFINI_OPS_NVIDIA_TAKE_KERNEL_H_ + +#include + +#include +#include +#include +#include + +#include "base/take.h" +#include "native/cuda/nvidia/runtime_.h" + +namespace infini::ops { +namespace nvidia_take_detail { +void Launch(const void* input, const std::int64_t* index, void* out, + std::size_t input_count, std::size_t output_count, + std::size_t element_size, cudaStream_t stream); +} // namespace nvidia_take_detail + +template <> +class Operator : public Take { + public: + static bool supports(const Tensor input, const Tensor index, Tensor out) { + if (input.device().type() != Device::Type::kNvidia || + input.device() != index.device() || input.device() != out.device() || + !CheckedContiguous(input) || !CheckedContiguous(index) || + !CheckedContiguous(out) || + index.dtype() != DataType::kInt64 || input.dtype() != out.dtype() || + index.shape() != out.shape()) { + return false; + } + int current_device = -1; + if (cudaGetDevice(¤t_device) != cudaSuccess || + current_device != input.device().index()) return false; + switch (input.dtype()) { + case DataType::kInt8: case DataType::kUInt8: + case DataType::kInt16: case DataType::kUInt16: + case DataType::kFloat16: case DataType::kBFloat16: + case DataType::kInt32: case DataType::kUInt32: + case DataType::kFloat32: + case DataType::kInt64: case DataType::kUInt64: + case DataType::kFloat64: + break; + default: + return false; + } + std::size_t input_count, output_count; + if (!CheckedNumel(input, &input_count) || + !CheckedNumel(index, &output_count) || + input_count > static_cast(std::numeric_limits::max()) || + (output_count && !input_count)) { + return false; + } + const auto width = input.element_size(); + if (!width || input_count > std::numeric_limits::max() / width || + output_count > std::numeric_limits::max() / width || + output_count > std::numeric_limits::max() / sizeof(std::int64_t)) { + return false; + } + if ((input_count && !input.data()) || + (output_count && (!index.data() || !out.data())) || + (input_count && reinterpret_cast(input.data()) % width) || + (output_count && (reinterpret_cast(out.data()) % width || + reinterpret_cast(index.data()) % alignof(std::int64_t)))) { + return false; + } + return !Overlaps(input.data(), input_count * width, index.data(), + output_count * sizeof(std::int64_t)) && + !Overlaps(input.data(), input_count * width, out.data(), + output_count * width) && + !Overlaps(index.data(), output_count * sizeof(std::int64_t), + out.data(), output_count * width); + } + + Operator(const Tensor input, const Tensor index, Tensor out) + : Take(input, index, out) { + Validate(input, index, out); + } + + bool capture_safe() const override { return true; } + + void operator()(const Tensor input, const Tensor index, Tensor out) const override { + if (input.shape() != input_shape_ || input.strides() != input_strides_ || + input.dtype() != input_type_ || index.shape() != index_shape_ || + index.strides() != index_strides_ || index.dtype() != index_type_ || + out.shape() != out_shape_ || out.strides() != out_strides_ || + out.dtype() != out_type_ || out.device().index() != device_index_) { + throw std::invalid_argument("NVIDIA Take call metadata changed after Make"); + } + Validate(input, index, out); + const auto output_count = out.numel(); + if (!output_count) return; + nvidia_take_detail::Launch(input.data(), + static_cast(index.data()), + out.data(), input.numel(), output_count, + input.element_size(), + static_cast(stream_ ? stream_ : 0)); + } + + private: + static bool CheckedContiguous(const Tensor tensor) { + const auto& shape = tensor.shape(); + const auto& strides = tensor.strides(); + if (shape.size() != strides.size()) return false; + std::size_t expected = 1; + for (std::size_t i = shape.size(); i > 0; --i) { + if (expected > static_cast( + std::numeric_limits::max()) || + strides[i - 1] != static_cast(expected)) return false; + if (shape[i - 1] && + expected > std::numeric_limits::max() / shape[i - 1]) + return false; + expected *= shape[i - 1]; + } + return true; + } + + static bool CheckedNumel(const Tensor tensor, std::size_t* result) { + *result = 1; + for (const auto dim : tensor.shape()) { + if (dim && *result > std::numeric_limits::max() / dim) return false; + *result *= dim; + } + return true; + } + + static bool Overlaps(const void* a, std::size_t a_bytes, + const void* b, std::size_t b_bytes) { + if (!a_bytes || !b_bytes) return false; + const auto start_a = reinterpret_cast(a); + const auto start_b = reinterpret_cast(b); + const auto max = std::numeric_limits::max(); + if (a_bytes > max - start_a || b_bytes > max - start_b) return true; + return start_a < start_b + b_bytes && start_b < start_a + a_bytes; + } + + static void Validate(const Tensor input, const Tensor index, Tensor out) { + if (!supports(input, index, out)) + throw std::invalid_argument("unsupported NVIDIA Take metadata, device or alias"); + } +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_NVIDIA_TAKE_KERNEL_H_ diff --git a/src/native/cuda/nvidia/ops/upsample_bilinear2d/kernel.cu b/src/native/cuda/nvidia/ops/upsample_bilinear2d/kernel.cu new file mode 100644 index 000000000..a88f7bc22 --- /dev/null +++ b/src/native/cuda/nvidia/ops/upsample_bilinear2d/kernel.cu @@ -0,0 +1,40 @@ +#include "native/cuda/nvidia/ops/upsample_bilinear2d/kernel.h" +#include "native/cuda/ops/upsample_bilinear2d/kernel.cuh" + +#include + +#include +#include + +namespace infini::ops::nvidia_upsample_bilinear2d_detail { +namespace { +template +void Submit(const void* input, void* out, Layout layout, + std::size_t output_count, cudaStream_t stream) { + constexpr unsigned kThreads = 256; + const auto blocks = static_cast( + std::min((output_count - 1) / kThreads + 1, 65535)); + UpsampleBilinear2dKernel<<>>( + static_cast(input), static_cast(out), layout, output_count); +} +} // namespace + +void Launch(const void* input, void* out, DataType dtype, Layout layout, + std::size_t output_count, cudaStream_t stream) { + switch (dtype) { + case DataType::kFloat16: + Submit<__half>(input, out, layout, output_count, stream); + break; + case DataType::kBFloat16: + Submit<__nv_bfloat16>(input, out, layout, output_count, stream); + break; + case DataType::kFloat32: + Submit(input, out, layout, output_count, stream); + break; + default: + throw std::invalid_argument("NVIDIA UpsampleBilinear2d unsupported dtype"); + } + const auto status = cudaPeekAtLastError(); + if (status != cudaSuccess) throw std::runtime_error(cudaGetErrorString(status)); +} +} // namespace infini::ops::nvidia_upsample_bilinear2d_detail diff --git a/src/native/cuda/nvidia/ops/upsample_bilinear2d/kernel.h b/src/native/cuda/nvidia/ops/upsample_bilinear2d/kernel.h new file mode 100644 index 000000000..9cd63f04d --- /dev/null +++ b/src/native/cuda/nvidia/ops/upsample_bilinear2d/kernel.h @@ -0,0 +1,184 @@ +#ifndef INFINI_OPS_NVIDIA_UPSAMPLE_BILINEAR2D_KERNEL_H_ +#define INFINI_OPS_NVIDIA_UPSAMPLE_BILINEAR2D_KERNEL_H_ + +#include + +#include +#include +#include +#include +#include +#include + +#include "base/upsample_bilinear2d.h" +#include "native/cuda/nvidia/runtime_.h" + +namespace infini::ops { +namespace nvidia_upsample_bilinear2d_detail { +struct Layout { + std::size_t channels; + std::size_t input_height; + std::size_t input_width; + std::size_t output_height; + std::size_t output_width; + std::size_t input_strides[4]; + float scale_height; + float scale_width; + bool align_corners; +}; +void Launch(const void* input, void* out, DataType dtype, Layout layout, + std::size_t output_count, cudaStream_t stream); +} // namespace nvidia_upsample_bilinear2d_detail + +template <> +class Operator + : public UpsampleBilinear2d { + public: + static bool supports(const Tensor input, + const std::vector output_size, + const bool align_corners, + const std::optional scales_h, + const std::optional scales_w, Tensor out) { + (void)align_corners; + if (input.device().type() != Device::Type::kNvidia || + input.device() != out.device() || + input.shape().size() != 4 || out.shape().size() != 4 || + input.strides().size() != 4 || out.strides().size() != 4 || + output_size.size() != 2 || output_size[0] <= 0 || output_size[1] <= 0 || + scales_h || scales_w || input.dtype() != out.dtype() || + (input.dtype() != DataType::kFloat16 && + input.dtype() != DataType::kBFloat16 && + input.dtype() != DataType::kFloat32) || + input.shape()[0] != out.shape()[0] || + input.shape()[1] != out.shape()[1]) return false; + int current_device = -1; + if (cudaGetDevice(¤t_device) != cudaSuccess || + current_device != input.device().index()) return false; + if (input.dtype() == DataType::kBFloat16) { + cudaDeviceProp properties{}; + if (cudaGetDeviceProperties(&properties, current_device) != cudaSuccess || + properties.major < 8) return false; + } + const auto max_spatial = static_cast( + std::numeric_limits::max()); + if (static_cast(output_size[0]) > max_spatial || + static_cast(output_size[1]) > max_spatial || + out.shape()[2] != static_cast(output_size[0]) || + out.shape()[3] != static_cast(output_size[1]) || + input.shape()[2] > max_spatial || input.shape()[3] > max_spatial) + return false; + + std::size_t input_span = 1; + std::size_t output_count = 1; + for (std::size_t d = 0; d < 4; ++d) { + const auto extent = input.shape()[d]; + const auto stride = input.strides()[d]; + if (!extent || !out.shape()[d] || stride < 0 || + (extent > 1 && stride == 0) || + static_cast(stride) > + static_cast(std::numeric_limits::max()) || + extent > static_cast(std::numeric_limits::max()) || + out.shape()[d] > static_cast(std::numeric_limits::max()) || + (extent - 1 > (std::numeric_limits::max() - input_span) / + static_cast(stride ? stride : 1)) || + output_count > std::numeric_limits::max() / out.shape()[d]) + return false; + input_span += (extent - 1) * static_cast(stride); + output_count *= out.shape()[d]; + } + const auto width = input.element_size(); + if (!width || input_span > std::numeric_limits::max() / width || + output_count > std::numeric_limits::max() / width || + !CheckedContiguous(out) || !input.data() || !out.data() || + reinterpret_cast(input.data()) % width || + reinterpret_cast(out.data()) % width) + return false; + return !Overlaps(input.data(), input_span * width, out.data(), + output_count * width); + } + + Operator(const Tensor input, const std::vector output_size, + const bool align_corners, const std::optional scales_h, + const std::optional scales_w, Tensor out) + : UpsampleBilinear2d(input, output_size, align_corners, scales_h, + scales_w, out) { + Validate(input, output_size, align_corners, scales_h, scales_w, out); + } + + bool capture_safe() const override { return true; } + + void operator()(const Tensor input, + const std::vector output_size, + const bool align_corners, + const std::optional scales_h, + const std::optional scales_w, Tensor out) const override { + if (input.shape() != input_shape_ || input.strides() != input_strides_ || + input.dtype() != input_type_ || out.shape() != out_shape_ || + out.strides() != out_strides_ || out.dtype() != out_type_ || + output_size != output_size_ || align_corners != align_corners_ || + scales_h != scales_h_ || scales_w != scales_w_ || + out.device().index() != device_index_) + throw std::invalid_argument( + "NVIDIA UpsampleBilinear2d metadata changed after Make"); + Validate(input, output_size, align_corners, scales_h, scales_w, out); + nvidia_upsample_bilinear2d_detail::Layout layout{}; + layout.channels = input.shape()[1]; + layout.input_height = input.shape()[2]; + layout.input_width = input.shape()[3]; + layout.output_height = out.shape()[2]; + layout.output_width = out.shape()[3]; + for (std::size_t d = 0; d < 4; ++d) + layout.input_strides[d] = static_cast(input.strides()[d]); + layout.align_corners = align_corners; + layout.scale_height = align_corners + ? (layout.output_height > 1 + ? static_cast(layout.input_height - 1) / + static_cast(layout.output_height - 1) + : 0.0F) + : static_cast(layout.input_height) / + static_cast(layout.output_height); + layout.scale_width = align_corners + ? (layout.output_width > 1 + ? static_cast(layout.input_width - 1) / + static_cast(layout.output_width - 1) + : 0.0F) + : static_cast(layout.input_width) / + static_cast(layout.output_width); + nvidia_upsample_bilinear2d_detail::Launch( + input.data(), out.data(), input.dtype(), layout, out.numel(), + static_cast(stream_ ? stream_ : 0)); + } + + private: + static bool CheckedContiguous(const Tensor tensor) { + std::size_t expected = 1; + for (std::size_t i = 4; i > 0; --i) { + if (expected > static_cast( + std::numeric_limits::max()) || + tensor.strides()[i - 1] != static_cast(expected) || + expected > std::numeric_limits::max() / + tensor.shape()[i - 1]) return false; + expected *= tensor.shape()[i - 1]; + } + return true; + } + static bool Overlaps(const void* a, std::size_t a_bytes, + const void* b, std::size_t b_bytes) { + const auto first = reinterpret_cast(a); + const auto second = reinterpret_cast(b); + const auto max = std::numeric_limits::max(); + return a_bytes > max - first || b_bytes > max - second || + (first < second + b_bytes && second < first + a_bytes); + } + static void Validate(const Tensor input, + const std::vector& output_size, + bool align_corners, std::optional scales_h, + std::optional scales_w, Tensor out) { + if (!supports(input, output_size, align_corners, scales_h, scales_w, out)) + throw std::invalid_argument( + "unsupported NVIDIA UpsampleBilinear2d metadata, device or alias"); + } +}; +} // namespace infini::ops + +#endif // INFINI_OPS_NVIDIA_UPSAMPLE_BILINEAR2D_KERNEL_H_ diff --git a/src/native/cuda/ops/convolution/kernel.h b/src/native/cuda/ops/convolution/kernel.h index b00565cff..288ea057c 100644 --- a/src/native/cuda/ops/convolution/kernel.h +++ b/src/native/cuda/ops/convolution/kernel.h @@ -5,6 +5,7 @@ #include #include #include +#include #include #include "base/conv1d.h" @@ -122,7 +123,14 @@ class CudaConv : public Op { bias_strides_size + 3 * attrs_size; std::vector host_metadata(metadata_size); - Backend::Malloc((void**)&d_metadata_, metadata_size); + if (Backend::Malloc((void**)&d_metadata_, metadata_size) != 0 || + d_metadata_ == nullptr) { + if (d_metadata_ != nullptr) { + Backend::Free(d_metadata_); + d_metadata_ = nullptr; + } + throw std::runtime_error("CudaConv failed to allocate device metadata"); + } size_t offset = 0; d_input_shape_ = reinterpret_cast(d_metadata_ + offset); @@ -175,8 +183,12 @@ class CudaConv : public Op { std::memcpy(host_metadata.data() + offset, metadata.dilation.data(), attrs_size); - Backend::Memcpy(d_metadata_, host_metadata.data(), metadata_size, - Backend::kMemcpyHostToDevice); + if (Backend::Memcpy(d_metadata_, host_metadata.data(), metadata_size, + Backend::kMemcpyHostToDevice) != 0) { + Backend::Free(d_metadata_); + d_metadata_ = nullptr; + throw std::runtime_error("CudaConv failed to upload device metadata"); + } } std::byte* d_metadata_{nullptr}; diff --git a/src/native/cuda/ops/gelu/kernel.h b/src/native/cuda/ops/gelu/kernel.h index 70a0a6cd8..0d93af460 100644 --- a/src/native/cuda/ops/gelu/kernel.h +++ b/src/native/cuda/ops/gelu/kernel.h @@ -5,6 +5,9 @@ #include #include #include +#include +#include +#include #include #include "base/gelu.h" @@ -31,41 +34,63 @@ class CudaGelu : public Gelu { "`Gelu` input and output devices must match"); assert(!out.HasBroadcastDim() && "`Gelu` output must not have broadcasted dimensions"); + if (Backend::GetDevice(&owner_device_) != Backend::kSuccess) + throw std::runtime_error("CudaGelu cannot query metadata owner device"); + if (output_size_ == 0 || (is_input_contiguous_ && is_out_contiguous_)) + return; const size_t shape_size = ndim_ * sizeof(*d_input_shape_); const size_t strides_size = ndim_ * sizeof(*d_input_strides_); const size_t metadata_size = 2 * (shape_size + strides_size); - if (metadata_size == 0) { - return; - } std::vector metadata(metadata_size); - Backend::Malloc(reinterpret_cast(&d_metadata_), metadata_size); - - size_t offset = 0; - d_input_shape_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, input_shape_.data(), shape_size); - offset += shape_size; - - d_out_shape_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, out_shape_.data(), shape_size); - offset += shape_size; - - d_input_strides_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, input_strides_.data(), strides_size); - offset += strides_size; - - d_out_strides_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, out_strides_.data(), strides_size); - - Backend::Memcpy(d_metadata_, metadata.data(), metadata_size, - Backend::kMemcpyHostToDevice); + void* allocation = nullptr; + const auto allocation_status = Backend::Malloc(&allocation, metadata_size); + if (allocation_status != Backend::kSuccess || !allocation) { + if (allocation && Backend::Free(allocation) != Backend::kSuccess) + std::terminate(); + throw std::bad_alloc(); + } + d_metadata_ = static_cast(allocation); + + try { + size_t offset = 0; + d_input_shape_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, input_shape_.data(), shape_size); + offset += shape_size; + + d_out_shape_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, out_shape_.data(), shape_size); + offset += shape_size; + + d_input_strides_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, input_strides_.data(), strides_size); + offset += strides_size; + + d_out_strides_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, out_strides_.data(), strides_size); + + if (Backend::Memcpy(d_metadata_, metadata.data(), metadata_size, + Backend::kMemcpyHostToDevice) != Backend::kSuccess) + throw std::runtime_error("CudaGelu metadata copy failed"); + } catch (...) { + if (Backend::Free(d_metadata_) != Backend::kSuccess) std::terminate(); + d_metadata_ = nullptr; + throw; + } } - ~CudaGelu() { - if (d_metadata_ != nullptr) { - Backend::Free(d_metadata_); - } + ~CudaGelu() override { + if (!d_metadata_) return; + int previous = -1; + if (Backend::GetDevice(&previous) != Backend::kSuccess) std::terminate(); + const bool switched = previous != owner_device_; + if (switched && Backend::SetDevice(owner_device_) != Backend::kSuccess) + std::terminate(); + const auto freed = Backend::Free(d_metadata_); + if (switched && Backend::SetDevice(previous) != Backend::kSuccess) + std::terminate(); + if (freed != Backend::kSuccess) std::terminate(); } void operator()(const Tensor input, const std::string approximate, @@ -112,6 +137,7 @@ class CudaGelu : public Gelu { private: std::byte* d_metadata_{nullptr}; + int owner_device_{-1}; Tensor::Size* d_input_shape_{nullptr}; diff --git a/src/native/cuda/ops/mrope/kernel.cuh b/src/native/cuda/ops/mrope/kernel.cuh new file mode 100644 index 000000000..23e3383e1 --- /dev/null +++ b/src/native/cuda/ops/mrope/kernel.cuh @@ -0,0 +1,42 @@ +#pragma once +#include +#include +#include +#include "native/cuda/caster.cuh" +namespace infini::ops { +// Read both pair members before stores: exact input/output alias is supported. +template +__device__ void MropeRotate(T *out, const T *in, const T *cos, const T *sin, const P *positions, + size_t heads, size_t head, size_t rot, ptrdiff_t os, ptrdiff_t is, + ptrdiff_t cs, ptrdiff_t ss, ptrdiff_t pa, ptrdiff_t pt, + size_t table_len, size_t st, size_t sh, size_t sw, bool axes, bool interleaved) { + using A = std::conditional_t, double, float>; + const size_t token = blockIdx.x, half = rot/2; + for (size_t pair = threadIdx.x; pair < heads*half; pair += blockDim.x) { + const size_t i = pair%half, h = pair/half; + const size_t axis = interleaved ? (i%3 == 1 && i < sh*3 ? 1 : i%3 == 2 && i < sw*3 ? 2 : 0) + : (i < st ? 0 : i < st+sh ? 1 : 2); + const int64_t raw = static_cast(positions[token*pt + (axes ? axis*pa : 0)]); + const size_t position = raw >= 0 && static_cast(raw) < table_len ? static_cast(raw) : 0; + const A c = Caster::template Cast(cos[position*cs+i]); + const A s = Caster::template Cast(sin[position*ss+i]); + const A x = Caster::template Cast(in[token*is+h*head+i]); + const A y = Caster::template Cast(in[token*is+h*head+i+half]); + out[token*os+h*head+i] = Caster::template Cast(x*c-y*s); + out[token*os+h*head+i+half] = Caster::template Cast(y*c+x*s); + } + for (size_t i = threadIdx.x; i < heads*(head-rot); i += blockDim.x) { + const size_t offset = i/(head-rot)*head + rot + i%(head-rot); + out[token*os+offset] = in[token*is+offset]; + } +} +template +__global__ void MropeKernel(T *qo, T *ko, const T *q, const T *k, const T *cos, const T *sin, const P *p, + size_t qheads, size_t kheads, size_t head, size_t rot, + ptrdiff_t qos, ptrdiff_t kos, ptrdiff_t qs, ptrdiff_t ks, + ptrdiff_t cs, ptrdiff_t ss, ptrdiff_t pa, ptrdiff_t pt, + size_t table_len, size_t st, size_t sh, size_t sw, bool axes, bool interleaved) { + MropeRotate(qo, q, cos, sin, p, qheads, head, rot, qos, qs, cs, ss, pa, pt, table_len, st, sh, sw, axes, interleaved); + MropeRotate(ko, k, cos, sin, p, kheads, head, rot, kos, ks, cs, ss, pa, pt, table_len, st, sh, sw, axes, interleaved); +} +} diff --git a/src/native/cuda/ops/mrope/kernel.h b/src/native/cuda/ops/mrope/kernel.h new file mode 100644 index 000000000..0a2391c8b --- /dev/null +++ b/src/native/cuda/ops/mrope/kernel.h @@ -0,0 +1,26 @@ +#pragma once +#include "base/mrope.h" +#include "dispatcher.h" +#include "native/cuda/ops/mrope/kernel.cuh" +namespace infini::ops { +template class CudaMrope : public Mrope { + public: + using Mrope::Mrope; + void operator()(const Tensor q, const Tensor k, const Tensor cos, const Tensor sin, const Tensor p, + int head, int rot, int st, int sh, int sw, bool interleaved, Tensor qo, Tensor ko) const override { + if (q.size(0) == 0) return; + auto stream = static_cast(stream_ ? stream_ : 0); + DispatchFunc, ReducedFloatTypes>, List>( + {static_cast(q.dtype()), static_cast(p.dtype())}, [&](auto tag) { + using T = TypeMapType(tag)>; + using P = TypeMapType(tag)>; + MropeKernel<<(q.size(0)), 256, 0, stream>>>( + reinterpret_cast(qo.data()), reinterpret_cast(ko.data()), + reinterpret_cast(q.data()), reinterpret_cast(k.data()), + reinterpret_cast(cos.data()), reinterpret_cast(sin.data()), reinterpret_cast(p.data()), + q.size(1)/head, k.size(1)/head, head, rot, qo.stride(0), ko.stride(0), q.stride(0), k.stride(0), cos.stride(0), sin.stride(0), + p.ndim() == 2 ? p.stride(0) : 0, p.stride(-1), cos.size(0), st, sh, sw, p.ndim() == 2, interleaved); + }, "CudaMrope::operator()"); + } +}; +} diff --git a/src/native/cuda/ops/separate_table_rope/kernel.h b/src/native/cuda/ops/separate_table_rope/kernel.h new file mode 100644 index 000000000..f26e03bf4 --- /dev/null +++ b/src/native/cuda/ops/separate_table_rope/kernel.h @@ -0,0 +1,68 @@ +#ifndef INFINI_OPS_CUDA_SEPARATE_TABLE_ROPE_KERNEL_H_ +#define INFINI_OPS_CUDA_SEPARATE_TABLE_ROPE_KERNEL_H_ +#include "base/separate_table_rope.h" +#include "layout.h" +#include "native/cuda/ops/rotary_embedding_infinilm/kernel.h" + +namespace infini::ops { +template +class CudaSeparateTableRope : public SeparateTableRope { + public: + static bool supports(const Tensor x, const Tensor p, const Tensor s, + const Tensor c, bool neox, const Tensor y) { + if (Backend::kDeviceType != Device::Type::kNvidia || + !SeparateTableRope::supports(x,p,s,c,neox,y)) return false; + return separate_rope_nvidia::layout_supported(x,p,s,c,neox,y); + } + CudaSeparateTableRope(const Tensor x, const Tensor p, const Tensor s, + const Tensor c, bool neox, Tensor y) + : SeparateTableRope(x,p,s,c,neox,y), + block_size_{prepare_block_size()} {} + void operator()(const Tensor x, const Tensor p, const Tensor s, + const Tensor c, bool neox, Tensor y) const override { + // Positions must be in [0, s.size(0)). + if (!separate_rope_nvidia::layout_supported(x,p,s,c,neox,y)) + throw std::invalid_argument("SeparateTableRope: NVIDIA unsupported rebind"); + if (x.numel() == 0) return; + const auto rank = x.ndim(); + const bool batched = rank == 4; + const bool batch_pos = p.ndim() == 2; + const auto batch = batched ? x.size(0) : 1; + dim3 grid(static_cast(x.size(rank-3)), + static_cast(x.size(rank-2)),static_cast(batch)); + const auto stream = static_cast(stream_ ? stream_ : 0); + using DataTypes = ConcatType, ReducedFloatTypes>; + using AngleTypes = ConcatType, ReducedFloatTypes>; + DispatchFunc>( + {static_cast(x.dtype()),static_cast(s.dtype()),block_size_}, + [&](auto tag) { + using Data = TypeMapType(tag)>; + using Angle = TypeMapType(tag)>; + constexpr int kBlock = ListGet<2>(tag); + // Existing kernel true = interleaved pair (GPT-J), false = split + // half (GPT-NeoX). Preserve Rs4's bool neox, not the old wrapper name. + auto launch = [&](auto interleaved) { + LaunchRopeKernel( + grid,kBlock,stream,reinterpret_cast(y.data()), + reinterpret_cast(x.data()),p.data(),p.dtype(), + reinterpret_cast(s.data()), + reinterpret_cast(c.data()),s.size(1), + batched ? y.stride(0) : 0,y.stride(rank-3),y.stride(rank-2), + batched ? x.stride(0) : 0,x.stride(rank-3),x.stride(rank-2), + batch_pos ? p.stride(0) : 0,batch_pos,batched); + }; + if (neox) launch(std::false_type{}); + else launch(std::true_type{}); + },"SeparateTableRope NVIDIA"); + } + private: + static int prepare_block_size() { + const int block = RuntimeUtils::GetOptimalBlockSize(); + if (block != 128 && block != 256 && block != 512 && block != 1024) + throw std::runtime_error("SeparateTableRope: unsupported NVIDIA block size"); + return block; + } + const int block_size_; +}; +} // namespace infini::ops +#endif diff --git a/src/native/cuda/ops/separate_table_rope/layout.h b/src/native/cuda/ops/separate_table_rope/layout.h new file mode 100644 index 000000000..a641fe5b3 --- /dev/null +++ b/src/native/cuda/ops/separate_table_rope/layout.h @@ -0,0 +1,42 @@ +#ifndef INFINI_OPS_CUDA_SEPARATE_TABLE_ROPE_LAYOUT_H_ +#define INFINI_OPS_CUDA_SEPARATE_TABLE_ROPE_LAYOUT_H_ +#include +#include +#include +#include + +namespace infini::ops::separate_rope_nvidia { +template +bool layout_supported(const Tensor& x, const Tensor& p, const Tensor& s, + const Tensor& c, bool neox, const Tensor& y) { + if (x.numel() == 0) return true; // operator returns without a launch. + const auto rank = x.ndim(); + const auto batch = rank == 4 ? x.size(0) : 1; + if (batch > 65535 || x.size(rank-2) > 65535 || + x.size(rank-3) > 2147483647u) return false; + using D = std::decay_t; + auto float_width = [](D dtype) -> std::size_t { + return dtype == D::kFloat32 ? 4 : 2; // base checks NVIDIA float domain. + }; + auto aligned = [](const Tensor& t, std::size_t width) { + // Pointer-free CanMake/Make metadata cannot establish address alignment; + // actual nonnull addresses are checked again at prepared invocation. + return !t.data() || reinterpret_cast(t.data()) % width == 0; + }; + if (!aligned(x,float_width(x.dtype())) || !aligned(y,float_width(y.dtype())) || + !aligned(s,float_width(s.dtype())) || !aligned(c,float_width(c.dtype())) || + !aligned(p,p.dtype() == D::kInt64 ? 8 : 4)) return false; + // Existing interleaved kernel reinterprets half/bfloat16 pairs as half2 / + // bfloat162. Every reachable row origin must be 4-byte aligned. Split-half + // uses scalar loads/stores: 2-byte origins and odd outer strides are valid. + if (!neox && x.dtype() != D::kFloat32) { + if (!aligned(x,4) || !aligned(y,4)) return false; + for (std::size_t axis = 0; axis + 1 < rank; ++axis) { + if ((x.size(axis) > 1 && x.stride(axis) % 2 != 0) || + (y.size(axis) > 1 && y.stride(axis) % 2 != 0)) return false; + } + } + return true; +} +} // namespace infini::ops::separate_rope_nvidia +#endif diff --git a/src/native/cuda/ops/sigmoid/kernel.h b/src/native/cuda/ops/sigmoid/kernel.h index 51e1f7644..b811a9c13 100644 --- a/src/native/cuda/ops/sigmoid/kernel.h +++ b/src/native/cuda/ops/sigmoid/kernel.h @@ -5,6 +5,9 @@ #include #include #include +#include +#include +#include #include #include "base/sigmoid.h" @@ -30,6 +33,11 @@ class CudaSigmoid : public Sigmoid { assert(!out.HasBroadcastDim() && "`Sigmoid` output must not have broadcasted dimensions"); + if (Backend::GetDevice(&owner_device_) != Backend::kSuccess) + throw std::runtime_error("CudaSigmoid cannot query metadata owner device"); + if (output_size_ == 0 || (is_input_contiguous_ && is_out_contiguous_)) + return; + const size_t shape_size = ndim_ * sizeof(*d_input_shape_); const size_t strides_size = ndim_ * sizeof(*d_input_strides_); const size_t metadata_size = 2 * (shape_size + strides_size); @@ -38,32 +46,53 @@ class CudaSigmoid : public Sigmoid { } std::vector metadata(metadata_size); - Backend::Malloc(reinterpret_cast(&d_metadata_), metadata_size); - - size_t offset = 0; - d_input_shape_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, input_shape_.data(), shape_size); - offset += shape_size; - - d_out_shape_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, out_shape_.data(), shape_size); - offset += shape_size; - - d_input_strides_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, input_strides_.data(), strides_size); - offset += strides_size; - - d_out_strides_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, out_strides_.data(), strides_size); - - Backend::Memcpy(d_metadata_, metadata.data(), metadata_size, - Backend::kMemcpyHostToDevice); + void* allocation = nullptr; + const auto allocation_status = Backend::Malloc(&allocation, metadata_size); + if (allocation_status != Backend::kSuccess || !allocation) { + if (allocation && Backend::Free(allocation) != Backend::kSuccess) + std::terminate(); + throw std::bad_alloc(); + } + d_metadata_ = static_cast(allocation); + + try { + size_t offset = 0; + d_input_shape_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, input_shape_.data(), shape_size); + offset += shape_size; + + d_out_shape_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, out_shape_.data(), shape_size); + offset += shape_size; + + d_input_strides_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, input_strides_.data(), strides_size); + offset += strides_size; + + d_out_strides_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, out_strides_.data(), strides_size); + + if (Backend::Memcpy(d_metadata_, metadata.data(), metadata_size, + Backend::kMemcpyHostToDevice) != Backend::kSuccess) + throw std::runtime_error("CudaSigmoid metadata copy failed"); + } catch (...) { + if (Backend::Free(d_metadata_) != Backend::kSuccess) std::terminate(); + d_metadata_ = nullptr; + throw; + } } ~CudaSigmoid() { - if (d_metadata_ != nullptr) { - Backend::Free(d_metadata_); - } + if (!d_metadata_) return; + int previous = -1; + if (Backend::GetDevice(&previous) != Backend::kSuccess) std::terminate(); + const bool switched = previous != owner_device_; + if (switched && Backend::SetDevice(owner_device_) != Backend::kSuccess) + std::terminate(); + const auto freed = Backend::Free(d_metadata_); + if (switched && Backend::SetDevice(previous) != Backend::kSuccess) + std::terminate(); + if (freed != Backend::kSuccess) std::terminate(); } void operator()(const Tensor input, Tensor out) const override { @@ -97,6 +126,7 @@ class CudaSigmoid : public Sigmoid { private: std::byte* d_metadata_{nullptr}; + int owner_device_{-1}; Tensor::Size* d_input_shape_{nullptr}; diff --git a/src/native/cuda/ops/sum/kernel.cuh b/src/native/cuda/ops/sum/kernel.cuh new file mode 100644 index 000000000..78032e4a4 --- /dev/null +++ b/src/native/cuda/ops/sum/kernel.cuh @@ -0,0 +1,57 @@ +#ifndef INFINI_OPS_CUDA_SUM_KERNEL_CUH_ +#define INFINI_OPS_CUDA_SUM_KERNEL_CUH_ + +#include +#include + +#include "native/cuda/nvidia/ops/sum/kernel.h" + +namespace infini::ops::nvidia_sum_detail { +template +__device__ float ToFloat(T value); +template <> +__device__ inline float ToFloat(float value) { return value; } +template <> +__device__ inline float ToFloat<__half>(__half value) { return __half2float(value); } +template <> +__device__ inline float ToFloat<__nv_bfloat16>(__nv_bfloat16 value) { + return __bfloat162float(value); +} + +template +__device__ T FromFloat(float value); +template <> +__device__ inline float FromFloat(float value) { return value; } +template <> +__device__ inline __half FromFloat<__half>(float value) { + return __float2half_rn(value); +} +template <> +__device__ inline __nv_bfloat16 FromFloat<__nv_bfloat16>(float value) { + return __float2bfloat16_rn(value); +} + +template +__global__ void SumKernel(const T* input, T* out, Layout layout, + std::size_t output_count) { + const auto step = static_cast(gridDim.x) * blockDim.x; + for (auto pos = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + pos < output_count; pos += step) { + auto remaining = pos; + std::int64_t source = 0; + for (int d = layout.rank - 1; d >= 0; --d) { + if (d == layout.axis) continue; + source += static_cast(remaining % layout.shape[d]) * + layout.stride[d]; + remaining /= layout.shape[d]; + } + float total = 0.0F; + for (std::int64_t r = 0; r < layout.reduce_length; ++r) { + total += ToFloat(input[source + r * layout.stride[layout.axis]]); + } + out[pos] = FromFloat(total); + } +} +} // namespace infini::ops::nvidia_sum_detail + +#endif // INFINI_OPS_CUDA_SUM_KERNEL_CUH_ diff --git a/src/native/cuda/ops/upsample_bilinear2d/kernel.cuh b/src/native/cuda/ops/upsample_bilinear2d/kernel.cuh new file mode 100644 index 000000000..c66c86152 --- /dev/null +++ b/src/native/cuda/ops/upsample_bilinear2d/kernel.cuh @@ -0,0 +1,86 @@ +#ifndef INFINI_OPS_CUDA_UPSAMPLE_BILINEAR2D_KERNEL_CUH_ +#define INFINI_OPS_CUDA_UPSAMPLE_BILINEAR2D_KERNEL_CUH_ + +#include +#include + +#include + +#include "native/cuda/nvidia/ops/upsample_bilinear2d/kernel.h" + +namespace infini::ops::nvidia_upsample_bilinear2d_detail { +template +__device__ float ToFloat(T value); +template <> +__device__ inline float ToFloat(float value) { return value; } +template <> +__device__ inline float ToFloat<__half>(__half value) { return __half2float(value); } +template <> +__device__ inline float ToFloat<__nv_bfloat16>(__nv_bfloat16 value) { + return __bfloat162float(value); +} + +template +__device__ T FromFloat(float value); +template <> +__device__ inline float FromFloat(float value) { return value; } +template <> +__device__ inline __half FromFloat<__half>(float value) { + return __float2half_rn(value); +} +template <> +__device__ inline __nv_bfloat16 FromFloat<__nv_bfloat16>(float value) { + return __float2bfloat16_rn(value); +} + +__device__ inline std::size_t Clamp(std::int64_t value, std::size_t maximum) { + if (value < 0) return 0; + const auto positive = static_cast(value); + return positive > maximum ? maximum : positive; +} + +template +__global__ void UpsampleBilinear2dKernel( + const T* input, T* out, Layout layout, std::size_t output_count) { + const auto step = static_cast(gridDim.x) * blockDim.x; + for (auto pos = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + pos < output_count; pos += step) { + auto remainder = pos; + const auto w = remainder % layout.output_width; + remainder /= layout.output_width; + const auto h = remainder % layout.output_height; + remainder /= layout.output_height; + const auto c = remainder % layout.channels; + const auto n = remainder / layout.channels; + + const float source_h = layout.align_corners + ? static_cast(h) * layout.scale_height + : (static_cast(h) + 0.5F) * layout.scale_height - 0.5F; + const float source_w = layout.align_corners + ? static_cast(w) * layout.scale_width + : (static_cast(w) + 0.5F) * layout.scale_width - 0.5F; + const auto h0_raw = static_cast(floorf(source_h)); + const auto w0_raw = static_cast(floorf(source_w)); + const float h1_weight = source_h - static_cast(h0_raw); + const float w1_weight = source_w - static_cast(w0_raw); + const auto h0 = Clamp(h0_raw, layout.input_height - 1); + const auto h1 = Clamp(h0_raw + 1, layout.input_height - 1); + const auto w0 = Clamp(w0_raw, layout.input_width - 1); + const auto w1 = Clamp(w0_raw + 1, layout.input_width - 1); + const auto base = n * layout.input_strides[0] + c * layout.input_strides[1]; + const float v00 = ToFloat(input[base + h0 * layout.input_strides[2] + + w0 * layout.input_strides[3]]); + const float v01 = ToFloat(input[base + h0 * layout.input_strides[2] + + w1 * layout.input_strides[3]]); + const float v10 = ToFloat(input[base + h1 * layout.input_strides[2] + + w0 * layout.input_strides[3]]); + const float v11 = ToFloat(input[base + h1 * layout.input_strides[2] + + w1 * layout.input_strides[3]]); + const float top = (1.0F - w1_weight) * v00 + w1_weight * v01; + const float bottom = (1.0F - w1_weight) * v10 + w1_weight * v11; + out[pos] = FromFloat((1.0F - h1_weight) * top + h1_weight * bottom); + } +} +} // namespace infini::ops::nvidia_upsample_bilinear2d_detail + +#endif // INFINI_OPS_CUDA_UPSAMPLE_BILINEAR2D_KERNEL_CUH_ From 22440eb762428a939e23efc8d3a1c153a41d81d3 Mon Sep 17 00:00:00 2001 From: whjthu Date: Tue, 13 Oct 2026 01:44:40 +0800 Subject: [PATCH 03/25] refactor: migrate expert routing and compute --- ..._ops_nvidia_dense_sampling_torch_free.json | 6 + src/base/deepseek_moe.h | 37 + src/base/moe_align_block_size.h | 8 +- src/base/moe_compute_common.h | 43 + src/base/moe_fused_dense.h | 46 + src/base/moe_fused_gate.h | 28 + src/base/topk_softmax.h | 53 + .../cuda/nvidia/ops/deepseek_moe/kernel.cu | 1087 +++++++++++++++++ .../cuda/nvidia/ops/deepseek_moe/kernel.h | 25 + .../nvidia/ops/moe_align_block_size/kernel.cu | 3 +- .../ops/moe_align_block_size/kernel.cuh | 13 +- .../nvidia/ops/moe_align_block_size/kernel.h | 2 + .../cuda/nvidia/ops/moe_fused_dense/kernel.cu | 695 +++++++++++ .../cuda/nvidia/ops/moe_fused_dense/kernel.h | 18 + .../ops/moe_fused_gate/device_guard.cuh | 15 + .../cuda/nvidia/ops/moe_fused_gate/kernel.cu | 234 ++++ .../cuda/nvidia/ops/moe_fused_gate/kernel.h | 14 + .../cuda/nvidia/ops/topk_softmax/kernel.cu | 2 +- .../cuda/nvidia/ops/topk_softmax/kernel.h | 17 + src/native/cuda/ops/topk_softmax/kernel.cuh | 23 +- src/native/cuda/ops/topk_softmax/kernel.h | 2 +- 21 files changed, 2353 insertions(+), 18 deletions(-) create mode 100644 src/base/deepseek_moe.h create mode 100644 src/base/moe_compute_common.h create mode 100644 src/base/moe_fused_dense.h create mode 100644 src/base/moe_fused_gate.h create mode 100644 src/native/cuda/nvidia/ops/deepseek_moe/kernel.cu create mode 100644 src/native/cuda/nvidia/ops/deepseek_moe/kernel.h create mode 100644 src/native/cuda/nvidia/ops/moe_fused_dense/kernel.cu create mode 100644 src/native/cuda/nvidia/ops/moe_fused_dense/kernel.h create mode 100644 src/native/cuda/nvidia/ops/moe_fused_gate/device_guard.cuh create mode 100644 src/native/cuda/nvidia/ops/moe_fused_gate/kernel.cu create mode 100644 src/native/cuda/nvidia/ops/moe_fused_gate/kernel.h diff --git a/scripts/configs/infiniops_ops_nvidia_dense_sampling_torch_free.json b/scripts/configs/infiniops_ops_nvidia_dense_sampling_torch_free.json index 0d442f542..9ff2ca830 100644 --- a/scripts/configs/infiniops_ops_nvidia_dense_sampling_torch_free.json +++ b/scripts/configs/infiniops_ops_nvidia_dense_sampling_torch_free.json @@ -5,6 +5,7 @@ "cat": {"implementations": [0]}, "convolution": {"implementations": [0]}, "copy": {"implementations": [0]}, + "deepseek_moe": {"implementations": [0]}, "embedding": {"implementations": [0]}, "fill": {"implementations": [0]}, "flash_attn_varlen_func": {"implementations": [0]}, @@ -13,6 +14,9 @@ "gelu": {"implementations": [0]}, "gemm": {"implementations": [0]}, "layer_norm": {"implementations": [0]}, + "moe_align_block_size": {"implementations": [0]}, + "moe_fused_dense": {"implementations": [0]}, + "moe_fused_gate": {"implementations": [0]}, "mrope": {"implementations": [0]}, "mul": {"implementations": [0]}, "mul_scalar": {"implementations": [0]}, @@ -31,6 +35,8 @@ "sum": {"implementations": [0]}, "take": {"implementations": [0]}, "top_k_top_p_sampling_from_logits": {"implementations": [16]}, + "topk_sigmoid": {"implementations": [0]}, + "topk_softmax": {"implementations": [0]}, "tril": {"implementations": [0]}, "triu": {"implementations": [0]}, "upsample_bilinear2d": {"implementations": [0]} diff --git a/src/base/deepseek_moe.h b/src/base/deepseek_moe.h new file mode 100644 index 000000000..26b5191b9 --- /dev/null +++ b/src/base/deepseek_moe.h @@ -0,0 +1,37 @@ +#pragma once +#include "base/moe_compute_common.h" +namespace infini::ops { +class DeepseekMoe : public Operator { + public: + + static std::vector Collect(Tensor o,Tensor h,Tensor i,Tensor w,const std::vector& g,const std::vector& u,const std::vector& d) { + std::vector v{o,h,i,w}; for(const auto& t:g) v.push_back(t); for(const auto& t:u) v.push_back(t); for(const auto& t:d) v.push_back(t); return v; + } + + DeepseekMoe(Tensor output, const Tensor hidden, const Tensor ids, const Tensor weights, const std::vector gate, const std::vector up, const std::vector down, size_t intermediate, size_t num_experts) : metadata_{Collect(output,hidden,ids,weights,gate,up,down)}, intermediate_(intermediate), num_experts_(num_experts) { Validate(output, hidden, ids, weights, gate, up, down, intermediate, num_experts); } + static bool supports(Tensor output, const Tensor hidden, const Tensor ids, const Tensor weights, const std::vector gate, const std::vector up, const std::vector down, size_t intermediate, size_t num_experts) { try { Validate(output, hidden, ids, weights, gate, up, down, intermediate, num_experts); return true; } catch(const std::invalid_argument&) { return false; } } + virtual void operator()(Tensor output, const Tensor hidden, const Tensor ids, const Tensor weights, const std::vector gate, const std::vector up, const std::vector down, size_t intermediate, size_t num_experts) const = 0; + protected: + void Invocation(Tensor output, const Tensor hidden, const Tensor ids, const Tensor weights, const std::vector gate, const std::vector up, const std::vector down, size_t intermediate, size_t num_experts) const { + moe_detail::Require(intermediate==intermediate_ && num_experts==num_experts_); moe_detail::Same(metadata_,Collect(output,hidden,ids,weights,gate,up,down)); Validate(output, hidden, ids, weights, gate, up, down, intermediate, num_experts); + } + static void Validate(Tensor output, const Tensor hidden, const Tensor ids, const Tensor weights, const std::vector gate, const std::vector up, const std::vector down, size_t intermediate, size_t num_experts) { using namespace moe_detail; + Require(hidden.ndim()==2 && output.ndim()==2 && ids.ndim()==2 && weights.ndim()==2); + Require(Reduced(hidden.dtype()) && output.dtype()==hidden.dtype() && ids.dtype()==DataType::kInt32 && weights.dtype()==DataType::kFloat32); + Require(intermediate && num_experts && ids.size(1) && ids.size(1)<=num_experts && hidden.size(1) && output.shape()==hidden.shape() && weights.shape()==ids.shape() && ids.size(0)==hidden.size(0)); + Bound(hidden.size(0)); Bound(hidden.size(1)); Bound(ids.size(1)); Bound(intermediate); Bound(num_experts); + Require(gate.size()==num_experts && up.size()==num_experts && down.size()==num_experts); + Bound(Product(hidden.size(0),ids.size(1))); Bound(Product(Product(hidden.size(0),ids.size(1)),intermediate)); Bound(Product(Product(hidden.size(0),ids.size(1)),hidden.size(1))); + for(size_t e=0;e(cos[position*cs+i]); - const A s = Caster::template Cast(sin[position*ss+i]); - const A x = Caster::template Cast(in[token*is+h*head+i]); - const A y = Caster::template Cast(in[token*is+h*head+i+half]); - out[token*os+h*head+i] = Caster::template Cast(x*c-y*s); - out[token*os+h*head+i+half] = Caster::template Cast(y*c+x*s); + const size_t token = blockIdx.x, half = rot / 2; + for (size_t pair = threadIdx.x; pair < heads * half; pair += blockDim.x) { + const size_t i = pair % half, h = pair / half; + const size_t axis = interleaved ? (i % 3 == 1 && i < sh * 3 ? 1 + : i % 3 == 2 && i < sw * 3 ? 2 + : 0) + : (i < st ? 0 + : i < st + sh ? 1 + : 2); + const int64_t raw = + static_cast(positions[token * pt + (axes ? axis * pa : 0)]); + const size_t position = raw >= 0 && static_cast(raw) < table_len + ? static_cast(raw) + : 0; + const A c = Caster::template Cast(cos[position * cs + i]); + const A s = Caster::template Cast(sin[position * ss + i]); + const A x = Caster::template Cast(in[token * is + h * head + i]); + const A y = + Caster::template Cast(in[token * is + h * head + i + half]); + out[token * os + h * head + i] = + Caster::template Cast(x * c - y * s); + out[token * os + h * head + i + half] = + Caster::template Cast(y * c + x * s); } - for (size_t i = threadIdx.x; i < heads*(head-rot); i += blockDim.x) { - const size_t offset = i/(head-rot)*head + rot + i%(head-rot); - out[token*os+offset] = in[token*is+offset]; + for (size_t i = threadIdx.x; i < heads * (head - rot); i += blockDim.x) { + const size_t offset = i / (head - rot) * head + rot + i % (head - rot); + out[token * os + offset] = in[token * is + offset]; } } -template -__global__ void MropeKernel(T *qo, T *ko, const T *q, const T *k, const T *cos, const T *sin, const P *p, - size_t qheads, size_t kheads, size_t head, size_t rot, - ptrdiff_t qos, ptrdiff_t kos, ptrdiff_t qs, ptrdiff_t ks, - ptrdiff_t cs, ptrdiff_t ss, ptrdiff_t pa, ptrdiff_t pt, - size_t table_len, size_t st, size_t sh, size_t sw, bool axes, bool interleaved) { - MropeRotate(qo, q, cos, sin, p, qheads, head, rot, qos, qs, cs, ss, pa, pt, table_len, st, sh, sw, axes, interleaved); - MropeRotate(ko, k, cos, sin, p, kheads, head, rot, kos, ks, cs, ss, pa, pt, table_len, st, sh, sw, axes, interleaved); -} +template +__global__ void MropeKernel(T* qo, T* ko, const T* q, const T* k, const T* cos, + const T* sin, const P* p, size_t qheads, + size_t kheads, size_t head, size_t rot, + ptrdiff_t qos, ptrdiff_t kos, ptrdiff_t qs, + ptrdiff_t ks, ptrdiff_t cs, ptrdiff_t ss, + ptrdiff_t pa, ptrdiff_t pt, size_t table_len, + size_t st, size_t sh, size_t sw, bool axes, + bool interleaved) { + MropeRotate(qo, q, cos, sin, p, qheads, head, rot, qos, qs, cs, ss, + pa, pt, table_len, st, sh, sw, axes, interleaved); + MropeRotate(ko, k, cos, sin, p, kheads, head, rot, kos, ks, cs, ss, + pa, pt, table_len, st, sh, sw, axes, interleaved); } +} // namespace infini::ops diff --git a/src/native/cuda/ops/mrope/kernel.h b/src/native/cuda/ops/mrope/kernel.h index 0a2391c8b..c9ca518e5 100644 --- a/src/native/cuda/ops/mrope/kernel.h +++ b/src/native/cuda/ops/mrope/kernel.h @@ -3,24 +3,38 @@ #include "dispatcher.h" #include "native/cuda/ops/mrope/kernel.cuh" namespace infini::ops { -template class CudaMrope : public Mrope { +template +class CudaMrope : public Mrope { public: using Mrope::Mrope; - void operator()(const Tensor q, const Tensor k, const Tensor cos, const Tensor sin, const Tensor p, - int head, int rot, int st, int sh, int sw, bool interleaved, Tensor qo, Tensor ko) const override { + void operator()(const Tensor q, const Tensor k, const Tensor cos, + const Tensor sin, const Tensor p, int head, int rot, int st, + int sh, int sw, bool interleaved, Tensor qo, + Tensor ko) const override { if (q.size(0) == 0) return; auto stream = static_cast(stream_ ? stream_ : 0); - DispatchFunc, ReducedFloatTypes>, List>( - {static_cast(q.dtype()), static_cast(p.dtype())}, [&](auto tag) { - using T = TypeMapType(tag)>; - using P = TypeMapType(tag)>; - MropeKernel<<(q.size(0)), 256, 0, stream>>>( - reinterpret_cast(qo.data()), reinterpret_cast(ko.data()), - reinterpret_cast(q.data()), reinterpret_cast(k.data()), - reinterpret_cast(cos.data()), reinterpret_cast(sin.data()), reinterpret_cast(p.data()), - q.size(1)/head, k.size(1)/head, head, rot, qo.stride(0), ko.stride(0), q.stride(0), k.stride(0), cos.stride(0), sin.stride(0), - p.ndim() == 2 ? p.stride(0) : 0, p.stride(-1), cos.size(0), st, sh, sw, p.ndim() == 2, interleaved); - }, "CudaMrope::operator()"); + DispatchFunc, + ReducedFloatTypes>, + List>( + {static_cast(q.dtype()), static_cast(p.dtype())}, + [&](auto tag) { + using T = TypeMapType(tag)>; + using P = TypeMapType(tag)>; + MropeKernel + <<(q.size(0)), 256, 0, stream>>>( + reinterpret_cast(qo.data()), + reinterpret_cast(ko.data()), + reinterpret_cast(q.data()), + reinterpret_cast(k.data()), + reinterpret_cast(cos.data()), + reinterpret_cast(sin.data()), + reinterpret_cast(p.data()), q.size(1) / head, + k.size(1) / head, head, rot, qo.stride(0), ko.stride(0), + q.stride(0), k.stride(0), cos.stride(0), sin.stride(0), + p.ndim() == 2 ? p.stride(0) : 0, p.stride(-1), cos.size(0), + st, sh, sw, p.ndim() == 2, interleaved); + }, + "CudaMrope::operator()"); } }; -} +} // namespace infini::ops diff --git a/src/native/cuda/ops/mul/kernel.h b/src/native/cuda/ops/mul/kernel.h index 4c41b6f5d..6f5e5a78f 100644 --- a/src/native/cuda/ops/mul/kernel.h +++ b/src/native/cuda/ops/mul/kernel.h @@ -56,12 +56,16 @@ class CudaMul : public Mul { std::memcpy(metadata.data() + offset, out_shape_.data(), shape_size); offset += shape_size; - d_input_strides_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, input_strides_.data(), strides_size); + d_input_strides_ = + reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, input_strides_.data(), + strides_size); offset += strides_size; - d_other_strides_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, other_strides_.data(), strides_size); + d_other_strides_ = + reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, other_strides_.data(), + strides_size); offset += strides_size; d_out_strides_ = reinterpret_cast(d_metadata_ + offset); diff --git a/src/native/cuda/ops/separate_table_rope/kernel.h b/src/native/cuda/ops/separate_table_rope/kernel.h index f26e03bf4..cbbcfd8b7 100644 --- a/src/native/cuda/ops/separate_table_rope/kernel.h +++ b/src/native/cuda/ops/separate_table_rope/kernel.h @@ -11,55 +11,68 @@ class CudaSeparateTableRope : public SeparateTableRope { static bool supports(const Tensor x, const Tensor p, const Tensor s, const Tensor c, bool neox, const Tensor y) { if (Backend::kDeviceType != Device::Type::kNvidia || - !SeparateTableRope::supports(x,p,s,c,neox,y)) return false; - return separate_rope_nvidia::layout_supported(x,p,s,c,neox,y); + !SeparateTableRope::supports(x, p, s, c, neox, y)) + return false; + return separate_rope_nvidia::layout_supported(x, p, s, c, neox, y); } CudaSeparateTableRope(const Tensor x, const Tensor p, const Tensor s, - const Tensor c, bool neox, Tensor y) - : SeparateTableRope(x,p,s,c,neox,y), + const Tensor c, bool neox, Tensor y) + : SeparateTableRope(x, p, s, c, neox, y), block_size_{prepare_block_size()} {} void operator()(const Tensor x, const Tensor p, const Tensor s, const Tensor c, bool neox, Tensor y) const override { // Positions must be in [0, s.size(0)). - if (!separate_rope_nvidia::layout_supported(x,p,s,c,neox,y)) - throw std::invalid_argument("SeparateTableRope: NVIDIA unsupported rebind"); + if (!separate_rope_nvidia::layout_supported(x, p, s, c, neox, y)) + throw std::invalid_argument( + "SeparateTableRope: NVIDIA unsupported rebind"); if (x.numel() == 0) return; const auto rank = x.ndim(); const bool batched = rank == 4; const bool batch_pos = p.ndim() == 2; const auto batch = batched ? x.size(0) : 1; - dim3 grid(static_cast(x.size(rank-3)), - static_cast(x.size(rank-2)),static_cast(batch)); - const auto stream = static_cast(stream_ ? stream_ : 0); + dim3 grid(static_cast(x.size(rank - 3)), + static_cast(x.size(rank - 2)), + static_cast(batch)); + const auto stream = + static_cast(stream_ ? stream_ : 0); using DataTypes = ConcatType, ReducedFloatTypes>; using AngleTypes = ConcatType, ReducedFloatTypes>; DispatchFunc>( - {static_cast(x.dtype()),static_cast(s.dtype()),block_size_}, + {static_cast(x.dtype()), static_cast(s.dtype()), + block_size_}, [&](auto tag) { - using Data = TypeMapType(tag)>; - using Angle = TypeMapType(tag)>; + using Data = TypeMapType(tag)>; + using Angle = TypeMapType(tag)>; constexpr int kBlock = ListGet<2>(tag); // Existing kernel true = interleaved pair (GPT-J), false = split - // half (GPT-NeoX). Preserve Rs4's bool neox, not the old wrapper name. + // half (GPT-NeoX). Preserve Rs4's bool neox, not the old wrapper + // name. auto launch = [&](auto interleaved) { - LaunchRopeKernel( - grid,kBlock,stream,reinterpret_cast(y.data()), - reinterpret_cast(x.data()),p.data(),p.dtype(), + LaunchRopeKernel( + grid, kBlock, stream, reinterpret_cast(y.data()), + reinterpret_cast(x.data()), p.data(), p.dtype(), reinterpret_cast(s.data()), - reinterpret_cast(c.data()),s.size(1), - batched ? y.stride(0) : 0,y.stride(rank-3),y.stride(rank-2), - batched ? x.stride(0) : 0,x.stride(rank-3),x.stride(rank-2), - batch_pos ? p.stride(0) : 0,batch_pos,batched); + reinterpret_cast(c.data()), s.size(1), + batched ? y.stride(0) : 0, y.stride(rank - 3), + y.stride(rank - 2), batched ? x.stride(0) : 0, + x.stride(rank - 3), x.stride(rank - 2), + batch_pos ? p.stride(0) : 0, batch_pos, batched); }; - if (neox) launch(std::false_type{}); - else launch(std::true_type{}); - },"SeparateTableRope NVIDIA"); + if (neox) + launch(std::false_type{}); + else + launch(std::true_type{}); + }, + "SeparateTableRope NVIDIA"); } + private: static int prepare_block_size() { const int block = RuntimeUtils::GetOptimalBlockSize(); if (block != 128 && block != 256 && block != 512 && block != 1024) - throw std::runtime_error("SeparateTableRope: unsupported NVIDIA block size"); + throw std::runtime_error( + "SeparateTableRope: unsupported NVIDIA block size"); return block; } const int block_size_; diff --git a/src/native/cuda/ops/separate_table_rope/layout.h b/src/native/cuda/ops/separate_table_rope/layout.h index a641fe5b3..9e2cfb0a1 100644 --- a/src/native/cuda/ops/separate_table_rope/layout.h +++ b/src/native/cuda/ops/separate_table_rope/layout.h @@ -12,8 +12,9 @@ bool layout_supported(const Tensor& x, const Tensor& p, const Tensor& s, if (x.numel() == 0) return true; // operator returns without a launch. const auto rank = x.ndim(); const auto batch = rank == 4 ? x.size(0) : 1; - if (batch > 65535 || x.size(rank-2) > 65535 || - x.size(rank-3) > 2147483647u) return false; + if (batch > 65535 || x.size(rank - 2) > 65535 || + x.size(rank - 3) > 2147483647u) + return false; using D = std::decay_t; auto float_width = [](D dtype) -> std::size_t { return dtype == D::kFloat32 ? 4 : 2; // base checks NVIDIA float domain. @@ -23,17 +24,21 @@ bool layout_supported(const Tensor& x, const Tensor& p, const Tensor& s, // actual nonnull addresses are checked again at prepared invocation. return !t.data() || reinterpret_cast(t.data()) % width == 0; }; - if (!aligned(x,float_width(x.dtype())) || !aligned(y,float_width(y.dtype())) || - !aligned(s,float_width(s.dtype())) || !aligned(c,float_width(c.dtype())) || - !aligned(p,p.dtype() == D::kInt64 ? 8 : 4)) return false; + if (!aligned(x, float_width(x.dtype())) || + !aligned(y, float_width(y.dtype())) || + !aligned(s, float_width(s.dtype())) || + !aligned(c, float_width(c.dtype())) || + !aligned(p, p.dtype() == D::kInt64 ? 8 : 4)) + return false; // Existing interleaved kernel reinterprets half/bfloat16 pairs as half2 / // bfloat162. Every reachable row origin must be 4-byte aligned. Split-half // uses scalar loads/stores: 2-byte origins and odd outer strides are valid. if (!neox && x.dtype() != D::kFloat32) { - if (!aligned(x,4) || !aligned(y,4)) return false; + if (!aligned(x, 4) || !aligned(y, 4)) return false; for (std::size_t axis = 0; axis + 1 < rank; ++axis) { if ((x.size(axis) > 1 && x.stride(axis) % 2 != 0) || - (y.size(axis) > 1 && y.stride(axis) % 2 != 0)) return false; + (y.size(axis) > 1 && y.stride(axis) % 2 != 0)) + return false; } } return true; diff --git a/src/native/cuda/ops/sigmoid/kernel.h b/src/native/cuda/ops/sigmoid/kernel.h index b811a9c13..f333cd595 100644 --- a/src/native/cuda/ops/sigmoid/kernel.h +++ b/src/native/cuda/ops/sigmoid/kernel.h @@ -34,7 +34,8 @@ class CudaSigmoid : public Sigmoid { "`Sigmoid` output must not have broadcasted dimensions"); if (Backend::GetDevice(&owner_device_) != Backend::kSuccess) - throw std::runtime_error("CudaSigmoid cannot query metadata owner device"); + throw std::runtime_error( + "CudaSigmoid cannot query metadata owner device"); if (output_size_ == 0 || (is_input_contiguous_ && is_out_contiguous_)) return; @@ -65,8 +66,10 @@ class CudaSigmoid : public Sigmoid { std::memcpy(metadata.data() + offset, out_shape_.data(), shape_size); offset += shape_size; - d_input_strides_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, input_strides_.data(), strides_size); + d_input_strides_ = + reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, input_strides_.data(), + strides_size); offset += strides_size; d_out_strides_ = reinterpret_cast(d_metadata_ + offset); diff --git a/src/native/cuda/ops/silu/kernel.h b/src/native/cuda/ops/silu/kernel.h index 95de3459d..b4ec1d9d9 100644 --- a/src/native/cuda/ops/silu/kernel.h +++ b/src/native/cuda/ops/silu/kernel.h @@ -52,8 +52,10 @@ class CudaSilu : public Silu { std::memcpy(metadata.data() + offset, out_shape_.data(), shape_size); offset += shape_size; - d_input_strides_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, input_strides_.data(), strides_size); + d_input_strides_ = + reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, input_strides_.data(), + strides_size); offset += strides_size; d_out_strides_ = reinterpret_cast(d_metadata_ + offset); diff --git a/src/native/cuda/ops/silu_and_mul/kernel.h b/src/native/cuda/ops/silu_and_mul/kernel.h index aa3dbb9c3..83d50c0af 100644 --- a/src/native/cuda/ops/silu_and_mul/kernel.h +++ b/src/native/cuda/ops/silu_and_mul/kernel.h @@ -21,7 +21,8 @@ class CudaSiluAndMul : public SiluAndMul { public: CudaSiluAndMul(const Tensor input, Tensor out) : SiluAndMul{input, out} { if (Backend::GetDevice(&owner_device_) != Backend::kSuccess) - throw std::runtime_error("CudaSiluAndMul cannot query metadata owner device"); + throw std::runtime_error( + "CudaSiluAndMul cannot query metadata owner device"); if (output_size_ == 0 || (is_input_contiguous_ && is_out_contiguous_)) return; @@ -50,8 +51,10 @@ class CudaSiluAndMul : public SiluAndMul { std::memcpy(metadata.data() + offset, out_shape_.data(), shape_size); offset += shape_size; - d_input_strides_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, input_strides_.data(), strides_size); + d_input_strides_ = + reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, input_strides_.data(), + strides_size); offset += strides_size; d_out_strides_ = reinterpret_cast(d_metadata_ + offset); diff --git a/src/native/cuda/ops/softmax/kernel.h b/src/native/cuda/ops/softmax/kernel.h index 131c91ca6..8bba2c0f7 100644 --- a/src/native/cuda/ops/softmax/kernel.h +++ b/src/native/cuda/ops/softmax/kernel.h @@ -40,7 +40,8 @@ class CudaSoftmax : public Softmax { "`Softmax` output must not have broadcasted dimensions"); if (Backend::GetDevice(&owner_device_) != Backend::kSuccess) - throw std::runtime_error("CudaSoftmax cannot query metadata owner device"); + throw std::runtime_error( + "CudaSoftmax cannot query metadata owner device"); block_size_ = std::min( RuntimeUtils::GetOptimalBlockSize(), 1024); if (row_count_ == 0 || dim_size_ == 0) return; @@ -68,8 +69,10 @@ class CudaSoftmax : public Softmax { std::memcpy(metadata.data() + offset, out_shape_.data(), shape_size); offset += shape_size; - d_input_strides_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, input_strides_.data(), strides_size); + d_input_strides_ = + reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, input_strides_.data(), + strides_size); offset += strides_size; d_out_strides_ = reinterpret_cast(d_metadata_ + offset); diff --git a/src/native/cuda/ops/sum/kernel.cuh b/src/native/cuda/ops/sum/kernel.cuh index 78032e4a4..aadb9f893 100644 --- a/src/native/cuda/ops/sum/kernel.cuh +++ b/src/native/cuda/ops/sum/kernel.cuh @@ -10,9 +10,13 @@ namespace infini::ops::nvidia_sum_detail { template __device__ float ToFloat(T value); template <> -__device__ inline float ToFloat(float value) { return value; } +__device__ inline float ToFloat(float value) { + return value; +} template <> -__device__ inline float ToFloat<__half>(__half value) { return __half2float(value); } +__device__ inline float ToFloat<__half>(__half value) { + return __half2float(value); +} template <> __device__ inline float ToFloat<__nv_bfloat16>(__nv_bfloat16 value) { return __bfloat162float(value); @@ -21,7 +25,9 @@ __device__ inline float ToFloat<__nv_bfloat16>(__nv_bfloat16 value) { template __device__ T FromFloat(float value); template <> -__device__ inline float FromFloat(float value) { return value; } +__device__ inline float FromFloat(float value) { + return value; +} template <> __device__ inline __half FromFloat<__half>(float value) { return __float2half_rn(value); @@ -35,7 +41,8 @@ template __global__ void SumKernel(const T* input, T* out, Layout layout, std::size_t output_count) { const auto step = static_cast(gridDim.x) * blockDim.x; - for (auto pos = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + for (auto pos = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; pos < output_count; pos += step) { auto remaining = pos; std::int64_t source = 0; diff --git a/src/native/cuda/ops/topk_softmax/kernel.cuh b/src/native/cuda/ops/topk_softmax/kernel.cuh index c6f6d2f56..d1f3152a9 100644 --- a/src/native/cuda/ops/topk_softmax/kernel.cuh +++ b/src/native/cuda/ops/topk_softmax/kernel.cuh @@ -105,8 +105,8 @@ __global__ void TopkSoftmaxKernel(const Input* __restrict__ gating_output, Index* __restrict__ topk_indices, int32_t* __restrict__ token_expert_indices, int32_t num_tokens, int32_t num_experts, - int32_t topk, bool renormalize, - float softcap, float min_sum = 0.0f) { + int32_t topk, bool renormalize, float softcap, + float min_sum = 0.0f) { const int32_t token = static_cast(blockIdx.x); const int64_t input_offset = static_cast(token) * num_experts; const int64_t output_offset = static_cast(token) * topk; @@ -114,15 +114,17 @@ __global__ void TopkSoftmaxKernel(const Input* __restrict__ gating_output, float thread_max = -FLT_MAX; for (int32_t expert = static_cast(threadIdx.x); expert < num_experts; expert += kBlockSize) { - thread_max = fmaxf(thread_max, - Score(gating_output[input_offset + expert], softcap)); + thread_max = + fmaxf(thread_max, + Score(gating_output[input_offset + expert], softcap)); } const float row_max = BlockMax(thread_max); float thread_sum = 0.0f; for (int32_t expert = static_cast(threadIdx.x); expert < num_experts; expert += kBlockSize) { - const float value = Score(gating_output[input_offset + expert], softcap); + const float value = + Score(gating_output[input_offset + expert], softcap); thread_sum += expf(value - row_max); } const float row_sum = BlockSum(thread_sum); diff --git a/src/native/cuda/ops/triangular/kernel.cuh b/src/native/cuda/ops/triangular/kernel.cuh index 14960e65a..f3f9d32e0 100644 --- a/src/native/cuda/ops/triangular/kernel.cuh +++ b/src/native/cuda/ops/triangular/kernel.cuh @@ -24,30 +24,28 @@ struct Metadata { template __global__ void TriangularKernel(const T* input, T* output, Metadata meta) { const std::size_t step = static_cast(gridDim.x) * blockDim.x; - for (std::size_t index = static_cast(blockIdx.x) * blockDim.x + - threadIdx.x; + for (std::size_t index = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; index < meta.elements; index += step) { const std::size_t column = index % meta.columns; const std::size_t row = index / meta.columns % meta.rows; std::int64_t input_offset = 0; std::int64_t output_offset = 0; if (meta.rank <= 3) { - const std::size_t batch = meta.rank == 3 - ? index / (meta.columns * meta.rows) - : 0; - input_offset = static_cast(row) * - meta.input_strides[meta.rank - 2] + - static_cast(column) * - meta.input_strides[meta.rank - 1]; - output_offset = static_cast(row) * - meta.output_strides[meta.rank - 2] + - static_cast(column) * - meta.output_strides[meta.rank - 1]; + const std::size_t batch = + meta.rank == 3 ? index / (meta.columns * meta.rows) : 0; + input_offset = + static_cast(row) * meta.input_strides[meta.rank - 2] + + static_cast(column) * meta.input_strides[meta.rank - 1]; + output_offset = + static_cast(row) * meta.output_strides[meta.rank - 2] + + static_cast(column) * + meta.output_strides[meta.rank - 1]; if (meta.rank == 3) { - input_offset += static_cast(batch) * - meta.input_strides[0]; - output_offset += static_cast(batch) * - meta.output_strides[0]; + input_offset += + static_cast(batch) * meta.input_strides[0]; + output_offset += + static_cast(batch) * meta.output_strides[0]; } } else { std::size_t remaining = index; @@ -60,10 +58,10 @@ __global__ void TriangularKernel(const T* input, T* output, Metadata meta) { meta.output_strides[dimension]; } } - const std::int64_t relative = static_cast(column) - - static_cast(row); - const bool keep = kUpper ? relative >= meta.diagonal - : relative <= meta.diagonal; + const std::int64_t relative = + static_cast(column) - static_cast(row); + const bool keep = + kUpper ? relative >= meta.diagonal : relative <= meta.diagonal; output[output_offset] = keep ? input[input_offset] : T{}; } } diff --git a/src/native/cuda/ops/triangular/kernel.h b/src/native/cuda/ops/triangular/kernel.h index 0905b8d82..d4544f869 100644 --- a/src/native/cuda/ops/triangular/kernel.h +++ b/src/native/cuda/ops/triangular/kernel.h @@ -20,9 +20,10 @@ namespace infini::ops::triangular_detail { inline std::size_t CheckedElementCount(const Tensor& tensor) { std::size_t count = 1; for (const auto extent : tensor.shape()) { - if (extent > static_cast( - std::numeric_limits::max())) { - throw std::invalid_argument("triangular shape exceeds signed indexing range"); + if (extent > + static_cast(std::numeric_limits::max())) { + throw std::invalid_argument( + "triangular shape exceeds signed indexing range"); } if (extent && count > std::numeric_limits::max() / extent) { throw std::invalid_argument("triangular element count overflows"); @@ -37,13 +38,15 @@ inline std::size_t MaxOffset(const Tensor& tensor) { for (std::size_t dim = 0; dim < tensor.ndim(); ++dim) { const auto stride = tensor.strides()[dim]; if (stride < 0) { - throw std::invalid_argument("triangular negative strides are unsupported"); + throw std::invalid_argument( + "triangular negative strides are unsupported"); } const std::size_t extent = tensor.shape()[dim]; const std::size_t positive_stride = static_cast(stride); - if (extent > 1 && positive_stride > - (static_cast( - std::numeric_limits::max()) - offset) / + if (extent > 1 && + positive_stride > (static_cast( + std::numeric_limits::max()) - + offset) / (extent - 1)) { throw std::invalid_argument("triangular stride span overflows"); } @@ -53,11 +56,9 @@ inline std::size_t MaxOffset(const Tensor& tensor) { } inline bool SameView(const Tensor& input, const Tensor& output) { - return input.data() == output.data() && - input.shape() == output.shape() && + return input.data() == output.data() && input.shape() == output.shape() && input.strides() == output.strides() && - input.dtype() == output.dtype() && - input.device() == output.device(); + input.dtype() == output.dtype() && input.device() == output.device(); } inline std::pair ByteSpan( @@ -65,8 +66,7 @@ inline std::pair ByteSpan( if (!elements) return {0, 0}; const auto maximum = MaxOffset(tensor); const auto bytes_per_element = tensor.element_size(); - if (maximum >= std::numeric_limits::max() / - bytes_per_element) { + if (maximum >= std::numeric_limits::max() / bytes_per_element) { throw std::invalid_argument("triangular byte span overflows"); } const auto bytes = (maximum + 1) * bytes_per_element; @@ -146,11 +146,11 @@ void Launch(const Tensor& input, Tensor output, Metadata meta, cudaStream_t stream) { constexpr unsigned kThreads = 256; const std::size_t blocks = 1 + (meta.elements - 1) / kThreads; - const unsigned grid = static_cast( - std::min(blocks, 65535)); - TriangularKernel<<>>( - static_cast(input.data()), static_cast(output.data()), - meta); + const unsigned grid = + static_cast(std::min(blocks, 65535)); + TriangularKernel + <<>>(static_cast(input.data()), + static_cast(output.data()), meta); const auto status = cudaPeekAtLastError(); if (status != cudaSuccess) { throw std::runtime_error(cudaGetErrorString(status)); diff --git a/src/native/cuda/ops/upsample_bilinear2d/kernel.cuh b/src/native/cuda/ops/upsample_bilinear2d/kernel.cuh index c66c86152..d8cb65ca0 100644 --- a/src/native/cuda/ops/upsample_bilinear2d/kernel.cuh +++ b/src/native/cuda/ops/upsample_bilinear2d/kernel.cuh @@ -12,9 +12,13 @@ namespace infini::ops::nvidia_upsample_bilinear2d_detail { template __device__ float ToFloat(T value); template <> -__device__ inline float ToFloat(float value) { return value; } +__device__ inline float ToFloat(float value) { + return value; +} template <> -__device__ inline float ToFloat<__half>(__half value) { return __half2float(value); } +__device__ inline float ToFloat<__half>(__half value) { + return __half2float(value); +} template <> __device__ inline float ToFloat<__nv_bfloat16>(__nv_bfloat16 value) { return __bfloat162float(value); @@ -23,7 +27,9 @@ __device__ inline float ToFloat<__nv_bfloat16>(__nv_bfloat16 value) { template __device__ T FromFloat(float value); template <> -__device__ inline float FromFloat(float value) { return value; } +__device__ inline float FromFloat(float value) { + return value; +} template <> __device__ inline __half FromFloat<__half>(float value) { return __float2half_rn(value); @@ -40,10 +46,11 @@ __device__ inline std::size_t Clamp(std::int64_t value, std::size_t maximum) { } template -__global__ void UpsampleBilinear2dKernel( - const T* input, T* out, Layout layout, std::size_t output_count) { +__global__ void UpsampleBilinear2dKernel(const T* input, T* out, Layout layout, + std::size_t output_count) { const auto step = static_cast(gridDim.x) * blockDim.x; - for (auto pos = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + for (auto pos = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; pos < output_count; pos += step) { auto remainder = pos; const auto w = remainder % layout.output_width; @@ -53,12 +60,14 @@ __global__ void UpsampleBilinear2dKernel( const auto c = remainder % layout.channels; const auto n = remainder / layout.channels; - const float source_h = layout.align_corners - ? static_cast(h) * layout.scale_height - : (static_cast(h) + 0.5F) * layout.scale_height - 0.5F; - const float source_w = layout.align_corners - ? static_cast(w) * layout.scale_width - : (static_cast(w) + 0.5F) * layout.scale_width - 0.5F; + const float source_h = + layout.align_corners + ? static_cast(h) * layout.scale_height + : (static_cast(h) + 0.5F) * layout.scale_height - 0.5F; + const float source_w = + layout.align_corners + ? static_cast(w) * layout.scale_width + : (static_cast(w) + 0.5F) * layout.scale_width - 0.5F; const auto h0_raw = static_cast(floorf(source_h)); const auto w0_raw = static_cast(floorf(source_w)); const float h1_weight = source_h - static_cast(h0_raw); diff --git a/src/operator.h b/src/operator.h index 85803b793..64175aab5 100644 --- a/src/operator.h +++ b/src/operator.h @@ -39,7 +39,8 @@ struct AutomaticDefault> : std::bool_constant {}; // Providers may reject unsupported metadata before construction or execution. -// Legacy providers retain their existing domain; this is not capture eligibility. +// Legacy providers retain their existing domain; this is not capture +// eligibility. template auto ProviderSupports(int, const Args&... args) -> decltype(Provider::supports(args...), bool{}) { @@ -54,9 +55,9 @@ bool ProviderSupports(long, const Args&...) { template class HasConfigAwareSupports { template - static auto Test(int) - -> decltype(P::supports(std::declval(), - std::declval()...), std::true_type{}); + static auto Test(int) -> decltype(P::supports(std::declval(), + std::declval()...), + std::true_type{}); template static std::false_type Test(long); @@ -390,13 +391,14 @@ class Operator : public OperatorBase { // Metadata eligibility only. This never constructs a provider and does not // predict allocation success or capture safety. No eligible provider returns // false; an inactive explicit implementation index throws invalid_argument. - // Providers without supports() retain their legacy, interface-only eligibility. + // Providers without supports() retain their legacy, interface-only + // eligibility. template static bool CanMake(const Config& config, const Tensor tensor, const Args&... args) { - return SelectWithDevice(config, tensor.device().type(), false, - [](auto, auto, auto, auto&...) {}, tensor, - args...); + return SelectWithDevice( + config, tensor.device().type(), false, + [](auto, auto, auto, auto&...) {}, tensor, args...); } template @@ -473,11 +475,13 @@ class Operator : public OperatorBase { } // An explicit workspace takes precedence. Preserve legacy optional // workspace behavior unless this provider declares it mandatory. - if (handle.workspace() != nullptr || handle.workspace_size_in_bytes() != 0) { + if (handle.workspace() != nullptr || + handle.workspace_size_in_bytes() != 0) { if (op->requires_workspace() && (handle.workspace() == nullptr || handle.workspace_size_in_bytes() < required_workspace)) { - throw std::invalid_argument("operator workspace is missing or too small"); + throw std::invalid_argument( + "operator workspace is missing or too small"); } return (*op)(handle, args...); } @@ -489,7 +493,8 @@ class Operator : public OperatorBase { return (*op)(handle, args...); } if (lease.size_in_bytes < required_workspace) { - throw std::invalid_argument("operator workspace lease is missing or too small"); + throw std::invalid_argument( + "operator workspace lease is missing or too small"); } auto invocation_handle = handle.Clone(); invocation_handle->set_workspace(lease.data.get()); @@ -623,7 +628,8 @@ class Operator : public OperatorBase { std::forward(args)...); if (!op_ptr) { - throw std::invalid_argument("no active operator implementation supports these parameters"); + throw std::invalid_argument( + "no active operator implementation supports these parameters"); } return op_ptr; @@ -631,9 +637,8 @@ class Operator : public OperatorBase { template static bool SelectWithDevice(const Config& config, - Device::Type dispatch_device_type, - bool strict, Selected&& selected, - Args&&... args) { + Device::Type dispatch_device_type, bool strict, + Selected&& selected, Args&&... args) { bool matched = false; auto cache_args = std::forward_as_tuple(args...); @@ -648,65 +653,75 @@ class Operator : public OperatorBase { dispatch_device_type, [&](auto device_tag) { constexpr Device::Type kDev = decltype(device_tag)::value; - using Implementations = typename ActiveImplementations::type; - const bool automatic = config.implementation_index() == - Config::kAutomaticImplementation; + using Implementations = + typename ActiveImplementations::type; + const bool automatic = + config.implementation_index() == Config::kAutomaticImplementation; const auto construct = [&](auto implementation_tag) { - constexpr std::size_t kImplementationIndex = - decltype(implementation_tag)::value; - using Provider = Operator; - constexpr bool kConfigAware = - std::is_constructible_v && - detail::HasConfigAwareSupports::value; - if constexpr (kConfigAware || - std::is_constructible_v) { - bool supported = false; - if constexpr (kConfigAware) { - supported = std::apply( - [&](const auto&... cached_args) { - return Provider::supports(config, cached_args...); - }, cache_args); - } else if (!config.requires_config_aware_provider()) { - supported = std::apply( - [](const auto&... cached_args) { - return detail::ProviderSupports(0, cached_args...); - }, cache_args); - } - if (!supported) { - if (automatic) return; - if (strict) { - throw std::invalid_argument( - "requested operator implementation does not support these parameters"); - } - return; - } - std::apply( - [&](auto&... cached_args) { - matched = true; - selected(device_tag, implementation_tag, - std::bool_constant{}, - cached_args...); - }, - cache_args); - } else { - if (!automatic && strict) { - throw std::invalid_argument("operator implementation has no matching interface"); - } + constexpr std::size_t kImplementationIndex = + decltype(implementation_tag)::value; + using Provider = Operator; + constexpr bool kConfigAware = + std::is_constructible_v && + detail::HasConfigAwareSupports::value; + if constexpr (kConfigAware || + std::is_constructible_v) { + bool supported = false; + if constexpr (kConfigAware) { + supported = std::apply( + [&](const auto&... cached_args) { + return Provider::supports(config, cached_args...); + }, + cache_args); + } else if (!config.requires_config_aware_provider()) { + supported = std::apply( + [](const auto&... cached_args) { + return detail::ProviderSupports(0, + cached_args...); + }, + cache_args); + } + if (!supported) { + if (automatic) return; + if (strict) { + throw std::invalid_argument( + "requested operator implementation does not support " + "these parameters"); } - }; + return; + } + std::apply( + [&](auto&... cached_args) { + matched = true; + selected(device_tag, implementation_tag, + std::bool_constant{}, + cached_args...); + }, + cache_args); + } else { + if (!automatic && strict) { + throw std::invalid_argument( + "operator implementation has no matching interface"); + } + } + }; if (automatic) { // Selection has no kernel side effects. A constructor failure is // propagated, never treated as permission to try another provider. for (const auto index : detail::ListToVector(Implementations{})) { detail::DispatchImplementation(index, construct, - "Operator::Make(automatic)", Implementations{}); + "Operator::Make(automatic)", + Implementations{}); if (matched) break; } } else { - if (!detail::ListContains(config.implementation_index(), Implementations{})) { - throw std::invalid_argument("requested operator implementation is not active"); + if (!detail::ListContains(config.implementation_index(), + Implementations{})) { + throw std::invalid_argument( + "requested operator implementation is not active"); } - detail::DispatchImplementation(config.implementation_index(), construct, + detail::DispatchImplementation( + config.implementation_index(), construct, "Operator::Make(implementation_index)", Implementations{}); } }, From 52d7b330f0faa300431a654d9a7aec486b735ab9 Mon Sep 17 00:00:00 2001 From: whjthu Date: Mon, 19 Oct 2026 08:34:16 +0800 Subject: [PATCH 16/25] build: preserve upstream source and redistribution notices --- LICENSES/CUTLASS-BSD-3-Clause.txt | 34 ++ LICENSES/FlashInfer-0.6.7/LICENSE.cutlass.txt | 27 ++ .../LICENSE.flashattention3.txt | 29 ++ LICENSES/FlashInfer-0.6.7/LICENSE.fmt.txt | 27 ++ LICENSES/FlashInfer-0.6.7/LICENSE.spdlog.txt | 25 + LICENSES/FlashInfer-0.6.7/LICENSE.txt | 229 +++++++++ LICENSES/FlashInfer-0.6.7/NOTICE.txt | 8 + LICENSES/TVM-FFI-0.1.10/LICENSE.dlpack.txt | 201 ++++++++ .../TVM-FFI-0.1.10/LICENSE.libbacktrace.txt | 29 ++ LICENSES/TVM-FFI-0.1.10/LICENSE.pytorch.txt | 84 ++++ LICENSES/TVM-FFI-0.1.10/LICENSE.txt | 217 +++++++++ LICENSES/TVM-FFI-0.1.10/NOTICE.pytorch.txt | 456 ++++++++++++++++++ LICENSES/TVM-FFI-0.1.10/NOTICE.txt | 5 + THIRD_PARTY.txt | 31 +- cmake/NativeFlashAttention.cmake | 9 + docs/build-nvidia-inference.md | 6 + scripts/prepare_flash_attention_inference.py | 20 + src/CMakeLists.txt | 5 + .../chunk_gated_delta_rule_legacy_kernel.cuh | 4 + .../cuda/nvidia/ops/deepseek_moe/kernel.cu | 6 +- .../cuda/nvidia/ops/moe_fused_dense/kernel.cu | 6 +- .../cuda/nvidia/ops/moe_fused_gate/kernel.cu | 4 + .../recurrent_delta_rule_common.cuh | 4 + ...current_gated_delta_rule_legacy_kernel.cuh | 4 + .../cuda/nvidia/ops/w4_marlin/parameters.cu | 8 +- .../vendor/awq_repack/cuda/kernel.cuh | 4 + .../vendor/gemm/core/scalar_type.hpp | 4 + .../vendor/gemm/core/source_location.h | 4 + .../ops/w4_marlin/vendor/gemm/core/utils.h | 4 + .../w4_marlin/vendor/gemm/marlin/dequant.h | 4 + .../ops/w4_marlin/vendor/gemm/marlin/kernel.h | 4 + .../w4_marlin/vendor/gemm/marlin/marlin.cuh | 4 + .../vendor/gemm/marlin/marlin_dtypes.cuh | 4 + .../w4_marlin/vendor/gemm/marlin/marlin_mma.h | 4 + .../vendor/gemm/marlin/marlin_template.h | 4 + .../w4_marlin/vendor/gemm/nvidia/kernel.cuh | 4 + .../vendor/gemm/nvidia/kernel_selector.h | 5 + .../sm80_kernel_bfloat16_u4_bfloat16.cu | 7 +- .../sm80_kernel_bfloat16_u4b8_bfloat16.cu | 7 +- .../nvidia/sm80_kernel_float16_u4_float16.cu | 7 +- .../sm80_kernel_float16_u4b8_float16.cu | 7 +- .../vendor/gptq_repack/cuda/kernel.cuh | 4 + .../cuda/ops/mamba_selective_scan/kernel.h | 4 + 43 files changed, 1554 insertions(+), 9 deletions(-) create mode 100644 LICENSES/CUTLASS-BSD-3-Clause.txt create mode 100644 LICENSES/FlashInfer-0.6.7/LICENSE.cutlass.txt create mode 100644 LICENSES/FlashInfer-0.6.7/LICENSE.flashattention3.txt create mode 100644 LICENSES/FlashInfer-0.6.7/LICENSE.fmt.txt create mode 100644 LICENSES/FlashInfer-0.6.7/LICENSE.spdlog.txt create mode 100644 LICENSES/FlashInfer-0.6.7/LICENSE.txt create mode 100644 LICENSES/FlashInfer-0.6.7/NOTICE.txt create mode 100644 LICENSES/TVM-FFI-0.1.10/LICENSE.dlpack.txt create mode 100644 LICENSES/TVM-FFI-0.1.10/LICENSE.libbacktrace.txt create mode 100644 LICENSES/TVM-FFI-0.1.10/LICENSE.pytorch.txt create mode 100644 LICENSES/TVM-FFI-0.1.10/LICENSE.txt create mode 100644 LICENSES/TVM-FFI-0.1.10/NOTICE.pytorch.txt create mode 100644 LICENSES/TVM-FFI-0.1.10/NOTICE.txt diff --git a/LICENSES/CUTLASS-BSD-3-Clause.txt b/LICENSES/CUTLASS-BSD-3-Clause.txt new file mode 100644 index 000000000..a7389952d --- /dev/null +++ b/LICENSES/CUTLASS-BSD-3-Clause.txt @@ -0,0 +1,34 @@ +Copyright (c) 2017 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +SPDX-License-Identifier: BSD-3-Clause + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + +1. Redistributions of source code must retain the above copyright notice, this +list of conditions and the following disclaimer. + +2. Redistributions in binary form must reproduce the above copyright notice, +this list of conditions and the following disclaimer in the documentation +and/or other materials provided with the distribution. + +3. Neither the name of the copyright holder nor the names of its +contributors may be used to endorse or promote products derived from +this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +Certain files within this repository are subject to separate licensing terms: + +- The files located in the `python/CuTeDSL` directory are licensed under the + NVIDIA End User License Agreement (EULA). Please refer to + https://docs.nvidia.com/cutlass/media/docs/pythonDSL/license.html + for the full terms. diff --git a/LICENSES/FlashInfer-0.6.7/LICENSE.cutlass.txt b/LICENSES/FlashInfer-0.6.7/LICENSE.cutlass.txt new file mode 100644 index 000000000..525500841 --- /dev/null +++ b/LICENSES/FlashInfer-0.6.7/LICENSE.cutlass.txt @@ -0,0 +1,27 @@ +Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +SPDX-License-Identifier: BSD-3-Clause + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + +1. Redistributions of source code must retain the above copyright notice, this +list of conditions and the following disclaimer. + +2. Redistributions in binary form must reproduce the above copyright notice, +this list of conditions and the following disclaimer in the documentation +and/or other materials provided with the distribution. + +3. Neither the name of the copyright holder nor the names of its +contributors may be used to endorse or promote products derived from +this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/LICENSES/FlashInfer-0.6.7/LICENSE.flashattention3.txt b/LICENSES/FlashInfer-0.6.7/LICENSE.flashattention3.txt new file mode 100644 index 000000000..5860e4b33 --- /dev/null +++ b/LICENSES/FlashInfer-0.6.7/LICENSE.flashattention3.txt @@ -0,0 +1,29 @@ +BSD 3-Clause License + +Copyright (c) 2022, the respective contributors, as shown by the AUTHORS file. +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + +* Redistributions of source code must retain the above copyright notice, this + list of conditions and the following disclaimer. + +* Redistributions in binary form must reproduce the above copyright notice, + this list of conditions and the following disclaimer in the documentation + and/or other materials provided with the distribution. + +* Neither the name of the copyright holder nor the names of its + contributors may be used to endorse or promote products derived from + this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/LICENSES/FlashInfer-0.6.7/LICENSE.fmt.txt b/LICENSES/FlashInfer-0.6.7/LICENSE.fmt.txt new file mode 100644 index 000000000..f0ec3db4d --- /dev/null +++ b/LICENSES/FlashInfer-0.6.7/LICENSE.fmt.txt @@ -0,0 +1,27 @@ +Copyright (c) 2012 - present, Victor Zverovich + +Permission is hereby granted, free of charge, to any person obtaining +a copy of this software and associated documentation files (the +"Software"), to deal in the Software without restriction, including +without limitation the rights to use, copy, modify, merge, publish, +distribute, sublicense, and/or sell copies of the Software, and to +permit persons to whom the Software is furnished to do so, subject to +the following conditions: + +The above copyright notice and this permission notice shall be +included in all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, +EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF +MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND +NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE +LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION +OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION +WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + +--- Optional exception to the license --- + +As an exception, if, as a result of your compiling your source code, portions +of this Software are embedded into a machine-executable object form of such +source code, you may redistribute such embedded portions in such object form +without including the above copyright and permission notices. diff --git a/LICENSES/FlashInfer-0.6.7/LICENSE.spdlog.txt b/LICENSES/FlashInfer-0.6.7/LICENSE.spdlog.txt new file mode 100644 index 000000000..c46826e51 --- /dev/null +++ b/LICENSES/FlashInfer-0.6.7/LICENSE.spdlog.txt @@ -0,0 +1,25 @@ +The MIT License (MIT) + +Copyright (c) 2016 Gabi Melman. + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in +all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +THE SOFTWARE. + +-- NOTE: Third party dependency used by this software -- +This software depends on the fmt lib (MIT License), +and users must comply to its license: https://raw.githubusercontent.com/fmtlib/fmt/master/LICENSE diff --git a/LICENSES/FlashInfer-0.6.7/LICENSE.txt b/LICENSES/FlashInfer-0.6.7/LICENSE.txt new file mode 100644 index 000000000..ba557f2bb --- /dev/null +++ b/LICENSES/FlashInfer-0.6.7/LICENSE.txt @@ -0,0 +1,229 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. + +------------------------------------------------------------------------------------------------- +Some of the code in this project are adapted from other open-source projects with different +licenses. This product also bundles some third-party components under other open source licenses. +This section summarizes those components and their licenses. +See licenses/ for text of these licenses. + +BSD 3-Clause License +-------------------- + +include/flashinfer/attention/hopper/epilogue.cuh +include/flashinfer/attention/hopper/mainloop.cuh +include/flashinfer/attention/hopper/kernel_traits.cuh +include/flashinfer/attention/hopper/named_barrier.cuh +include/flashinfer/attention/hopper/tile_scheduler.cuh +include/flashinfer/attention/hopper/utils.cuh + +BSD 3-Clause "New" License +-------------------------- + +3rdparty/cutlass +include/flashinfer/attention/hopper/block_sparse_gather.cuh + +MIT License +----------- + +3rdparty/spdlog +3rdparty/spdlog/include/spdlog/fmt/bundled (fmt library) diff --git a/LICENSES/FlashInfer-0.6.7/NOTICE.txt b/LICENSES/FlashInfer-0.6.7/NOTICE.txt new file mode 100644 index 000000000..f7a1d2da4 --- /dev/null +++ b/LICENSES/FlashInfer-0.6.7/NOTICE.txt @@ -0,0 +1,8 @@ +FlashInfer +Copyright 2025-2026 NVIDIA +Copyright 2023-2026 FlashInfer community (https://flashinfer.ai/) + +------------------------------------------------------------------------------------------------- +Some of the code in this project are adapted from other open-source projects with different +licenses. This product also bundles some third-party components under other open source licenses. +See licenses/ for text of these licenses. diff --git a/LICENSES/TVM-FFI-0.1.10/LICENSE.dlpack.txt b/LICENSES/TVM-FFI-0.1.10/LICENSE.dlpack.txt new file mode 100644 index 000000000..20a9c8a7b --- /dev/null +++ b/LICENSES/TVM-FFI-0.1.10/LICENSE.dlpack.txt @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "{}" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright 2017 by Contributors + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/LICENSES/TVM-FFI-0.1.10/LICENSE.libbacktrace.txt b/LICENSES/TVM-FFI-0.1.10/LICENSE.libbacktrace.txt new file mode 100644 index 000000000..e9e256244 --- /dev/null +++ b/LICENSES/TVM-FFI-0.1.10/LICENSE.libbacktrace.txt @@ -0,0 +1,29 @@ +# Copyright (C) 2012-2016 Free Software Foundation, Inc. + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are +# met: + +# (1) Redistributions of source code must retain the above copyright +# notice, this list of conditions and the following disclaimer. + +# (2) Redistributions in binary form must reproduce the above copyright +# notice, this list of conditions and the following disclaimer in +# the documentation and/or other materials provided with the +# distribution. + +# (3) The name of the author may not be used to +# endorse or promote products derived from this software without +# specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR +# IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED +# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT, +# INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES +# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) +# HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, +# STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING +# IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE +# POSSIBILITY OF SUCH DAMAGE. diff --git a/LICENSES/TVM-FFI-0.1.10/LICENSE.pytorch.txt b/LICENSES/TVM-FFI-0.1.10/LICENSE.pytorch.txt new file mode 100644 index 000000000..966a609b6 --- /dev/null +++ b/LICENSES/TVM-FFI-0.1.10/LICENSE.pytorch.txt @@ -0,0 +1,84 @@ +From PyTorch: + +Copyright (c) 2016- Facebook, Inc (Adam Paszke) +Copyright (c) 2014- Facebook, Inc (Soumith Chintala) +Copyright (c) 2011-2014 Idiap Research Institute (Ronan Collobert) +Copyright (c) 2012-2014 Deepmind Technologies (Koray Kavukcuoglu) +Copyright (c) 2011-2012 NEC Laboratories America (Koray Kavukcuoglu) +Copyright (c) 2011-2013 NYU (Clement Farabet) +Copyright (c) 2006-2010 NEC Laboratories America (Ronan Collobert, Leon Bottou, Iain Melvin, Jason Weston) +Copyright (c) 2006 Idiap Research Institute (Samy Bengio) +Copyright (c) 2001-2004 Idiap Research Institute (Ronan Collobert, Samy Bengio, Johnny Mariethoz) + +From Caffe2: + +Copyright (c) 2016-present, Facebook Inc. All rights reserved. + +All contributions by Facebook: +Copyright (c) 2016 Facebook Inc. + +All contributions by Google: +Copyright (c) 2015 Google Inc. +All rights reserved. + +All contributions by Yangqing Jia: +Copyright (c) 2015 Yangqing Jia +All rights reserved. + +All contributions by Kakao Brain: +Copyright 2019-2020 Kakao Brain + +All contributions by Cruise LLC: +Copyright (c) 2022 Cruise LLC. +All rights reserved. + +All contributions by Tri Dao: +Copyright (c) 2024 Tri Dao. +All rights reserved. + +All contributions by Arm: +Copyright (c) 2021, 2023-2024 Arm Limited and/or its affiliates + +All contributions from Caffe: +Copyright(c) 2013, 2014, 2015, the respective contributors +All rights reserved. + +All other contributions: +Copyright(c) 2015, 2016 the respective contributors +All rights reserved. + +Caffe2 uses a copyright model similar to Caffe: each contributor holds +copyright over their contributions to Caffe2. The project versioning records +all such contribution and copyright details. If a contributor wants to further +mark their specific copyright on a particular contribution, they should +indicate their copyright solely in the commit message of the change when it is +committed. + +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + +1. Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + +2. Redistributions in binary form must reproduce the above copyright + notice, this list of conditions and the following disclaimer in the + documentation and/or other materials provided with the distribution. + +3. Neither the names of Facebook, Deepmind Technologies, NYU, NEC Laboratories America + and IDIAP Research Institute nor the names of its contributors may be + used to endorse or promote products derived from this software without + specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE +LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR +CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF +SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN +CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) +ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE +POSSIBILITY OF SUCH DAMAGE. diff --git a/LICENSES/TVM-FFI-0.1.10/LICENSE.txt b/LICENSES/TVM-FFI-0.1.10/LICENSE.txt new file mode 100644 index 000000000..49ac4b3df --- /dev/null +++ b/LICENSES/TVM-FFI-0.1.10/LICENSE.txt @@ -0,0 +1,217 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "{}" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright {yyyy} {name of copyright owner} + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. + +------------------------------------------------------------------------------------ +This product bundles various third-party components under other open source licenses. +This section summarizes those components and their licenses. See licenses/ +for text of these licenses. + + +Apache Software Foundation License 2.0 +-------------------------------------- + +3rdparty/dlpack + +BSD 3-Clause "New" or "Revised" License +--------------------------------------- + +3rdparty/libbacktrace diff --git a/LICENSES/TVM-FFI-0.1.10/NOTICE.pytorch.txt b/LICENSES/TVM-FFI-0.1.10/NOTICE.pytorch.txt new file mode 100644 index 000000000..6effb8b5d --- /dev/null +++ b/LICENSES/TVM-FFI-0.1.10/NOTICE.pytorch.txt @@ -0,0 +1,456 @@ +======================================================================= +Software under third_party +======================================================================= +Software libraries under third_party are provided as github submodule +links, and their content is not part of the Caffe2 codebase. Their +licences can be found under the respective software repositories. + +======================================================================= +Earlier BSD License +======================================================================= +Early development of Caffe2 in 2015 and early 2016 is licensed under the +BSD license. The license is attached below: + +All contributions by Facebook: +Copyright (c) 2016 Facebook Inc. + +All contributions by Google: +Copyright (c) 2015 Google Inc. +All rights reserved. + +All contributions by Yangqing Jia: +Copyright (c) 2015 Yangqing Jia +All rights reserved. + +All contributions by Kakao Brain: +Copyright 2019-2020 Kakao Brain + +All other contributions: +Copyright(c) 2015, 2016 the respective contributors +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + +1. Redistributions of source code must retain the above copyright notice, this + list of conditions and the following disclaimer. +2. Redistributions in binary form must reproduce the above copyright notice, + this list of conditions and the following disclaimer in the documentation + and/or other materials provided with the distribution. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND +ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED +WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR +ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES +(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; +LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND +ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + + +======================================================================= +Caffe's BSD License +======================================================================= +Some parts of the caffe2 code is derived from the original Caffe code, which is +created by Yangqing Jia and is now a BSD-licensed open-source project. The Caffe +license is as follows: + +COPYRIGHT + +All contributions by the University of California: +Copyright (c) 2014, The Regents of the University of California (Regents) +All rights reserved. + +All other contributions: +Copyright (c) 2014, the respective contributors +All rights reserved. + +Caffe uses a shared copyright model: each contributor holds copyright over +their contributions to Caffe. The project versioning records all such +contribution and copyright details. If a contributor wants to further mark +their specific copyright on a particular contribution, they should indicate +their copyright solely in the commit message of the change when it is +committed. + +LICENSE + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + +1. Redistributions of source code must retain the above copyright notice, this + list of conditions and the following disclaimer. +2. Redistributions in binary form must reproduce the above copyright notice, + this list of conditions and the following disclaimer in the documentation + and/or other materials provided with the distribution. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND +ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED +WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR +ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES +(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; +LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND +ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +CONTRIBUTION AGREEMENT + +By contributing to the BVLC/caffe repository through pull-request, comment, +or otherwise, the contributor releases their content to the +license and copyright terms herein. + +======================================================================= +Caffe2's Apache License +======================================================================= + +This repo contains Caffe2 code, which was previously licensed under +Apache License Version 2.0: + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + +======================================================================= +Cephes's 3-Clause BSD License +======================================================================= + +Code derived from implementations in the Cephes Math Library should mention +its derivation and reference the following license: + + 3-Clause BSD License for the Cephes Math Library + Copyright (c) 2018, Steven Moshier + All rights reserved. + + Redistribution and use in source and binary forms, with or without + modification, are permitted provided that the following conditions are met: + + * Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + + * Redistributions in binary form must reproduce the above copyright + notice, this list of conditions and the following disclaimer in the + documentation and/or other materials provided with the distribution. + + * Neither the name of the nor the + names of its contributors may be used to endorse or promote products + derived from this software without specific prior written permission. + + THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND + ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED + WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + DISCLAIMED. IN NO EVENT SHALL Steven Moshier BE LIABLE FOR ANY + DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES + (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; + LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND + ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT + (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS + SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + + +======================================================================= +SciPy's 3-Clause BSD License +======================================================================= + +Code derived from implementations in SciPy should mention its derivation +and reference the following license: + + Copyright (c) 2001-2002 Enthought, Inc. 2003-2019, SciPy Developers. + All rights reserved. + + Redistribution and use in source and binary forms, with or without + modification, are permitted provided that the following conditions + are met: + + 1. Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + + 2. Redistributions in binary form must reproduce the above + copyright notice, this list of conditions and the following + disclaimer in the documentation and/or other materials provided + with the distribution. + + 3. Neither the name of the copyright holder nor the names of its + contributors may be used to endorse or promote products derived + from this software without specific prior written permission. + + THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS + "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT + LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR + A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT + OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, + SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT + LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, + DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY + THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT + (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +======================================================================= +Boost's 1.0 Software License +======================================================================= + +Code derived from implementations in Boost 1.0 should mention its +derivation and reference the following license: + + Boost Software License - Version 1.0 - August 17th, 2003 + + Permission is hereby granted, free of charge, to any person or organization + obtaining a copy of the software and accompanying documentation covered by + this license (the "Software") to use, reproduce, display, distribute, + execute, and transmit the Software, and to prepare derivative works of the + Software, and to permit third-parties to whom the Software is furnished to + do so, all subject to the following: + + The copyright notices in the Software and this entire statement, including + the above license grant, this restriction and the following disclaimer, + must be included in all copies of the Software, in whole or in part, and + all derivative works of the Software, unless such copies or derivative + works are solely in the form of machine-executable object code generated by + a source language processor. + + THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + FITNESS FOR A PARTICULAR PURPOSE, TITLE AND NON-INFRINGEMENT. IN NO EVENT + SHALL THE COPYRIGHT HOLDERS OR ANYONE DISTRIBUTING THE SOFTWARE BE LIABLE + FOR ANY DAMAGES OR OTHER LIABILITY, WHETHER IN CONTRACT, TORT OR OTHERWISE, + ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER + DEALINGS IN THE SOFTWARE. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. + +======================================================================= +PILLOW-SIMD Software License +======================================================================= + +Code derived from implementations in PILLOW-SIMD should mention its derivation +and reference the following license: + + The Python Imaging Library (PIL) is + + Copyright © 1997-2011 by Secret Labs AB + Copyright © 1995-2011 by Fredrik Lundh + + Pillow is the friendly PIL fork. It is + + Copyright © 2010-2022 by Alex Clark and contributors + + Like PIL, Pillow is licensed under the open source HPND License: + + By obtaining, using, and/or copying this software and/or its associated + documentation, you agree that you have read, understood, and will comply + with the following terms and conditions: + + Permission to use, copy, modify, and distribute this software and its + associated documentation for any purpose and without fee is hereby granted, + provided that the above copyright notice appears in all copies, and that + both that copyright notice and this permission notice appear in supporting + documentation, and that the name of Secret Labs AB or the author not be + used in advertising or publicity pertaining to distribution of the software + without specific, written prior permission. + + SECRET LABS AB AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH REGARD TO THIS + SOFTWARE, INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS. + IN NO EVENT SHALL SECRET LABS AB OR THE AUTHOR BE LIABLE FOR ANY SPECIAL, + INDIRECT OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM + LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE + OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR + PERFORMANCE OF THIS SOFTWARE. diff --git a/LICENSES/TVM-FFI-0.1.10/NOTICE.txt b/LICENSES/TVM-FFI-0.1.10/NOTICE.txt new file mode 100644 index 000000000..160794b12 --- /dev/null +++ b/LICENSES/TVM-FFI-0.1.10/NOTICE.txt @@ -0,0 +1,5 @@ +Apache TVM FFI +Copyright 2024-present The Apache Software Foundation + +This product includes software developed at +The Apache Software Foundation (http://www.apache.org/). diff --git a/THIRD_PARTY.txt b/THIRD_PARTY.txt index 6a14c700b..6f172541b 100644 --- a/THIRD_PARTY.txt +++ b/THIRD_PARTY.txt @@ -17,7 +17,8 @@ Upstream projects: https://github.com/IST-DASLab/marlin and https://github.com/v Marlin notices retain Copyright 2024 Elias Frantar, modifications by Neural Magic, Apache-2.0. The original generator identifies Apache-2.0 / Copyright contributors to the vLLM project. Marlin original Apache headers remain in marlin_template.h and nvidia/kernel.cuh; Apache text: LICENSE. -Target <- direct donor; COPY means unchanged bytes, ADAPTED/GENERATED identifies modifications: +Target <- direct donor; COPY retains the donor implementation (apart from added attribution +headers); ADAPTED/PORTED/GENERATED identify changes described below. Original notices remain. ADAPTED W/awq_repack/cuda/kernel.cuh <- C/awq_marlin_repack/cuda/kernel.cuh (include-path-only; kernel body unchanged) COPY W/gemm/core/scalar_type.hpp <- C/awq_marlin_gemm/core/scalar_type.hpp COPY W/gemm/core/source_location.h <- C/awq_marlin_gemm/core/source_location.h @@ -35,6 +36,9 @@ GENERATED W/gemm/nvidia/sm80_kernel_bfloat16_u4b8_bfloat16.cu <- C/awq_marlin_ge GENERATED W/gemm/nvidia/sm80_kernel_float16_u4_float16.cu <- C/awq_marlin_gemm/nvidia/generate_kernels.py GENERATED W/gemm/nvidia/sm80_kernel_float16_u4b8_float16.cu <- C/awq_marlin_gemm/nvidia/generate_kernels.py ADAPTED W/gptq_repack/cuda/kernel.cuh <- C/gptq_marlin_repack/cuda/kernel.cuh (include-path-only; kernel body unchanged) +ADAPTED P/w4_marlin/parameters.cu implements the scale/zero permutations from +vLLM's vllm/model_executor/layers/quantization/utils/marlin_utils.py; +the fixed reference revision and loading-only CUDA adaptation are noted in that file. MoE / recurrent-state ports ADAPTED P/deepseek_moe/kernel.cu <- C/deepseek_moe/nvidia/deepseek_moe_nvidia.cu @@ -44,7 +48,7 @@ SGLang project: https://github.com/sgl-project/sglang; earlier upstream path: sg The gate keeps the original Copyright 2025 SGLang Team / Apache-2.0 notice; Apache text: LICENSE. PORTED P/chunk_gated_delta_rule/chunk_gated_delta_rule_legacy_kernel.cuh <- C/chunk_gated_delta_rule/cuda/kernel.cuh Only the active helpers/state-workspace kernel are ported; the disabled donor branch is not copied. -COPY P/recurrent_gated_delta_rule/recurrent_gated_delta_rule_legacy_kernel.cuh <- C/recurrent_gated_delta_rule/cuda/kernel.cuh +ADAPTED P/recurrent_gated_delta_rule/recurrent_gated_delta_rule_legacy_kernel.cuh <- C/recurrent_gated_delta_rule/cuda/kernel.cuh (formatting and source notice only) PORTED P/recurrent_gated_delta_rule/recurrent_delta_rule_common.cuh <- C/recurrent_gated_delta_rule/cuda/recurrent_delta_rule_common.cuh ADAPTED src/native/cuda/ops/mamba_selective_scan/kernel.h <- C/mamba_selective_scan/nvidia/mamba_selective_scan_nvidia.cu Namespace, layout/backend binding and resource adaptations are not described as source-identical copies. @@ -57,3 +61,26 @@ It preserves upstream header notices and copies the archive LICENSE and AUTHORS The adapted scheduling heuristic comes from csrc/flash_attn/flash_api.cpp and retains Copyright 2023 Tri Dao. The original archive/CUTLASS notices govern those portions; the Core MIT notice does not relicense them. Native launch templates, descriptors and workspace/lifetime integration are glue around the imported forward mathematics. + +Additional vLLM / TensorRT-LLM / Ascend adaptations +Apache-2.0 text: LICENSE; per-file source paths, revisions and original notices are retained: +src/native/cuda/ops/topk_softmax/kernel.cuh <- vLLM 2f75e7f712fb2a013ce05ff357d94135231c8ae2 +P/topk_sigmoid/kernel.cuh <- vLLM 88402a41c4ab272ebbbd33f4a77fbbac0431cbb9 +P/grouped_topk/kernel.cuh <- vLLM ffc4f08c8ee130d4ea6347c1bf31ffd4f8af28ab +P/get_cutlass_moe_mm_data/kernel.cuh <- vLLM 2f75e7f712fb2a013ce05ff357d94135231c8ae2 +P/moe_wna16_gemm/kernel.cu and kernel.cuh <- vLLM ffc4f08c8ee130d4ea6347c1bf31ffd4f8af28ab +src/native/ascend/custom/{ops.h,torch_binding.cpp,utils/torch_kernel_helper.h} +retain their Huawei/InfiniTensor notices and vLLM-Ascend source links. + +Distribution +SDK installation includes LICENSE, this attribution list and LICENSES under +share/licenses/InfiniOps. Native FlashAttention installation also includes the exact +archive LICENSE and AUTHORS and the bundled CUTLASS header's complete license block. +CUTLASS-BSD-3-Clause.txt records the separate CMake dependency at +087c84df83d254b5fb295a7a408f1a1d554085cf. CUTLASS notices apply separately to its compiled headers. +Linked libraries are not relicensed by Ops: distributions bundling sampling.so or +libtvm_ffi.so must retain the resolved FlashInfer/TVM-FFI licenses, NOTICE files and +applicable third-party notices. See docs/build-nvidia-inference.md for the pinned recipe. +LICENSES/FlashInfer-0.6.7 and LICENSES/TVM-FFI-0.1.10 retain the corresponding +upstream release LICENSE, NOTICE and licenses/ materials for that recipe verbatim. +They are not a substitute for the notices of a different resolved binary version. diff --git a/cmake/NativeFlashAttention.cmake b/cmake/NativeFlashAttention.cmake index 7cdef6024..b1ca0bdea 100644 --- a/cmake/NativeFlashAttention.cmake +++ b/cmake/NativeFlashAttention.cmake @@ -49,4 +49,13 @@ function(infiniops_add_native_flash target) # Match the pinned upstream setup.py numerical compilation mode. "$<$:--expt-relaxed-constexpr;--expt-extended-lambda;--use_fast_math>") target_link_libraries(${target} PRIVATE CUDA::cudart) + install(FILES "${_prepared}/LICENSE" + DESTINATION "${CMAKE_INSTALL_DATADIR}/licenses/InfiniOps" + RENAME FlashAttention-LICENSE.txt) + install(FILES "${_prepared}/AUTHORS" + DESTINATION "${CMAKE_INSTALL_DATADIR}/licenses/InfiniOps" + RENAME FlashAttention-AUTHORS.txt) + install(FILES "${_prepared}/CUTLASS-LICENSE.txt" + DESTINATION "${CMAKE_INSTALL_DATADIR}/licenses/InfiniOps" + RENAME FlashAttention-CUTLASS-LICENSE.txt) endfunction() diff --git a/docs/build-nvidia-inference.md b/docs/build-nvidia-inference.md index 1c8e664f8..80856218a 100644 --- a/docs/build-nvidia-inference.md +++ b/docs/build-nvidia-inference.md @@ -26,6 +26,12 @@ Installed library hashes for the two wheels above: The YAML version ranges describe the provider's accepted packages, not a lock for a release. A release using a different wheel must record its resolved version/hash and run the corresponding direct operator tests again. +`LICENSES/FlashInfer-0.6.7` and `LICENSES/TVM-FFI-0.1.10` retain the original +release licenses, notices and third-party license texts for these exact inputs. +If replacing either wheel, update those redistribution materials from the +actual release as well. SDK installation puts these materials, the native +FlashAttention notices and the Ops attribution list in `share/licenses/InfiniOps`; +downstream binary packages must retain that directory's contents. ## Acquire and install the fixed inputs diff --git a/scripts/prepare_flash_attention_inference.py b/scripts/prepare_flash_attention_inference.py index 0b55faff6..6b155c0de 100644 --- a/scripts/prepare_flash_attention_inference.py +++ b/scripts/prepare_flash_attention_inference.py @@ -69,6 +69,17 @@ def adapt(name, text): text = replace_once( text, " exit(1);", " INFINI_FLASH_CUDA_CHECK(status_);" ) + if name in ( + "flash.h", + "flash_fwd_kernel.h", + "flash_fwd_launch_template.h", + "hardware_info.h", + ): + text = ( + "// Modified by InfiniOps from FlashAttention 2.8.3: inference-only native runtime integration.\n" + "// Upstream: https://github.com/Dao-AILab/flash-attention/tree/v2.8.3/csrc/flash_attn/src\n" + "// Original copyright and license notices below are retained.\n" + text + ) return text @@ -117,6 +128,8 @@ def read(name): ) headers["num_splits.h"] = ( "// Copyright (c) 2023, Tri Dao. See accompanying LICENSE.\n" + "// Extracted from flash_attn-2.8.3/csrc/flash_attn/flash_api.cpp;\n" + "// modified by InfiniOps to remove Torch dependencies and use a private namespace.\n" "#pragma once\n#include \n#include \n#include \n" "namespace infini_flash {\n" + api[begin:end] + "\n}\n" ) @@ -135,6 +148,13 @@ def read(name): write_if_changed(path, read(member.name)) for name in ("LICENSE", "AUTHORS"): write_if_changed(output / name, read(ROOT + name)) + # The sdist omits CUTLASS's root LICENSE.txt; preserve the full license + # block from its actual bundled header, rather than substituting main. + cutlass_header = read(ROOT + "csrc/cutlass/include/cutlass/cutlass.h").decode() + cutlass_notice = re.match(r"\s*(/\*.*?\*/)", cutlass_header, re.DOTALL) + if cutlass_notice is None or "Redistribution" not in cutlass_notice.group(1): + raise ValueError("missing pinned CUTLASS copyright/license block") + write_if_changed(output / "CUTLASS-LICENSE.txt", cutlass_notice.group(1) + "\n") manifest = { "upstream": "flash-attn", "version": "2.8.3", diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 9ecfee576..9a30c7337 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -1690,6 +1690,11 @@ install(TARGETS infiniops RUNTIME DESTINATION ${CMAKE_INSTALL_BINDIR} ) +install(FILES "${PROJECT_SOURCE_DIR}/LICENSE" "${PROJECT_SOURCE_DIR}/THIRD_PARTY.txt" + DESTINATION "${CMAKE_INSTALL_DATADIR}/licenses/InfiniOps") +install(DIRECTORY "${PROJECT_SOURCE_DIR}/LICENSES/" + DESTINATION "${CMAKE_INSTALL_DATADIR}/licenses/InfiniOps") + set(INFINI_OPS_INSTALL_CMAKEDIR "${CMAKE_INSTALL_LIBDIR}/cmake/InfiniOps") install(EXPORT InfiniOpsTargets diff --git a/src/native/cuda/nvidia/ops/chunk_gated_delta_rule/chunk_gated_delta_rule_legacy_kernel.cuh b/src/native/cuda/nvidia/ops/chunk_gated_delta_rule/chunk_gated_delta_rule_legacy_kernel.cuh index 8a8bdcfce..733b71c1c 100644 --- a/src/native/cuda/nvidia/ops/chunk_gated_delta_rule/chunk_gated_delta_rule_legacy_kernel.cuh +++ b/src/native/cuda/nvidia/ops/chunk_gated_delta_rule/chunk_gated_delta_rule_legacy_kernel.cuh @@ -1,3 +1,7 @@ +// Ported from InfiniCore 1ab85ef1d3c2bce49edf1fc73a5b48becaf7552e: +// src/infiniop/ops/chunk_gated_delta_rule/cuda/kernel.cuh. +// Modified for InfiniOps: active helpers/state kernel and namespace adaptation. +// Copyright 2025 InfiniTensor; MIT notice: LICENSES/InfiniCore-MIT.txt. #pragma once #include // Exact baseline helpers and active state-workspace kernel. The baseline diff --git a/src/native/cuda/nvidia/ops/deepseek_moe/kernel.cu b/src/native/cuda/nvidia/ops/deepseek_moe/kernel.cu index e149f42ad..a76794471 100644 --- a/src/native/cuda/nvidia/ops/deepseek_moe/kernel.cu +++ b/src/native/cuda/nvidia/ops/deepseek_moe/kernel.cu @@ -1,4 +1,8 @@ -// Actual fused decode, batched cuBLAS, expert grouped/chunked CUDA bodies ported from frozen InfiniCore 0.2.9. +// Adapted from InfiniCore 1ab85ef1d3c2bce49edf1fc73a5b48becaf7552e: +// src/infiniop/ops/deepseek_moe/nvidia/deepseek_moe_nvidia.cu. +// Modified for InfiniOps: backend, layout and resource integration. +// Copyright 2025 InfiniTensor; MIT notice: LICENSES/InfiniCore-MIT.txt. +// Additional upstream notices are retained; see repository THIRD_PARTY.txt. #include "native/cuda/nvidia/ops/deepseek_moe/kernel.h" #include "native/cuda/nvidia/ops/moe_fused_gate/device_guard.cuh" #include diff --git a/src/native/cuda/nvidia/ops/moe_fused_dense/kernel.cu b/src/native/cuda/nvidia/ops/moe_fused_dense/kernel.cu index b7171cc4b..f80fb5ce8 100644 --- a/src/native/cuda/nvidia/ops/moe_fused_dense/kernel.cu +++ b/src/native/cuda/nvidia/ops/moe_fused_dense/kernel.cu @@ -1,4 +1,8 @@ -// CUDA/CUTLASS body ported from frozen InfiniCore 0.2.9. See THIRD_PARTY.txt for source and license details. +// Adapted from InfiniCore 1ab85ef1d3c2bce49edf1fc73a5b48becaf7552e: +// src/infiniop/ops/moe_fused_dense/nvidia/moe_fused_dense_nvidia.cu. +// Modified for InfiniOps: backend, layout and resource integration. +// Copyright 2025 InfiniTensor; MIT notice: LICENSES/InfiniCore-MIT.txt. +// CUTLASS headers retain their own license; see repository THIRD_PARTY.txt. #include "native/cuda/nvidia/ops/moe_fused_dense/kernel.h" #include "native/cuda/nvidia/ops/moe_fused_gate/device_guard.cuh" #include diff --git a/src/native/cuda/nvidia/ops/moe_fused_gate/kernel.cu b/src/native/cuda/nvidia/ops/moe_fused_gate/kernel.cu index e0dca49b6..f6f49f021 100644 --- a/src/native/cuda/nvidia/ops/moe_fused_gate/kernel.cu +++ b/src/native/cuda/nvidia/ops/moe_fused_gate/kernel.cu @@ -1,3 +1,7 @@ +// Adapted via InfiniCore 1ab85ef1d3c2bce49edf1fc73a5b48becaf7552e: +// src/infiniop/ops/moe_fused_gate/nvidia/moe_fused_gate_nvidia.cu. +// Modified for InfiniOps: native backend and resource integration. +// Original SGLang notice below applies; see repository THIRD_PARTY.txt. /* * Portions of the CUDA kernels in this file are adapted from SGLang: * /sgl-kernel/csrc/moe/moe_fused_gate.cu diff --git a/src/native/cuda/nvidia/ops/recurrent_gated_delta_rule/recurrent_delta_rule_common.cuh b/src/native/cuda/nvidia/ops/recurrent_gated_delta_rule/recurrent_delta_rule_common.cuh index cb6421be8..e5855f8ab 100644 --- a/src/native/cuda/nvidia/ops/recurrent_gated_delta_rule/recurrent_delta_rule_common.cuh +++ b/src/native/cuda/nvidia/ops/recurrent_gated_delta_rule/recurrent_delta_rule_common.cuh @@ -1,3 +1,7 @@ +// Ported from InfiniCore 1ab85ef1d3c2bce49edf1fc73a5b48becaf7552e: +// src/infiniop/ops/recurrent_gated_delta_rule/cuda/recurrent_delta_rule_common.cuh. +// Modified for InfiniOps: NVIDIA-only warp definitions and backend integration. +// Copyright 2025 InfiniTensor; MIT notice: LICENSES/InfiniCore-MIT.txt. #ifndef __RECURRENT_DELTA_RULE_COMMON_CUH__ #define __RECURRENT_DELTA_RULE_COMMON_CUH__ diff --git a/src/native/cuda/nvidia/ops/recurrent_gated_delta_rule/recurrent_gated_delta_rule_legacy_kernel.cuh b/src/native/cuda/nvidia/ops/recurrent_gated_delta_rule/recurrent_gated_delta_rule_legacy_kernel.cuh index 138678422..a50ceaa1f 100644 --- a/src/native/cuda/nvidia/ops/recurrent_gated_delta_rule/recurrent_gated_delta_rule_legacy_kernel.cuh +++ b/src/native/cuda/nvidia/ops/recurrent_gated_delta_rule/recurrent_gated_delta_rule_legacy_kernel.cuh @@ -1,3 +1,7 @@ +// Copied from InfiniCore 1ab85ef1d3c2bce49edf1fc73a5b48becaf7552e: +// src/infiniop/ops/recurrent_gated_delta_rule/cuda/kernel.cuh. +// Modified for InfiniOps: formatting and this source notice only. +// Copyright 2025 InfiniTensor; MIT notice: LICENSES/InfiniCore-MIT.txt. #ifndef __RECURRENT_GATED_DELTA_RULE_KERNEL_CUH__ #define __RECURRENT_GATED_DELTA_RULE_KERNEL_CUH__ diff --git a/src/native/cuda/nvidia/ops/w4_marlin/parameters.cu b/src/native/cuda/nvidia/ops/w4_marlin/parameters.cu index ad22e7ba3..c6ce73586 100644 --- a/src/native/cuda/nvidia/ops/w4_marlin/parameters.cu +++ b/src/native/cuda/nvidia/ops/w4_marlin/parameters.cu @@ -1,8 +1,14 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright contributors to the vLLM project. +// Scale/zero permutations follow vLLM's Marlin utilities: +// https://github.com/vllm-project/vllm/blob/ffc4f08c8ee130d4ea6347c1bf31ffd4f8af28ab/vllm/model_executor/layers/quantization/utils/marlin_utils.py +// Modified for InfiniOps: loading-only CUDA kernels, bit-preserving FP16/BF16 output. +// See repository THIRD_PARTY.txt and LICENSE. #include "kernel.h" #include "launch.h" namespace infini::ops::native_w4 { -// These are the byte permutations from the original marlin_utils.cpp; the +// These are the byte permutations from the vLLM Marlin utilities cited above; the // arithmetic keeps FP16/BF16 bit patterns unchanged. __device__ int scale_source(int i, bool single) { if (!single) return (i / 8) + 8 * (i % 8); diff --git a/src/native/cuda/nvidia/ops/w4_marlin/vendor/awq_repack/cuda/kernel.cuh b/src/native/cuda/nvidia/ops/w4_marlin/vendor/awq_repack/cuda/kernel.cuh index a448c54e4..c05c3a5ff 100644 --- a/src/native/cuda/nvidia/ops/w4_marlin/vendor/awq_repack/cuda/kernel.cuh +++ b/src/native/cuda/nvidia/ops/w4_marlin/vendor/awq_repack/cuda/kernel.cuh @@ -1,3 +1,7 @@ +// Adapted from InfiniCore 1ab85ef1d3c2bce49edf1fc73a5b48becaf7552e: +// src/infiniop/ops/awq_marlin_repack/cuda/kernel.cuh (Marlin/vLLM). +// Modified for InfiniOps: relative include paths; kernel mathematics retained. +// Original licenses and copyright notices: repository THIRD_PARTY.txt and LICENSE. #include "../../gemm/marlin/marlin.cuh" namespace marlin { diff --git a/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/core/scalar_type.hpp b/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/core/scalar_type.hpp index 7cc0aa09d..ba45d7245 100644 --- a/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/core/scalar_type.hpp +++ b/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/core/scalar_type.hpp @@ -1,3 +1,7 @@ +// Source: InfiniCore 1ab85ef1d3c2bce49edf1fc73a5b48becaf7552e, +// src/infiniop/ops/awq_marlin_gemm/core/scalar_type.hpp. +// Donor implementation retained; this attribution notice was added by InfiniOps. +// Marlin/vLLM and direct-donor licenses: repository THIRD_PARTY.txt. #pragma once #include diff --git a/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/core/source_location.h b/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/core/source_location.h index 9a06fb380..74ef9c422 100644 --- a/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/core/source_location.h +++ b/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/core/source_location.h @@ -1,3 +1,7 @@ +// Source: InfiniCore 1ab85ef1d3c2bce49edf1fc73a5b48becaf7552e, +// src/infiniop/ops/awq_marlin_gemm/core/source_location.h. +// Donor implementation retained; this attribution notice was added by InfiniOps. +// Marlin/vLLM and direct-donor licenses: repository THIRD_PARTY.txt. /// \file source_location.h /// \brief Portable `source_location` wrapper. /// diff --git a/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/core/utils.h b/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/core/utils.h index a6b6cddfd..f1497cd60 100644 --- a/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/core/utils.h +++ b/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/core/utils.h @@ -1,3 +1,7 @@ +// Adapted from InfiniCore 1ab85ef1d3c2bce49edf1fc73a5b48becaf7552e: +// src/infiniop/ops/awq_marlin_gemm/core/utils.h. +// Modified for InfiniOps: unused pointer, range and dtype helpers removed. +// See repository THIRD_PARTY.txt and LICENSES/InfiniCore-MIT.txt. /// \file utils.h /// \brief Host-side C++ utilities used by JIT kernel wrappers. /// diff --git a/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/marlin/dequant.h b/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/marlin/dequant.h index 1340b6bbe..b4597ea90 100644 --- a/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/marlin/dequant.h +++ b/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/marlin/dequant.h @@ -1,3 +1,7 @@ +// Source: InfiniCore 1ab85ef1d3c2bce49edf1fc73a5b48becaf7552e, +// src/infiniop/ops/awq_marlin_gemm/marlin/dequant.h. +// Donor implementation retained; this attribution notice was added by InfiniOps. +// Marlin/vLLM and direct-donor licenses: repository THIRD_PARTY.txt. /* Fast Dequantization (Converting INT4/INT8/FP4/FP8 to FP16/BF16) diff --git a/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/marlin/kernel.h b/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/marlin/kernel.h index b1fabe323..bcfac3e5a 100644 --- a/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/marlin/kernel.h +++ b/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/marlin/kernel.h @@ -1,4 +1,8 @@ +// Source: InfiniCore 1ab85ef1d3c2bce49edf1fc73a5b48becaf7552e, +// src/infiniop/ops/awq_marlin_gemm/marlin/kernel.h. +// Donor implementation retained; this attribution notice was added by InfiniOps. +// Marlin/vLLM and direct-donor licenses: repository THIRD_PARTY.txt. #ifndef MARLIN_NAMESPACE_NAME #define MARLIN_NAMESPACE_NAME marlin #endif diff --git a/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/marlin/marlin.cuh b/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/marlin/marlin.cuh index f3d897d27..a752519b2 100644 --- a/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/marlin/marlin.cuh +++ b/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/marlin/marlin.cuh @@ -1,3 +1,7 @@ +// Source: InfiniCore 1ab85ef1d3c2bce49edf1fc73a5b48becaf7552e, +// src/infiniop/ops/awq_marlin_gemm/marlin/marlin.cuh. +// Donor implementation retained; this attribution notice was added by InfiniOps. +// Marlin/vLLM and direct-donor licenses: repository THIRD_PARTY.txt. #pragma once #ifndef _marlin_cuh diff --git a/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/marlin/marlin_dtypes.cuh b/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/marlin/marlin_dtypes.cuh index 3ee06faef..98340b53f 100644 --- a/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/marlin/marlin_dtypes.cuh +++ b/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/marlin/marlin_dtypes.cuh @@ -1,4 +1,8 @@ +// Source: InfiniCore 1ab85ef1d3c2bce49edf1fc73a5b48becaf7552e, +// src/infiniop/ops/awq_marlin_gemm/marlin/marlin_dtypes.cuh. +// Donor implementation retained; this attribution notice was added by InfiniOps. +// Marlin/vLLM and direct-donor licenses: repository THIRD_PARTY.txt. #ifndef _data_types_cuh #define _data_types_cuh #include "../core/scalar_type.hpp" diff --git a/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/marlin/marlin_mma.h b/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/marlin/marlin_mma.h index d8fb2946e..ca9ca73fb 100644 --- a/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/marlin/marlin_mma.h +++ b/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/marlin/marlin_mma.h @@ -1,4 +1,8 @@ +// Source: InfiniCore 1ab85ef1d3c2bce49edf1fc73a5b48becaf7552e, +// src/infiniop/ops/awq_marlin_gemm/marlin/marlin_mma.h. +// Donor implementation retained; this attribution notice was added by InfiniOps. +// Marlin/vLLM and direct-donor licenses: repository THIRD_PARTY.txt. #include "marlin_dtypes.cuh" namespace MARLIN_NAMESPACE_NAME { diff --git a/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/marlin/marlin_template.h b/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/marlin/marlin_template.h index f78d883c5..7415dd971 100644 --- a/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/marlin/marlin_template.h +++ b/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/marlin/marlin_template.h @@ -1,3 +1,7 @@ +// Adapted from InfiniCore 1ab85ef1d3c2bce49edf1fc73a5b48becaf7552e: +// src/infiniop/ops/awq_marlin_gemm/marlin/marlin_template.h. +// Modified for InfiniOps: native launcher integration; original notices retained. +// See repository THIRD_PARTY.txt and LICENSE (Apache-2.0). /* * Modified by Neural Magic * Copyright (C) Marlin.2024 Elias Frantar diff --git a/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/nvidia/kernel.cuh b/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/nvidia/kernel.cuh index 80b9fb5a3..21e5367b8 100644 --- a/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/nvidia/kernel.cuh +++ b/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/nvidia/kernel.cuh @@ -1,3 +1,7 @@ +// Source: InfiniCore 1ab85ef1d3c2bce49edf1fc73a5b48becaf7552e, +// src/infiniop/ops/awq_marlin_gemm/nvidia/kernel.cuh. +// Donor implementation retained; this attribution notice was added by InfiniOps. +// Marlin/vLLM and direct-donor licenses: repository THIRD_PARTY.txt. /* * Modified by Neural Magic * Copyright (C) Marlin.2024 Elias Frantar diff --git a/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/nvidia/kernel_selector.h b/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/nvidia/kernel_selector.h index 56acf17ba..5a07f4219 100644 --- a/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/nvidia/kernel_selector.h +++ b/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/nvidia/kernel_selector.h @@ -1,3 +1,8 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright contributors to the vLLM project. +// Generated from InfiniCore 1ab85ef1d3c2bce49edf1fc73a5b48becaf7552e, +// src/infiniop/ops/awq_marlin_gemm/nvidia/generate_kernels.py. +// Modified for InfiniOps: W4A16-only dispatch. See THIRD_PARTY.txt and LICENSE. if (a_type == vllm::kFloat16 && b_type == vllm::kU4 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == -1 && is_zp_float == false) kernel = Marlin; else if (a_type == vllm::kFloat16 && b_type == vllm::kU4 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == true && stages == 4 && group_blocks == -1 && is_zp_float == false) diff --git a/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/nvidia/sm80_kernel_bfloat16_u4_bfloat16.cu b/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/nvidia/sm80_kernel_bfloat16_u4_bfloat16.cu index 5cd175e45..c9f344cc8 100644 --- a/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/nvidia/sm80_kernel_bfloat16_u4_bfloat16.cu +++ b/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/nvidia/sm80_kernel_bfloat16_u4_bfloat16.cu @@ -1,4 +1,9 @@ -// Derived from fixed original generate_kernels.py; W4A16 only. +// SPDX-License-Identifier: Apache-2.0 +// Copyright contributors to the vLLM project. +// Generated from InfiniCore 1ab85ef1d3c2bce49edf1fc73a5b48becaf7552e, +// src/infiniop/ops/awq_marlin_gemm/nvidia/generate_kernels.py. +// Modified for InfiniOps: W4A16-only instantiations and relative includes. +// See repository THIRD_PARTY.txt and LICENSE. #include "../marlin/kernel.h" #include "../marlin/marlin_template.h" namespace MARLIN_NAMESPACE_NAME { diff --git a/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/nvidia/sm80_kernel_bfloat16_u4b8_bfloat16.cu b/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/nvidia/sm80_kernel_bfloat16_u4b8_bfloat16.cu index 316c4d20f..f3af34444 100644 --- a/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/nvidia/sm80_kernel_bfloat16_u4b8_bfloat16.cu +++ b/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/nvidia/sm80_kernel_bfloat16_u4b8_bfloat16.cu @@ -1,4 +1,9 @@ -// Derived from fixed original generate_kernels.py; W4A16 only. +// SPDX-License-Identifier: Apache-2.0 +// Copyright contributors to the vLLM project. +// Generated from InfiniCore 1ab85ef1d3c2bce49edf1fc73a5b48becaf7552e, +// src/infiniop/ops/awq_marlin_gemm/nvidia/generate_kernels.py. +// Modified for InfiniOps: W4A16-only instantiations and relative includes. +// See repository THIRD_PARTY.txt and LICENSE. #include "../marlin/kernel.h" #include "../marlin/marlin_template.h" namespace MARLIN_NAMESPACE_NAME { diff --git a/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/nvidia/sm80_kernel_float16_u4_float16.cu b/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/nvidia/sm80_kernel_float16_u4_float16.cu index 06b7faa04..bf51abb2c 100644 --- a/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/nvidia/sm80_kernel_float16_u4_float16.cu +++ b/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/nvidia/sm80_kernel_float16_u4_float16.cu @@ -1,4 +1,9 @@ -// Derived from fixed original generate_kernels.py; W4A16 only. +// SPDX-License-Identifier: Apache-2.0 +// Copyright contributors to the vLLM project. +// Generated from InfiniCore 1ab85ef1d3c2bce49edf1fc73a5b48becaf7552e, +// src/infiniop/ops/awq_marlin_gemm/nvidia/generate_kernels.py. +// Modified for InfiniOps: W4A16-only instantiations and relative includes. +// See repository THIRD_PARTY.txt and LICENSE. #include "../marlin/kernel.h" #include "../marlin/marlin_template.h" namespace MARLIN_NAMESPACE_NAME { diff --git a/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/nvidia/sm80_kernel_float16_u4b8_float16.cu b/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/nvidia/sm80_kernel_float16_u4b8_float16.cu index 519f7082b..d35826b7c 100644 --- a/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/nvidia/sm80_kernel_float16_u4b8_float16.cu +++ b/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/nvidia/sm80_kernel_float16_u4b8_float16.cu @@ -1,4 +1,9 @@ -// Derived from fixed original generate_kernels.py; W4A16 only. +// SPDX-License-Identifier: Apache-2.0 +// Copyright contributors to the vLLM project. +// Generated from InfiniCore 1ab85ef1d3c2bce49edf1fc73a5b48becaf7552e, +// src/infiniop/ops/awq_marlin_gemm/nvidia/generate_kernels.py. +// Modified for InfiniOps: W4A16-only instantiations and relative includes. +// See repository THIRD_PARTY.txt and LICENSE. #include "../marlin/kernel.h" #include "../marlin/marlin_template.h" namespace MARLIN_NAMESPACE_NAME { diff --git a/src/native/cuda/nvidia/ops/w4_marlin/vendor/gptq_repack/cuda/kernel.cuh b/src/native/cuda/nvidia/ops/w4_marlin/vendor/gptq_repack/cuda/kernel.cuh index f1e623ea4..8b9bed21f 100644 --- a/src/native/cuda/nvidia/ops/w4_marlin/vendor/gptq_repack/cuda/kernel.cuh +++ b/src/native/cuda/nvidia/ops/w4_marlin/vendor/gptq_repack/cuda/kernel.cuh @@ -1,3 +1,7 @@ +// Adapted from InfiniCore 1ab85ef1d3c2bce49edf1fc73a5b48becaf7552e: +// src/infiniop/ops/gptq_marlin_repack/cuda/kernel.cuh (Marlin/vLLM). +// Modified for InfiniOps: relative include paths; kernel mathematics retained. +// Original licenses and copyright notices: repository THIRD_PARTY.txt and LICENSE. #include "../../gemm/marlin/marlin.cuh" namespace marlin { diff --git a/src/native/cuda/ops/mamba_selective_scan/kernel.h b/src/native/cuda/ops/mamba_selective_scan/kernel.h index 9eb0b5cb0..ed6bdbbda 100644 --- a/src/native/cuda/ops/mamba_selective_scan/kernel.h +++ b/src/native/cuda/ops/mamba_selective_scan/kernel.h @@ -1,3 +1,7 @@ +// Adapted from InfiniCore 1ab85ef1d3c2bce49edf1fc73a5b48becaf7552e: +// src/infiniop/ops/mamba_selective_scan/nvidia/mamba_selective_scan_nvidia.cu. +// Modified for InfiniOps: layout, type dispatch and runtime integration. +// Copyright 2025 InfiniTensor; MIT notice: LICENSES/InfiniCore-MIT.txt. #ifndef INFINI_OPS_CUDA_MAMBA_SELECTIVE_SCAN_KERNEL_H_ #define INFINI_OPS_CUDA_MAMBA_SELECTIVE_SCAN_KERNEL_H_ #include From 254c4343b569ec49837a13ad5830f945d6337f10 Mon Sep 17 00:00:00 2001 From: whjthu Date: Mon, 19 Oct 2026 10:56:06 +0800 Subject: [PATCH 17/25] fix: retain optional Marlin public operator interfaces --- scripts/generate_wrappers.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/scripts/generate_wrappers.py b/scripts/generate_wrappers.py index 22f2643b6..975390ee0 100644 --- a/scripts/generate_wrappers.py +++ b/scripts/generate_wrappers.py @@ -2016,7 +2016,11 @@ def _get_all_ops( impl_headers_by_operator.get(_op_relative_type(op_name), ()) ) - if not implementations: + if not implementations and op_name not in { + "w4_marlin_gemm", + "w4_marlin_repack", + "w4_marlin_parameters", + }: continue ops[op_name] = implementations @@ -2246,6 +2250,11 @@ def _dispatch_gen_batch_size(): if config_path is not None: config = ops_config.load_ops_config(config_path) + # Optional Marlin providers do not remove the public checkpoint API. + # Reuse the existing empty-provider generation path when none is built. + for name in ("w4_marlin_gemm", "w4_marlin_repack", "w4_marlin_parameters"): + if name not in config or (name in ops and not ops[name]): + config[name] = {"public_api_only": True} ops = _select_ops_from_config(ops, config, config_path) ops = _filter_ops( From 945184c02d767f9dfc343ba7a17f3409f111f32c Mon Sep 17 00:00:00 2001 From: whjthu Date: Mon, 19 Oct 2026 23:03:24 +0800 Subject: [PATCH 18/25] refactor: simplify migrated providers and pin upstream kernel sources --- .gitmodules | 6 + CMakeLists.txt | 1 - LICENSES/Apache-2.0-components.txt | 82 + LICENSES/FasterTransformer-BSD-3-Clause.txt | 27 + LICENSES/FlashInfer-0.6.7/LICENSE.txt | 229 -- LICENSES/TVM-FFI-0.1.10/LICENSE.dlpack.txt | 201 -- LICENSES/TVM-FFI-0.1.10/LICENSE.txt | 217 -- THIRD_PARTY.txt | 64 +- cmake/NativeFlashAttention.cmake | 35 +- cmake/NativeW4Marlin.cmake | 38 +- docs/build-nvidia-inference.md | 66 +- patches/flash-attention/inference.patch | 148 ++ patches/marlin/w4a16.patch | 420 ++++ pyproject.toml | 19 +- scripts/generate_public_operator_header.py | 11 + scripts/generate_wrappers.py | 44 +- scripts/prepare_flash_attention_inference.py | 251 +-- scripts/prepare_marlin.py | 88 + scripts/resolve_linked_ops.py | 41 + src/CMakeLists.txt | 27 +- src/base/argmax.h | 2 + src/base/causal_conv1d.h | 11 +- src/base/chunk_gated_delta_rule.h | 11 +- src/base/gemm.h | 22 +- src/base/mamba_selective_scan.h | 6 +- src/base/recurrent_gated_delta_rule.h | 14 +- src/base/sum.h | 2 + src/base/take.h | 2 + src/base/upsample_bilinear2d.h | 2 + src/native/cuda/device_buffer.h | 56 + src/native/cuda/nvidia/ops/argmax/kernel.h | 48 +- .../cuda/nvidia/ops/causal_conv1d/kernel.h | 16 + .../ops/chunk_gated_delta_rule/kernel.h | 3 + .../ops/flash_attn_varlen_func/kernel.h | 2 +- .../ops/flash_attn_with_kvcache/kernel.cc | 23 +- .../ops/flash_attn_with_kvcache/kernel.h | 16 +- src/native/cuda/nvidia/ops/flash_head_dim.h | 22 +- src/native/cuda/nvidia/ops/gemm/cublas.h | 17 + src/native/cuda/nvidia/ops/gemm/cublaslt.h | 17 + .../cuda/nvidia/ops/layer_norm/kernel.h | 67 +- .../cuda/nvidia/ops/pad_constant/kernel.h | 16 - .../ops/recurrent_gated_delta_rule/kernel.h | 3 + src/native/cuda/nvidia/ops/sum/kernel.h | 27 +- src/native/cuda/nvidia/ops/take/kernel.h | 18 +- .../nvidia/ops/upsample_bilinear2d/kernel.h | 31 +- .../vendor/awq_repack/cuda/kernel.cuh | 201 -- .../w4_marlin/vendor/gemm/marlin/dequant.h | 605 ----- .../ops/w4_marlin/vendor/gemm/marlin/kernel.h | 47 - .../w4_marlin/vendor/gemm/marlin/marlin.cuh | 182 -- .../vendor/gemm/marlin/marlin_dtypes.cuh | 159 -- .../w4_marlin/vendor/gemm/marlin/marlin_mma.h | 271 --- .../vendor/gemm/marlin/marlin_template.h | 1992 ----------------- .../w4_marlin/vendor/gemm/nvidia/kernel.cuh | 2 +- .../vendor/gemm/nvidia/kernel_selector.h | 545 ----- .../sm80_kernel_bfloat16_u4_bfloat16.cu | 70 - .../sm80_kernel_bfloat16_u4b8_bfloat16.cu | 85 - .../nvidia/sm80_kernel_float16_u4_float16.cu | 70 - .../sm80_kernel_float16_u4b8_float16.cu | 85 - .../vendor/gptq_repack/cuda/kernel.cuh | 256 --- src/native/cuda/ops/add/kernel.h | 103 +- src/native/cuda/ops/copy/kernel.h | 86 +- src/native/cuda/ops/embedding/kernel.h | 115 +- src/native/cuda/ops/fill/kernel.h | 67 +- src/native/cuda/ops/gelu/kernel.h | 84 +- .../cuda/ops/mamba_selective_scan/kernel.h | 31 +- src/native/cuda/ops/mul/kernel.h | 101 +- src/native/cuda/ops/sigmoid/kernel.h | 85 +- src/native/cuda/ops/silu/kernel.h | 83 +- src/native/cuda/ops/silu_and_mul/kernel.h | 84 +- src/native/cuda/ops/softmax/kernel.h | 69 +- src/operator.h | 9 + tests/test_generate_wrappers.py | 50 +- third_party/flash-attention | 1 + third_party/vllm | 1 + 74 files changed, 1783 insertions(+), 6225 deletions(-) create mode 100644 LICENSES/Apache-2.0-components.txt create mode 100644 LICENSES/FasterTransformer-BSD-3-Clause.txt delete mode 100644 LICENSES/FlashInfer-0.6.7/LICENSE.txt delete mode 100644 LICENSES/TVM-FFI-0.1.10/LICENSE.dlpack.txt delete mode 100644 LICENSES/TVM-FFI-0.1.10/LICENSE.txt create mode 100644 patches/flash-attention/inference.patch create mode 100644 patches/marlin/w4a16.patch create mode 100644 scripts/prepare_marlin.py create mode 100644 src/native/cuda/device_buffer.h delete mode 100644 src/native/cuda/nvidia/ops/w4_marlin/vendor/awq_repack/cuda/kernel.cuh delete mode 100644 src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/marlin/dequant.h delete mode 100644 src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/marlin/kernel.h delete mode 100644 src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/marlin/marlin.cuh delete mode 100644 src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/marlin/marlin_dtypes.cuh delete mode 100644 src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/marlin/marlin_mma.h delete mode 100644 src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/marlin/marlin_template.h delete mode 100644 src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/nvidia/kernel_selector.h delete mode 100644 src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/nvidia/sm80_kernel_bfloat16_u4_bfloat16.cu delete mode 100644 src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/nvidia/sm80_kernel_bfloat16_u4b8_bfloat16.cu delete mode 100644 src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/nvidia/sm80_kernel_float16_u4_float16.cu delete mode 100644 src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/nvidia/sm80_kernel_float16_u4b8_float16.cu delete mode 100644 src/native/cuda/nvidia/ops/w4_marlin/vendor/gptq_repack/cuda/kernel.cuh create mode 160000 third_party/flash-attention create mode 160000 third_party/vllm diff --git a/.gitmodules b/.gitmodules index be99e8a8f..77094cdd7 100644 --- a/.gitmodules +++ b/.gitmodules @@ -1,3 +1,9 @@ [submodule ".ci"] path = .ci url = https://github.com/InfiniTensor/ci.git +[submodule "third_party/flash-attention"] + path = third_party/flash-attention + url = https://github.com/Dao-AILab/flash-attention.git +[submodule "third_party/vllm"] + path = third_party/vllm + url = https://github.com/vllm-project/vllm.git diff --git a/CMakeLists.txt b/CMakeLists.txt index 63f505d6d..b4cdcb9f1 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -28,7 +28,6 @@ option(WITH_LINKED "Enable linked third-party operator backends" OFF) option(BUILD_NVIDIA_FLASH_INFERENCE "Build pinned Torch-free NVIDIA inference Attention" OFF) option(BUILD_NVIDIA_W4_MARLIN "Build NVIDIA AWQ/GPTQ Marlin providers (CUDA >= 12.3)" OFF) -set(FLASH_ATTN_SOURCE_ARCHIVE "" CACHE FILEPATH "Exact flash_attn-2.8.3 source archive") if(BUILD_NVIDIA_FLASH_INFERENCE AND NOT WITH_NVIDIA) message(FATAL_ERROR "Native FlashAttention requires WITH_NVIDIA") endif() diff --git a/LICENSES/Apache-2.0-components.txt b/LICENSES/Apache-2.0-components.txt new file mode 100644 index 000000000..9b7be775a --- /dev/null +++ b/LICENSES/Apache-2.0-components.txt @@ -0,0 +1,82 @@ +Apache-2.0 components + +The complete common Apache License, Version 2.0, including sections 1-9 and +the application appendix, is in the repository LICENSE, installed beside +this file as share/licenses/InfiniOps/LICENSE. Distribute that file together +with this component list and the accompanying notices. + +This list consolidates the common Apache terms from the fixed releases below. +Their distinct copyrights, NOTICE files and secondary licenses remain +applicable. It is not a license inventory for arbitrary resolved versions. + +FlashInfer 0.6.7 +https://github.com/flashinfer-ai/flashinfer/tree/v0.6.7 +Copyright 2025-2026 NVIDIA +Copyright 2023-2026 FlashInfer community (https://flashinfer.ai/) +License: Apache-2.0, complete terms in LICENSE. +Original NOTICE retained in FlashInfer-0.6.7/NOTICE.txt. + +The following component list is retained from the release LICENSE: +------------------------------------------------------------------------------------------------- +Some of the code in this project are adapted from other open-source projects with different +licenses. This product also bundles some third-party components under other open source licenses. +This section summarizes those components and their licenses. +See licenses/ for text of these licenses. + +BSD 3-Clause License +-------------------- + +include/flashinfer/attention/hopper/epilogue.cuh +include/flashinfer/attention/hopper/mainloop.cuh +include/flashinfer/attention/hopper/kernel_traits.cuh +include/flashinfer/attention/hopper/named_barrier.cuh +include/flashinfer/attention/hopper/tile_scheduler.cuh +include/flashinfer/attention/hopper/utils.cuh + +BSD 3-Clause "New" License +-------------------------- + +3rdparty/cutlass +include/flashinfer/attention/hopper/block_sparse_gather.cuh + +MIT License +----------- + +3rdparty/spdlog +3rdparty/spdlog/include/spdlog/fmt/bundled (fmt library) + +The release's licenses/ materials are retained here in FlashInfer-0.6.7/: +LICENSE.flashattention3.txt, LICENSE.cutlass.txt, LICENSE.spdlog.txt, +LICENSE.fmt.txt. These keep each original copyright and complete terms. + +Apache TVM FFI 0.1.10 +https://github.com/apache/tvm-ffi/tree/v0.1.10 +Copyright 2024-present The Apache Software Foundation +License: Apache-2.0, complete terms in LICENSE. +Original NOTICE retained in TVM-FFI-0.1.10/NOTICE.txt. + +The following component list is retained from the release LICENSE: +------------------------------------------------------------------------------------ +This product bundles various third-party components under other open source licenses. +This section summarizes those components and their licenses. See licenses/ +for text of these licenses. + + +Apache Software Foundation License 2.0 +-------------------------------------- + +3rdparty/dlpack + +BSD 3-Clause "New" or "Revised" License +--------------------------------------- + +3rdparty/libbacktrace + +DLPack: Copyright 2017 by Contributors +License: Apache-2.0, complete terms in LICENSE. + +The release's other licenses/ materials are retained in TVM-FFI-0.1.10/: +LICENSE.libbacktrace.txt, LICENSE.pytorch.txt, NOTICE.pytorch.txt. +Their complete original terms and copyrights are retained, including the +PyTorch secondary notices. This consolidation does not assert that every +component in the full releases is compiled into sampling.so or libtvm_ffi.so. diff --git a/LICENSES/FasterTransformer-BSD-3-Clause.txt b/LICENSES/FasterTransformer-BSD-3-Clause.txt new file mode 100644 index 000000000..d9219ec9b --- /dev/null +++ b/LICENSES/FasterTransformer-BSD-3-Clause.txt @@ -0,0 +1,27 @@ +Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +SPDX-License-Identifier: BSD-3-Clause + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + +1. Redistributions of source code must retain the above copyright notice, this +list of conditions and the following disclaimer. + +2. Redistributions in binary form must reproduce the above copyright notice, +this list of conditions and the following disclaimer in the documentation +and/or other materials provided with the distribution. + +3. Neither the name of the copyright holder nor the names of its +contributors may be used to endorse or promote products derived from +this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/LICENSES/FlashInfer-0.6.7/LICENSE.txt b/LICENSES/FlashInfer-0.6.7/LICENSE.txt deleted file mode 100644 index ba557f2bb..000000000 --- a/LICENSES/FlashInfer-0.6.7/LICENSE.txt +++ /dev/null @@ -1,229 +0,0 @@ - Apache License - Version 2.0, January 2004 - http://www.apache.org/licenses/ - - TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION - - 1. Definitions. - - "License" shall mean the terms and conditions for use, reproduction, - and distribution as defined by Sections 1 through 9 of this document. - - "Licensor" shall mean the copyright owner or entity authorized by - the copyright owner that is granting the License. - - "Legal Entity" shall mean the union of the acting entity and all - other entities that control, are controlled by, or are under common - control with that entity. For the purposes of this definition, - "control" means (i) the power, direct or indirect, to cause the - direction or management of such entity, whether by contract or - otherwise, or (ii) ownership of fifty percent (50%) or more of the - outstanding shares, or (iii) beneficial ownership of such entity. - - "You" (or "Your") shall mean an individual or Legal Entity - exercising permissions granted by this License. - - "Source" form shall mean the preferred form for making modifications, - including but not limited to software source code, documentation - source, and configuration files. - - "Object" form shall mean any form resulting from mechanical - transformation or translation of a Source form, including but - not limited to compiled object code, generated documentation, - and conversions to other media types. - - "Work" shall mean the work of authorship, whether in Source or - Object form, made available under the License, as indicated by a - copyright notice that is included in or attached to the work - (an example is provided in the Appendix below). - - "Derivative Works" shall mean any work, whether in Source or Object - form, that is based on (or derived from) the Work and for which the - editorial revisions, annotations, elaborations, or other modifications - represent, as a whole, an original work of authorship. For the purposes - of this License, Derivative Works shall not include works that remain - separable from, or merely link (or bind by name) to the interfaces of, - the Work and Derivative Works thereof. - - "Contribution" shall mean any work of authorship, including - the original version of the Work and any modifications or additions - to that Work or Derivative Works thereof, that is intentionally - submitted to Licensor for inclusion in the Work by the copyright owner - or by an individual or Legal Entity authorized to submit on behalf of - the copyright owner. For the purposes of this definition, "submitted" - means any form of electronic, verbal, or written communication sent - to the Licensor or its representatives, including but not limited to - communication on electronic mailing lists, source code control systems, - and issue tracking systems that are managed by, or on behalf of, the - Licensor for the purpose of discussing and improving the Work, but - excluding communication that is conspicuously marked or otherwise - designated in writing by the copyright owner as "Not a Contribution." - - "Contributor" shall mean Licensor and any individual or Legal Entity - on behalf of whom a Contribution has been received by Licensor and - subsequently incorporated within the Work. - - 2. Grant of Copyright License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - copyright license to reproduce, prepare Derivative Works of, - publicly display, publicly perform, sublicense, and distribute the - Work and such Derivative Works in Source or Object form. - - 3. Grant of Patent License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - (except as stated in this section) patent license to make, have made, - use, offer to sell, sell, import, and otherwise transfer the Work, - where such license applies only to those patent claims licensable - by such Contributor that are necessarily infringed by their - Contribution(s) alone or by combination of their Contribution(s) - with the Work to which such Contribution(s) was submitted. If You - institute patent litigation against any entity (including a - cross-claim or counterclaim in a lawsuit) alleging that the Work - or a Contribution incorporated within the Work constitutes direct - or contributory patent infringement, then any patent licenses - granted to You under this License for that Work shall terminate - as of the date such litigation is filed. - - 4. Redistribution. You may reproduce and distribute copies of the - Work or Derivative Works thereof in any medium, with or without - modifications, and in Source or Object form, provided that You - meet the following conditions: - - (a) You must give any other recipients of the Work or - Derivative Works a copy of this License; and - - (b) You must cause any modified files to carry prominent notices - stating that You changed the files; and - - (c) You must retain, in the Source form of any Derivative Works - that You distribute, all copyright, patent, trademark, and - attribution notices from the Source form of the Work, - excluding those notices that do not pertain to any part of - the Derivative Works; and - - (d) If the Work includes a "NOTICE" text file as part of its - distribution, then any Derivative Works that You distribute must - include a readable copy of the attribution notices contained - within such NOTICE file, excluding those notices that do not - pertain to any part of the Derivative Works, in at least one - of the following places: within a NOTICE text file distributed - as part of the Derivative Works; within the Source form or - documentation, if provided along with the Derivative Works; or, - within a display generated by the Derivative Works, if and - wherever such third-party notices normally appear. The contents - of the NOTICE file are for informational purposes only and - do not modify the License. You may add Your own attribution - notices within Derivative Works that You distribute, alongside - or as an addendum to the NOTICE text from the Work, provided - that such additional attribution notices cannot be construed - as modifying the License. - - You may add Your own copyright statement to Your modifications and - may provide additional or different license terms and conditions - for use, reproduction, or distribution of Your modifications, or - for any such Derivative Works as a whole, provided Your use, - reproduction, and distribution of the Work otherwise complies with - the conditions stated in this License. - - 5. Submission of Contributions. Unless You explicitly state otherwise, - any Contribution intentionally submitted for inclusion in the Work - by You to the Licensor shall be under the terms and conditions of - this License, without any additional terms or conditions. - Notwithstanding the above, nothing herein shall supersede or modify - the terms of any separate license agreement you may have executed - with Licensor regarding such Contributions. - - 6. Trademarks. This License does not grant permission to use the trade - names, trademarks, service marks, or product names of the Licensor, - except as required for reasonable and customary use in describing the - origin of the Work and reproducing the content of the NOTICE file. - - 7. Disclaimer of Warranty. Unless required by applicable law or - agreed to in writing, Licensor provides the Work (and each - Contributor provides its Contributions) on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or - implied, including, without limitation, any warranties or conditions - of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A - PARTICULAR PURPOSE. You are solely responsible for determining the - appropriateness of using or redistributing the Work and assume any - risks associated with Your exercise of permissions under this License. - - 8. Limitation of Liability. In no event and under no legal theory, - whether in tort (including negligence), contract, or otherwise, - unless required by applicable law (such as deliberate and grossly - negligent acts) or agreed to in writing, shall any Contributor be - liable to You for damages, including any direct, indirect, special, - incidental, or consequential damages of any character arising as a - result of this License or out of the use or inability to use the - Work (including but not limited to damages for loss of goodwill, - work stoppage, computer failure or malfunction, or any and all - other commercial damages or losses), even if such Contributor - has been advised of the possibility of such damages. - - 9. Accepting Warranty or Additional Liability. While redistributing - the Work or Derivative Works thereof, You may choose to offer, - and charge a fee for, acceptance of support, warranty, indemnity, - or other liability obligations and/or rights consistent with this - License. However, in accepting such obligations, You may act only - on Your own behalf and on Your sole responsibility, not on behalf - of any other Contributor, and only if You agree to indemnify, - defend, and hold each Contributor harmless for any liability - incurred by, or claims asserted against, such Contributor by reason - of your accepting any such warranty or additional liability. - - END OF TERMS AND CONDITIONS - - APPENDIX: How to apply the Apache License to your work. - - To apply the Apache License to your work, attach the following - boilerplate notice, with the fields enclosed by brackets "[]" - replaced with your own identifying information. (Don't include - the brackets!) The text should be enclosed in the appropriate - comment syntax for the file format. We also recommend that a - file or class name and description of purpose be included on the - same "printed page" as the copyright notice for easier - identification within third-party archives. - - Copyright [yyyy] [name of copyright owner] - - Licensed under the Apache License, Version 2.0 (the "License"); - you may not use this file except in compliance with the License. - You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - - Unless required by applicable law or agreed to in writing, software - distributed under the License is distributed on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - See the License for the specific language governing permissions and - limitations under the License. - -------------------------------------------------------------------------------------------------- -Some of the code in this project are adapted from other open-source projects with different -licenses. This product also bundles some third-party components under other open source licenses. -This section summarizes those components and their licenses. -See licenses/ for text of these licenses. - -BSD 3-Clause License --------------------- - -include/flashinfer/attention/hopper/epilogue.cuh -include/flashinfer/attention/hopper/mainloop.cuh -include/flashinfer/attention/hopper/kernel_traits.cuh -include/flashinfer/attention/hopper/named_barrier.cuh -include/flashinfer/attention/hopper/tile_scheduler.cuh -include/flashinfer/attention/hopper/utils.cuh - -BSD 3-Clause "New" License --------------------------- - -3rdparty/cutlass -include/flashinfer/attention/hopper/block_sparse_gather.cuh - -MIT License ------------ - -3rdparty/spdlog -3rdparty/spdlog/include/spdlog/fmt/bundled (fmt library) diff --git a/LICENSES/TVM-FFI-0.1.10/LICENSE.dlpack.txt b/LICENSES/TVM-FFI-0.1.10/LICENSE.dlpack.txt deleted file mode 100644 index 20a9c8a7b..000000000 --- a/LICENSES/TVM-FFI-0.1.10/LICENSE.dlpack.txt +++ /dev/null @@ -1,201 +0,0 @@ - Apache License - Version 2.0, January 2004 - http://www.apache.org/licenses/ - - TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION - - 1. Definitions. - - "License" shall mean the terms and conditions for use, reproduction, - and distribution as defined by Sections 1 through 9 of this document. - - "Licensor" shall mean the copyright owner or entity authorized by - the copyright owner that is granting the License. - - "Legal Entity" shall mean the union of the acting entity and all - other entities that control, are controlled by, or are under common - control with that entity. For the purposes of this definition, - "control" means (i) the power, direct or indirect, to cause the - direction or management of such entity, whether by contract or - otherwise, or (ii) ownership of fifty percent (50%) or more of the - outstanding shares, or (iii) beneficial ownership of such entity. - - "You" (or "Your") shall mean an individual or Legal Entity - exercising permissions granted by this License. - - "Source" form shall mean the preferred form for making modifications, - including but not limited to software source code, documentation - source, and configuration files. - - "Object" form shall mean any form resulting from mechanical - transformation or translation of a Source form, including but - not limited to compiled object code, generated documentation, - and conversions to other media types. - - "Work" shall mean the work of authorship, whether in Source or - Object form, made available under the License, as indicated by a - copyright notice that is included in or attached to the work - (an example is provided in the Appendix below). - - "Derivative Works" shall mean any work, whether in Source or Object - form, that is based on (or derived from) the Work and for which the - editorial revisions, annotations, elaborations, or other modifications - represent, as a whole, an original work of authorship. For the purposes - of this License, Derivative Works shall not include works that remain - separable from, or merely link (or bind by name) to the interfaces of, - the Work and Derivative Works thereof. - - "Contribution" shall mean any work of authorship, including - the original version of the Work and any modifications or additions - to that Work or Derivative Works thereof, that is intentionally - submitted to Licensor for inclusion in the Work by the copyright owner - or by an individual or Legal Entity authorized to submit on behalf of - the copyright owner. For the purposes of this definition, "submitted" - means any form of electronic, verbal, or written communication sent - to the Licensor or its representatives, including but not limited to - communication on electronic mailing lists, source code control systems, - and issue tracking systems that are managed by, or on behalf of, the - Licensor for the purpose of discussing and improving the Work, but - excluding communication that is conspicuously marked or otherwise - designated in writing by the copyright owner as "Not a Contribution." - - "Contributor" shall mean Licensor and any individual or Legal Entity - on behalf of whom a Contribution has been received by Licensor and - subsequently incorporated within the Work. - - 2. Grant of Copyright License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - copyright license to reproduce, prepare Derivative Works of, - publicly display, publicly perform, sublicense, and distribute the - Work and such Derivative Works in Source or Object form. - - 3. Grant of Patent License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - (except as stated in this section) patent license to make, have made, - use, offer to sell, sell, import, and otherwise transfer the Work, - where such license applies only to those patent claims licensable - by such Contributor that are necessarily infringed by their - Contribution(s) alone or by combination of their Contribution(s) - with the Work to which such Contribution(s) was submitted. If You - institute patent litigation against any entity (including a - cross-claim or counterclaim in a lawsuit) alleging that the Work - or a Contribution incorporated within the Work constitutes direct - or contributory patent infringement, then any patent licenses - granted to You under this License for that Work shall terminate - as of the date such litigation is filed. - - 4. Redistribution. You may reproduce and distribute copies of the - Work or Derivative Works thereof in any medium, with or without - modifications, and in Source or Object form, provided that You - meet the following conditions: - - (a) You must give any other recipients of the Work or - Derivative Works a copy of this License; and - - (b) You must cause any modified files to carry prominent notices - stating that You changed the files; and - - (c) You must retain, in the Source form of any Derivative Works - that You distribute, all copyright, patent, trademark, and - attribution notices from the Source form of the Work, - excluding those notices that do not pertain to any part of - the Derivative Works; and - - (d) If the Work includes a "NOTICE" text file as part of its - distribution, then any Derivative Works that You distribute must - include a readable copy of the attribution notices contained - within such NOTICE file, excluding those notices that do not - pertain to any part of the Derivative Works, in at least one - of the following places: within a NOTICE text file distributed - as part of the Derivative Works; within the Source form or - documentation, if provided along with the Derivative Works; or, - within a display generated by the Derivative Works, if and - wherever such third-party notices normally appear. The contents - of the NOTICE file are for informational purposes only and - do not modify the License. You may add Your own attribution - notices within Derivative Works that You distribute, alongside - or as an addendum to the NOTICE text from the Work, provided - that such additional attribution notices cannot be construed - as modifying the License. - - You may add Your own copyright statement to Your modifications and - may provide additional or different license terms and conditions - for use, reproduction, or distribution of Your modifications, or - for any such Derivative Works as a whole, provided Your use, - reproduction, and distribution of the Work otherwise complies with - the conditions stated in this License. - - 5. Submission of Contributions. Unless You explicitly state otherwise, - any Contribution intentionally submitted for inclusion in the Work - by You to the Licensor shall be under the terms and conditions of - this License, without any additional terms or conditions. - Notwithstanding the above, nothing herein shall supersede or modify - the terms of any separate license agreement you may have executed - with Licensor regarding such Contributions. - - 6. Trademarks. This License does not grant permission to use the trade - names, trademarks, service marks, or product names of the Licensor, - except as required for reasonable and customary use in describing the - origin of the Work and reproducing the content of the NOTICE file. - - 7. Disclaimer of Warranty. Unless required by applicable law or - agreed to in writing, Licensor provides the Work (and each - Contributor provides its Contributions) on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or - implied, including, without limitation, any warranties or conditions - of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A - PARTICULAR PURPOSE. You are solely responsible for determining the - appropriateness of using or redistributing the Work and assume any - risks associated with Your exercise of permissions under this License. - - 8. Limitation of Liability. In no event and under no legal theory, - whether in tort (including negligence), contract, or otherwise, - unless required by applicable law (such as deliberate and grossly - negligent acts) or agreed to in writing, shall any Contributor be - liable to You for damages, including any direct, indirect, special, - incidental, or consequential damages of any character arising as a - result of this License or out of the use or inability to use the - Work (including but not limited to damages for loss of goodwill, - work stoppage, computer failure or malfunction, or any and all - other commercial damages or losses), even if such Contributor - has been advised of the possibility of such damages. - - 9. Accepting Warranty or Additional Liability. While redistributing - the Work or Derivative Works thereof, You may choose to offer, - and charge a fee for, acceptance of support, warranty, indemnity, - or other liability obligations and/or rights consistent with this - License. However, in accepting such obligations, You may act only - on Your own behalf and on Your sole responsibility, not on behalf - of any other Contributor, and only if You agree to indemnify, - defend, and hold each Contributor harmless for any liability - incurred by, or claims asserted against, such Contributor by reason - of your accepting any such warranty or additional liability. - - END OF TERMS AND CONDITIONS - - APPENDIX: How to apply the Apache License to your work. - - To apply the Apache License to your work, attach the following - boilerplate notice, with the fields enclosed by brackets "{}" - replaced with your own identifying information. (Don't include - the brackets!) The text should be enclosed in the appropriate - comment syntax for the file format. We also recommend that a - file or class name and description of purpose be included on the - same "printed page" as the copyright notice for easier - identification within third-party archives. - - Copyright 2017 by Contributors - - Licensed under the Apache License, Version 2.0 (the "License"); - you may not use this file except in compliance with the License. - You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - - Unless required by applicable law or agreed to in writing, software - distributed under the License is distributed on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - See the License for the specific language governing permissions and - limitations under the License. diff --git a/LICENSES/TVM-FFI-0.1.10/LICENSE.txt b/LICENSES/TVM-FFI-0.1.10/LICENSE.txt deleted file mode 100644 index 49ac4b3df..000000000 --- a/LICENSES/TVM-FFI-0.1.10/LICENSE.txt +++ /dev/null @@ -1,217 +0,0 @@ - Apache License - Version 2.0, January 2004 - http://www.apache.org/licenses/ - - TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION - - 1. Definitions. - - "License" shall mean the terms and conditions for use, reproduction, - and distribution as defined by Sections 1 through 9 of this document. - - "Licensor" shall mean the copyright owner or entity authorized by - the copyright owner that is granting the License. - - "Legal Entity" shall mean the union of the acting entity and all - other entities that control, are controlled by, or are under common - control with that entity. For the purposes of this definition, - "control" means (i) the power, direct or indirect, to cause the - direction or management of such entity, whether by contract or - otherwise, or (ii) ownership of fifty percent (50%) or more of the - outstanding shares, or (iii) beneficial ownership of such entity. - - "You" (or "Your") shall mean an individual or Legal Entity - exercising permissions granted by this License. - - "Source" form shall mean the preferred form for making modifications, - including but not limited to software source code, documentation - source, and configuration files. - - "Object" form shall mean any form resulting from mechanical - transformation or translation of a Source form, including but - not limited to compiled object code, generated documentation, - and conversions to other media types. - - "Work" shall mean the work of authorship, whether in Source or - Object form, made available under the License, as indicated by a - copyright notice that is included in or attached to the work - (an example is provided in the Appendix below). - - "Derivative Works" shall mean any work, whether in Source or Object - form, that is based on (or derived from) the Work and for which the - editorial revisions, annotations, elaborations, or other modifications - represent, as a whole, an original work of authorship. For the purposes - of this License, Derivative Works shall not include works that remain - separable from, or merely link (or bind by name) to the interfaces of, - the Work and Derivative Works thereof. - - "Contribution" shall mean any work of authorship, including - the original version of the Work and any modifications or additions - to that Work or Derivative Works thereof, that is intentionally - submitted to Licensor for inclusion in the Work by the copyright owner - or by an individual or Legal Entity authorized to submit on behalf of - the copyright owner. For the purposes of this definition, "submitted" - means any form of electronic, verbal, or written communication sent - to the Licensor or its representatives, including but not limited to - communication on electronic mailing lists, source code control systems, - and issue tracking systems that are managed by, or on behalf of, the - Licensor for the purpose of discussing and improving the Work, but - excluding communication that is conspicuously marked or otherwise - designated in writing by the copyright owner as "Not a Contribution." - - "Contributor" shall mean Licensor and any individual or Legal Entity - on behalf of whom a Contribution has been received by Licensor and - subsequently incorporated within the Work. - - 2. Grant of Copyright License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - copyright license to reproduce, prepare Derivative Works of, - publicly display, publicly perform, sublicense, and distribute the - Work and such Derivative Works in Source or Object form. - - 3. Grant of Patent License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - (except as stated in this section) patent license to make, have made, - use, offer to sell, sell, import, and otherwise transfer the Work, - where such license applies only to those patent claims licensable - by such Contributor that are necessarily infringed by their - Contribution(s) alone or by combination of their Contribution(s) - with the Work to which such Contribution(s) was submitted. If You - institute patent litigation against any entity (including a - cross-claim or counterclaim in a lawsuit) alleging that the Work - or a Contribution incorporated within the Work constitutes direct - or contributory patent infringement, then any patent licenses - granted to You under this License for that Work shall terminate - as of the date such litigation is filed. - - 4. Redistribution. You may reproduce and distribute copies of the - Work or Derivative Works thereof in any medium, with or without - modifications, and in Source or Object form, provided that You - meet the following conditions: - - (a) You must give any other recipients of the Work or - Derivative Works a copy of this License; and - - (b) You must cause any modified files to carry prominent notices - stating that You changed the files; and - - (c) You must retain, in the Source form of any Derivative Works - that You distribute, all copyright, patent, trademark, and - attribution notices from the Source form of the Work, - excluding those notices that do not pertain to any part of - the Derivative Works; and - - (d) If the Work includes a "NOTICE" text file as part of its - distribution, then any Derivative Works that You distribute must - include a readable copy of the attribution notices contained - within such NOTICE file, excluding those notices that do not - pertain to any part of the Derivative Works, in at least one - of the following places: within a NOTICE text file distributed - as part of the Derivative Works; within the Source form or - documentation, if provided along with the Derivative Works; or, - within a display generated by the Derivative Works, if and - wherever such third-party notices normally appear. The contents - of the NOTICE file are for informational purposes only and - do not modify the License. You may add Your own attribution - notices within Derivative Works that You distribute, alongside - or as an addendum to the NOTICE text from the Work, provided - that such additional attribution notices cannot be construed - as modifying the License. - - You may add Your own copyright statement to Your modifications and - may provide additional or different license terms and conditions - for use, reproduction, or distribution of Your modifications, or - for any such Derivative Works as a whole, provided Your use, - reproduction, and distribution of the Work otherwise complies with - the conditions stated in this License. - - 5. Submission of Contributions. Unless You explicitly state otherwise, - any Contribution intentionally submitted for inclusion in the Work - by You to the Licensor shall be under the terms and conditions of - this License, without any additional terms or conditions. - Notwithstanding the above, nothing herein shall supersede or modify - the terms of any separate license agreement you may have executed - with Licensor regarding such Contributions. - - 6. Trademarks. This License does not grant permission to use the trade - names, trademarks, service marks, or product names of the Licensor, - except as required for reasonable and customary use in describing the - origin of the Work and reproducing the content of the NOTICE file. - - 7. Disclaimer of Warranty. Unless required by applicable law or - agreed to in writing, Licensor provides the Work (and each - Contributor provides its Contributions) on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or - implied, including, without limitation, any warranties or conditions - of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A - PARTICULAR PURPOSE. You are solely responsible for determining the - appropriateness of using or redistributing the Work and assume any - risks associated with Your exercise of permissions under this License. - - 8. Limitation of Liability. In no event and under no legal theory, - whether in tort (including negligence), contract, or otherwise, - unless required by applicable law (such as deliberate and grossly - negligent acts) or agreed to in writing, shall any Contributor be - liable to You for damages, including any direct, indirect, special, - incidental, or consequential damages of any character arising as a - result of this License or out of the use or inability to use the - Work (including but not limited to damages for loss of goodwill, - work stoppage, computer failure or malfunction, or any and all - other commercial damages or losses), even if such Contributor - has been advised of the possibility of such damages. - - 9. Accepting Warranty or Additional Liability. While redistributing - the Work or Derivative Works thereof, You may choose to offer, - and charge a fee for, acceptance of support, warranty, indemnity, - or other liability obligations and/or rights consistent with this - License. However, in accepting such obligations, You may act only - on Your own behalf and on Your sole responsibility, not on behalf - of any other Contributor, and only if You agree to indemnify, - defend, and hold each Contributor harmless for any liability - incurred by, or claims asserted against, such Contributor by reason - of your accepting any such warranty or additional liability. - - END OF TERMS AND CONDITIONS - - APPENDIX: How to apply the Apache License to your work. - - To apply the Apache License to your work, attach the following - boilerplate notice, with the fields enclosed by brackets "{}" - replaced with your own identifying information. (Don't include - the brackets!) The text should be enclosed in the appropriate - comment syntax for the file format. We also recommend that a - file or class name and description of purpose be included on the - same "printed page" as the copyright notice for easier - identification within third-party archives. - - Copyright {yyyy} {name of copyright owner} - - Licensed under the Apache License, Version 2.0 (the "License"); - you may not use this file except in compliance with the License. - You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - - Unless required by applicable law or agreed to in writing, software - distributed under the License is distributed on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - See the License for the specific language governing permissions and - limitations under the License. - ------------------------------------------------------------------------------------- -This product bundles various third-party components under other open source licenses. -This section summarizes those components and their licenses. See licenses/ -for text of these licenses. - - -Apache Software Foundation License 2.0 --------------------------------------- - -3rdparty/dlpack - -BSD 3-Clause "New" or "Revised" License ---------------------------------------- - -3rdparty/libbacktrace diff --git a/THIRD_PARTY.txt b/THIRD_PARTY.txt index 6f172541b..f9965ac8a 100644 --- a/THIRD_PARTY.txt +++ b/THIRD_PARTY.txt @@ -14,28 +14,34 @@ W = P/w4_marlin/vendor/ W4 / Marlin Upstream projects: https://github.com/IST-DASLab/marlin and https://github.com/vllm-project/vllm +Current kernel source: third_party/vllm, pinned gitlink +ffc4f08c8ee130d4ea6347c1bf31ffd4f8af28ab, +csrc/libtorch_stable/quantization/marlin. Only its source kernels and generator +are consumed; the vLLM/Torch runtime is not a dependency of this native target. +The original Ops port came through the fixed Core donor above; this selected +vLLM revision is not asserted to be the unique historical origin of that port. +Preparation applies patches/marlin to build-directory copies, retaining the +existing Torch-free includes, repack bindings and activation-order index behavior. +CUDA template instances and their selector are generated together in the build +directory. The four host/type adaptations below remain in Ops. Marlin notices retain Copyright 2024 Elias Frantar, modifications by Neural Magic, Apache-2.0. The original generator identifies Apache-2.0 / Copyright contributors to the vLLM project. -Marlin original Apache headers remain in marlin_template.h and nvidia/kernel.cuh; Apache text: LICENSE. -Target <- direct donor; COPY retains the donor implementation (apart from added attribution -headers); ADAPTED/PORTED/GENERATED identify changes described below. Original notices remain. -ADAPTED W/awq_repack/cuda/kernel.cuh <- C/awq_marlin_repack/cuda/kernel.cuh (include-path-only; kernel body unchanged) +Marlin original Apache headers remain in the prepared source and nvidia/kernel.cuh; +Apache text: LICENSE. Original source notices remain in prepared files. +Remaining adaptations <- direct Core donor: COPY W/gemm/core/scalar_type.hpp <- C/awq_marlin_gemm/core/scalar_type.hpp COPY W/gemm/core/source_location.h <- C/awq_marlin_gemm/core/source_location.h ADAPTED W/gemm/core/utils.h <- C/awq_marlin_gemm/core/utils.h (unused pointer/range/dtype helpers removed; runtime checks unchanged) -COPY W/gemm/marlin/dequant.h <- C/awq_marlin_gemm/marlin/dequant.h -COPY W/gemm/marlin/kernel.h <- C/awq_marlin_gemm/marlin/kernel.h -COPY W/gemm/marlin/marlin.cuh <- C/awq_marlin_gemm/marlin/marlin.cuh -COPY W/gemm/marlin/marlin_dtypes.cuh <- C/awq_marlin_gemm/marlin/marlin_dtypes.cuh -COPY W/gemm/marlin/marlin_mma.h <- C/awq_marlin_gemm/marlin/marlin_mma.h -ADAPTED W/gemm/marlin/marlin_template.h <- C/awq_marlin_gemm/marlin/marlin_template.h +The prepared marlin/dequant.h INT4/INT8 conversion portions also derive from NVIDIA FasterTransformer: +https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h +The dequant.h references map INT4 FP16/BF16 to upstream lines 215-287/327-385 +and INT8 FP16/BF16 to upstream lines 53-85/125-175; packing, bias and scalar +specializations are adapted by the Marlin donor, not source-identical copies. +Original Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES and the +complete BSD-3-Clause terms are retained in the prepared dequant.h and +LICENSES/FasterTransformer-BSD-3-Clause.txt. The separate CUTLASS dependency's +2017-2026 notice does not replace this original copyright notice. COPY W/gemm/nvidia/kernel.cuh <- C/awq_marlin_gemm/nvidia/kernel.cuh -GENERATED W/gemm/nvidia/kernel_selector.h <- C/awq_marlin_gemm/nvidia/generate_kernels.py -GENERATED W/gemm/nvidia/sm80_kernel_bfloat16_u4_bfloat16.cu <- C/awq_marlin_gemm/nvidia/generate_kernels.py -GENERATED W/gemm/nvidia/sm80_kernel_bfloat16_u4b8_bfloat16.cu <- C/awq_marlin_gemm/nvidia/generate_kernels.py -GENERATED W/gemm/nvidia/sm80_kernel_float16_u4_float16.cu <- C/awq_marlin_gemm/nvidia/generate_kernels.py -GENERATED W/gemm/nvidia/sm80_kernel_float16_u4b8_float16.cu <- C/awq_marlin_gemm/nvidia/generate_kernels.py -ADAPTED W/gptq_repack/cuda/kernel.cuh <- C/gptq_marlin_repack/cuda/kernel.cuh (include-path-only; kernel body unchanged) ADAPTED P/w4_marlin/parameters.cu implements the scale/zero permutations from vLLM's vllm/model_executor/layers/quantization/utils/marlin_utils.py; the fixed reference revision and loading-only CUDA adaptation are noted in that file. @@ -55,11 +61,15 @@ Namespace, layout/backend binding and resource adaptations are not described as FlashAttention / CUTLASS Projects: https://github.com/Dao-AILab/flash-attention and https://github.com/NVIDIA/cutlass -Direct archive: flash_attn-2.8.3, SHA256 1e71dd64a9e0280e0447b8a0c2541bad4bf6ac65bdeaa2f90e51a9e57de0370d. -scripts/prepare_flash_attention_inference.py imports csrc/flash_attn/src forward headers and csrc/cutlass/include. -It preserves upstream header notices and copies the archive LICENSE and AUTHORS verbatim to the prepared build directory. +Source: third_party/flash-attention, pinned gitlink +060c9188beec3a8b62b33a3bfa6d5d2d44975fab (v2.8.3). +Its csrc/cutlass gitlink is dc4817921edda44a549197ff3a9dcf5df0636e7b. +scripts/prepare_flash_attention_inference.py prepares csrc/flash_attn/src +forward headers with patches/flash-attention, preserving source notices, and +copies upstream LICENSE and AUTHORS. CUTLASS headers are included directly +from the nested source dependency; its LICENSE.txt is copied for distribution. The adapted scheduling heuristic comes from csrc/flash_attn/flash_api.cpp and retains Copyright 2023 Tri Dao. -The original archive/CUTLASS notices govern those portions; the Core MIT notice does not relicense them. +The original FlashAttention/CUTLASS notices govern those portions; the Core MIT notice does not relicense them. Native launch templates, descriptors and workspace/lifetime integration are glue around the imported forward mathematics. Additional vLLM / TensorRT-LLM / Ascend adaptations @@ -75,12 +85,20 @@ retain their Huawei/InfiniTensor notices and vLLM-Ascend source links. Distribution SDK installation includes LICENSE, this attribution list and LICENSES under share/licenses/InfiniOps. Native FlashAttention installation also includes the exact -archive LICENSE and AUTHORS and the bundled CUTLASS header's complete license block. +upstream LICENSE and AUTHORS and the nested CUTLASS LICENSE.txt. CUTLASS-BSD-3-Clause.txt records the separate CMake dependency at 087c84df83d254b5fb295a7a408f1a1d554085cf. CUTLASS notices apply separately to its compiled headers. Linked libraries are not relicensed by Ops: distributions bundling sampling.so or libtvm_ffi.so must retain the resolved FlashInfer/TVM-FFI licenses, NOTICE files and applicable third-party notices. See docs/build-nvidia-inference.md for the pinned recipe. -LICENSES/FlashInfer-0.6.7 and LICENSES/TVM-FFI-0.1.10 retain the corresponding -upstream release LICENSE, NOTICE and licenses/ materials for that recipe verbatim. +LICENSES/Apache-2.0-components.txt maps the corresponding release components +and copyrights to the complete common Apache terms in LICENSE, and preserves +their original secondary-component lists. LICENSES/FlashInfer-0.6.7 and +LICENSES/TVM-FFI-0.1.10 retain the upstream NOTICE and secondary-license texts +verbatim; the repeated common Apache terms are consolidated, not omitted. They are not a substitute for the notices of a different resolved binary version. +These release bundles conservatively retain upstream secondary notices, including +PyTorch-derived TVM-FFI code; not importing Torch does not remove source-derived +obligations. The full Python wheels contain more components than sampling.so and +libtvm_ffi.so alone. Do not infer a native-only notice subset from Python imports +or from an absent wheel license directory; check the compiled source provenance. diff --git a/cmake/NativeFlashAttention.cmake b/cmake/NativeFlashAttention.cmake index b1ca0bdea..c0db8ad8f 100644 --- a/cmake/NativeFlashAttention.cmake +++ b/cmake/NativeFlashAttention.cmake @@ -3,19 +3,40 @@ function(infiniops_add_native_flash target) find_package(Python3 REQUIRED COMPONENTS Interpreter) find_package(CUDAToolkit REQUIRED) - if(NOT EXISTS "${FLASH_ATTN_SOURCE_ARCHIVE}") - message(FATAL_ERROR "Native FlashAttention requires the pinned source archive") - endif() + find_program(_flash_patch NAMES patch REQUIRED) set(_provider "${CMAKE_CURRENT_FUNCTION_LIST_DIR}/../src/native/cuda/nvidia/ops/flash_attn_with_kvcache") set(_varlen_provider "${CMAKE_CURRENT_FUNCTION_LIST_DIR}/../src/native/cuda/nvidia/ops/flash_attn_varlen_func") set(_source "${CMAKE_CURRENT_FUNCTION_LIST_DIR}/..") + set(_upstream "${_source}/third_party/flash-attention") + set(_kernel_source "${_upstream}/csrc/flash_attn/src") + set(_cutlass_include "${_upstream}/csrc/cutlass/include") set(_prepared "${CMAKE_CURRENT_BINARY_DIR}/native-flash-inference") + set(_inputs "${_source}/scripts/prepare_flash_attention_inference.py" + "${_source}/patches/flash-attention/inference.patch" + "${_upstream}/csrc/flash_attn/flash_api.cpp" + "${_upstream}/LICENSE" "${_upstream}/AUTHORS" + "${_upstream}/csrc/cutlass/LICENSE.txt" + "${_cutlass_include}/cutlass/cutlass.h") + # Fixed forward closure: no recursive source scanning or copied CUTLASS tree. + foreach(_header block_info.h dropout.h flash.h flash_fwd_kernel.h + flash_fwd_launch_template.h hardware_info.h kernel_traits.h mask.h + namespace_config.h philox.cuh rotary.h softmax.h static_switch.h utils.h) + list(APPEND _inputs "${_kernel_source}/${_header}") + endforeach() + foreach(_input IN LISTS _inputs) + if(NOT EXISTS "${_input}") + message(FATAL_ERROR + "Native FlashAttention is missing ${_input}. Initialize third_party/flash-attention and its CUTLASS submodule.") + endif() + endforeach() + set_property(DIRECTORY APPEND PROPERTY CMAKE_CONFIGURE_DEPENDS ${_inputs}) execute_process(COMMAND "${Python3_EXECUTABLE}" "${_source}/scripts/prepare_flash_attention_inference.py" - --archive "${FLASH_ATTN_SOURCE_ARCHIVE}" --output "${_prepared}" - RESULT_VARIABLE _status) + --source-dir "${_upstream}" --output "${_prepared}" + --patch-executable "${_flash_patch}" + RESULT_VARIABLE _status ERROR_VARIABLE _prepare_error) if(NOT _status EQUAL 0) - message(FATAL_ERROR "Pinned FlashAttention preparation failed") + message(FATAL_ERROR "Pinned FlashAttention preparation failed: ${_prepare_error}") endif() set(_launches) foreach(_flash_kind fp16 bf16) @@ -43,7 +64,7 @@ function(infiniops_add_native_flash target) FLASH_NAMESPACE=infini_flash FLASHATTENTION_DISABLE_DROPOUT) target_include_directories(${target} PRIVATE "${_source}/src" "${_source}/include" "${_prepared}" - "${_prepared}/cutlass" "${_provider}" "${_varlen_provider}" + "${_cutlass_include}" "${_provider}" "${_varlen_provider}" ${INFINI_RT_INCLUDE_DIRS}) target_compile_options(${target} PRIVATE # Match the pinned upstream setup.py numerical compilation mode. diff --git a/cmake/NativeW4Marlin.cmake b/cmake/NativeW4Marlin.cmake index 2bb08936f..ebc27c1c7 100644 --- a/cmake/NativeW4Marlin.cmake +++ b/cmake/NativeW4Marlin.cmake @@ -2,22 +2,52 @@ if(CMAKE_CUDA_COMPILER_VERSION VERSION_LESS 12.3) message(FATAL_ERROR "NVIDIA Marlin requires CUDA >=12.3 and C++20") endif() set(_w4_root "${PROJECT_SOURCE_DIR}/src/native/cuda/nvidia/ops/w4_marlin") +find_package(Python3 REQUIRED COMPONENTS Interpreter) +find_program(_w4_patch NAMES patch REQUIRED) +set(_w4_upstream "${PROJECT_SOURCE_DIR}/third_party/vllm") +set(_w4_upstream_marlin "${_w4_upstream}/csrc/libtorch_stable/quantization/marlin") +set(_w4_prepared "${CMAKE_CURRENT_BINARY_DIR}/native-w4-marlin") +set(_w4_inputs + "${PROJECT_SOURCE_DIR}/scripts/prepare_marlin.py" + "${PROJECT_SOURCE_DIR}/patches/marlin/w4a16.patch") +foreach(_w4_input dequant.h kernel.h marlin.cuh marlin_dtypes.cuh + marlin_mma.h marlin_template.h awq_marlin_repack.cu + gptq_marlin_repack.cu generate_kernels.py) + list(APPEND _w4_inputs "${_w4_upstream_marlin}/${_w4_input}") +endforeach() +foreach(_w4_input IN LISTS _w4_inputs) + if(NOT EXISTS "${_w4_input}") + message(FATAL_ERROR + "Native Marlin is missing ${_w4_input}. Initialize third_party/vllm.") + endif() +endforeach() +set_property(DIRECTORY APPEND PROPERTY CMAKE_CONFIGURE_DEPENDS ${_w4_inputs}) +execute_process( + COMMAND "${Python3_EXECUTABLE}" "${PROJECT_SOURCE_DIR}/scripts/prepare_marlin.py" + --source-dir "${_w4_upstream}" --output "${_w4_prepared}" + --patch-executable "${_w4_patch}" + RESULT_VARIABLE _w4_status ERROR_VARIABLE _w4_prepare_error) +if(NOT _w4_status EQUAL 0) + message(FATAL_ERROR "Pinned Marlin preparation failed: ${_w4_prepare_error}") +endif() add_library(infiniops_w4_marlin OBJECT "${_w4_root}/repack.cc" "${_w4_root}/repack_awq.cu" "${_w4_root}/repack_gptq.cu" "${_w4_root}/parameters.cu" "${_w4_root}/gemm.cu" - "${_w4_root}/vendor/gemm/nvidia/sm80_kernel_bfloat16_u4_bfloat16.cu" - "${_w4_root}/vendor/gemm/nvidia/sm80_kernel_bfloat16_u4b8_bfloat16.cu" - "${_w4_root}/vendor/gemm/nvidia/sm80_kernel_float16_u4_float16.cu" - "${_w4_root}/vendor/gemm/nvidia/sm80_kernel_float16_u4b8_float16.cu") + "${_w4_prepared}/vendor/gemm/nvidia/sm80_kernel_bfloat16_u4_bfloat16.cu" + "${_w4_prepared}/vendor/gemm/nvidia/sm80_kernel_bfloat16_u4b8_bfloat16.cu" + "${_w4_prepared}/vendor/gemm/nvidia/sm80_kernel_float16_u4_float16.cu" + "${_w4_prepared}/vendor/gemm/nvidia/sm80_kernel_float16_u4b8_float16.cu") set_target_properties(infiniops_w4_marlin PROPERTIES POSITION_INDEPENDENT_CODE ON CXX_STANDARD 20 CXX_STANDARD_REQUIRED ON CUDA_STANDARD 20 CUDA_STANDARD_REQUIRED ON CUDA_ARCHITECTURES 80) target_include_directories(infiniops_w4_marlin PRIVATE + "${_w4_prepared}" "${_w4_prepared}/vendor/gemm" + "${_w4_prepared}/vendor/gemm/nvidia" "${_w4_root}/vendor/gemm" "${PROJECT_SOURCE_DIR}/src" "${PROJECT_SOURCE_DIR}/include" ${INFINI_RT_INCLUDE_DIRS}) target_link_libraries(infiniops_w4_marlin PRIVATE infinirt CUDA::cudart) diff --git a/docs/build-nvidia-inference.md b/docs/build-nvidia-inference.md index 80856218a..e979a0c82 100644 --- a/docs/build-nvidia-inference.md +++ b/docs/build-nvidia-inference.md @@ -13,9 +13,21 @@ execute their Python inference APIs. No `flashinfer-python` or Torch is required |---|---|---| | Sampling | `flashinfer_jit_cache-0.6.7+cu128-cp39-abi3-manylinux_2_28_x86_64.whl` | `40139f1aab5435d9cb5ac0261883bf26740ba26680a9e2c9d7703168a6b4648d` | | TVM FFI | `apache_tvm_ffi-0.1.10-cp312-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl` | `73d478395a8625dd92fde7b7fd92b4719f18f480b78336e422cb66cc7985213d` | -| FlashAttention | `flash_attn-2.8.3.tar.gz` | `1e71dd64a9e0280e0447b8a0c2541bad4bf6ac65bdeaa2f90e51a9e57de0370d` | | CUTLASS | commit `087c84df83d254b5fb295a7a408f1a1d554085cf` archive | `99340f7a8c87f5295b91aecf8295df6693ad954807672c3869a7fb3d0084f709` | +Native Attention and W4 kernels use the fixed submodule revisions below. They +are source inputs only; neither the Torch extension nor the vLLM Python package +is built or installed. + +| Source | Fixed commit | +|---|---| +| `third_party/flash-attention` | `060c9188beec3a8b62b33a3bfa6d5d2d44975fab` | +| Its `csrc/cutlass` submodule | `dc4817921edda44a549197ff3a9dcf5df0636e7b` | +| `third_party/vllm` (Marlin only) | `ffc4f08c8ee130d4ea6347c1bf31ffd4f8af28ab` | + +The other CUTLASS input above remains the existing dependency of ordinary Ops +providers. Its version is not substituted for FlashAttention's private headers. + Installed library hashes for the two wheels above: ```text @@ -26,8 +38,10 @@ Installed library hashes for the two wheels above: The YAML version ranges describe the provider's accepted packages, not a lock for a release. A release using a different wheel must record its resolved version/hash and run the corresponding direct operator tests again. +`LICENSES/Apache-2.0-components.txt` records the release component copyrights +and maps their shared Apache terms to `LICENSE`. `LICENSES/FlashInfer-0.6.7` and `LICENSES/TVM-FFI-0.1.10` retain the original -release licenses, notices and third-party license texts for these exact inputs. +release NOTICE and secondary-license texts for these exact inputs. If replacing either wheel, update those redistribution materials from the actual release as well. SDK installation puts these materials, the native FlashAttention notices and the Ops attribution list in `share/licenses/InfiniOps`; @@ -44,17 +58,19 @@ export CORE=/absolute/path/to/InfiniCore export INFINI_ROOT="$PWD/sdk" export CUDA_HOME=/usr/local/cuda export INFINI_RT_ROOT="$INFINI_ROOT" +git -C "$CORE/submodules/InfiniOps" submodule update --init \ + third_party/flash-attention third_party/vllm +git -C "$CORE/submodules/InfiniOps/third_party/flash-attention" \ + submodule update --init csrc/cutlass mkdir artifacts cd artifacts curl -fL -o 'flashinfer_jit_cache-0.6.7+cu128-cp39-abi3-manylinux_2_28_x86_64.whl' \ 'https://github.com/flashinfer-ai/flashinfer/releases/download/v0.6.7/flashinfer_jit_cache-0.6.7%2Bcu128-cp39-abi3-manylinux_2_28_x86_64.whl' curl -fL -O 'https://files.pythonhosted.org/packages/50/6f/23740f06829030704e6f8f1f7093a06b7a68f904baa40053a5f594705bae/apache_tvm_ffi-0.1.10-cp312-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl' -curl -fL -O 'https://files.pythonhosted.org/packages/3b/b2/8d76c41ad7974ee264754709c22963447f7f8134613fd9ce80984ed0dab7/flash_attn-2.8.3.tar.gz' curl -fL -o cutlass.tar.gz 'https://github.com/NVIDIA/cutlass/archive/087c84df83d254b5fb295a7a408f1a1d554085cf.tar.gz' sha256sum -c <<'SHA256' 40139f1aab5435d9cb5ac0261883bf26740ba26680a9e2c9d7703168a6b4648d flashinfer_jit_cache-0.6.7+cu128-cp39-abi3-manylinux_2_28_x86_64.whl 73d478395a8625dd92fde7b7fd92b4719f18f480b78336e422cb66cc7985213d apache_tvm_ffi-0.1.10-cp312-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl -1e71dd64a9e0280e0447b8a0c2541bad4bf6ac65bdeaa2f90e51a9e57de0370d flash_attn-2.8.3.tar.gz 99340f7a8c87f5295b91aecf8295df6693ad954807672c3869a7fb3d0084f709 cutlass.tar.gz SHA256 tar -xzf cutlass.tar.gz @@ -62,7 +78,7 @@ export INFERENCE_ARTIFACTS="$PWD" cd .. python3.12 -m venv build-env source build-env/bin/activate -python -m pip install cmake ninja libclang numpy packaging scikit-build-core pybind11 pyyaml +python -m pip install cmake ninja libclang numpy packaging scikit-build-core pybind11 pyyaml jinja2 python -m pip install --no-index --no-deps --target "$INFINI_ROOT/python" \ "$INFERENCE_ARTIFACTS/flashinfer_jit_cache-0.6.7+cu128-cp39-abi3-manylinux_2_28_x86_64.whl" \ "$INFERENCE_ARTIFACTS/apache_tvm_ffi-0.1.10-cp312-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl" @@ -96,7 +112,6 @@ cmake -S "$CORE/submodules/InfiniOps" -B build/ops \ -DAUTO_DETECT_DEVICES=OFF -DAUTO_DETECT_BACKENDS=OFF \ -DGENERATE_OPERATOR_CALL_INSTANTIATIONS=ON -DGENERATE_PYTHON_BINDINGS=OFF \ -DBUILD_NVIDIA_FLASH_INFERENCE=ON -DBUILD_NVIDIA_W4_MARLIN=ON \ - -DFLASH_ATTN_SOURCE_ARCHIVE="$INFERENCE_ARTIFACTS/flash_attn-2.8.3.tar.gz" \ -DFETCHCONTENT_SOURCE_DIR_CUTLASS="$INFERENCE_ARTIFACTS/cutlass-087c84df83d254b5fb295a7a408f1a1d554085cf" \ -DINFINI_OPS_OPS="$CORE/submodules/InfiniOps/scripts/configs/infiniops_ops_nvidia_dense_sampling_torch_free.json" \ -DCMAKE_CUDA_ARCHITECTURES=80 \ @@ -110,28 +125,33 @@ cmake --install build/ops Marlin is enabled here because the LM profile explicitly contains W4 operators. A dense-only SDK can set it `OFF` and use an allowlist without `w4_marlin_*`; then its CUDA >=12.3/C++20 check and provider registration are both omitted. +Both native components require the standard `patch` command when enabled. +Their preparation scripts patch build-directory copies, leaving submodules +unchanged. Jinja2 is used only to generate Marlin instances at build time. +Source distributions include the consumed upstream files and can be built +without Git metadata or network access. Component-disabled builds do not need +these CUDA source submodules. An ordinary Git source archive does not include +submodule contents; use the complete source distribution for offline builds. Build CCL separately using the Core instructions and the same ABI/prefix before building LM. Pin and record the installed NCCL version as part of that deployment. ## Third-party notices and validation -Keep the TVM wheel's `apache_tvm_ffi-0.1.10.dist-info/licenses/LICENSE` and -`NOTICE`. The selected FlashInfer wheel omits a license directory: obtain the -Apache-2.0 `LICENSE` from its matching upstream release, not from a floating -branch. For example: - -```bash -install -d "$INFINI_ROOT/share/licenses/native-sampling" -curl -fL 'https://raw.githubusercontent.com/flashinfer-ai/flashinfer/v0.6.7/LICENSE' \ - -o "$INFINI_ROOT/share/licenses/native-sampling/FlashInfer-LICENSE.txt" -install -m644 "$INFINI_ROOT/python/apache_tvm_ffi-0.1.10.dist-info/licenses/LICENSE" \ - "$INFINI_ROOT/share/licenses/native-sampling/TVM-FFI-LICENSE.txt" -install -m644 "$INFINI_ROOT/python/apache_tvm_ffi-0.1.10.dist-info/licenses/NOTICE" \ - "$INFINI_ROOT/share/licenses/native-sampling/TVM-FFI-NOTICE.txt" -# When bundling these DSOs into an LM wheel, include the SDK notices too. -export LM=/absolute/path/to/InfiniLM -cp "$INFINI_ROOT/share/licenses/native-sampling/"*.txt "$LM/python/infinilm/licenses/" -``` +`cmake --install build/ops` installs the checked-in notices for the fixed +FlashInfer 0.6.7 and TVM-FFI 0.1.10 inputs into +`$INFINI_ROOT/share/licenses/InfiniOps`. The consuming LM setup collects this +SDK license directory when bundling the native libraries. No additional +download or manual copy into the LM source tree is needed for these inputs. + +The checked-in release notices cover a conservative superset of the two native +libraries. `sampling.so` is one JIT-cache artifact, not the entire FlashInfer +wheel; `libtvm_ffi.so` is the TVM-FFI runtime, not its full Python package. +In particular, TVM-FFI contains code adapted from PyTorch, and can compile +libbacktrace and DLPack interfaces into its runtime. A Torch-free execution path +does not establish that these secondary notices can be removed. Retain the +release bundles unless a source-to-binary audit establishes a smaller set for +the exact library hashes. Distributing the complete wheels requires their +complete applicable notices, independently of the Ops SDK. Preserve the separate FlashAttention/CUTLASS and Marlin notices described in `THIRD_PARTY.txt` when distributing those compiled providers. Do not treat this diff --git a/patches/flash-attention/inference.patch b/patches/flash-attention/inference.patch new file mode 100644 index 000000000..8e25d67d3 --- /dev/null +++ b/patches/flash-attention/inference.patch @@ -0,0 +1,148 @@ +--- a/flash.h ++++ b/flash.h +@@ -1,3 +1,6 @@ ++// Modified by InfiniOps from FlashAttention 2.8.3: inference-only native runtime integration. ++// Upstream: https://github.com/Dao-AILab/flash-attention/tree/v2.8.3/csrc/flash_attn/src ++// Original copyright and license notices below are retained. + /****************************************************************************** + * Copyright (c) 2023, Tri Dao. + ******************************************************************************/ +@@ -9,7 +12,7 @@ + #include + #include + +-#include // For at::Generator and at::PhiloxCudaState ++ + + namespace FLASH_NAMESPACE { + constexpr int TOTAL_DIM = 0; +@@ -119,7 +122,7 @@ + float softcap; + + // Random state. +- at::PhiloxCudaState philox_args; ++ // Inference-only build: dropout and its Torch RNG state are absent. + + // Pointer to the RNG seed (idx 0) and offset (idx 1). + uint64_t * rng_state; +--- a/flash_fwd_kernel.h ++++ b/flash_fwd_kernel.h +@@ -1,3 +1,6 @@ ++// Modified by InfiniOps from FlashAttention 2.8.3: inference-only native runtime integration. ++// Upstream: https://github.com/Dao-AILab/flash-attention/tree/v2.8.3/csrc/flash_attn/src ++// Original copyright and license notices below are retained. + /****************************************************************************** + * Copyright (c) 2024, Tri Dao. + ******************************************************************************/ +@@ -5,7 +8,7 @@ + #pragma once + + #include "namespace_config.h" +-#include "philox_unpack.cuh" // For at::cuda::philox::unpack ++ + + #include + +@@ -66,7 +69,8 @@ + constexpr int kHeadDim = Kernel_traits::kHeadDim; + constexpr int kNWarps = Kernel_traits::kNWarps; + +- auto seed_offset = at::cuda::philox::unpack(params.philox_args); ++ static_assert(!Is_dropout, "InfiniOps inference build does not support dropout"); ++ const auto seed_offset = std::make_tuple(uint64_t{0}, uint64_t{0}); + FLASH_NAMESPACE::Dropout dropout(std::get<0>(seed_offset), std::get<1>(seed_offset), params.p_dropout_in_uint8_t, + bidb, bidh, tidx, params.h); + +--- a/flash_fwd_launch_template.h ++++ b/flash_fwd_launch_template.h +@@ -1,10 +1,13 @@ ++// Modified by InfiniOps from FlashAttention 2.8.3: inference-only native runtime integration. ++// Upstream: https://github.com/Dao-AILab/flash-attention/tree/v2.8.3/csrc/flash_attn/src ++// Original copyright and license notices below are retained. + /****************************************************************************** + * Copyright (c) 2023, Tri Dao. + ******************************************************************************/ + + #pragma once + #include "namespace_config.h" +-#include // For C10_CUDA_CHECK and C10_CUDA_KERNEL_LAUNCH_CHECK ++#include "infini_flash_cuda.h" + + #include "static_switch.h" + #include "hardware_info.h" +@@ -81,7 +84,7 @@ + // printf("IsEvenMNConst = %d, IsEvenKConst = %d, Is_local = %d, Is_causal = %d, ReturnSoftmaxConst = %d, Is_dropout = %d\n", int(IsEvenMNConst), int(IsEvenKConst), int(Is_local), int(Is_causal), int(ReturnSoftmaxConst), int(Is_dropout)); + // auto kernel = &flash_fwd_kernel; + if (smem_size >= 48 * 1024) { +- C10_CUDA_CHECK(cudaFuncSetAttribute( ++ INFINI_FLASH_CUDA_CHECK(cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + } + // int ctas_per_sm; +@@ -89,7 +92,7 @@ + // &ctas_per_sm, kernel, Kernel_traits::kNThreads, smem_size); + // printf("smem_size = %d, CTAs per SM = %d\n", int(smem_size), ctas_per_sm); + kernel<<>>(params); +- C10_CUDA_KERNEL_LAUNCH_CHECK(); ++ INFINI_FLASH_KERNEL_CHECK(); + }); + }); + }); +@@ -121,11 +124,11 @@ + // auto kernel = &flash_fwd_splitkv_kernel; + // auto kernel = &flash_fwd_splitkv_kernel; + if (smem_size >= 48 * 1024) { +- C10_CUDA_CHECK(cudaFuncSetAttribute( ++ INFINI_FLASH_CUDA_CHECK(cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + } + kernel<<>>(params); +- C10_CUDA_KERNEL_LAUNCH_CHECK(); ++ INFINI_FLASH_KERNEL_CHECK(); + }); + }); + }); +@@ -155,7 +158,7 @@ + } else if (params.num_splits <= 128) { + flash_fwd_splitkv_combine_kernel<<>>(params); + } +- C10_CUDA_KERNEL_LAUNCH_CHECK(); ++ INFINI_FLASH_KERNEL_CHECK(); + }); + } + } +@@ -284,7 +287,7 @@ + status_ = cudaDeviceGetAttribute( + &max_smem_per_block, cudaDevAttrMaxSharedMemoryPerBlockOptin, device); + if (status_ != cudaSuccess) { +- C10_CUDA_CHECK(status_); ++ INFINI_FLASH_CUDA_CHECK(status_); + } + // printf("max_smem_per_sm = %d, max_smem_per_block = %d\n", max_smem_per_sm, max_smem_per_block); + DROPOUT_SWITCH(params.p_dropout < 1.f, Is_dropout, [&] { +--- a/hardware_info.h ++++ b/hardware_info.h +@@ -1,3 +1,6 @@ ++// Modified by InfiniOps from FlashAttention 2.8.3: inference-only native runtime integration. ++// Upstream: https://github.com/Dao-AILab/flash-attention/tree/v2.8.3/csrc/flash_attn/src ++// Original copyright and license notices below are retained. + /****************************************************************************** + * Copyright (c) 2024, Tri Dao. + ******************************************************************************/ +@@ -5,6 +8,7 @@ + #pragma once + + #include ++#include "infini_flash_cuda.h" + + #if !defined(__CUDACC_RTC__) + #include "cuda_runtime.h" +@@ -16,7 +20,7 @@ + if (status_ != cudaSuccess) { \ + fprintf(stderr, "CUDA error (%s:%d): %s\n", __FILE__, __LINE__, \ + cudaGetErrorString(status_)); \ +- exit(1); \ ++ INFINI_FLASH_CUDA_CHECK(status_); \ + } \ + } while (0) + diff --git a/patches/marlin/w4a16.patch b/patches/marlin/w4a16.patch new file mode 100644 index 000000000..5d4f18805 --- /dev/null +++ b/patches/marlin/w4a16.patch @@ -0,0 +1,420 @@ +--- a/gemm/marlin/dequant.h ++++ b/gemm/marlin/dequant.h +@@ -1,3 +1,39 @@ ++// Source: InfiniCore 1ab85ef1d3c2bce49edf1fc73a5b48becaf7552e, ++// src/infiniop/ops/awq_marlin_gemm/marlin/dequant.h. ++// Donor implementation retained; this attribution notice was added by InfiniOps. ++// Marlin/vLLM and direct-donor licenses: repository THIRD_PARTY.txt. ++// INT4/INT8 conversion portions derive from NVIDIA FasterTransformer, ++// release/v5.3_tag, src/fastertransformer/cutlass_extensions/include/ ++// cutlass_extensions/interleaved_numeric_conversion.h (references below). ++/* ++Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ++SPDX-License-Identifier: BSD-3-Clause ++ ++Redistribution and use in source and binary forms, with or without ++modification, are permitted provided that the following conditions are met: ++ ++1. Redistributions of source code must retain the above copyright notice, this ++list of conditions and the following disclaimer. ++ ++2. Redistributions in binary form must reproduce the above copyright notice, ++this list of conditions and the following disclaimer in the documentation ++and/or other materials provided with the distribution. ++ ++3. Neither the name of the copyright holder nor the names of its ++contributors may be used to endorse or promote products derived from ++this software without specific prior written permission. ++ ++THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" ++AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE ++IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE ++DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE ++FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL ++DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR ++SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER ++CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, ++OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE ++OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. ++*/ + /* + Fast Dequantization (Converting INT4/INT8/FP4/FP8 to FP16/BF16) + +--- a/gemm/marlin/marlin.cuh ++++ b/gemm/marlin/marlin.cuh +@@ -10,6 +10,16 @@ + #ifndef MARLIN_NAMESPACE_NAME + #define MARLIN_NAMESPACE_NAME marlin + #endif ++ ++template ++__device__ __forceinline__ uint32_t __cvta_generic_to_shared(T *ptr) { ++ size_t smem_addr; ++ asm volatile( ++ "cvta.to.shared.u64 %0, %1;" ++ : "=l"(smem_addr) ++ : "l"(ptr)); ++ return static_cast(smem_addr); ++} + + namespace MARLIN_NAMESPACE_NAME { + +--- a/gemm/marlin/marlin_template.h ++++ b/gemm/marlin/marlin_template.h +@@ -327,9 +327,6 @@ + if constexpr (b_type == vllm::kFE2M1f) { + static_assert(s_type == vllm::kFE4M3fn && group_blocks == 1 || + s_type == vllm::kFE8M0fnu && group_blocks == 2); +- } else if constexpr (s_type == vllm::kFE8M0fnu) { +- // MXFP8: FP8 weights with e8m0 microscaling block scales +- static_assert(b_type == vllm::kFE4M3fn && group_blocks == 2); + } else if constexpr (std::is_same::value) { + static_assert(s_type == vllm::kBFloat16); + } else if constexpr (std::is_same::value) { +@@ -337,7 +334,6 @@ + } + + constexpr bool is_a_8bit = a_type.size_bits() == 8; +- constexpr bool is_8bit_scale = s_type.size_bits() == 8; + if constexpr (!is_a_8bit) { + static_assert(std::is_same::value); + } +@@ -347,7 +343,7 @@ + b_type == vllm::kU4B8 || b_type == vllm::kU8B128; + // see comments of dequant.h for more details + constexpr bool dequant_skip_flop = +- is_a_8bit || (b_type == vllm::kFE4M3fn && !(s_type == vllm::kFE8M0fnu)) || ++ is_a_8bit || b_type == vllm::kFE4M3fn || + b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn || + has_zp && !is_zp_float && !std::is_same::value || + has_zp && !is_zp_float && !(b_type == vllm::kU8); +@@ -559,8 +555,8 @@ + constexpr int b_sh_wr_iters = b_sh_stage / b_sh_wr_delta; + + // Scale sizes/strides without act_order +- int s_gl_stride = prob_n / (is_8bit_scale ? 16 : 8); +- constexpr int s_sh_stride = 16 * thread_n_blocks / (is_8bit_scale ? 16 : 8); ++ int s_gl_stride = prob_n / (b_type == vllm::kFE2M1f ? 16 : 8); ++ constexpr int s_sh_stride = 16 * thread_n_blocks / (b_type == vllm::kFE2M1f ? 16 : 8); + constexpr int s_tb_groups = + !has_act_order && group_blocks != -1 && group_blocks < thread_k_blocks + ? thread_k_blocks / group_blocks +@@ -1000,7 +996,7 @@ + + int4* sh_s_stage = sh_s + s_sh_stage * pipe; + +- if constexpr (!is_8bit_scale) { ++ if constexpr (b_type_id != vllm::kFE2M1f.id()) { + reinterpret_cast(&frag_s[k % 2])[0] = + sh_s_stage[s_sh_rd + cur_group_id * s_sh_stride]; + } else { +@@ -1009,7 +1005,7 @@ + sh_s_stage)[s_sh_rd + cur_group_id * (2 * s_sh_stride)]; + } + } else if (group_blocks >= b_sh_wr_iters) { +- if constexpr (!is_8bit_scale) { ++ if constexpr (b_type_id != vllm::kFE2M1f.id()) { + reinterpret_cast(&frag_s[1])[0] = + reinterpret_cast(&frag_s[0])[0]; + } else { +@@ -1035,7 +1031,7 @@ + cur_k = 0; + + // Progress to current iteration +- cur_k += k % b_sh_wr_iters; ++ cur_k += 16 * (k % b_sh_wr_iters); + + // Determine "position" inside the thread-block (based on warp and + // thread-id) +@@ -1210,7 +1206,7 @@ + } + } + +- if constexpr (s_type == vllm::kFE4M3fn || s_type == vllm::kFE8M0fnu) { ++ if constexpr (b_type == vllm::kFE2M1f) { + int s_quant_0 = reinterpret_cast(frag_s[k2])[0]; + int s_quant_1 = reinterpret_cast(frag_s[k2])[1]; + +--- a/awq_repack/cuda/kernel.cuh ++++ b/awq_repack/cuda/kernel.cuh +@@ -1,18 +1,9 @@ +-#include "marlin.cuh" +- +-#include +-#include +-#include +-#include +-#include +-#include +- +-#include "libtorch_stable/torch_utils.h" ++#include "../../gemm/marlin/marlin.cuh" + + namespace marlin { + + template +-__global__ void awq_marlin_repack_kernel( ++__device__ void awq_marlin_repack_kernel( + uint32_t const* __restrict__ b_q_weight_ptr, uint32_t* __restrict__ out_ptr, + int size_k, int size_n) { + constexpr int pack_factor = 32 / num_bits; +@@ -212,83 +203,3 @@ + } + + } // namespace marlin +- +-#define CALL_IF(NUM_BITS, IS_A_8BIT) \ +- else if (num_bits == NUM_BITS && is_a_8bit == IS_A_8BIT) { \ +- cudaFuncSetAttribute( \ +- marlin::awq_marlin_repack_kernel, \ +- cudaFuncAttributeMaxDynamicSharedMemorySize, max_shared_mem); \ +- marlin::awq_marlin_repack_kernel \ +- <<>>( \ +- b_q_weight_ptr, out_ptr, size_k, size_n); \ +- } +- +-torch::stable::Tensor awq_marlin_repack(torch::stable::Tensor& b_q_weight, +- int64_t size_k, int64_t size_n, +- int64_t num_bits, bool is_a_8bit) { +- // Verify compatibility with marlin tile of 16x64 +- STD_TORCH_CHECK(size_k % marlin::tile_k_size == 0, "size_k = ", size_k, +- " is not divisible by tile_k_size = ", marlin::tile_k_size); +- STD_TORCH_CHECK(size_n % marlin::tile_n_size == 0, "size_n = ", size_n, +- " is not divisible by tile_n_size = ", marlin::tile_n_size); +- +- STD_TORCH_CHECK(num_bits == 4 || num_bits == 8, +- "num_bits must be 4 or 8. Got = ", num_bits); +- int const pack_factor = 32 / num_bits; +- +- // Verify B +- STD_TORCH_CHECK(b_q_weight.size(0) == size_k, +- "b_q_weight.size(0) = ", b_q_weight.size(0), +- " is not size_k = ", size_k); +- STD_TORCH_CHECK((size_n / pack_factor) == b_q_weight.size(1), +- "Shape mismatch: b_q_weight.size(1) = ", b_q_weight.size(1), +- ", size_n = ", size_n, ", pack_factor = ", pack_factor); +- +- // Verify device and strides +- STD_TORCH_CHECK(b_q_weight.is_cuda(), "b_q_weight is not on GPU"); +- STD_TORCH_CHECK(b_q_weight.is_contiguous(), "b_q_weight is not contiguous"); +- STD_TORCH_CHECK( +- b_q_weight.scalar_type() == torch::headeronly::ScalarType::Int, +- "b_q_weight type is not kInt"); +- +- const int32_t device_index = b_q_weight.get_device_index(); +- torch::stable::accelerator::DeviceGuard device_guard(device_index); +- const cudaStream_t stream = get_current_cuda_stream(device_index); +- +- // Alloc buffers +- torch::stable::Tensor out = torch::stable::empty( +- {size_k / marlin::tile_size, size_n * marlin::tile_size / pack_factor}, +- b_q_weight.scalar_type(), std::nullopt, b_q_weight.device()); +- +- // Get ptrs +- uint32_t const* b_q_weight_ptr = +- reinterpret_cast(b_q_weight.const_data_ptr()); +- uint32_t* out_ptr = reinterpret_cast(out.mutable_data_ptr()); +- +- int blocks; +- cudaDeviceGetAttribute(&blocks, cudaDevAttrMultiProcessorCount, device_index); +- +- int max_shared_mem = 0; +- cudaDeviceGetAttribute(&max_shared_mem, +- cudaDevAttrMaxSharedMemoryPerBlockOptin, device_index); +- STD_TORCH_CHECK(max_shared_mem > 0); +- +- if (false) { +- } +- CALL_IF(4, false) +- CALL_IF(8, false) +- CALL_IF(4, true) +- CALL_IF(8, true) +- else { +- STD_TORCH_CHECK(false, "Unsupported repack config: num_bits = ", num_bits, +- ", is_a_8bit = ", is_a_8bit); +- } +- +- return out; +-} +- +-STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) { +- m.impl("awq_marlin_repack", TORCH_BOX(&awq_marlin_repack)); +-} +--- a/gptq_repack/cuda/kernel.cuh ++++ b/gptq_repack/cuda/kernel.cuh +@@ -1,19 +1,10 @@ +-#include "marlin.cuh" +- +-#include +-#include +-#include +-#include +-#include +-#include +- +-#include "libtorch_stable/torch_utils.h" ++#include "../../gemm/marlin/marlin.cuh" + + namespace marlin { + + template +-__global__ void gptq_marlin_repack_kernel( ++__device__ void gptq_marlin_repack_kernel( + uint32_t const* __restrict__ b_q_weight_ptr, + uint32_t const* __restrict__ perm_ptr, uint32_t* __restrict__ out_ptr, + int size_k, int size_n) { +@@ -268,99 +259,3 @@ + } + + } // namespace marlin +- +-#define CALL_IF(NUM_BITS, HAS_PERM, IS_A_8BIT) \ +- else if (num_bits == NUM_BITS && has_perm == HAS_PERM && \ +- is_a_8bit == IS_A_8BIT) { \ +- cudaFuncSetAttribute( \ +- marlin::gptq_marlin_repack_kernel, \ +- cudaFuncAttributeMaxDynamicSharedMemorySize, max_shared_mem); \ +- marlin::gptq_marlin_repack_kernel \ +- <<>>( \ +- b_q_weight_ptr, perm_ptr, out_ptr, size_k, size_n); \ +- } +- +-torch::stable::Tensor gptq_marlin_repack(torch::stable::Tensor& b_q_weight, +- torch::stable::Tensor& perm, +- int64_t size_k, int64_t size_n, +- int64_t num_bits, bool is_a_8bit) { +- // Verify compatibility with marlin tile of 16x64 +- STD_TORCH_CHECK(size_k % marlin::tile_k_size == 0, "size_k = ", size_k, +- " is not divisible by tile_k_size = ", marlin::tile_k_size); +- STD_TORCH_CHECK(size_n % marlin::tile_n_size == 0, "size_n = ", size_n, +- " is not divisible by tile_n_size = ", marlin::tile_n_size); +- +- STD_TORCH_CHECK(num_bits == 4 || num_bits == 8, +- "num_bits must be 4 or 8. Got = ", num_bits); +- int const pack_factor = 32 / num_bits; +- +- // Verify B +- STD_TORCH_CHECK((size_k / pack_factor) == b_q_weight.size(0), +- "Shape mismatch: b_q_weight.size(0) = ", b_q_weight.size(0), +- ", size_k = ", size_k, ", pack_factor = ", pack_factor); +- STD_TORCH_CHECK(b_q_weight.size(1) == size_n, +- "b_q_weight.size(1) = ", b_q_weight.size(1), +- " is not size_n = ", size_n); +- +- // Verify device and strides +- STD_TORCH_CHECK(b_q_weight.is_cuda(), "b_q_weight is not on GPU"); +- STD_TORCH_CHECK(b_q_weight.is_contiguous(), "b_q_weight is not contiguous"); +- STD_TORCH_CHECK( +- b_q_weight.scalar_type() == torch::headeronly::ScalarType::Int, +- "b_q_weight type is not kInt"); +- +- STD_TORCH_CHECK(perm.is_cuda(), "perm is not on GPU"); +- STD_TORCH_CHECK(perm.is_contiguous(), "perm is not contiguous"); +- STD_TORCH_CHECK(perm.scalar_type() == torch::headeronly::ScalarType::Int, +- "perm type is not at::kInt"); +- +- const int32_t device_index = b_q_weight.get_device_index(); +- torch::stable::accelerator::DeviceGuard device_guard(device_index); +- const cudaStream_t stream = get_current_cuda_stream(device_index); +- +- // Alloc buffers +- torch::stable::Tensor out = torch::stable::empty( +- {size_k / marlin::tile_size, size_n * marlin::tile_size / pack_factor}, +- b_q_weight.scalar_type(), std::nullopt, b_q_weight.device()); +- +- // Detect if there is act_order +- bool has_perm = perm.size(0) != 0; +- +- // Get ptrs +- uint32_t const* b_q_weight_ptr = +- reinterpret_cast(b_q_weight.const_data_ptr()); +- uint32_t const* perm_ptr = +- reinterpret_cast(perm.const_data_ptr()); +- uint32_t* out_ptr = reinterpret_cast(out.mutable_data_ptr()); +- +- int blocks; +- cudaDeviceGetAttribute(&blocks, cudaDevAttrMultiProcessorCount, device_index); +- +- int max_shared_mem = 0; +- cudaDeviceGetAttribute(&max_shared_mem, +- cudaDevAttrMaxSharedMemoryPerBlockOptin, device_index); +- STD_TORCH_CHECK(max_shared_mem > 0); +- +- if (false) { +- } +- CALL_IF(4, false, false) +- CALL_IF(4, true, false) +- CALL_IF(8, false, false) +- CALL_IF(8, true, false) +- +- CALL_IF(4, false, true) +- CALL_IF(8, false, true) +- +- else { +- STD_TORCH_CHECK(false, "Unsupported repack config: num_bits = ", num_bits, +- ", has_perm = ", has_perm, ", is_a_8bit = ", is_a_8bit); +- } +- +- return out; +-} +- +-STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) { +- m.impl("gptq_marlin_repack", TORCH_BOX(&gptq_marlin_repack)); +-} +--- a/gemm/nvidia/generate_kernels.py ++++ b/gemm/nvidia/generate_kernels.py +@@ -27,15 +27,17 @@ + SUPPORT_SM75 = True + + FILE_HEAD_COMMENT = """ +-// auto generated by generate_kernels.py ++// SPDX-License-Identifier: Apache-2.0 ++// Copyright contributors to the vLLM project. ++// Generated from pinned vLLM by prepare_marlin.py (W4A16 subset). + // clang-format off + """.lstrip() + + FILE_HEAD = ( + FILE_HEAD_COMMENT + + """ +-#include "kernel.h" +-#include "marlin_template.h" ++#include "../marlin/kernel.h" ++#include "../marlin/marlin_template.h" + + namespace MARLIN_NAMESPACE_NAME { + """ +@@ -162,6 +164,10 @@ + ] + + ++# InfiniOps retains the existing AWQ/GPTQ W4A16 instance set. ++QUANT_CONFIGS = QUANT_CONFIGS[:2] ++ ++ + def remove_old_kernels(): + for filename in glob.glob(os.path.dirname(__file__) + "/*kernel_*.cu"): + subprocess.call(["rm", "-f", filename]) +@@ -300,17 +306,9 @@ + with open(os.path.join(os.path.dirname(__file__), filename), "w") as f: + f.write(file_content) + +- if not SUPPORT_FP8 and kernel_selector_str != FILE_HEAD_COMMENT: +- kernel_selector_str += ( +- "else if (a_type == vllm::kFE4M3fn)\n" +- " STD_TORCH_CHECK(false, " +- '"marlin kernel with fp8 activation is not built.");' +- ) +- + with open(os.path.join(os.path.dirname(__file__), "kernel_selector.h"), "w") as f: + f.write(kernel_selector_str) + + + if __name__ == "__main__": +- remove_old_kernels() + generate_new_kernels() diff --git a/pyproject.toml b/pyproject.toml index 2dd4ba2db..68a883585 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,5 +1,5 @@ [build-system] -requires = ["scikit-build-core", "pybind11", "libclang", "packaging", "pyyaml"] +requires = ["scikit-build-core", "pybind11", "libclang", "packaging", "pyyaml", "jinja2"] build-backend = "scikit_build_core.build" [project] @@ -12,6 +12,23 @@ dev = ["pytest", "pytest-cov", "pytest-xdist", "ruff==0.15.22", "torch", "pyyaml [tool.scikit-build.wheel] install-dir = "infini" +[tool.scikit-build.sdist] +exclude = ["third_party/flash-attention/**", "third_party/vllm/**"] +include = [ + "third_party/flash-attention/LICENSE", + "third_party/flash-attention/AUTHORS", + "third_party/flash-attention/csrc/flash_attn/src/*.h", + "third_party/flash-attention/csrc/flash_attn/src/*.cuh", + "third_party/flash-attention/csrc/flash_attn/flash_api.cpp", + "third_party/flash-attention/csrc/cutlass/include/**", + "third_party/flash-attention/csrc/cutlass/LICENSE.txt", + "third_party/vllm/LICENSE", + "third_party/vllm/csrc/libtorch_stable/quantization/marlin/*.h", + "third_party/vllm/csrc/libtorch_stable/quantization/marlin/*.cuh", + "third_party/vllm/csrc/libtorch_stable/quantization/marlin/*repack.cu", + "third_party/vllm/csrc/libtorch_stable/quantization/marlin/generate_kernels.py", +] + [tool.scikit-build.cmake.define] AUTO_DETECT_DEVICES = "ON" AUTO_DETECT_BACKENDS = "ON" diff --git a/scripts/generate_public_operator_header.py b/scripts/generate_public_operator_header.py index 622ef0027..9dad53155 100644 --- a/scripts/generate_public_operator_header.py +++ b/scripts/generate_public_operator_header.py @@ -20,6 +20,17 @@ # lookup and construction bodies must not be optimized in a downstream TU, # even when that TU sees an `extern template` declaration. _PROVIDER_SIGNATURES = ( + ( + "Operator::CanMake(vector)", + re.compile( + r"(?P^[ \t]*)template\s*\s*\n" + r"(?P=indent)static\s+bool\s+CanMake\s*\(\s*" + r"const\s+Config\s*&\s*config\s*,\s*" + r"const\s+std::vector\s+tensors\s*,\s*" + r"const\s+Args\s*&\s*\.\.\.\s*args\s*\)\s*", + re.MULTILINE, + ), + ), ( "Operator::CanMake(Tensor)", re.compile( diff --git a/scripts/generate_wrappers.py b/scripts/generate_wrappers.py index 975390ee0..31b272525 100644 --- a/scripts/generate_wrappers.py +++ b/scripts/generate_wrappers.py @@ -1678,15 +1678,10 @@ def _append_unique(declaration, definition): f"extern template std::unique_ptr> {make_instantiation};", f"template std::unique_ptr> {make_instantiation};", ) - if operator.name in ( - "flash_attn_with_kvcache", - "reshape_and_cache_flash", - "flash_attn_varlen_func", - ): - _append_unique( - f"extern template bool {can_make_instantiation};", - f"template bool {can_make_instantiation};", - ) + _append_unique( + f"extern template bool {can_make_instantiation};", + f"template bool {can_make_instantiation};", + ) _append_unique( f"extern template void {operator_instantiation};", f"template void {operator_instantiation};", @@ -1904,9 +1899,20 @@ def _select_ops_from_config(ops, config, config_path): headers = selection["headers"] if headers is not None: - selected[op_name] = [ - _implementation_from_json(header) for header in headers - ] + requested = [_implementation_from_json(header) for header in headers] + active = {impl.path.resolve() for impl in ops.get(op_name, [])} + for impl in requested: + # Built-in headers still require their compiled backend. + # Preserve explicit external-provider header support. + if ( + impl.path.resolve().is_relative_to(_SRC_DIR.resolve()) + and impl.path.resolve() not in active + ): + raise ValueError( + f"{config_path}: provider {impl.path} for {op_name!r} " + "is not enabled by this build" + ) + selected[op_name] = requested continue if op_name not in ops: @@ -1985,19 +1991,25 @@ def _get_all_ops( if not with_native_flash: excluded_headers.update( { - "src/native/cuda/nvidia/ops/flash_attn_with_kvcache/kernel.h", - "src/native/cuda/nvidia/ops/flash_attn_varlen_func/kernel.h", + ( + _SRC_DIR / "native/cuda/nvidia/ops/flash_attn_with_kvcache/kernel.h" + ).resolve(), + ( + _SRC_DIR / "native/cuda/nvidia/ops/flash_attn_varlen_func/kernel.h" + ).resolve(), } ) if not with_native_w4: - excluded_headers.add("src/native/cuda/nvidia/ops/w4_marlin/kernel.h") + excluded_headers.add( + (_SRC_DIR / "native/cuda/nvidia/ops/w4_marlin/kernel.h").resolve() + ) if excluded_headers: # Optional providers are emitted only with their compiled kernel target. for name, headers in impl_headers_by_operator.items(): impl_headers_by_operator[name] = [ header for header in headers - if header.path.as_posix() not in excluded_headers + if header.path.resolve() not in excluded_headers ] for base_dir in base_dirs: diff --git a/scripts/prepare_flash_attention_inference.py b/scripts/prepare_flash_attention_inference.py index 6b155c0de..395e4f325 100644 --- a/scripts/prepare_flash_attention_inference.py +++ b/scripts/prepare_flash_attention_inference.py @@ -1,181 +1,116 @@ -"""Prepare pinned upstream forward kernels without a Torch runtime dependency. +"""Prepare the pinned FlashAttention forward headers for native inference. -This is a build-time, inference-only adaptation, not another Attention algorithm. -Do not link these altered parameter layouts to the original Torch extension. +Apply the reviewed adaptation only in build staging, never in the submodule. +The prepared parameter layout must not be linked to the upstream Torch extension. """ import argparse -import hashlib -import json import pathlib -import re -import tarfile - - -ARCHIVE_SHA256 = "1e71dd64a9e0280e0447b8a0c2541bad4bf6ac65bdeaa2f90e51a9e57de0370d" -ROOT = "flash_attn-2.8.3/" -KERNEL_ROOT = ROOT + "csrc/flash_attn/src/" - - -def replace_once(text, before, after): - if text.count(before) != 1: - raise ValueError(f"pinned FlashAttention source contract changed: {before!r}") - return text.replace(before, after, 1) +import subprocess +import tempfile + + +# Explicit forward closure for the pinned source. num_splits.h is extracted +# separately from flash_api.cpp; CUTLASS headers are included directly by CMake. +HEADERS = ( + "block_info.h", + "dropout.h", + "flash.h", + "flash_fwd_kernel.h", + "flash_fwd_launch_template.h", + "hardware_info.h", + "kernel_traits.h", + "mask.h", + "namespace_config.h", + "philox.cuh", + "rotary.h", + "softmax.h", + "static_switch.h", + "utils.h", +) +PATCH = pathlib.Path(__file__).resolve().parents[1] / "patches/flash-attention/inference.patch" def write_if_changed(path, content): - # CMake may prepare the same pinned closure again during an incremental - # build. Preserve timestamps so hundreds of CUDA kernels are not rebuilt - # merely because configure ran; changed content still invalidates them. + # Preserve timestamps across reconfiguration when prepared bytes agree. content = content.encode("utf-8") if isinstance(content, str) else content if not path.exists() or path.read_bytes() != content: path.parent.mkdir(parents=True, exist_ok=True) path.write_bytes(content) -def adapt(name, text): - if name == "flash.h": - text = replace_once( - text, - "#include // For at::Generator and at::PhiloxCudaState", - "", - ) - text = replace_once( - text, - " at::PhiloxCudaState philox_args;", - " // Inference-only build: dropout and its Torch RNG state are absent.", - ) - elif name == "flash_fwd_kernel.h": - text = replace_once( - text, '#include "philox_unpack.cuh" // For at::cuda::philox::unpack', "" - ) - text = replace_once( - text, - " auto seed_offset = at::cuda::philox::unpack(params.philox_args);", - ' static_assert(!Is_dropout, "InfiniOps inference build does not support dropout");\n const auto seed_offset = std::make_tuple(uint64_t{0}, uint64_t{0});', - ) - elif name == "flash_fwd_launch_template.h": - text = replace_once( - text, - "#include // For C10_CUDA_CHECK and C10_CUDA_KERNEL_LAUNCH_CHECK", - '#include "infini_flash_cuda.h"', - ) - text = text.replace("C10_CUDA_KERNEL_LAUNCH_CHECK", "INFINI_FLASH_KERNEL_CHECK") - text = text.replace("C10_CUDA_CHECK", "INFINI_FLASH_CUDA_CHECK") - elif name == "hardware_info.h": - text = replace_once( - text, "#include ", '#include \n#include "infini_flash_cuda.h"' - ) - text = replace_once( - text, " exit(1);", " INFINI_FLASH_CUDA_CHECK(status_);" - ) - if name in ( - "flash.h", - "flash_fwd_kernel.h", - "flash_fwd_launch_template.h", - "hardware_info.h", - ): - text = ( - "// Modified by InfiniOps from FlashAttention 2.8.3: inference-only native runtime integration.\n" - "// Upstream: https://github.com/Dao-AILab/flash-attention/tree/v2.8.3/csrc/flash_attn/src\n" - "// Original copyright and license notices below are retained.\n" + text - ) - return text - - -def prepare(archive, output): - archive, output = pathlib.Path(archive), pathlib.Path(output) - digest = hashlib.sha256(archive.read_bytes()).hexdigest() - if digest != ARCHIVE_SHA256: +def read_source(path): + if not path.is_file(): raise ValueError( - f"expected pinned flash_attn-2.8.3 archive {ARCHIVE_SHA256}, got {digest}" - ) - with tarfile.open(archive, "r:gz") as source: - members = {member.name: member for member in source.getmembers()} - - def read(name): - member = members.get(name) - if member is None or not member.isfile(): - raise ValueError(f"missing regular pinned source file: {name}") - return source.extractfile(member).read() - - # Only the transitive forward include closure is needed. No backward, - # ATen wrapper, Python extension or Torch allocator is compiled. - pending = ["flash_fwd_launch_template.h"] - headers = {} - while pending: - name = pending.pop() - if name in headers: - continue - if pathlib.PurePosixPath(name).name != name: - raise ValueError(f"unexpected upstream include: {name}") - text = adapt(name, read(KERNEL_ROOT + name).decode()) - if re.search( - r"#\s*include\s*[<\"](?:ATen|c10|torch)/|\bat::|\bc10::", text - ): - raise ValueError(f"unremoved Torch dependency in {name}") - headers[name] = text - for include in re.findall(r'#\s*include\s*"([^"]+)"', text): - if KERNEL_ROOT + include in members: - pending.append(include) - - # Reuse the pinned host scheduling heuristic as well as the kernels; - # do not invent a second split-KV tuning policy in LM. - api = read(ROOT + "csrc/flash_attn/flash_api.cpp").decode() - begin = api.index("inline int num_splits_heuristic(") - end = api.index( - "\nstd::tuple set_params_splitkv", begin - ) - headers["num_splits.h"] = ( - "// Copyright (c) 2023, Tri Dao. See accompanying LICENSE.\n" - "// Extracted from flash_attn-2.8.3/csrc/flash_attn/flash_api.cpp;\n" - "// modified by InfiniOps to remove Torch dependencies and use a private namespace.\n" - "#pragma once\n#include \n#include \n#include \n" - "namespace infini_flash {\n" + api[begin:end] + "\n}\n" + f"missing FlashAttention source dependency: {path}; " + "initialize third_party/flash-attention and its CUTLASS submodule" ) - - output.mkdir(parents=True, exist_ok=True) + return path.read_text(encoding="utf-8") + + +def prepare(source, output, patch_executable="patch"): + source, output = pathlib.Path(source).resolve(), pathlib.Path(output).resolve() + if output == source or source in output.parents: + raise ValueError("FlashAttention preparation output must be outside its source tree") + kernel_source = source / "csrc/flash_attn/src" + headers = {name: read_source(kernel_source / name) for name in HEADERS} + api = read_source(source / "csrc/flash_attn/flash_api.cpp") + notices = {name: read_source(source / name) for name in ("LICENSE", "AUTHORS")} + notices["CUTLASS-LICENSE.txt"] = read_source(source / "csrc/cutlass/LICENSE.txt") + if not PATCH.is_file(): + raise ValueError(f"missing FlashAttention inference adaptation: {PATCH}") + + output.parent.mkdir(parents=True, exist_ok=True) + with tempfile.TemporaryDirectory(prefix="flash-inference-", dir=output.parent) as temp: + staging = pathlib.Path(temp) for name, text in headers.items(): - write_if_changed(output / name, text) - for member in source.getmembers(): - prefix = ROOT + "csrc/cutlass/include/" - if member.name.startswith(prefix) and member.isfile(): - relative = pathlib.PurePosixPath(member.name[len(prefix) :]) - if relative.is_absolute() or ".." in relative.parts: - raise ValueError("unsafe CUTLASS archive path") - path = output / "cutlass" / pathlib.Path(*relative.parts) - path.parent.mkdir(parents=True, exist_ok=True) - write_if_changed(path, read(member.name)) - for name in ("LICENSE", "AUTHORS"): - write_if_changed(output / name, read(ROOT + name)) - # The sdist omits CUTLASS's root LICENSE.txt; preserve the full license - # block from its actual bundled header, rather than substituting main. - cutlass_header = read(ROOT + "csrc/cutlass/include/cutlass/cutlass.h").decode() - cutlass_notice = re.match(r"\s*(/\*.*?\*/)", cutlass_header, re.DOTALL) - if cutlass_notice is None or "Redistribution" not in cutlass_notice.group(1): - raise ValueError("missing pinned CUTLASS copyright/license block") - write_if_changed(output / "CUTLASS-LICENSE.txt", cutlass_notice.group(1) + "\n") - manifest = { - "upstream": "flash-attn", - "version": "2.8.3", - "archive_sha256": digest, - "namespace": "infini_flash", - "inference_only": True, - "dropout": False, - "headers_sha256": { - name: hashlib.sha256(text.encode()).hexdigest() - for name, text in sorted(headers.items()) - }, - } - write_if_changed( - output / "source-manifest.json", json.dumps(manifest, indent=2) + "\n" - ) - return manifest + write_if_changed(staging / name, text) + try: + applied = subprocess.run( + [ + str(patch_executable), "--batch", "--forward", "--fuzz=0", + "-p1", "--input", str(PATCH), + ], + cwd=staging, + text=True, + capture_output=True, + check=False, + timeout=30, + ) + except FileNotFoundError as exc: + raise ValueError( + "native FlashAttention preparation requires the patch executable" + ) from exc + except subprocess.TimeoutExpired as exc: + raise ValueError("FlashAttention inference patch timed out") from exc + if applied.returncode: + raise ValueError( + "pinned FlashAttention inference patch failed:\n" + + applied.stdout + applied.stderr + ) + headers = {name: (staging / name).read_text(encoding="utf-8") for name in HEADERS} + + # Keep the upstream scheduling heuristic, without introducing another policy. + begin = api.index("inline int num_splits_heuristic(") + end = api.index("\nstd::tuple set_params_splitkv", begin) + headers["num_splits.h"] = ( + "// Copyright (c) 2023, Tri Dao. See accompanying LICENSE.\n" + "// Extracted from flash_attn-2.8.3/csrc/flash_attn/flash_api.cpp;\n" + "// modified by InfiniOps to remove Torch dependencies and use a private namespace.\n" + "#pragma once\n#include \n#include \n#include \n" + "namespace infini_flash {\n" + api[begin:end] + "\n}\n" + ) + for name, text in {**headers, **notices}.items(): + write_if_changed(output / name, text) if __name__ == "__main__": parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--archive", required=True) + parser.add_argument("--source-dir", required=True) parser.add_argument("--output", required=True) + parser.add_argument("--patch-executable", default="patch") args = parser.parse_args() - prepare(args.archive, args.output) + try: + prepare(args.source_dir, args.output, args.patch_executable) + except (ValueError, OSError) as exc: + parser.exit(1, f"FlashAttention preparation failed: {exc}\n") diff --git a/scripts/prepare_marlin.py b/scripts/prepare_marlin.py new file mode 100644 index 000000000..e121decd6 --- /dev/null +++ b/scripts/prepare_marlin.py @@ -0,0 +1,88 @@ +#!/usr/bin/env python3 +"""Prepare the pinned vLLM Marlin W4A16 sources in the build directory.""" + +import argparse +import pathlib +import subprocess +import sys +import tempfile + + +REPO = pathlib.Path(__file__).resolve().parents[1] +UPSTREAM = pathlib.Path("csrc/libtorch_stable/quantization/marlin") +SOURCES = { + pathlib.Path("gemm/marlin") / name: UPSTREAM / name + for name in ( + "dequant.h", "kernel.h", "marlin.cuh", "marlin_dtypes.cuh", + "marlin_mma.h", "marlin_template.h", + ) +} +SOURCES.update({ + pathlib.Path("awq_repack/cuda/kernel.cuh"): UPSTREAM / "awq_marlin_repack.cu", + pathlib.Path("gptq_repack/cuda/kernel.cuh"): UPSTREAM / "gptq_marlin_repack.cu", + pathlib.Path("gemm/nvidia/generate_kernels.py"): UPSTREAM / "generate_kernels.py", +}) +GENERATED = ( + "kernel_selector.h", + "sm80_kernel_float16_u4_float16.cu", + "sm80_kernel_float16_u4b8_float16.cu", + "sm80_kernel_bfloat16_u4_bfloat16.cu", + "sm80_kernel_bfloat16_u4b8_bfloat16.cu", +) + + +def write_if_changed(path, content): + path.parent.mkdir(parents=True, exist_ok=True) + if not path.exists() or path.read_bytes() != content: + path.write_bytes(content) + + +def prepare(source, output, patch_executable): + if output.is_relative_to(source): + raise ValueError("Marlin output must be outside the upstream source tree") + output.parent.mkdir(parents=True, exist_ok=True) + with tempfile.TemporaryDirectory(prefix="marlin-", dir=output.parent) as temporary: + staging = pathlib.Path(temporary) + for destination, origin in SOURCES.items(): + upstream = source / origin + if not upstream.is_file(): + raise ValueError( + f"Missing pinned vLLM Marlin source: {upstream}; " + "initialize third_party/vllm (or include it in the source archive)" + ) + target = staging / destination + target.parent.mkdir(parents=True, exist_ok=True) + target.write_text(upstream.read_text(encoding="utf-8"), encoding="utf-8", newline="\n") + applied = subprocess.run( + [str(patch_executable), "--batch", "--forward", "--fuzz=0", "-p1", + "--input", str(REPO / "patches/marlin/w4a16.patch")], + cwd=staging, text=True, capture_output=True, check=False, timeout=60, + ) + if applied.returncode: + raise ValueError("Pinned Marlin adaptation failed:\n" + applied.stdout + applied.stderr) + # The upstream generator's single result/config list emits both CUDA + # instantiations and the selector. The patch retains only its W4A16 set. + subprocess.run( + [sys.executable, str(staging / "gemm/nvidia/generate_kernels.py"), "8.0"], + cwd=staging, check=True, + ) + prepared = [path for path in SOURCES if path.suffix != ".py"] + prepared.extend(pathlib.Path("gemm/nvidia") / name for name in GENERATED) + for destination in prepared: + write_if_changed(output / "vendor" / destination, (staging / destination).read_bytes()) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--source-dir", type=pathlib.Path, required=True) + parser.add_argument("--output", type=pathlib.Path, required=True) + parser.add_argument("--patch-executable", default="patch") + args = parser.parse_args() + try: + prepare(args.source_dir.resolve(), args.output.resolve(), args.patch_executable) + except (ValueError, OSError, subprocess.SubprocessError) as error: + parser.exit(1, f"{error}\n") + + +if __name__ == "__main__": + main() diff --git a/scripts/resolve_linked_ops.py b/scripts/resolve_linked_ops.py index 1541b1011..3de837db9 100644 --- a/scripts/resolve_linked_ops.py +++ b/scripts/resolve_linked_ops.py @@ -387,6 +387,39 @@ def _load_distribution(config): return distribution +def _locate_distribution_notices(config): + try: + distribution = importlib.metadata.distribution( + config.python_distribution_package + ) + except importlib.metadata.PackageNotFoundError: + # A custom library locator can supply a DSO without package metadata. + # Retain the static conservative redistribution notices in that case. + return [] + distribution_root = pathlib.Path(distribution.locate_file("")).resolve() + notices = set() + # Use only this resolved package's inventory. Sparse RECORDs and editable + # packages may omit notices; absence must not gate operator availability. + for entry in distribution.files or (): + relative = pathlib.PurePosixPath(str(entry).replace("\\", "/")) + if relative.is_absolute() or ".." in relative.parts: + continue + name = relative.name.upper() + in_license_directory = any( + part.lower() in {"license", "licenses"} for part in relative.parts[:-1] + ) + is_notice = any( + name == prefix or name.startswith((prefix + ".", prefix + "-")) + for prefix in ("LICENSE", "NOTICE", "COPYING") + ) + if not (in_license_directory or is_notice): + continue + candidate = pathlib.Path(distribution.locate_file(entry)).resolve() + if candidate.is_file() and candidate.is_relative_to(distribution_root): + notices.add(candidate) + return sorted(notices) + + def _locate_distribution_library(config): distribution = _load_distribution(config) @@ -637,6 +670,11 @@ def _render_cmake_manifest(payload): "INFINI_OPS_LINKED_LIBRARIES": [ library["path"] for library in payload["libraries"] ], + "INFINI_OPS_LINKED_NOTICE_FILES": [ + path + for library in payload["libraries"] + for path in library.get("notice_files", ()) + ], "INFINI_OPS_LINKED_FORCE_LOAD_LIBRARIES": [ library["path"] for library in payload["libraries"] if library["force_load"] ], @@ -819,6 +857,9 @@ def resolve_linked_ops( "name": config.name, "path": str(library_path), "python_distribution_package": (config.python_distribution_package), + "notice_files": [ + str(path) for path in _locate_distribution_notices(config) + ], "runtime_dir": str(library_path.parent), "transport": config.transport, } diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 9a30c7337..258c8e375 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -1692,8 +1692,33 @@ install(TARGETS infiniops install(FILES "${PROJECT_SOURCE_DIR}/LICENSE" "${PROJECT_SOURCE_DIR}/THIRD_PARTY.txt" DESTINATION "${CMAKE_INSTALL_DATADIR}/licenses/InfiniOps") -install(DIRECTORY "${PROJECT_SOURCE_DIR}/LICENSES/" +install(FILES "${PROJECT_SOURCE_DIR}/LICENSES/InfiniCore-MIT.txt" DESTINATION "${CMAKE_INSTALL_DATADIR}/licenses/InfiniOps") +if(BUILD_NVIDIA_W4_MARLIN) + install(FILES + "${PROJECT_SOURCE_DIR}/LICENSES/FasterTransformer-BSD-3-Clause.txt" + DESTINATION "${CMAKE_INSTALL_DATADIR}/licenses/InfiniOps") +endif() +if(WITH_NVIDIA) + install(FILES "${PROJECT_SOURCE_DIR}/LICENSES/CUTLASS-BSD-3-Clause.txt" + DESTINATION "${CMAKE_INSTALL_DATADIR}/licenses/InfiniOps") +endif() +if(WITH_LINKED) + # Keep the fixed recipe's secondary notices as a conservative fallback; + # additionally preserve the inventory of the distributions actually used. + install(FILES "${PROJECT_SOURCE_DIR}/LICENSES/Apache-2.0-components.txt" + DESTINATION "${CMAKE_INSTALL_DATADIR}/licenses/InfiniOps") + install(DIRECTORY + "${PROJECT_SOURCE_DIR}/LICENSES/FlashInfer-0.6.7" + "${PROJECT_SOURCE_DIR}/LICENSES/TVM-FFI-0.1.10" + DESTINATION "${CMAKE_INSTALL_DATADIR}/licenses/InfiniOps") + foreach(_notice IN LISTS INFINI_OPS_LINKED_NOTICE_FILES) + get_filename_component(_notice_dir "${_notice}" DIRECTORY) + string(SHA256 _notice_group "${_notice_dir}") + install(FILES "${_notice}" + DESTINATION "${CMAKE_INSTALL_DATADIR}/licenses/InfiniOps/resolved/${_notice_group}") + endforeach() +endif() set(INFINI_OPS_INSTALL_CMAKEDIR "${CMAKE_INSTALL_LIBDIR}/cmake/InfiniOps") diff --git a/src/base/argmax.h b/src/base/argmax.h index 904d66c3f..816d232a6 100644 --- a/src/base/argmax.h +++ b/src/base/argmax.h @@ -9,6 +9,8 @@ namespace infini::ops { class Argmax : public Operator { public: + static constexpr bool kAutomaticDefault = true; + Argmax(const Tensor input, const std::optional dim, const bool keepdim, Tensor out) : input_shape_{input.shape()}, diff --git a/src/base/causal_conv1d.h b/src/base/causal_conv1d.h index d8202d93c..cd339364f 100644 --- a/src/base/causal_conv1d.h +++ b/src/base/causal_conv1d.h @@ -75,9 +75,9 @@ class CausalConv1d : public Operator { out.shape() != q.shape()) throw std::invalid_argument("causal conv rank/shape mismatch"); const size_t channels = q.shape()[2]; - if (w.shape() != Tensor::Shape{channels, 1, 4} || - state.shape()[1] != channels || state.shape()[2] != 3) - throw std::invalid_argument("causal conv requires K4 and state_len3"); + if (w.shape()[0] != channels || w.shape()[1] != 1 || w.shape()[2] == 0 || + state.shape()[1] != channels || state.shape()[2] != w.shape()[2] - 1) + throw std::invalid_argument("causal conv weight/state shape mismatch"); if (v[6].has_value() == v[8].has_value()) throw std::invalid_argument( "causal conv requires exactly one final state destination mode"); @@ -90,7 +90,8 @@ class CausalConv1d : public Operator { } if (!v[5] && !v[6] && state.shape()[0] != requests) throw std::invalid_argument("causal conv state batch mismatch"); - if (v[8] && v[8]->shape() != Tensor::Shape{requests, channels, 3}) + if (v[8] && + v[8]->shape() != Tensor::Shape{requests, channels, state.shape()[2]}) throw std::invalid_argument("causal conv final state shape mismatch"); if (v[3] && v[3]->shape() != Tensor::Shape{channels}) throw std::invalid_argument("causal conv bias shape mismatch"); @@ -117,7 +118,7 @@ class CausalConv1d : public Operator { "causal conv broadcast destination unsupported"); } if (q.shape()[1] > - static_cast(std::numeric_limits::max() - 3) || + static_cast(std::numeric_limits::max()) || state.shape()[0] > static_cast(std::numeric_limits::max())) throw std::overflow_error("causal conv index extent exceeds int64"); diff --git a/src/base/chunk_gated_delta_rule.h b/src/base/chunk_gated_delta_rule.h index 308052a42..220fda061 100644 --- a/src/base/chunk_gated_delta_rule.h +++ b/src/base/chunk_gated_delta_rule.h @@ -88,7 +88,7 @@ class ChunkGatedDeltaRule : public Operator { beta.shape() != g.shape()) throw std::invalid_argument("chunk delta rank/shape mismatch"); const auto physical_B = q.shape()[0], T = q.shape()[1], Hk = q.shape()[2], - Hv = v.shape()[2], D = q.shape()[3]; + Hv = v.shape()[2], D = q.shape()[3], Dv = v.shape()[3]; size_t B = physical_B; if (a[6]) { if (physical_B != 1 || a[6]->ndim() != 1 || a[6]->numel() < 2) @@ -97,16 +97,15 @@ class ChunkGatedDeltaRule : public Operator { B = a[6]->numel() - 1; } if (v.shape()[0] != physical_B || v.shape()[1] != T || !Hk || !Hv || - Hv % Hk || v.shape()[3] != D || (D != 64 && D != 128) || - g.shape() != Tensor::Shape{physical_B, T, Hv}) + Hv % Hk || !D || !Dv || g.shape() != Tensor::Shape{physical_B, T, Hv}) throw std::invalid_argument( - "chunk delta requires Dk=Dv64/128 and Hv multiple of Hk"); + "chunk delta data/gate shape or head grouping mismatch"); if (a[8].has_value() == a[10].has_value()) throw std::invalid_argument( "chunk delta requires exactly one final destination mode"); - if (state.shape()[1] != Hv || state.shape()[2] != D || + if (state.shape()[1] != Hv || state.shape()[2] != Dv || state.shape()[3] != D || (!a[7] && !a[8] && state.shape()[0] != B) || - (a[10] && a[10]->shape() != Tensor::Shape{B, Hv, D, D})) + (a[10] && a[10]->shape() != Tensor::Shape{B, Hv, Dv, D})) throw std::invalid_argument( "chunk delta state layout must be pool/B,Hv,Dv,Dk"); for (size_t i = 0; i < a.size(); ++i) { diff --git a/src/base/gemm.h b/src/base/gemm.h index ecbbc8d76..0bfc7c730 100644 --- a/src/base/gemm.h +++ b/src/base/gemm.h @@ -39,36 +39,20 @@ class GemmConfig : public Config { return precision; } - static bool SupportsFp32Precision(Device::Type device, - std::size_t implementation_index, - Fp32Precision precision) { - if (!IsValid(precision)) return false; - return precision == Fp32Precision::DefaultFast || - (device == Device::Type::kNvidia && - (implementation_index == 0 || implementation_index == 1)); - } - - static void RequireSupported(const Config& config, Device::Type device) { - if (!SupportsFp32Precision(device, config.implementation_index(), - PrecisionFrom(config))) { - throw std::invalid_argument( - "FullFp32 GEMM requires NVIDIA implementation 0 or 1"); - } + bool requires_config_aware_provider() const override { + return fp32_precision_ == Fp32Precision::FullFp32; } static bool SupportsMetadata(const Config& config, const Tensor& a, const Tensor& b, const Tensor& y) { const auto precision = PrecisionFrom(config); if (precision == Fp32Precision::DefaultFast) return true; - return SupportsFp32Precision(a.device().type(), - config.implementation_index(), precision) && - a.dtype() == DataType::kFloat32 && b.dtype() == DataType::kFloat32 && + return a.dtype() == DataType::kFloat32 && b.dtype() == DataType::kFloat32 && y.dtype() == DataType::kFloat32; } static void RequireSupported(const Config& config, const Tensor& a, const Tensor& b, const Tensor& y) { - RequireSupported(config, a.device().type()); if (!SupportsMetadata(config, a, b, y)) { throw std::invalid_argument( "FullFp32 GEMM requires float32 A, B, and output"); diff --git a/src/base/mamba_selective_scan.h b/src/base/mamba_selective_scan.h index 0781a3176..3f5ae4ec6 100644 --- a/src/base/mamba_selective_scan.h +++ b/src/base/mamba_selective_scan.h @@ -40,11 +40,9 @@ class MambaSelectiveScan : public Operator { "MambaSelectiveScan: expected F32/F16/BF16 inputs"); const auto parameter_dtype = ts[5].dtype(); if (ts[6].dtype() != parameter_dtype || - (parameter_dtype != dtype && parameter_dtype != DataType::kFloat32) || - (parameter_dtype != dtype && - ts[1].device().type() != Device::Type::kNvidia)) + (parameter_dtype != dtype && parameter_dtype != DataType::kFloat32)) throw std::invalid_argument( - "MambaSelectiveScan: A_log/D must share input dtype or NVIDIA F32"); + "MambaSelectiveScan: A_log/D must share input dtype or F32"); constexpr int ranks[10] = {3, 3, 3, 3, 3, 2, 1, 3, 1, 3}; for (int i = 0; i < 10; ++i) { if (ts[i].ndim() != static_cast(ranks[i]) || diff --git a/src/base/recurrent_gated_delta_rule.h b/src/base/recurrent_gated_delta_rule.h index 245cf3014..61c991420 100644 --- a/src/base/recurrent_gated_delta_rule.h +++ b/src/base/recurrent_gated_delta_rule.h @@ -7,7 +7,6 @@ #include "operator.h" namespace infini::ops { -// Candidate src/base/recurrent_gated_delta_rule.h. Decode-only T=1. class RecurrentGatedDeltaRule : public Operator { public: RecurrentGatedDeltaRule(const Tensor q, const Tensor k, const Tensor v, @@ -84,18 +83,17 @@ class RecurrentGatedDeltaRule : public Operator { beta.shape() != g.shape()) throw std::invalid_argument("recurrent delta rank/shape mismatch"); const auto B = q.shape()[0], Hk = q.shape()[2], Hv = v.shape()[2], - D = q.shape()[3]; - if (q.shape()[1] != 1 || v.shape()[0] != B || v.shape()[1] != 1 || !Hk || - !Hv || Hv % Hk || v.shape()[3] != D || (D != 64 && D != 128) || - g.shape() != Tensor::Shape{B, 1, Hv}) + D = q.shape()[3], Dv = v.shape()[3], T = q.shape()[1]; + if (v.shape()[0] != B || v.shape()[1] != T || !Hk || !Hv || Hv % Hk || !D || + !Dv || g.shape() != Tensor::Shape{B, T, Hv}) throw std::invalid_argument( - "recurrent delta requires T1, Dk=Dv64/128 and Hv multiple of Hk"); + "recurrent delta data/gate shape or head grouping mismatch"); if (a[7].has_value() == a[9].has_value()) throw std::invalid_argument( "recurrent delta requires exactly one final destination mode"); - if (state.shape()[1] != Hv || state.shape()[2] != D || + if (state.shape()[1] != Hv || state.shape()[2] != Dv || state.shape()[3] != D || (!a[6] && !a[7] && state.shape()[0] != B) || - (a[9] && a[9]->shape() != Tensor::Shape{B, Hv, D, D})) + (a[9] && a[9]->shape() != Tensor::Shape{B, Hv, Dv, D})) throw std::invalid_argument( "recurrent delta state layout must be pool/B,Hv,Dv,Dk"); for (size_t i = 0; i < a.size(); ++i) { diff --git a/src/base/sum.h b/src/base/sum.h index cfad95839..46eafa2d2 100644 --- a/src/base/sum.h +++ b/src/base/sum.h @@ -10,6 +10,8 @@ namespace infini::ops { class Sum : public Operator { public: + static constexpr bool kAutomaticDefault = true; + Sum(const Tensor input, const std::optional> dim, const bool keepdim, const std::optional dtype, Tensor out) : input_shape_{input.shape()}, diff --git a/src/base/take.h b/src/base/take.h index 24339abc5..99e76bdc8 100644 --- a/src/base/take.h +++ b/src/base/take.h @@ -7,6 +7,8 @@ namespace infini::ops { class Take : public Operator { public: + static constexpr bool kAutomaticDefault = true; + Take(const Tensor input, const Tensor index, Tensor out) : input_shape_{input.shape()}, input_strides_{input.strides()}, diff --git a/src/base/upsample_bilinear2d.h b/src/base/upsample_bilinear2d.h index e279a6fda..b81616a3a 100644 --- a/src/base/upsample_bilinear2d.h +++ b/src/base/upsample_bilinear2d.h @@ -10,6 +10,8 @@ namespace infini::ops { class UpsampleBilinear2d : public Operator { public: + static constexpr bool kAutomaticDefault = true; + UpsampleBilinear2d(const Tensor input, const std::vector output_size, const bool align_corners, const std::optional scales_h, diff --git a/src/native/cuda/device_buffer.h b/src/native/cuda/device_buffer.h new file mode 100644 index 000000000..58134b03c --- /dev/null +++ b/src/native/cuda/device_buffer.h @@ -0,0 +1,56 @@ +#ifndef INFINI_OPS_CUDA_DEVICE_BUFFER_H_ +#define INFINI_OPS_CUDA_DEVICE_BUFFER_H_ + +#include +#include +#include +#include + +namespace infini::ops { + +// Private ownership for descriptor metadata and scratch allocated at Make. +// Backend keeps this usable by every CUDA-like provider without a CUDA API tie. +template +class CudaDeviceBuffer { + public: + CudaDeviceBuffer() = default; + CudaDeviceBuffer(const CudaDeviceBuffer&) = delete; + CudaDeviceBuffer& operator=(const CudaDeviceBuffer&) = delete; + + std::byte* Allocate(std::size_t bytes) { + if (!bytes) return nullptr; + if (data_) throw std::logic_error("device buffer already allocated"); + if (Backend::GetDevice(&owner_device_) != Backend::kSuccess) + throw std::runtime_error("cannot query allocation owner device"); + void* allocation = nullptr; + const auto status = Backend::Malloc(&allocation, bytes); + if (status != Backend::kSuccess || !allocation) { + if (allocation && Backend::Free(allocation) != Backend::kSuccess) + std::terminate(); + throw std::bad_alloc(); + } + data_ = allocation; + return static_cast(data_); + } + + ~CudaDeviceBuffer() { + if (!data_) return; + int previous = -1; + if (Backend::GetDevice(&previous) != Backend::kSuccess) std::terminate(); + const bool switched = previous != owner_device_; + if (switched && Backend::SetDevice(owner_device_) != Backend::kSuccess) + std::terminate(); + const auto freed = Backend::Free(data_); + if (switched && Backend::SetDevice(previous) != Backend::kSuccess) + std::terminate(); + if (freed != Backend::kSuccess) std::terminate(); + } + + private: + void* data_{nullptr}; + int owner_device_{-1}; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/nvidia/ops/argmax/kernel.h b/src/native/cuda/nvidia/ops/argmax/kernel.h index c8a8b4e48..32f461d67 100644 --- a/src/native/cuda/nvidia/ops/argmax/kernel.h +++ b/src/native/cuda/nvidia/ops/argmax/kernel.h @@ -19,6 +19,11 @@ namespace infini::ops { template <> class Operator : public Argmax { public: + static bool supports(const Tensor input, const std::optional dim, + bool keepdim, const Tensor out) { + return Supports(input, dim, keepdim, out, false); + } + Operator(const Tensor input, const std::optional dim, const bool keepdim, Tensor out) : Argmax(input, dim, keepdim, out) { @@ -27,10 +32,6 @@ class Operator : public Argmax { const auto width = input.shape().back(); if (width > nvidia_argmax_detail::kChunkWidth) { const auto chunks = (width - 1) / nvidia_argmax_detail::kChunkWidth + 1; - if (rows > std::numeric_limits::max() / chunks / - sizeof(nvidia_argmax_detail::Candidate)) { - throw std::invalid_argument("NVIDIA Argmax workspace size overflows"); - } workspace_bytes_ = rows * chunks * sizeof(nvidia_argmax_detail::Candidate); } @@ -102,8 +103,8 @@ class Operator : public Argmax { } } - static void Validate(const Tensor input, const std::optional dim, - bool keepdim, const Tensor out) { + static bool Supports(const Tensor input, const std::optional dim, + bool keepdim, const Tensor out, bool invocation) { if (input.device().type() != Device::Type::kNvidia || input.device() != out.device() || (input.ndim() != 1 && input.ndim() != 2) || keepdim || @@ -113,23 +114,33 @@ class Operator : public Argmax { input.dtype() != DataType::kBFloat16 && input.dtype() != DataType::kFloat32) || out.dtype() != DataType::kInt64) { - throw std::invalid_argument("unsupported NVIDIA Argmax metadata"); + return false; } const auto rows = input.ndim() == 2 ? input.shape()[0] : 1; const auto width = input.shape().back(); if (!rows || !width || rows > static_cast(INT_MAX) || width > static_cast(INT_MAX) || (input.ndim() == 1 ? out.ndim() != 0 - : out.ndim() != 1 || out.shape()[0] != rows) || - !input.data() || !out.data()) { - throw std::invalid_argument("NVIDIA Argmax has empty or invalid storage"); + : out.ndim() != 1 || out.shape()[0] != rows)) { + return false; } if (rows > std::numeric_limits::max() / width || rows * width > std::numeric_limits::max() / input.element_size() || rows > std::numeric_limits::max() / sizeof(std::int64_t)) { - throw std::invalid_argument("NVIDIA Argmax storage size overflows"); + return false; } + if (width > nvidia_argmax_detail::kChunkWidth) { + const auto chunks = (width - 1) / nvidia_argmax_detail::kChunkWidth + 1; + if (rows > std::numeric_limits::max() / chunks / + sizeof(nvidia_argmax_detail::Candidate)) + return false; + } + if (!invocation) return true; + if (!input.data() || !out.data() || + reinterpret_cast(input.data()) % input.element_size() || + reinterpret_cast(out.data()) % alignof(std::int64_t)) + return false; const auto input_bytes = input.numel() * input.element_size(); const auto output_bytes = out.numel() * out.element_size(); const auto in_begin = reinterpret_cast(input.data()); @@ -138,17 +149,18 @@ class Operator : public Argmax { output_bytes > std::numeric_limits::max() - out_begin || (in_begin < out_begin + output_bytes && out_begin < in_begin + input_bytes)) { - throw std::invalid_argument( - "NVIDIA Argmax input/output storage overlaps"); + return false; } int current_device = -1; const auto status = cudaGetDevice(¤t_device); - if (status != cudaSuccess) - throw std::runtime_error(cudaGetErrorString(status)); - if (current_device != input.device().index()) { + return status == cudaSuccess && current_device == input.device().index(); + } + + static void Validate(const Tensor input, const std::optional dim, + bool keepdim, const Tensor out) { + if (!Supports(input, dim, keepdim, out, true)) throw std::invalid_argument( - "NVIDIA Argmax tensor device is not current device"); - } + "unsupported NVIDIA Argmax metadata, storage or current device"); } template diff --git a/src/native/cuda/nvidia/ops/causal_conv1d/kernel.h b/src/native/cuda/nvidia/ops/causal_conv1d/kernel.h index ba4ac4c55..9353b11c6 100644 --- a/src/native/cuda/nvidia/ops/causal_conv1d/kernel.h +++ b/src/native/cuda/nvidia/ops/causal_conv1d/kernel.h @@ -109,6 +109,12 @@ class Operator : public CausalConv1d { final_conv_state) { if (device_.type() != Device::Type::kNvidia) throw std::invalid_argument("causal conv requires NVIDIA"); + if (weight.shape()[2] != 4 || conv_state.shape()[2] != 3) + throw std::invalid_argument( + "NVIDIA causal conv kernel requires K4 and state_len3"); + if (tokens_ > static_cast(std::numeric_limits::max() - 3)) + throw std::overflow_error( + "NVIDIA causal conv history index exceeds int64"); check_device(); // Baseline uses int request/channel variables and CUDA's grid.y dimension. if (requests_ >= static_cast(INT_MAX) || channels_ > 65535) @@ -142,6 +148,16 @@ class Operator : public CausalConv1d { final_state_indices, out, final_conv_state}; + for (size_t i = 0; i < views.size(); ++i) { + const auto& t = views[i]; + if (t.has_value() != metadata_[i].present || + (t && + (t->device() != device_ || t->shape() != metadata_[i].shape || + t->strides() != metadata_[i].strides || + t->dtype() != metadata_[i].dtype || (t->numel() && !t->data())))) + throw std::invalid_argument( + "causal conv metadata changed or nonempty storage is null"); + } check_device(); if (!requests_ || !channels_) return; switch (dtype_) { diff --git a/src/native/cuda/nvidia/ops/chunk_gated_delta_rule/kernel.h b/src/native/cuda/nvidia/ops/chunk_gated_delta_rule/kernel.h index aba4f0247..451443059 100644 --- a/src/native/cuda/nvidia/ops/chunk_gated_delta_rule/kernel.h +++ b/src/native/cuda/nvidia/ops/chunk_gated_delta_rule/kernel.h @@ -48,6 +48,9 @@ class Operator final_state, use_qk_l2norm, chunk_size) { if (device_.type() != Device::Type::kNvidia) throw std::invalid_argument("chunk delta requires NVIDIA"); + if (v.shape()[3] != D_ || (D_ != 64 && D_ != 128)) + throw std::invalid_argument( + "NVIDIA chunk delta kernels require Dk=Dv64/128"); CheckDevice(); cudaDeviceProp property{}; if (cudaGetDeviceProperties(&property, device_.index()) != cudaSuccess || diff --git a/src/native/cuda/nvidia/ops/flash_attn_varlen_func/kernel.h b/src/native/cuda/nvidia/ops/flash_attn_varlen_func/kernel.h index 3b7c4886b..8bb148462 100644 --- a/src/native/cuda/nvidia/ops/flash_attn_varlen_func/kernel.h +++ b/src/native/cuda/nvidia/ops/flash_attn_varlen_func/kernel.h @@ -21,7 +21,7 @@ namespace infini::ops { -// Inference-only qualified head dimensions are checked by supports(). +// Inference-only rounded widths and vector layouts are checked by supports(). // Automatic selection leaves other configurations to an existing provider; // construction failures propagate. template <> diff --git a/src/native/cuda/nvidia/ops/flash_attn_with_kvcache/kernel.cc b/src/native/cuda/nvidia/ops/flash_attn_with_kvcache/kernel.cc index 0cea771f6..8697c3b79 100644 --- a/src/native/cuda/nvidia/ops/flash_attn_with_kvcache/kernel.cc +++ b/src/native/cuda/nvidia/ops/flash_attn_with_kvcache/kernel.cc @@ -51,15 +51,6 @@ struct DeviceScope { ~DeviceScope() { cudaSetDevice(previous); } }; -// Candidate ordinary MHA-only extension. Keep true d/strides, use existing 96 -// specialization's Is_even_K=false masking. No physical padding or new -// algorithm. -int ordinary_head_dim_bin(int dimension) { - return dimension == 72 || dimension == 80 - ? 96 - : native_flash::head_dim_bin(dimension); -} - bool plain_layout(Tensor t, int dimensions) { if (dimensions < 1 || dimensions > 4 || t.ndim() != dimensions || t.stride(-1) != 1) @@ -121,14 +112,9 @@ bool Operator::supported_layout( const Tensor q, Tensor k, Tensor v, std::optional lengths, std::optional blocks, Tensor out) { if (!plain_layout(q, 4) || !plain_layout(k, 4) || !plain_layout(v, 4) || - !plain_layout(out, 4) || !ordinary_head_dim_bin(q.size(3)) || + !plain_layout(out, 4) || !native_flash::head_dim_bin(q.size(3)) || q.device().type() != Device::Type::kNvidia) return false; - // New dimensions are restricted to ordinary dense MHA, not decode/GQA, - // length-masked or paged cache domains that require separate qualification. - if ((q.size(3) == 72 || q.size(3) == 80) && - (lengths || blocks || q.size(0) != k.size(0) || q.size(2) != k.size(2))) - return false; // Upstream vectorized loads/stores require 16-byte-aligned rows and heads. // Metadata-only descriptor views may have null data until invocation. for (auto tensor : {q, k, v, out}) { @@ -142,7 +128,10 @@ bool Operator::supported_layout( (!plain_layout(*blocks, 2) || blocks->size(1) > (std::numeric_limits::max() - 127) / k.size(1)))) return false; - return true; + int major = 0; + return cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, + q.device().index()) == cudaSuccess && + major >= 8; } void Operator::prepare( @@ -169,7 +158,7 @@ void Operator::prepare( p.h = q.size(2); p.h_k = cache.size(2); p.d = q.size(3); - p.d_rounded = ordinary_head_dim_bin(p.d); + p.d_rounded = native_flash::head_dim_bin(p.d); p.seqlen_q = q.size(1); p.seqlen_k = logical_bound > 0 ? logical_bound diff --git a/src/native/cuda/nvidia/ops/flash_attn_with_kvcache/kernel.h b/src/native/cuda/nvidia/ops/flash_attn_with_kvcache/kernel.h index b6c42dc79..8d6ac332d 100644 --- a/src/native/cuda/nvidia/ops/flash_attn_with_kvcache/kernel.h +++ b/src/native/cuda/nvidia/ops/flash_attn_with_kvcache/kernel.h @@ -29,10 +29,8 @@ namespace infini::ops { -// Native inference provider: qualified heads, read-only cache, global or finite -// left-window/right-zero attention, no ALiBi/softcap. -// Other existing configurations remain with their previous providers. The -// explicit build option is not enabled in the accepted LM candidate yet. +// Native inference provider: rounded kernel widths, read-only cache, global or +// finite left-window/right-zero attention, no ALiBi/softcap. template <> class Operator : public FlashAttnWithKvcache { @@ -51,9 +49,7 @@ class Operator window_size[1] == 0)) && softcap == 0.0 && num_splits <= 128 && !return_softmax_lse && supported_layout(q, k_cache, v_cache, cache_seqlens, block_table, - out) && - ((q.size(3) != 72 && q.size(3) != 80) || - (!causal && window_size[0] == -1)); + out); } } @@ -151,11 +147,7 @@ class Operator } // Graph callers must supply stable external workspace for the graph lifetime. - // The qualified native domain is enforced by supports(), not the GPU type. - bool capture_safe() const override { - // New ordinary head dimensions are SOURCE_ONLY, never capture-qualified. - return q_shape_[3] != 72 && q_shape_[3] != 80; - } + bool capture_safe() const override { return true; } std::size_t workspace_size_in_bytes() const override; private: diff --git a/src/native/cuda/nvidia/ops/flash_head_dim.h b/src/native/cuda/nvidia/ops/flash_head_dim.h index 3554a191b..99990830a 100644 --- a/src/native/cuda/nvidia/ops/flash_head_dim.h +++ b/src/native/cuda/nvidia/ops/flash_head_dim.h @@ -2,24 +2,12 @@ namespace infini::ops::native_flash { -// This intentionally names only the head dimensions qualified by this native -// provider. The pinned FA 2.8.3 kernels round 160 and 224 to the next 64. +// The pinned FA 2.8.3 launchers specialize these rounded widths. Kernels keep +// the logical d and use Is_even_K=false masks for shorter vector-aligned rows. inline int head_dim_bin(int dimension) { - switch (dimension) { - case 32: - case 64: - case 96: - case 128: - case 192: - case 256: - return dimension; - case 160: - return 192; - case 224: - return 256; - default: - return 0; - } + if (dimension <= 0 || dimension > 256 || dimension % 8) return 0; + if (dimension <= 128) return ((dimension + 31) / 32) * 32; + return ((dimension + 63) / 64) * 64; } inline int split_block_n(int bin) { diff --git a/src/native/cuda/nvidia/ops/gemm/cublas.h b/src/native/cuda/nvidia/ops/gemm/cublas.h index 6c48234d4..feeedb9e5 100644 --- a/src/native/cuda/nvidia/ops/gemm/cublas.h +++ b/src/native/cuda/nvidia/ops/gemm/cublas.h @@ -11,6 +11,23 @@ template <> class Operator : public BlasGemm> { public: + Operator(const Config&, const Tensor a, const Tensor b, + const std::optional c, std::optional alpha, + std::optional beta, std::optional trans_a, + std::optional trans_b, Tensor y) + : Operator(a, b, c, alpha, beta, trans_a, trans_b, y) {} + + Operator(const Config&, const Tensor a, const Tensor b, Tensor y) + : Operator(a, b, y) {} + + template + static bool supports(const Config& config, const Tensor a, const Tensor b, + const Tail&... tail) { + static_assert(sizeof...(Tail) == 1 || sizeof...(Tail) == 6); + const auto& y = std::get(std::tie(tail...)); + return GemmConfig::SupportsMetadata(config, a, b, y); + } + Operator(const Tensor a, const Tensor b, const std::optional c, std::optional alpha, std::optional beta, std::optional trans_a, std::optional trans_b, Tensor y) diff --git a/src/native/cuda/nvidia/ops/gemm/cublaslt.h b/src/native/cuda/nvidia/ops/gemm/cublaslt.h index 677abdf96..4ae15fd18 100644 --- a/src/native/cuda/nvidia/ops/gemm/cublaslt.h +++ b/src/native/cuda/nvidia/ops/gemm/cublaslt.h @@ -18,6 +18,23 @@ namespace infini::ops { template <> class Operator : public Gemm { public: + Operator(const Config&, const Tensor a, const Tensor b, + const std::optional c, std::optional alpha, + std::optional beta, std::optional trans_a, + std::optional trans_b, Tensor y) + : Operator(a, b, c, alpha, beta, trans_a, trans_b, y) {} + + Operator(const Config&, const Tensor a, const Tensor b, Tensor y) + : Operator(a, b, y) {} + + template + static bool supports(const Config& config, const Tensor a, const Tensor b, + const Tail&... tail) { + static_assert(sizeof...(Tail) == 1 || sizeof...(Tail) == 6); + const auto& y = std::get(std::tie(tail...)); + return GemmConfig::SupportsMetadata(config, a, b, y); + } + Operator(const Tensor a, const Tensor b, const std::optional c, std::optional alpha, std::optional beta, std::optional trans_a, std::optional trans_b, Tensor y) diff --git a/src/native/cuda/nvidia/ops/layer_norm/kernel.h b/src/native/cuda/nvidia/ops/layer_norm/kernel.h index f3f83ea12..51ac98b64 100644 --- a/src/native/cuda/nvidia/ops/layer_norm/kernel.h +++ b/src/native/cuda/nvidia/ops/layer_norm/kernel.h @@ -102,7 +102,25 @@ class Operator : public LayerNorm { Operator(const Tensor input, const Tensor weight, const std::optional bias, float eps, Tensor out) : LayerNorm(checked(input, weight, bias, eps, out), weight, bias, eps, - out) {} + out) { + using namespace nvidia_binary_view; + SpanBytes(input, &input_bytes_); + SpanBytes(weight, &weight_bytes_); + SpanBytes(out, &out_bytes_); + if (bias) SpanBytes(*bias, &bias_bytes_); + // An identical input layout inherits the validated output non-overlap. + exact_inplace_layout_ = input_strides_ == out_strides_; + layout_.rank = static_cast(input.ndim()); + layout_.dim = input.size(-1); + layout_.weight_stride = static_cast(weight.stride(0)); + layout_.bias_stride = bias ? static_cast(bias->stride(0)) : 0; + for (std::size_t i = 0; i < input.ndim(); ++i) { + layout_.shape[i] = input.shape()[i]; + layout_.input_strides[i] = static_cast(input.strides()[i]); + layout_.out_strides[i] = static_cast(out.strides()[i]); + if (i + 1 < input.ndim()) rows_ *= input.shape()[i]; + } + } // Replay submits a kernel with by-value layout and no device allocation. bool capture_safe() const override { return true; } @@ -115,34 +133,47 @@ class Operator : public LayerNorm { weight.strides() != weight_strides_ || out.shape() != out_shape_ || out.strides() != out_strides_ || input.dtype() != dtype_ || weight.dtype() != dtype_ || out.dtype() != dtype_ || - input.device() != device_ || bias.has_value() != has_bias_ || + input.device() != device_ || weight.device() != device_ || + out.device() != device_ || bias.has_value() != has_bias_ || (bias && (bias->shape() != bias_shape_ || bias->strides() != bias_strides_ || - bias->dtype() != dtype_)) || - bits(eps) != bits(eps_) || !supports(input, weight, bias, eps, out)) + bias->dtype() != dtype_ || bias->device() != device_)) || + bits(eps) != bits(eps_)) throw std::invalid_argument( "NVIDIA LayerNorm call metadata or alias changed after Make"); - std::size_t rows = 1; - nvidia_layer_norm_detail::Layout layout{}; - layout.rank = static_cast(input.ndim()); - layout.dim = input.size(-1); - layout.weight_stride = static_cast(weight.stride(0)); - layout.bias_stride = bias ? static_cast(bias->stride(0)) : 0; - for (std::size_t i = 0; i < input.ndim(); ++i) { - layout.shape[i] = input.shape()[i]; - layout.input_strides[i] = static_cast(input.strides()[i]); - layout.out_strides[i] = static_cast(out.strides()[i]); - if (i + 1 < input.ndim()) rows *= input.shape()[i]; - } - if (!rows) return; + using namespace nvidia_binary_view; + const auto bound = [](const Tensor tensor, std::size_t bytes) { + return !bytes || + (tensor.data() && reinterpret_cast(tensor.data()) % + tensor.element_size() == + 0); + }; + int current = -1; + if (cudaGetDevice(¤t) != cudaSuccess || current != device_.index() || + !bound(input, input_bytes_) || !bound(weight, weight_bytes_) || + !bound(out, out_bytes_) || (bias && !bound(*bias, bias_bytes_)) || + Overlaps(out, out_bytes_, weight, weight_bytes_) || + (bias && Overlaps(out, out_bytes_, *bias, bias_bytes_)) || + (Overlaps(out, out_bytes_, input, input_bytes_) && + !(input.data() == out.data() && exact_inplace_layout_))) + throw std::invalid_argument( + "NVIDIA LayerNorm current device, pointer or alias changed after " + "Make"); + if (!rows_) return; nvidia_layer_norm_detail::Launch( input.data(), weight.data(), bias ? bias->data() : nullptr, out.data(), - dtype_, layout, rows, eps, + dtype_, layout_, rows_, eps_, static_cast(stream_ ? stream_ : 0)); } private: + nvidia_layer_norm_detail::Layout layout_{}; + std::size_t rows_ = 1; + std::size_t input_bytes_ = 0, weight_bytes_ = 0, out_bytes_ = 0, + bias_bytes_ = 0; + bool exact_inplace_layout_ = false; + static std::uint32_t bits(float value) { std::uint32_t result; std::memcpy(&result, &value, sizeof(result)); diff --git a/src/native/cuda/nvidia/ops/pad_constant/kernel.h b/src/native/cuda/nvidia/ops/pad_constant/kernel.h index bc13ae64b..b90142374 100644 --- a/src/native/cuda/nvidia/ops/pad_constant/kernel.h +++ b/src/native/cuda/nvidia/ops/pad_constant/kernel.h @@ -51,22 +51,6 @@ class Operator : public PadConstant { if (cudaGetDevice(¤t_device) != cudaSuccess || current_device != input.device().index()) return false; - if (input.dtype() == DataType::kBFloat16) { - // Device architecture is immutable, while supports() is on the Call hot - // path. Keep the metadata-only qualification local to each host thread. - static thread_local int cached_device = -1; - static thread_local bool cached_bf16 = false; - if (cached_device != current_device) { - int major = 0; - if (cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, - current_device) != cudaSuccess) - return false; - cached_bf16 = major >= 8; - cached_device = current_device; - } - if (!cached_bf16) return false; - } - const auto max_i64 = static_cast(std::numeric_limits::max()); std::size_t input_count = 1, output_count = 1; diff --git a/src/native/cuda/nvidia/ops/recurrent_gated_delta_rule/kernel.h b/src/native/cuda/nvidia/ops/recurrent_gated_delta_rule/kernel.h index 165795ccd..c687268a1 100644 --- a/src/native/cuda/nvidia/ops/recurrent_gated_delta_rule/kernel.h +++ b/src/native/cuda/nvidia/ops/recurrent_gated_delta_rule/kernel.h @@ -46,6 +46,9 @@ class Operator final_state, use_qk_l2norm) { if (device_.type() != Device::Type::kNvidia) throw std::invalid_argument("recurrent delta requires NVIDIA"); + if (q.shape()[1] != 1 || v.shape()[3] != D_ || (D_ != 64 && D_ != 128)) + throw std::invalid_argument( + "NVIDIA recurrent delta kernels require T1 and Dk=Dv64/128"); CheckDevice(); cudaDeviceProp property{}; if (cudaGetDeviceProperties(&property, device_.index()) != cudaSuccess || diff --git a/src/native/cuda/nvidia/ops/sum/kernel.h b/src/native/cuda/nvidia/ops/sum/kernel.h index 3b3337992..c5a3d1eb4 100644 --- a/src/native/cuda/nvidia/ops/sum/kernel.h +++ b/src/native/cuda/nvidia/ops/sum/kernel.h @@ -33,6 +33,14 @@ class Operator : public Sum { const std::optional> dim, const bool keepdim, const std::optional dtype, Tensor out) { + return Supports(input, dim, keepdim, dtype, out, false); + } + + private: + static bool Supports(const Tensor input, + const std::optional> dim, + const bool keepdim, const std::optional dtype, + Tensor out, bool invocation) { if (input.device().type() != Device::Type::kNvidia || input.device() != out.device() || input.dtype() != out.dtype() || !dim || dim->size() != 1 || keepdim || dtype || @@ -43,13 +51,10 @@ class Operator : public Sum { input.dtype() != DataType::kBFloat16 && input.dtype() != DataType::kFloat32) return false; - int current_device = -1; - if (cudaGetDevice(¤t_device) != cudaSuccess || - current_device != input.device().index()) - return false; if (input.dtype() == DataType::kBFloat16) { cudaDeviceProp properties{}; - if (cudaGetDeviceProperties(&properties, current_device) != cudaSuccess || + if (cudaGetDeviceProperties(&properties, input.device().index()) != + cudaSuccess || properties.major < 8) return false; } @@ -93,14 +98,20 @@ class Operator : public Sum { last_element + 1 > std::numeric_limits::max() / width || output_count > std::numeric_limits::max() || output_count > std::numeric_limits::max() / width || - !CheckedContiguous(out) || !input.data() || !out.data() || - reinterpret_cast(input.data()) % width || + !CheckedContiguous(out)) + return false; + if (!invocation) return true; + int current_device = -1; + if (cudaGetDevice(¤t_device) != cudaSuccess || + current_device != input.device().index() || !input.data() || + !out.data() || reinterpret_cast(input.data()) % width || reinterpret_cast(out.data()) % width) return false; return !Overlaps(input.data(), (last_element + 1) * width, out.data(), output_count * width); } + public: Operator(const Tensor input, const std::optional> dim, const bool keepdim, const std::optional dtype, Tensor out) @@ -164,7 +175,7 @@ class Operator : public Sum { const std::optional> dim, const bool keepdim, const std::optional dtype, Tensor out) { - if (!supports(input, dim, keepdim, dtype, out)) { + if (!Supports(input, dim, keepdim, dtype, out, true)) { throw std::invalid_argument( "unsupported NVIDIA Sum metadata, device or alias"); } diff --git a/src/native/cuda/nvidia/ops/take/kernel.h b/src/native/cuda/nvidia/ops/take/kernel.h index b5b1177d7..4a6c8b995 100644 --- a/src/native/cuda/nvidia/ops/take/kernel.h +++ b/src/native/cuda/nvidia/ops/take/kernel.h @@ -22,6 +22,12 @@ template <> class Operator : public Take { public: static bool supports(const Tensor input, const Tensor index, Tensor out) { + return Supports(input, index, out, false); + } + + private: + static bool Supports(const Tensor input, const Tensor index, Tensor out, + bool invocation) { if (input.device().type() != Device::Type::kNvidia || input.device() != index.device() || input.device() != out.device() || !CheckedContiguous(input) || !CheckedContiguous(index) || @@ -29,10 +35,6 @@ class Operator : public Take { input.dtype() != out.dtype() || index.shape() != out.shape()) { return false; } - int current_device = -1; - if (cudaGetDevice(¤t_device) != cudaSuccess || - current_device != input.device().index()) - return false; switch (input.dtype()) { case DataType::kInt8: case DataType::kUInt8: @@ -66,6 +68,11 @@ class Operator : public Take { std::numeric_limits::max() / sizeof(std::int64_t)) { return false; } + if (!invocation) return true; + int current_device = -1; + if (cudaGetDevice(¤t_device) != cudaSuccess || + current_device != input.device().index()) + return false; if ((input_count && !input.data()) || (output_count && (!index.data() || !out.data())) || (input_count && @@ -84,6 +91,7 @@ class Operator : public Take { out.data(), output_count * width); } + public: Operator(const Tensor input, const Tensor index, Tensor out) : Take(input, index, out) { Validate(input, index, out); @@ -150,7 +158,7 @@ class Operator : public Take { } static void Validate(const Tensor input, const Tensor index, Tensor out) { - if (!supports(input, index, out)) + if (!Supports(input, index, out, true)) throw std::invalid_argument( "unsupported NVIDIA Take metadata, device or alias"); } diff --git a/src/native/cuda/nvidia/ops/upsample_bilinear2d/kernel.h b/src/native/cuda/nvidia/ops/upsample_bilinear2d/kernel.h index aeb634985..4b3eadc51 100644 --- a/src/native/cuda/nvidia/ops/upsample_bilinear2d/kernel.h +++ b/src/native/cuda/nvidia/ops/upsample_bilinear2d/kernel.h @@ -39,6 +39,17 @@ class Operator const bool align_corners, const std::optional scales_h, const std::optional scales_w, Tensor out) { + return Supports(input, output_size, align_corners, scales_h, scales_w, out, + false); + } + + private: + static bool Supports(const Tensor input, + const std::vector output_size, + const bool align_corners, + const std::optional scales_h, + const std::optional scales_w, Tensor out, + bool invocation) { (void)align_corners; if (input.device().type() != Device::Type::kNvidia || input.device() != out.device() || input.shape().size() != 4 || @@ -52,13 +63,10 @@ class Operator input.shape()[0] != out.shape()[0] || input.shape()[1] != out.shape()[1]) return false; - int current_device = -1; - if (cudaGetDevice(¤t_device) != cudaSuccess || - current_device != input.device().index()) - return false; if (input.dtype() == DataType::kBFloat16) { cudaDeviceProp properties{}; - if (cudaGetDeviceProperties(&properties, current_device) != cudaSuccess || + if (cudaGetDeviceProperties(&properties, input.device().index()) != + cudaSuccess || properties.major < 8) return false; } @@ -97,14 +105,20 @@ class Operator if (!width || input_span > std::numeric_limits::max() / width || output_count > std::numeric_limits::max() / width || - !CheckedContiguous(out) || !input.data() || !out.data() || - reinterpret_cast(input.data()) % width || + !CheckedContiguous(out)) + return false; + if (!invocation) return true; + int current_device = -1; + if (cudaGetDevice(¤t_device) != cudaSuccess || + current_device != input.device().index() || !input.data() || + !out.data() || reinterpret_cast(input.data()) % width || reinterpret_cast(out.data()) % width) return false; return !Overlaps(input.data(), input_span * width, out.data(), output_count * width); } + public: Operator(const Tensor input, const std::vector output_size, const bool align_corners, const std::optional scales_h, const std::optional scales_w, Tensor out) @@ -184,7 +198,8 @@ class Operator const std::vector& output_size, bool align_corners, std::optional scales_h, std::optional scales_w, Tensor out) { - if (!supports(input, output_size, align_corners, scales_h, scales_w, out)) + if (!Supports(input, output_size, align_corners, scales_h, scales_w, out, + true)) throw std::invalid_argument( "unsupported NVIDIA UpsampleBilinear2d metadata, device or alias"); } diff --git a/src/native/cuda/nvidia/ops/w4_marlin/vendor/awq_repack/cuda/kernel.cuh b/src/native/cuda/nvidia/ops/w4_marlin/vendor/awq_repack/cuda/kernel.cuh deleted file mode 100644 index c05c3a5ff..000000000 --- a/src/native/cuda/nvidia/ops/w4_marlin/vendor/awq_repack/cuda/kernel.cuh +++ /dev/null @@ -1,201 +0,0 @@ -// Adapted from InfiniCore 1ab85ef1d3c2bce49edf1fc73a5b48becaf7552e: -// src/infiniop/ops/awq_marlin_repack/cuda/kernel.cuh (Marlin/vLLM). -// Modified for InfiniOps: relative include paths; kernel mathematics retained. -// Original licenses and copyright notices: repository THIRD_PARTY.txt and LICENSE. -#include "../../gemm/marlin/marlin.cuh" - -namespace marlin { - -template -__device__ void awq_marlin_repack_kernel( - uint32_t const *__restrict__ b_q_weight_ptr, uint32_t *__restrict__ out_ptr, - int size_k, int size_n) { - constexpr int pack_factor = 32 / num_bits; - - constexpr int target_tile_n_size = tile_n_size / (is_a_8bit ? 2 : 1); - constexpr int target_tile_k_size = tile_k_size * (is_a_8bit ? 2 : 1); - int k_tiles = size_k / target_tile_k_size; - int n_tiles = size_n / target_tile_n_size; - int block_k_tiles = div_ceil(k_tiles, gridDim.x); - - auto start_k_tile = blockIdx.x * block_k_tiles; - if (start_k_tile >= k_tiles) { - return; - } - - int finish_k_tile = min(start_k_tile + block_k_tiles, k_tiles); - - // Wait until the next thread tile has been loaded to shared memory. - auto wait_for_stage = [&]() { - // We only have `stages - 2` active fetches since we are double buffering - // and can only issue the next fetch when it is guaranteed that the previous - // shared memory load is fully complete (as it may otherwise be - // overwritten). - cp_async_wait(); - __syncthreads(); - }; - - extern __shared__ int4 sh[]; - - constexpr int tile_n_ints = target_tile_n_size / pack_factor; - - constexpr int stage_n_threads = tile_n_ints / 4; - constexpr int stage_k_threads = target_tile_k_size; - constexpr int stage_size = stage_k_threads * stage_n_threads; - - auto fetch_to_shared = [&](int pipe, int k_tile_id, int n_tile_id) { - if (n_tile_id >= n_tiles) { - cp_async_fence(); - return; - } - - int first_n = n_tile_id * target_tile_n_size; - int first_n_packed = first_n / pack_factor; - - int4 *sh_ptr = sh + stage_size * pipe; - - if (threadIdx.x < stage_size) { - auto k_id = threadIdx.x / stage_n_threads; - auto n_id = threadIdx.x % stage_n_threads; - - int first_k = k_tile_id * target_tile_k_size; - - cp_async4(&sh_ptr[k_id * stage_n_threads + n_id], - reinterpret_cast( - &(b_q_weight_ptr[(first_k + k_id) * (size_n / pack_factor) + first_n_packed + (n_id * 4)]))); - } - - cp_async_fence(); - }; - - auto repack_tile = [&](int pipe, int k_tile_id, int n_tile_id) { - if (n_tile_id >= n_tiles) { - return; - } - - auto warp_id = threadIdx.x / 32; - auto th_id = threadIdx.x % 32; - - if (warp_id >= 4) { - return; - } - - int tc_col = th_id / 4; - int tc_row = (th_id % 4) * (is_a_8bit ? 4 : 2); - - constexpr int tc_offsets[4] = {0, 1, 8, 9}; - - int cur_n = (warp_id / (is_a_8bit ? 2 : 1)) * 16 + tc_col; - int cur_n_packed = cur_n / pack_factor; - int cur_n_pos = cur_n % pack_factor; - - constexpr int sh_stride = tile_n_ints; - constexpr uint32_t mask = (1 << num_bits) - 1; - - int4 *sh_stage_ptr = sh + stage_size * pipe; - uint32_t *sh_stage_int_ptr = reinterpret_cast(sh_stage_ptr); - - // Undo interleaving - int cur_n_pos_unpacked; - if constexpr (num_bits == 4) { - constexpr int undo_pack[8] = {0, 4, 1, 5, 2, 6, 3, 7}; - cur_n_pos_unpacked = undo_pack[cur_n_pos]; - } else { - constexpr int undo_pack[4] = {0, 2, 1, 3}; - cur_n_pos_unpacked = undo_pack[cur_n_pos]; - } - - uint32_t vals[8]; -#pragma unroll - for (int i = 0; i < 4; i++) { - if constexpr (is_a_8bit) { - int cur_elem = tc_row + i; - - int packed_src_0 = sh_stage_int_ptr[cur_n_packed + (8 / pack_factor) * (warp_id % 2) + sh_stride * cur_elem]; - int packed_src_1 = sh_stage_int_ptr[cur_n_packed + (8 / pack_factor) * (warp_id % 2) + sh_stride * (cur_elem + 16)]; - - vals[i] = (packed_src_0 >> (cur_n_pos_unpacked * num_bits)) & mask; - vals[4 + i] = (packed_src_1 >> (cur_n_pos_unpacked * num_bits)) & mask; - } else { - int cur_elem = tc_row + tc_offsets[i]; - - int packed_src_0 = sh_stage_int_ptr[cur_n_packed + sh_stride * cur_elem]; - int packed_src_1 = sh_stage_int_ptr[cur_n_packed + (8 / pack_factor) + sh_stride * cur_elem]; - - vals[i] = (packed_src_0 >> (cur_n_pos_unpacked * num_bits)) & mask; - vals[4 + i] = (packed_src_1 >> (cur_n_pos_unpacked * num_bits)) & mask; - } - } - - constexpr int tile_size = target_tile_k_size * target_tile_n_size / pack_factor; - int out_offset = (k_tile_id * n_tiles + n_tile_id) * tile_size; - - // Result of: - // https://github.com/NVIDIA/FasterTransformer/blob/main/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h - if constexpr (!is_a_8bit && num_bits == 4) { - int pack_idx[8] = {0, 2, 4, 6, 1, 3, 5, 7}; - - uint32_t res = 0; -#pragma unroll - for (int i = 0; i < 8; i++) { - res |= vals[pack_idx[i]] << (i * 4); - } - - out_ptr[out_offset + th_id * 4 + warp_id] = res; - - } else if constexpr (is_a_8bit && num_bits == 4) { - int pack_idx[8] = {0, 4, 1, 5, 2, 6, 3, 7}; - - uint32_t res = 0; -#pragma unroll - for (int i = 0; i < 8; i++) { - res |= vals[pack_idx[i]] << (i * 4); - } - - out_ptr[out_offset + th_id * 4 + warp_id] = res; - - } else { - constexpr int pack_idx[4] = {0, 2, 1, 3}; - - uint32_t res1 = 0; - uint32_t res2 = 0; -#pragma unroll - for (int i = 0; i < 4; i++) { - const int ii = is_a_8bit ? i : pack_idx[i]; - res1 |= vals[ii] << (i * 8); - res2 |= vals[4 + ii] << (i * 8); - } - - out_ptr[out_offset + th_id * 8 + (warp_id * 2) + 0] = res1; - out_ptr[out_offset + th_id * 8 + (warp_id * 2) + 1] = res2; - } - }; - - auto start_pipes = [&](int k_tile_id, int n_tile_id) { -#pragma unroll - for (int pipe = 0; pipe < repack_stages - 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; - - start_pipes(k_tile_id, n_tile_id); - - while (n_tile_id < n_tiles) { -#pragma unroll - for (int pipe = 0; pipe < repack_stages; pipe++) { - fetch_to_shared((pipe + repack_stages - 1) % repack_stages, k_tile_id, - n_tile_id + pipe + repack_stages - 1); - repack_tile(pipe, k_tile_id, n_tile_id + pipe); - wait_for_stage(); - } - n_tile_id += repack_stages; - } - } -} - -} // namespace marlin diff --git a/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/marlin/dequant.h b/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/marlin/dequant.h deleted file mode 100644 index b4597ea90..000000000 --- a/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/marlin/dequant.h +++ /dev/null @@ -1,605 +0,0 @@ -// Source: InfiniCore 1ab85ef1d3c2bce49edf1fc73a5b48becaf7552e, -// src/infiniop/ops/awq_marlin_gemm/marlin/dequant.h. -// Donor implementation retained; this attribution notice was added by InfiniOps. -// Marlin/vLLM and direct-donor licenses: repository THIRD_PARTY.txt. -/* -Fast Dequantization (Converting INT4/INT8/FP4/FP8 to FP16/BF16) - -The process of fast dequantization can be summarized as a combination -of bitwise operations and floating-point computations: - -weight =>(bit_op / bitwise operations)=> -f16_value =>(flop / floating-point computation)=> -dequantized_weight - -Since the dequantized weights typically require subtracting the zero point and -applying a scale factor, the floating-point computation step can be fused with -the zero-point subtraction and scaling operations. - -The following are the parts that need to be modified for the fused operation -of zero-point subtraction and scaling. - -## INT4 => FP16/BF16 or INT8 => FP16 - -The floating-point computation is `__hsub2` - -If has zero points: - - flop(bit_op(weight)) - flop(bit_op(zp)) - = sub(bit_op(weight), bias) - sub(bit_op(zp), bias) - = bit_op(weight) - bit_op(zp) - -so we don't need additional modification. - -If has float zero points: - - flop(bit_op(weight)) - fzp - = sub(bit_op(weight), bias) - fzp - = bit_op(weight) - (fzp + bias) - -where the `fzp + bias` can be computed at weight loading. But this -may have accuracy issue, so we should not use this in most cases. - -If has not zero points: - - scale(flop(bit_op(weight))) - = scale(sub(bit_op(weight), bias)) - = scale(bit_op(weight)) - scale(bias) - = fma(bit_op(weight), scale_factor, scale(bias)) - -where the `scale(bias)` can be cached. But this may have accuracy issue, -so we should not use this in most cases. - - -## INT8 => BF16 - -INT8 => BF16 is a special case, it use byte_perm instead of flop. -We cannot fused byte_perm with scaling. - - -## FP4/FP8 => FP16/BF16 - - scale(flop(bit_op(weight))) - = scale(mul(bit_op(weight), multiplier)) - = mul(bit_op(weight), scale_factor * multiplier) - -where `scale_factor * multiplier` can be computed at weight loading. - -*/ - -#include "marlin_dtypes.cuh" - -namespace MARLIN_NAMESPACE_NAME { - -#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 750 -// Lookup-table based 3-input logical operation; explicitly used for -// dequantization as the compiler does not seem to automatically recognize it in -// all cases. -template -__device__ inline int lop3(int a, int b, int c) { - int res; - asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" - : "=r"(res) - : "r"(a), "r"(b), "r"(c), "n"(lut)); - return res; -} - -// Constructs destination register by taking bytes from 2 sources (based on -// mask) -template -__device__ inline uint32_t prmt(uint32_t a) { - uint32_t res; - asm volatile("prmt.b32 %0, %1, %2, %3;\n" - : "=r"(res) - : "r"(a), "n"(start_byte), "n"(mask)); - return res; -} - -template -__device__ inline void dequant(int q, scalar_t2 *frag_b); - -// -// Efficiently dequantize 4bit values packed in an int32 value into a full -// B-fragment of 4 fp16 values. We mostly follow the strategy in the link below, -// with some small changes: -// - FP16: -// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L215-L287 -// - BF16: -// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L327-L385 -// -template <> -__device__ inline void dequant(int q, - half2 *frag_b) { - const int MASK = 0x000f000f; - const int EX = 0x64006400; - // Guarantee that the `(a & b) | c` operations are LOP3s. - int lo = lop3<(0xf0 & 0xcc) | 0xaa>(q, MASK, EX); - q >>= 4; - int hi = lop3<(0xf0 & 0xcc) | 0xaa>(q, MASK, EX); - - frag_b[0] = *reinterpret_cast(&lo); - frag_b[1] = *reinterpret_cast(&hi); -} - -template <> -__device__ inline void dequant(int q, - half2 *frag_b) { - const int LO = 0x000f000f; - const int HI = 0x00f000f0; - const int EX = 0x64006400; - // Guarantee that the `(a & b) | c` operations are LOP3s. - // clang-format off - int lo = lop3<(0xf0 & 0xcc) | 0xaa>(q, LO, EX); - int hi = lop3<(0xf0 & 0xcc) | 0xaa>(q, HI, EX); - // clang-format on - // We want signed int4 outputs, hence we fuse the `-8` symmetric zero point - // directly into `SUB` and `ADD`. - const int SUB = 0x64086408; - const int MUL = 0x2c002c00; - const int ADD = 0xd480d480; - frag_b[0] = __hsub2(*reinterpret_cast(&lo), - *reinterpret_cast(&SUB)); - frag_b[1] = __hfma2(*reinterpret_cast(&hi), - *reinterpret_cast(&MUL), - *reinterpret_cast(&ADD)); -} - -template <> -__device__ inline void dequant(int q, - half2 *frag_b) { - dequant(q, frag_b); -} - -template <> -__device__ inline void dequant(int q, - half2 *frag_b) { - const int LO = 0x000f000f; - const int HI = 0x00f000f0; - const int EX = 0x64006400; - // Guarantee that the `(a & b) | c` operations are LOP3s. - // clang-format off - int lo = lop3<(0xf0 & 0xcc) | 0xaa>(q, LO, EX); - int hi = lop3<(0xf0 & 0xcc) | 0xaa>(q, HI, EX); - // clang-format on - // We want signed int4 outputs, hence we fuse the `-8` symmetric zero point - // directly into `SUB` and `ADD`. - const int SUB = 0x64006400; - const int MUL = 0x2c002c00; - const int ADD = 0xd400d400; - frag_b[0] = __hsub2(*reinterpret_cast(&lo), - *reinterpret_cast(&SUB)); - frag_b[1] = __hfma2(*reinterpret_cast(&hi), - *reinterpret_cast(&MUL), - *reinterpret_cast(&ADD)); -} - -template <> -__device__ inline void dequant( - int q, nv_bfloat162 *frag_b) { - static constexpr uint32_t MASK = 0x000f000f; - static constexpr uint32_t EX = 0x43004300; - - // Guarantee that the `(a & b) | c` operations are LOP3s. - // clang-format off - int lo = lop3<(0xf0 & 0xcc) | 0xaa>(q, MASK, EX); - q >>= 4; - int hi = lop3<(0xf0 & 0xcc) | 0xaa>(q, MASK, EX); - // clang-format on - - frag_b[0] = *reinterpret_cast(&lo); - frag_b[1] = *reinterpret_cast(&hi); -} - -template <> -__device__ inline void dequant( - int q, nv_bfloat162 *frag_b) { - dequant(q, frag_b); - - static constexpr uint32_t SUB = 0x43084308; - - frag_b[0] = __hsub2(frag_b[0], *reinterpret_cast(&SUB)); - frag_b[1] = __hsub2(frag_b[1], *reinterpret_cast(&SUB)); -} - -template <> -__device__ inline void dequant( - int q, nv_bfloat162 *frag_b) { - dequant(q, frag_b); -} - -template <> -__device__ inline void dequant( - int q, nv_bfloat162 *frag_b) { - dequant(q, frag_b); - - static constexpr uint32_t SUB = 0x43004300; - - frag_b[0] = __hsub2(frag_b[0], *reinterpret_cast(&SUB)); - frag_b[1] = __hsub2(frag_b[1], *reinterpret_cast(&SUB)); -} - -// -// Fast Int8ToFp16/Int8ToBf16: Efficiently dequantize 8bit int values to fp16 or -// bf16 Reference: -// - FP16: -// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L53-L85 -// - BF16: -// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L125-L175 -// -template <> -__device__ inline void dequant(int q, - half2 *frag_b) { - static constexpr uint32_t mask_for_elt_01 = 0x5250; - static constexpr uint32_t mask_for_elt_23 = 0x5351; - static constexpr uint32_t start_byte_for_fp16 = 0x64646464; - - uint32_t lo = prmt(q); - uint32_t hi = prmt(q); - - frag_b[0] = *reinterpret_cast(&lo); - frag_b[1] = *reinterpret_cast(&hi); -} - -template <> -__device__ inline void dequant( - int q, half2 *frag_b) { - dequant(q, frag_b); - - static constexpr uint32_t I8s_TO_F16s_MAGIC_NUM = 0x64806480; - frag_b[0] = __hsub2(frag_b[0], - *reinterpret_cast(&I8s_TO_F16s_MAGIC_NUM)); - frag_b[1] = __hsub2(frag_b[1], - *reinterpret_cast(&I8s_TO_F16s_MAGIC_NUM)); -} - -template <> -__device__ inline void dequant(int q, - half2 *frag_b) { - dequant(q, frag_b); -} - -template <> -__device__ inline void dequant(int q, - half2 *frag_b) { - dequant(q, frag_b); - - static constexpr uint32_t I8s_TO_F16s_MAGIC_NUM = 0x64006400; - frag_b[0] = __hsub2(frag_b[0], - *reinterpret_cast(&I8s_TO_F16s_MAGIC_NUM)); - frag_b[1] = __hsub2(frag_b[1], - *reinterpret_cast(&I8s_TO_F16s_MAGIC_NUM)); -} - -template <> -__device__ inline void dequant( - int q, nv_bfloat162 *frag_b) { - float fp32_intermediates[4]; - uint32_t *fp32_intermediates_casted = reinterpret_cast(fp32_intermediates); - - static constexpr uint32_t fp32_base = 0x4B000000; - fp32_intermediates_casted[0] = __byte_perm(q, fp32_base, 0x7650); - fp32_intermediates_casted[1] = __byte_perm(q, fp32_base, 0x7652); - fp32_intermediates_casted[2] = __byte_perm(q, fp32_base, 0x7651); - fp32_intermediates_casted[3] = __byte_perm(q, fp32_base, 0x7653); - - fp32_intermediates[0] -= 8388736.f; - fp32_intermediates[1] -= 8388736.f; - fp32_intermediates[2] -= 8388736.f; - fp32_intermediates[3] -= 8388736.f; - - uint32_t *bf16_result_ptr = reinterpret_cast(frag_b); - bf16_result_ptr[0] = __byte_perm(fp32_intermediates_casted[0], - fp32_intermediates_casted[1], 0x7632); - bf16_result_ptr[1] = __byte_perm(fp32_intermediates_casted[2], - fp32_intermediates_casted[3], 0x7632); -} - -template <> -__device__ inline void dequant( - int q, nv_bfloat162 *frag_b) { - float fp32_intermediates[4]; - uint32_t *fp32_intermediates_casted = reinterpret_cast(fp32_intermediates); - - static constexpr uint32_t fp32_base = 0x4B000000; - fp32_intermediates_casted[0] = __byte_perm(q, fp32_base, 0x7650); - fp32_intermediates_casted[1] = __byte_perm(q, fp32_base, 0x7652); - fp32_intermediates_casted[2] = __byte_perm(q, fp32_base, 0x7651); - fp32_intermediates_casted[3] = __byte_perm(q, fp32_base, 0x7653); - - fp32_intermediates[0] -= 8388608.f; - fp32_intermediates[1] -= 8388608.f; - fp32_intermediates[2] -= 8388608.f; - fp32_intermediates[3] -= 8388608.f; - - uint32_t *bf16_result_ptr = reinterpret_cast(frag_b); - bf16_result_ptr[0] = __byte_perm(fp32_intermediates_casted[0], - fp32_intermediates_casted[1], 0x7632); - bf16_result_ptr[1] = __byte_perm(fp32_intermediates_casted[2], - fp32_intermediates_casted[3], 0x7632); -} - -template <> -__device__ inline void dequant( - int q, half2 *frag_b) { - // Constants for FP8 (E4M3) and FP16 formats - constexpr int FP8_EXPONENT = 4, FP16_EXPONENT = 5; - constexpr int RIGHT_SHIFT = FP16_EXPONENT - FP8_EXPONENT; - constexpr int MASK = 0x7F007F00; - - // Extract and shift FP8 values to FP16 format - int Out1 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT); - q <<= 8; - int Out2 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT); - - // Note: reverse indexing is intentional because weights are permuted - frag_b[1] = *reinterpret_cast(&Out1); - frag_b[0] = *reinterpret_cast(&Out2); -} - -template <> -__device__ inline void dequant( - int q, half2 *frag_b) { - dequant(q, frag_b); - - // Constants for FP8 (E4M3) and FP16 formats - constexpr int FP8_EXPONENT = 4, FP16_EXPONENT = 5; - - // Construct and apply exponent bias - constexpr int BIAS_OFFSET = (1 << (FP16_EXPONENT - 1)) - (1 << (FP8_EXPONENT - 1)); - const half2 bias_reg = __float2half2_rn(float(1 << BIAS_OFFSET)); - - // Convert to half2 and apply bias - frag_b[1] = __hmul2(frag_b[1], bias_reg); - frag_b[0] = __hmul2(frag_b[0], bias_reg); -} - -template <> -__device__ inline void dequant( - int q, nv_bfloat162 *frag_b) { - // Constants for FP8 (E4M3) and BF16 formats - constexpr int FP8_EXPONENT = 4, BF16_EXPONENT = 8; - constexpr int RIGHT_SHIFT = BF16_EXPONENT - FP8_EXPONENT; - - constexpr int MASK = 0x7F007F00; - - // Extract and shift FP8 values to BF16 format - int Out1 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT); - q <<= 8; - int Out2 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT); - - // Note: reverse indexing is intentional because weights are permuted - frag_b[1] = *reinterpret_cast(&Out1); - frag_b[0] = *reinterpret_cast(&Out2); -} - -template <> -__device__ inline void dequant( - int q, nv_bfloat162 *frag_b) { - dequant(q, frag_b); - - // Constants for FP8 (E4M3) and BF16 formats - constexpr int FP8_EXPONENT = 4, BF16_EXPONENT = 8; - - // Construct and apply exponent bias - constexpr int BIAS_OFFSET = (1 << (BF16_EXPONENT - 1)) - (1 << (FP8_EXPONENT - 1)); - // Add 127 (float exponent bias) to BIAS_OFFSET and shift to float exponent - // position - constexpr uint32_t BIAS = (BIAS_OFFSET + 127) << 23; - const nv_bfloat162 bias_reg = __float2bfloat162_rn(*reinterpret_cast(&BIAS)); - - // Convert to bfloat162 and apply bias - frag_b[1] = __hmul2(frag_b[1], bias_reg); - frag_b[0] = __hmul2(frag_b[0], bias_reg); -} - -template <> -__device__ inline void dequant(int q, - half2 *frag_b) { - // Constants for FP4 (E2M1) and FP16 formats - constexpr int FP4_EXPONENT = 2, FP16_EXPONENT = 5; - constexpr int RIGHT_SHIFT = FP16_EXPONENT - FP4_EXPONENT; - constexpr int MASK = 0x70007000; - - // Extract and shift FP4 values to FP16 format - int Out1 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT); - q <<= 4; - int Out2 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT); - - // Note: reverse indexing is intentional because weights are permuted - frag_b[1] = *reinterpret_cast(&Out1); - frag_b[0] = *reinterpret_cast(&Out2); -} - -template <> -__device__ inline void dequant( - int q, half2 *frag_b) { - dequant(q, frag_b); - - // Constants for FP4 (E2M1) and FP16 formats - constexpr int FP4_EXPONENT = 2, FP16_EXPONENT = 5; - - // Construct and apply exponent bias - constexpr int BIAS_OFFSET = (1 << (FP16_EXPONENT - 1)) - (1 << (FP4_EXPONENT - 1)); - const half2 bias_reg = __float2half2_rn(float(1 << BIAS_OFFSET)); - - // Convert to half2 and apply bias - frag_b[1] = __hmul2(frag_b[1], bias_reg); - frag_b[0] = __hmul2(frag_b[0], bias_reg); -} - -template <> -__device__ inline void dequant( - int q, nv_bfloat162 *frag_b) { - // Constants for FP4 (E2M1) and FP16 formats - constexpr int FP4_EXPONENT = 2, BF16_EXPONENT = 8; - constexpr int RIGHT_SHIFT = BF16_EXPONENT - FP4_EXPONENT; - constexpr int MASK = 0x70007000; - - // Extract and shift FP4 values to FP16 format - int Out1 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT); - q <<= 4; - int Out2 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT); - - // Note: reverse indexing is intentional because weights are permuted - frag_b[1] = *reinterpret_cast(&Out1); - frag_b[0] = *reinterpret_cast(&Out2); -} - -template <> -__device__ inline void dequant( - int q, nv_bfloat162 *frag_b) { - dequant(q, frag_b); - - // Constants for FP4 (E2M1) and BF16 formats - constexpr int FP4_EXPONENT = 2, BF16_EXPONENT = 8; - - // Construct and apply exponent bias - constexpr int BIAS_OFFSET = (1 << (BF16_EXPONENT - 1)) - (1 << (FP4_EXPONENT - 1)); - // Add 127 (float exponent bias) to BIAS_OFFSET and shift to float exponent - // position - constexpr uint32_t BIAS = (BIAS_OFFSET + 127) << 23; - const nv_bfloat162 bias_reg = __float2bfloat162_rn(*reinterpret_cast(&BIAS)); - - // Convert to half2 and apply bias - frag_b[1] = __hmul2(frag_b[1], bias_reg); - frag_b[0] = __hmul2(frag_b[0], bias_reg); -} - -template <> -__device__ inline void dequant<__nv_fp8x4_e4m3, vllm::kFE2M1f.id(), true>( - int q, __nv_fp8x4_e4m3 *frag_b) { - // Constants for FP4 (E2M1) and FP16 formats - constexpr int FP4_EXPONENT = 2, FP8_EXPONENT = 4; - constexpr int RIGHT_SHIFT = FP8_EXPONENT - FP4_EXPONENT; - constexpr int MASK = 0x70707070; - - // Extract and shift FP4 values to FP16 format - int Out1 = (q & 0x80808080) | ((q & MASK) >> RIGHT_SHIFT); - q <<= 4; - int Out2 = (q & 0x80808080) | ((q & MASK) >> RIGHT_SHIFT); - - // Note1: reverse indexing is intentional because weights are permuted - // Note2: when dequant to 8bit type, we write to `frag_b[2]` instead of - // `frag_b[1]` to fit the layout of tensorcore - frag_b[1] = *reinterpret_cast(&Out1); - frag_b[0] = *reinterpret_cast(&Out2); -} - -template <> -__device__ inline void dequant( - int q, int32_t *frag_b) { - constexpr int repeated_zp = 0x08080808; - constexpr int MASK = 0x80808080; - - frag_b[0] = ((q & 0x0F0F0F0F | MASK) - repeated_zp) ^ MASK; - q >>= 4; - frag_b[1] = ((q & 0x0F0F0F0F | MASK) - repeated_zp) ^ MASK; -} - -template <> -__device__ inline void dequant<__nv_fp8x4_e4m3, vllm::kU4B8.id(), true>( - int q, __nv_fp8x4_e4m3 *frag_b) { - int s = q & 0x08080808; - int Out1 = ((q & 0x07070707) | (s << 4)) + (s >> 3); - q >>= 4; - s = q & 0x08080808; - int Out2 = ((q & 0x07070707) | (s << 4)) + (s >> 3); - - frag_b[0] = *reinterpret_cast(&Out1); - frag_b[1] = *reinterpret_cast(&Out2); -} - -template -__device__ inline void dequant_fp8_scales(int q, scalar_t2 *frag_b); - -template <> -__device__ inline void dequant_fp8_scales( - int q, half2 *frag_b) { - int Out1 = (q & 0xFF00FF00) >> 1; - ; - q <<= 8; - int Out2 = (q & 0xFF00FF00) >> 1; - - // Note: reverse indexing is intentional because weights are permuted - frag_b[1] = *reinterpret_cast(&Out1); - frag_b[0] = *reinterpret_cast(&Out2); -}; - -template <> -__device__ inline void dequant_fp8_scales( - int q, nv_bfloat162 *frag_b) { - constexpr int FP8_EXPONENT = 4, BF16_EXPONENT = 8; - constexpr int RIGHT_SHIFT = BF16_EXPONENT - FP8_EXPONENT; - constexpr int MASK = 0x7F007F00; - - // Extract and shift FP8 values to BF16 format - int Out1 = ((q & 0x80008000) >> 1) | ((q & MASK) >> RIGHT_SHIFT); - q <<= 8; - int Out2 = ((q & 0x80008000) >> 1) | ((q & MASK) >> RIGHT_SHIFT); - - // Note: reverse indexing is intentional because weights are permuted - frag_b[1] = *reinterpret_cast(&Out1); - frag_b[0] = *reinterpret_cast(&Out2); -} - -template <> -__device__ inline void dequant_fp8_scales( - int q, nv_bfloat162 *frag_b) { - // In this conversion, 2 ** -127 in FP8E8M0 would become 0 in BF16, - // but we assume that such a extreme value would not occur in real models. - int Out1 = (q & 0xFF00FF00) >> 1; - q <<= 7; - int Out2 = q & 0x7F807F80; - - // Note: reverse indexing is intentional because weights are permuted - frag_b[1] = *reinterpret_cast(&Out1); - frag_b[0] = *reinterpret_cast(&Out2); -}; - -// subtract zero point in quanted format and then dequant -template -__device__ inline void sub_zp_and_dequant(int q, scalar_t2 *frag_b, int zp); - -template <> -__device__ inline void sub_zp_and_dequant( - int q, int32_t *frag_b, int zp) { - // INT4 with zp -> INT8 - // see https://github.com/vllm-project/vllm/pull/24722 - int repeated_zp = 0x01010101 * zp; - int MASK = 0x80808080; - - frag_b[0] = ((q & 0x0F0F0F0F | MASK) - repeated_zp) ^ MASK; - q >>= 4; - frag_b[1] = ((q & 0x0F0F0F0F | MASK) - repeated_zp) ^ MASK; -} - -template <> -__device__ inline void sub_zp_and_dequant<__nv_fp8x4_e4m3, vllm::kU4.id(), - true>(int q, __nv_fp8x4_e4m3 *frag_b, - int zp) { - // INT4 with zp -> FP8 - // see https://github.com/vllm-project/vllm/pull/24722 - uint32_t u_q = *reinterpret_cast(&q); - uint32_t u_zp = *reinterpret_cast(&zp); - uint32_t u_zp1 = u_zp + 1; - uint32_t repeated_zp = 0x01010101 * u_zp; - - uint32_t q0, s; - q0 = (u_q & 0x0F0F0F0F) | 0x70707070; - s = (q0 + repeated_zp) & 0x80808080; - uint32_t Out1 = (q0 + (s >> 7) * u_zp1) & 0x0F0F0F0F | s; - - u_q >>= 4; - q0 = (u_q & 0x0F0F0F0F) | 0x70707070; - s = (q0 + repeated_zp) & 0x80808080; - uint32_t Out2 = (q0 + (s >> 7) * u_zp1) & 0x0F0F0F0F | s; - - frag_b[0] = *reinterpret_cast(&Out1); - frag_b[1] = *reinterpret_cast(&Out2); -} - -#endif - -} // namespace MARLIN_NAMESPACE_NAME diff --git a/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/marlin/kernel.h b/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/marlin/kernel.h deleted file mode 100644 index bcfac3e5a..000000000 --- a/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/marlin/kernel.h +++ /dev/null @@ -1,47 +0,0 @@ - -// Source: InfiniCore 1ab85ef1d3c2bce49edf1fc73a5b48becaf7552e, -// src/infiniop/ops/awq_marlin_gemm/marlin/kernel.h. -// Donor implementation retained; this attribution notice was added by InfiniOps. -// Marlin/vLLM and direct-donor licenses: repository THIRD_PARTY.txt. -#ifndef MARLIN_NAMESPACE_NAME -#define MARLIN_NAMESPACE_NAME marlin -#endif - -#include "../core/scalar_type.hpp" -#include "marlin.cuh" -#include "marlin_dtypes.cuh" - -#define MARLIN_KERNEL_PARAMS \ - const int4 *__restrict__ A, const int4 *__restrict__ B, \ - int4 *__restrict__ C, int4 *__restrict__ C_tmp, \ - const int4 *__restrict__ b_bias_ptr, \ - const float *__restrict__ a_scales_ptr, \ - const int4 *__restrict__ scales_ptr, \ - const float *__restrict__ global_scale_ptr, \ - const int4 *__restrict__ zp_ptr, const int *__restrict__ g_idx, \ - int num_groups, int prob_m, int prob_n, int prob_k, int lda, int *locks, \ - bool has_bias, bool use_atomic_add, bool use_fp32_reduce, \ - int max_shared_mem - -namespace MARLIN_NAMESPACE_NAME { -template shared - // fetch pipeline - const int group_blocks, // number of consecutive 16x16 blocks - // with a separate quantization scale - const bool is_zp_float // is zero point of float16 type? - > -__global__ void Marlin(MARLIN_KERNEL_PARAMS); - -} diff --git a/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/marlin/marlin.cuh b/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/marlin/marlin.cuh deleted file mode 100644 index a752519b2..000000000 --- a/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/marlin/marlin.cuh +++ /dev/null @@ -1,182 +0,0 @@ -// Source: InfiniCore 1ab85ef1d3c2bce49edf1fc73a5b48becaf7552e, -// src/infiniop/ops/awq_marlin_gemm/marlin/marlin.cuh. -// Donor implementation retained; this attribution notice was added by InfiniOps. -// Marlin/vLLM and direct-donor licenses: repository THIRD_PARTY.txt. -#pragma once - -#ifndef _marlin_cuh -#define _marlin_cuh - -#include -#include -#include -#include - -#ifndef MARLIN_NAMESPACE_NAME -#define MARLIN_NAMESPACE_NAME marlin -#endif - -template -__device__ __forceinline__ uint32_t __cvta_generic_to_shared(T *ptr) { - size_t smem_addr; - asm volatile( - "cvta.to.shared.u64 %0, %1;" - : "=l"(smem_addr) - : "l"(ptr)); - return static_cast(smem_addr); -} - -namespace MARLIN_NAMESPACE_NAME { - -// Marlin params - -// 8 warps are a good choice since every SM has 4 schedulers and having more -// than 1 warp per schedule allows some more latency hiding. At the same time, -// we want relatively few warps to have many registers per warp and small tiles. -static constexpr int default_threads = 256; - -static constexpr int pipe_stages = 4; // 4 pipeline stages fit into shared memory - -static constexpr int min_thread_n = 64; -static constexpr int min_thread_k = 64; -static constexpr int max_thread_n = 256; - -static constexpr int tile_size = 16; -static constexpr int max_par = 16; - -// Repack params -static constexpr int repack_stages = 8; - -static constexpr int repack_threads = 256; - -static constexpr int tile_k_size = tile_size; -static constexpr int tile_n_size = tile_k_size * 4; - -// Helpers -template -struct Vec { - T elems[n]; - __device__ T &operator[](int i) { return elems[i]; } -}; - -using I4 = Vec; - -constexpr int div_ceil(int a, int b) { return (a + b - 1) / b; } - -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 - -__device__ inline void cp_async1_ca_pred(void *smem_ptr, const void *glob_ptr, - bool pred = true) { - if (pred) { - reinterpret_cast(smem_ptr)[0] = reinterpret_cast(glob_ptr)[0]; - } -} - -__device__ inline void cp_async2_ca_pred(void *smem_ptr, const void *glob_ptr, - bool pred = true) { - if (pred) { - reinterpret_cast(smem_ptr)[0] = reinterpret_cast(glob_ptr)[0]; - } -} - -__device__ inline void cp_async4_ca_pred(void *smem_ptr, const void *glob_ptr, - bool pred = true) { - if (pred) { - reinterpret_cast(smem_ptr)[0] = reinterpret_cast(glob_ptr)[0]; - } -} - -__device__ inline void cp_async4_pred(void *smem_ptr, const void *glob_ptr, - bool pred = true) { - if (pred) { - reinterpret_cast(smem_ptr)[0] = reinterpret_cast(glob_ptr)[0]; - } -} - -__device__ inline void cp_async4(void *smem_ptr, const void *glob_ptr) { - reinterpret_cast(smem_ptr)[0] = reinterpret_cast(glob_ptr)[0]; -} - -__device__ inline void cp_async_fence() {} - -template -__device__ inline void cp_async_wait() {} - -#else - -__device__ inline void cp_async1_ca_pred(void *smem_ptr, const void *glob_ptr, - bool pred = true) { - const int BYTES = 4; - uint32_t smem = static_cast(__cvta_generic_to_shared(smem_ptr)); - asm volatile( - "{\n" - " .reg .pred p;\n" - " setp.ne.b32 p, %0, 0;\n" - " @p cp.async.ca.shared.global [%1], [%2], %3;\n" - "}\n" ::"r"((int)pred), - "r"(smem), "l"(glob_ptr), "n"(BYTES)); -} - -__device__ inline void cp_async2_ca_pred(void *smem_ptr, const void *glob_ptr, - bool pred = true) { - const int BYTES = 8; - uint32_t smem = static_cast(__cvta_generic_to_shared(smem_ptr)); - asm volatile( - "{\n" - " .reg .pred p;\n" - " setp.ne.b32 p, %0, 0;\n" - " @p cp.async.ca.shared.global [%1], [%2], %3;\n" - "}\n" ::"r"((int)pred), - "r"(smem), "l"(glob_ptr), "n"(BYTES)); -} - -__device__ inline void cp_async4_ca_pred(void *smem_ptr, const void *glob_ptr, - bool pred = true) { - const int BYTES = 16; - uint32_t smem = static_cast(__cvta_generic_to_shared(smem_ptr)); - asm volatile( - "{\n" - " .reg .pred p;\n" - " setp.ne.b32 p, %0, 0;\n" - " @p cp.async.ca.shared.global [%1], [%2], %3;\n" - "}\n" ::"r"((int)pred), - "r"(smem), "l"(glob_ptr), "n"(BYTES)); -} - -__device__ inline void cp_async4_pred(void *smem_ptr, const void *glob_ptr, - bool pred = true) { - const int BYTES = 16; - uint32_t smem = static_cast(__cvta_generic_to_shared(smem_ptr)); - asm volatile( - "{\n" - " .reg .pred p;\n" - " setp.ne.b32 p, %0, 0;\n" - " @p cp.async.cg.shared.global [%1], [%2], %3;\n" - "}\n" ::"r"((int)pred), - "r"(smem), "l"(glob_ptr), "n"(BYTES)); -} - -__device__ inline void cp_async4(void *smem_ptr, const void *glob_ptr) { - const int BYTES = 16; - uint32_t smem = static_cast(__cvta_generic_to_shared(smem_ptr)); - asm volatile( - "{\n" - " cp.async.cg.shared.global [%0], [%1], %2;\n" - "}\n" ::"r"(smem), - "l"(glob_ptr), "n"(BYTES)); -} - -__device__ inline void cp_async_fence() { - asm volatile("cp.async.commit_group;\n" ::); -} - -template -__device__ inline void cp_async_wait() { - asm volatile("cp.async.wait_group %0;\n" ::"n"(n)); -} - -#endif - -} // namespace MARLIN_NAMESPACE_NAME - -#endif diff --git a/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/marlin/marlin_dtypes.cuh b/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/marlin/marlin_dtypes.cuh deleted file mode 100644 index 98340b53f..000000000 --- a/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/marlin/marlin_dtypes.cuh +++ /dev/null @@ -1,159 +0,0 @@ - -// Source: InfiniCore 1ab85ef1d3c2bce49edf1fc73a5b48becaf7552e, -// src/infiniop/ops/awq_marlin_gemm/marlin/marlin_dtypes.cuh. -// Donor implementation retained; this attribution notice was added by InfiniOps. -// Marlin/vLLM and direct-donor licenses: repository THIRD_PARTY.txt. -#ifndef _data_types_cuh -#define _data_types_cuh -#include "../core/scalar_type.hpp" -#include "marlin.cuh" -#include -#include -#include - -#ifndef MARLIN_NAMESPACE_NAME -#define MARLIN_NAMESPACE_NAME marlin -#endif - -namespace MARLIN_NAMESPACE_NAME { - -template -class MarlinScalarType { -}; - -template <> -class MarlinScalarType { -public: - using scalar_t = half; - using scalar_t2 = half2; - using scalar_t4 = half2; - using scalar_32bit_t = half2; - - // Matrix fragments for tensor core instructions; their precise layout is - // documented here: - // https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#matrix-fragments-for-mma-m16n8k16-with-floating-point-type - using FragA = Vec; - using FragB = Vec; - using FragC = Vec; - using FragS = Vec; - using FragS0 = Vec<__nv_fp8x2_e4m3, 1>; - using FragZP = Vec; - - static __device__ float inline num2float(const half x) { - return __half2float(x); - } - - static __device__ half2 inline num2num2(const half x) { - return __half2half2(x); - } - - static __device__ half2 inline nums2num2(const half x1, const half x2) { - return __halves2half2(x1, x2); - } - - static __host__ __device__ half inline float2num(const float x) { - return __float2half(x); - } - - static __host__ __device__ float2 inline num22float2(const half2 x) { - return __half22float2(x); - } -}; - -template <> -class MarlinScalarType { -public: - using scalar_t = nv_bfloat16; - using scalar_t2 = nv_bfloat162; - using scalar_t4 = nv_bfloat162; - using scalar_32bit_t = nv_bfloat162; - - using FragA = Vec; - using FragB = Vec; - using FragC = Vec; - using FragS = Vec; - using FragS0 = Vec<__nv_fp8x2_e4m3, 1>; - using FragZP = Vec; - -#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 800 - static __device__ float inline num2float(const nv_bfloat16 x) { - return __bfloat162float(x); - } - - static __device__ nv_bfloat162 inline num2num2(const nv_bfloat16 x) { - return __bfloat162bfloat162(x); - } - - static __device__ nv_bfloat162 inline nums2num2(const nv_bfloat16 x1, - const nv_bfloat16 x2) { - return __halves2bfloat162(x1, x2); - } - - static __host__ __device__ nv_bfloat16 inline float2num(const float x) { - return __float2bfloat16(x); - } - - static __host__ __device__ float2 inline num22float2(const nv_bfloat162 x) { - return __bfloat1622float2(x); - } -#endif -}; - -template <> -class MarlinScalarType { -public: - using scalar_t = __nv_fp8_e4m3; - using scalar_t2 = __nv_fp8x2_e4m3; - using scalar_t4 = __nv_fp8x4_e4m3; - using scalar_32bit_t = __nv_fp8x4_e4m3; - - using FragA = Vec<__nv_fp8x4_e4m3, 4>; - using FragB = Vec<__nv_fp8x4_e4m3, 2>; - using FragC = Vec; - using FragZP = Vec<__nv_fp8x2_e4m3, 4>; - - static __host__ __device__ - float2 inline num22float2(const __nv_fp8x2_e4m3 x) { - return (float2)x; - } -}; - -template <> -class MarlinScalarType { -public: - using scalar_t = int8_t; - using scalar_t2 = int16_t; - using scalar_t4 = int32_t; - using scalar_32bit_t = int32_t; - - using FragA = Vec; - using FragB = Vec; - using FragC = Vec; - using FragZP = Vec; -}; - -template -class MarlinScalarType2 { -}; - -template <> -class MarlinScalarType2 : public MarlinScalarType { -}; - -template <> -class MarlinScalarType2 - : public MarlinScalarType { -}; - -template <> -class MarlinScalarType2<__nv_fp8_e4m3> - : public MarlinScalarType { -}; - -template <> -class MarlinScalarType2 : public MarlinScalarType { -}; - -} // namespace MARLIN_NAMESPACE_NAME - -#endif diff --git a/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/marlin/marlin_mma.h b/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/marlin/marlin_mma.h deleted file mode 100644 index ca9ca73fb..000000000 --- a/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/marlin/marlin_mma.h +++ /dev/null @@ -1,271 +0,0 @@ - -// Source: InfiniCore 1ab85ef1d3c2bce49edf1fc73a5b48becaf7552e, -// src/infiniop/ops/awq_marlin_gemm/marlin/marlin_mma.h. -// Donor implementation retained; this attribution notice was added by InfiniOps. -// Marlin/vLLM and direct-donor licenses: repository THIRD_PARTY.txt. -#include "marlin_dtypes.cuh" - -namespace MARLIN_NAMESPACE_NAME { - -// m16n8k16 tensor core mma instruction with fp16 inputs and fp32 -// output/accumulation. -template -__device__ inline void mma( - const typename MarlinScalarType::FragA &a_frag, - const typename MarlinScalarType::FragB &frag_b, - typename MarlinScalarType::FragC &frag_c, int idx = 0) { - const uint32_t *a = reinterpret_cast(&a_frag); - const uint32_t *b = reinterpret_cast(&frag_b); - using scalar_t = typename MarlinScalarType::scalar_t; - if constexpr (!std::is_same::value || k_size != 16) { - static_assert(!use_fp16_accum); - } - - if constexpr (k_size == 16) { - if constexpr (std::is_same::value && !use_fp16_accum) { -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 750 - float *c = reinterpret_cast(&frag_c); - asm volatile( - "mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 " - "{%0,%1,%2,%3}, {%4,%5}, {%6}, {%7,%8,%9,%10};\n" - : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3]) - : "r"(a[0]), "r"(a[1]), "r"(b[0]), "f"(c[0]), "f"(c[1]), "f"(c[2]), - "f"(c[3])); - asm volatile( - "mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 " - "{%0,%1,%2,%3}, {%4,%5}, {%6}, {%7,%8,%9,%10};\n" - : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3]) - : "r"(a[2]), "r"(a[3]), "r"(b[1]), "f"(c[0]), "f"(c[1]), "f"(c[2]), - "f"(c[3])); -#else - float *c = reinterpret_cast(&frag_c); - asm volatile( - "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 " - "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n" - : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3]) - : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]), - "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3])); -#endif - } else if constexpr (std::is_same::value && use_fp16_accum) { -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 750 - uint32_t *c = reinterpret_cast(&frag_c); - asm volatile( - "mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16 " - "{%0,%1}, {%2,%3}, {%4}, {%5,%6};\n" - : "=r"(c[0]), "=r"(c[1]) - : "r"(a[0]), "r"(a[1]), "r"(b[0]), "r"(c[0]), "r"(c[1])); - asm volatile( - "mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16 " - "{%0,%1}, {%2,%3}, {%4}, {%5,%6};\n" - : "=r"(c[0]), "=r"(c[1]) - : "r"(a[2]), "r"(a[3]), "r"(b[1]), "r"(c[0]), "r"(c[1])); -#else - uint32_t *c = reinterpret_cast(&frag_c); - asm volatile( - "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 " - "{%0,%1}, {%2,%3,%4,%5}, {%6,%7}, {%8,%9};\n" - : "=r"(c[0]), "=r"(c[1]) - : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]), - "r"(c[0]), "r"(c[1])); -#endif - } else if constexpr (std::is_same::value) { - float *c = reinterpret_cast(&frag_c); - asm volatile( - "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 " - "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n" - : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3]) - : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]), - "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3])); - } else if constexpr (std::is_same::value) { - float *c = reinterpret_cast(&frag_c); - asm volatile( - "mma.sync.aligned.m16n8k16.row.col.f32.e4m3.e4m3.f32 " - "{%0,%1,%2,%3}, {%4,%5}, {%6}, {%7,%8,%9,%10};\n" - : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3]) - : "r"(a[idx * 2]), "r"(a[idx * 2 + 1]), "r"(b[idx]), "f"(c[0]), - "f"(c[1]), "f"(c[2]), "f"(c[3])); - } else if constexpr (std::is_same::value) { - int32_t *c = reinterpret_cast(&frag_c); - asm volatile( - "mma.sync.aligned.m16n8k16.row.col.s32.s8.s8.s32.satfinite " - "{%0,%1,%2,%3}, {%4,%5}, {%6}, {%7,%8,%9,%10};\n" - : "=r"(c[0]), "=r"(c[1]), "=r"(c[2]), "=r"(c[3]) - : "r"(a[idx * 2]), "r"(a[idx * 2 + 1]), "r"(b[idx]), "r"(c[0]), - "r"(c[1]), "r"(c[2]), "r"(c[3])); - } - } else if (k_size == 32) { - if constexpr (std::is_same::value) { - float *c = reinterpret_cast(&frag_c); - asm volatile( - "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 " - "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n" - : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3]) - : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]), - "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3])); - } else if constexpr (std::is_same::value) { - int32_t *c = reinterpret_cast(&frag_c); -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 750 - asm volatile( - "mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32.satfinite " - "{%0,%1}, {%2}, {%3}, {%4,%5};\n" - : "=r"(c[0]), "=r"(c[1]) - : "r"(a[0]), "r"(b[0]), "r"(c[0]), "r"(c[1])); - asm volatile( - "mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32.satfinite " - "{%0,%1}, {%2}, {%3}, {%4,%5};\n" - : "=r"(c[2]), "=r"(c[3]) - : "r"(a[1]), "r"(b[0]), "r"(c[2]), "r"(c[3])); - asm volatile( - "mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32.satfinite " - "{%0,%1}, {%2}, {%3}, {%4,%5};\n" - : "=r"(c[0]), "=r"(c[1]) - : "r"(a[2]), "r"(b[1]), "r"(c[0]), "r"(c[1])); - asm volatile( - "mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32.satfinite " - "{%0,%1}, {%2}, {%3}, {%4,%5};\n" - : "=r"(c[2]), "=r"(c[3]) - : "r"(a[3]), "r"(b[1]), "r"(c[2]), "r"(c[3])); -#else - asm volatile( - "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32.satfinite " - "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n" - : "=r"(c[0]), "=r"(c[1]), "=r"(c[2]), "=r"(c[3]) - : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]), - "r"(c[0]), "r"(c[1]), "r"(c[2]), "r"(c[3])); -#endif - } - } -} - -template -__device__ inline void mma_trans( - const typename MarlinScalarType::FragA &a_frag, - const typename MarlinScalarType::FragB &frag_b, - const typename MarlinScalarType::FragB &frag_b2, - typename MarlinScalarType::FragC &frag_c) { - const uint32_t *a = reinterpret_cast(&a_frag); - const uint32_t *b = reinterpret_cast(&frag_b); - const uint32_t *b2 = reinterpret_cast(&frag_b2); - float *c = reinterpret_cast(&frag_c); - using scalar_t = typename MarlinScalarType::scalar_t; - if constexpr (!std::is_same::value || k_size != 16) { - static_assert(!use_fp16_accum); - } - - if constexpr (k_size == 16) { - if constexpr (std::is_same::value && !use_fp16_accum) { -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 750 - float *c = reinterpret_cast(&frag_c); - asm volatile( - "mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 " - "{%0,%1,%2,%3}, {%4,%5}, {%6}, {%7,%8,%9,%10};\n" - : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3]) - : "r"(b[0]), "r"(b2[0]), "r"(a[0]), "f"(c[0]), "f"(c[1]), "f"(c[2]), - "f"(c[3])); - asm volatile( - "mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 " - "{%0,%1,%2,%3}, {%4,%5}, {%6}, {%7,%8,%9,%10};\n" - : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3]) - : "r"(b[1]), "r"(b2[1]), "r"(a[1]), "f"(c[0]), "f"(c[1]), "f"(c[2]), - "f"(c[3])); -#else - float *c = reinterpret_cast(&frag_c); - asm volatile( - "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 " - "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n" - : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3]) - : "r"(b[0]), "r"(b2[0]), "r"(b[1]), "r"(b2[1]), "r"(a[0]), "r"(a[1]), - "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3])); -#endif - } else if constexpr (std::is_same::value && use_fp16_accum) { -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 750 - uint32_t *c = reinterpret_cast(&frag_c); - asm volatile( - "mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16 " - "{%0,%1}, {%2,%3}, {%4}, {%5,%6};\n" - : "=r"(c[0]), "=r"(c[1]) - : "r"(b[0]), "r"(b2[0]), "r"(a[0]), "r"(c[0]), "r"(c[1])); - asm volatile( - "mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16 " - "{%0,%1}, {%2,%3}, {%4}, {%5,%6};\n" - : "=r"(c[0]), "=r"(c[1]) - : "r"(b[1]), "r"(b2[1]), "r"(a[1]), "r"(c[0]), "r"(c[1])); -#else - uint32_t *c = reinterpret_cast(&frag_c); - asm volatile( - "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 " - "{%0,%1}, {%2,%3,%4,%5}, {%6,%7}, {%8,%9};\n" - : "=r"(c[0]), "=r"(c[1]) - : "r"(b[0]), "r"(b2[0]), "r"(b[1]), "r"(b2[1]), "r"(a[0]), "r"(a[1]), - "r"(c[0]), "r"(c[1])); -#endif - } else if constexpr (std::is_same::value) { - float *c = reinterpret_cast(&frag_c); - asm volatile( - "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 " - "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n" - : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3]) - : "r"(b[0]), "r"(b2[0]), "r"(b[1]), "r"(b2[1]), "r"(a[0]), "r"(a[1]), - "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3])); - } else if constexpr (std::is_same::value) { - float *c = reinterpret_cast(&frag_c); - asm volatile( - "mma.sync.aligned.m16n8k16.row.col.f32.e4m3.e4m3.f32 " - "{%0,%1,%2,%3}, {%4,%5}, {%6}, {%7,%8,%9,%10};\n" - : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3]) - : "r"(b[0]), "r"(b2[0]), "r"(a[0]), "f"(c[0]), "f"(c[1]), "f"(c[2]), - "f"(c[3])); - } else if constexpr (std::is_same::value) { - int32_t *c = reinterpret_cast(&frag_c); - asm volatile( - "mma.sync.aligned.m16n8k16.row.col.s32.s8.s8.s32.satfinite " - "{%0,%1,%2,%3}, {%4,%5}, {%6}, {%7,%8,%9,%10};\n" - : "=r"(c[0]), "=r"(c[1]), "=r"(c[2]), "=r"(c[3]) - : "r"(b[0]), "r"(b2[0]), "r"(a[0]), "r"(c[0]), "r"(c[1]), "r"(c[2]), - "r"(c[3])); - } - } else { - if constexpr (std::is_same::value) { - float *c = reinterpret_cast(&frag_c); - asm volatile( - "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 " - "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n" - : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3]) - : "r"(b[0]), "r"(b2[0]), "r"(b[1]), "r"(b2[1]), "r"(a[0]), "r"(a[1]), - "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3])); - } else if constexpr (std::is_same::value) { - int32_t *c = reinterpret_cast(&frag_c); -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 750 - asm volatile( - "mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32.satfinite " - "{%0,%1}, {%2}, {%3}, {%4,%5};\n" - : "=r"(c[0]), "=r"(c[1]) - : "r"(b[0]), "r"(a[0]), "r"(c[0]), "r"(c[1])); - asm volatile( - "mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32.satfinite " - "{%0,%1}, {%2}, {%3}, {%4,%5};\n" - : "=r"(c[2]), "=r"(c[3]) - : "r"(b2[1]), "r"(a[0]), "r"(c[2]), "r"(c[3])); - asm volatile( - "mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32.satfinite " - "{%0,%1}, {%2}, {%3}, {%4,%5};\n" - : "=r"(c[0]), "=r"(c[1]) - : "r"(b[0]), "r"(a[1]), "r"(c[0]), "r"(c[1])); - asm volatile( - "mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32.satfinite " - "{%0,%1}, {%2}, {%3}, {%4,%5};\n" - : "=r"(c[2]), "=r"(c[3]) - : "r"(b2[1]), "r"(a[1]), "r"(c[2]), "r"(c[3])); -#else - asm volatile( - "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32.satfinite " - "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n" - : "=r"(c[0]), "=r"(c[1]), "=r"(c[2]), "=r"(c[3]) - : "r"(b[0]), "r"(b2[0]), "r"(b[1]), "r"(b2[1]), "r"(a[0]), "r"(a[1]), - "r"(c[0]), "r"(c[1]), "r"(c[2]), "r"(c[3])); -#endif - } - } -} - -} // namespace MARLIN_NAMESPACE_NAME diff --git a/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/marlin/marlin_template.h b/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/marlin/marlin_template.h deleted file mode 100644 index 7415dd971..000000000 --- a/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/marlin/marlin_template.h +++ /dev/null @@ -1,1992 +0,0 @@ -// Adapted from InfiniCore 1ab85ef1d3c2bce49edf1fc73a5b48becaf7552e: -// src/infiniop/ops/awq_marlin_gemm/marlin/marlin_template.h. -// Modified for InfiniOps: native launcher integration; original notices retained. -// See repository THIRD_PARTY.txt and LICENSE (Apache-2.0). -/* - * Modified by Neural Magic - * Copyright (C) Marlin.2024 Elias Frantar - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -/* - * Adapted from https://github.com/IST-DASLab/marlin - */ - -#ifndef MARLIN_NAMESPACE_NAME -#define MARLIN_NAMESPACE_NAME marlin -#endif - -#include "../core/scalar_type.hpp" -#include "dequant.h" -#include "marlin.cuh" -#include "marlin_dtypes.cuh" -#include "marlin_mma.h" - -#define STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t) \ - static_assert(std::is_same::value || std::is_same::value, \ - "only float16 and bfloat16 is supported"); - -namespace MARLIN_NAMESPACE_NAME { - -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 750 - -template shared - // fetch pipeline - const bool has_act_order, // whether act_order is enabled - const int group_blocks, // number of consecutive 16x16 blocks - // with a separate quantization scale - const bool is_zp_float // is zero point of float16 type? - > -__global__ void Marlin( - const int4 *__restrict__ A, // fp16 input matrix of shape mxk - const int4 *__restrict__ B, // 4bit quantized weight matrix of shape kxn - int4 *__restrict__ C, // fp16 output buffer of shape mxn - int4 *__restrict__ C_tmp, // fp32 tmp output buffer (for reduce) - const int4 *__restrict__ scales_ptr, // fp16 quantization scales of shape - // (k/groupsize)xn - const int *__restrict__ g_idx, // int32 group indices of shape k - int num_groups, // number of scale groups per output channel - int prob_m, // batch dimension m - int prob_n, // output dimension n - int prob_k, // reduction dimension k - int *locks, // extra global storage for barrier synchronization - bool use_fp32_reduce // whether to use fp32 global reduce -) { -} - -} // namespace marlin - -#else - -// Instruction for loading a full 16x16 matrix fragment of operand A from shared -// memory, directly in tensor core layout. -template -__device__ inline void ldsm(typename MarlinScalarType::FragA &frag_a, - const void *smem_ptr) { - uint32_t *a = reinterpret_cast(&frag_a); - uint32_t smem = static_cast(__cvta_generic_to_shared(smem_ptr)); - if constexpr (count == 4) { - asm volatile( - "ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];\n" - : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) - : "r"(smem)); - } else if constexpr (count == 2) { - asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0,%1}, [%2];\n" - : "=r"(a[0]), "=r"(a[1]) - : "r"(smem)); - } else if constexpr (count == 1) { - asm volatile("ldmatrix.sync.aligned.m8n8.x1.shared.b16 {%0}, [%1];\n" - : "=r"(a[0]) - : "r"(smem)); - } else { - static_assert(count == 1 || count == 2 || count == 4, "invalid count"); - } -} - -// Multiply dequantized values by the corresponding quantization scale; used -// only for grouped quantization. -template -__device__ inline void scale(typename MarlinScalarType::FragB &frag_b, - typename MarlinScalarType::FragS &frag_s, - int i) { - using scalar_t = typename MarlinScalarType::scalar_t; - using scalar_t2 = typename MarlinScalarType::scalar_t2; - scalar_t2 s = MarlinScalarType::num2num2( - reinterpret_cast(&frag_s)[i]); - frag_b[0] = __hmul2(frag_b[0], s); - frag_b[1] = __hmul2(frag_b[1], s); -} - -template -__device__ inline void scale_and_sub( - typename MarlinScalarType::FragB &frag_b, - typename MarlinScalarType::scalar_t s, - typename MarlinScalarType::scalar_t zp) { - using scalar_t = typename MarlinScalarType::scalar_t; - using scalar_t2 = typename MarlinScalarType::scalar_t2; - scalar_t2 s2 = MarlinScalarType::num2num2(s); - scalar_t2 zp2 = MarlinScalarType::num2num2(zp); - frag_b[0] = __hfma2(frag_b[0], s2, __hneg2(zp2)); - frag_b[1] = __hfma2(frag_b[1], s2, __hneg2(zp2)); -} - -template -__device__ inline void sub_zp( - typename MarlinScalarType::FragB &frag_b, - typename MarlinScalarType::scalar_t2 &frag_zp, int i) { - using scalar_t = typename MarlinScalarType::scalar_t; - using scalar_t2 = typename MarlinScalarType::scalar_t2; - scalar_t2 zp = MarlinScalarType::num2num2( - reinterpret_cast(&frag_zp)[i]); - frag_b[0] = __hsub2(frag_b[0], zp); - frag_b[1] = __hsub2(frag_b[1], zp); -} - -// Same as above, but for act_order (each K is multiplied individually) -template -__device__ inline void scale4( - typename MarlinScalarType::FragB &frag_b, - typename MarlinScalarType::FragS &frag_s_1, - typename MarlinScalarType::FragS &frag_s_2, - typename MarlinScalarType::FragS &frag_s_3, - typename MarlinScalarType::FragS &frag_s_4, int i) { - using scalar_t = typename MarlinScalarType::scalar_t; - using scalar_t2 = typename MarlinScalarType::scalar_t2; - - scalar_t2 s_val_1_2; - s_val_1_2.x = reinterpret_cast(&frag_s_1)[i]; - s_val_1_2.y = reinterpret_cast(&frag_s_2)[i]; - - scalar_t2 s_val_3_4; - s_val_3_4.x = reinterpret_cast(&frag_s_3)[i]; - s_val_3_4.y = reinterpret_cast(&frag_s_4)[i]; - - frag_b[0] = __hmul2(frag_b[0], s_val_1_2); - frag_b[1] = __hmul2(frag_b[1], s_val_3_4); -} - -// Given 2 floats multiply by 2 scales (halves) -template -__device__ inline void scale_float( - float *c, typename MarlinScalarType::FragS &s) { - using scalar_t = typename MarlinScalarType::scalar_t; - scalar_t *s_ptr = reinterpret_cast(&s); - c[0] = __fmul_rn(c[0], MarlinScalarType::num2float(s_ptr[0])); - c[1] = __fmul_rn(c[1], MarlinScalarType::num2float(s_ptr[1])); -} - -// Wait until barrier reaches `count`, then lock for current threadblock. -__device__ inline void barrier_acquire(int *lock, int count) { - if (threadIdx.x == 0) { - int state = -1; - do { - // Guarantee that subsequent writes by this threadblock will be visible - // globally. - asm volatile("ld.global.acquire.gpu.b32 %0, [%1];\n" - : "=r"(state) - : "l"(lock)); - } while (state != count); - } - __syncthreads(); -} - -// Release barrier and increment visitation count. -__device__ inline void barrier_release(int *lock, bool reset = false) { - __syncthreads(); - if (threadIdx.x == 0) { - if (reset) { - lock[0] = 0; - return; - } - int val = 1; - // Make sure that all writes since acquiring this barrier are visible - // globally, while releasing the barrier. - asm volatile("fence.acq_rel.gpu;\n"); - asm volatile("red.relaxed.gpu.global.add.s32 [%0], %1;\n" - : - : "l"(lock), "r"(val)); - } -} - -// Wait until value of lock to be negative, and then add 1 -__device__ inline void wait_negative_and_add(int *lock) { - if (threadIdx.x == 0) { - int state = 0; - do { - // Guarantee that subsequent writes by this threadblock will be visible - // globally. - asm volatile("ld.global.acquire.gpu.b32 %0, [%1];\n" - : "=r"(state) - : "l"(lock)); - } while (state >= 0); - atomicAdd(lock, 1); - } - __syncthreads(); -} - -template shared - // fetch pipeline - const int group_blocks, // number of consecutive 16x16 blocks - // with a separate quantization scale - const bool is_zp_float // is zero point of float16 type? - > -__global__ void Marlin( - const int4 *__restrict__ A0, // fp16 input matrix of shape mxk - const int4 *__restrict__ B, // 4bit quantized weight matrix of shape kxn - int4 *__restrict__ C0, // fp16 output buffer of shape mxn - int4 *__restrict__ C_tmp, // fp32 tmp output buffer (for reduce) - const int4 *__restrict__ b_bias_ptr, - // float scales of input matrix, only used when is_a_8bit == true. - // shape (m,) - const float *__restrict__ a_scales_ptr, - // fp16 quantization scales. shape (k/groupsize, n) - const int4 *__restrict__ scales_ptr, - // float global scale (for nvfp4// only) - const float *__restrict__ global_scale_ptr, - // 4bit packed zero-points of shape - // (k/groupsize, n/pack_factor) - const int4 *__restrict__ zp_ptr, - // int32 group indices of shape k - const int *__restrict__ g_idx, - int num_groups, // number of scale groups per output channel - int prob_m, // batch dimension m - int prob_n, // output dimension n - int prob_k, // reduction dimension k - int lda, // A.stride(0), equal to prob_k is A is contiguous - int *locks, // extra global storage for barrier synchronization - bool has_bias, - bool use_atomic_add, // whether to use atomic add to reduce - bool use_fp32_reduce, // whether to use fp32 global reduce - int max_shared_mem) { - // Each threadblock processes one "stripe" of the B matrix with (roughly) the - // same size, which might involve multiple column "slices" (of width 16 * - // `thread_n_blocks`). Stripes are defined as shown in the 3x3 matrix 5 SM - // example: - // 0 1 3 - // 0 2 3 - // 1 2 4 - // While this kind of partitioning makes things somewhat more complicated, it - // ensures good utilization of all SMs for many kinds of shape and GPU - // configurations, while requiring as few slow global cross-threadblock - // reductions as possible. - -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 890 - // FP8 computation is only supported for Ada Lovelace or newer architectures. - if constexpr (a_type_id == vllm::kFE4M3fn.id()) { - return; - } -#endif - -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 750 - // Turing TensorCore only supports fp16 and int8 - if constexpr (a_type_id != vllm::kFloat16.id() && a_type_id != vllm::kS8.id()) { - return; - } -#endif - -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 750 - constexpr auto num_bits = vllm::ScalarType::from_id(b_type_id).size_bits(); - // Disable use_fp16_accum for NVFP4 and cases when group_size == -1 && - // num_bits == 4 - constexpr bool use_fp16_accum = a_type_id == vllm::kFloat16.id() && (!(b_type_id == vllm::kFE2M1f.id() && s_type_id == vllm::kFE4M3fn.id()) && !(group_blocks == -1 && num_bits == 4)); -#else - constexpr bool use_fp16_accum = false; -#endif - using Adtype = MarlinScalarType; - using Cdtype = MarlinScalarType; - const int4 *A = A0; - int4 *C = C0; - - using scalar_t = typename MarlinScalarType::scalar_t; - using scalar_t2 = typename MarlinScalarType::scalar_t2; - using scalar_32bit_t = typename MarlinScalarType::scalar_32bit_t; - - using c_scalar_t = typename MarlinScalarType::scalar_t; - using c_scalar_t2 = typename MarlinScalarType::scalar_t2; - - using FragA = typename MarlinScalarType::FragA; - using FragB = typename MarlinScalarType::FragB; - using FragC = typename MarlinScalarType::FragC; - using FragS = typename MarlinScalarType::FragS; - using FragZP = typename MarlinScalarType::FragZP; - - static constexpr auto a_type = vllm::ScalarType::from_id(a_type_id); - static constexpr auto b_type = vllm::ScalarType::from_id(b_type_id); - static constexpr auto c_type = vllm::ScalarType::from_id(c_type_id); - static constexpr auto s_type = vllm::ScalarType::from_id(s_type_id); - if constexpr (b_type == vllm::kFE2M1f) { - static_assert(s_type == vllm::kFE4M3fn && group_blocks == 1 || s_type == vllm::kFE8M0fnu && group_blocks == 2); - } else if constexpr (std::is_same::value) { - static_assert(s_type == vllm::kBFloat16); - } else if constexpr (std::is_same::value) { - static_assert(s_type == vllm::kFloat16); - } - - constexpr bool is_a_8bit = a_type.size_bits() == 8; - if constexpr (!is_a_8bit) { - static_assert(std::is_same::value); - } - constexpr bool has_zp = b_type == vllm::kU4 || b_type == vllm::kU8; - constexpr bool is_int_type = b_type == vllm::kU4 || b_type == vllm::kU8 || b_type == vllm::kS4 || b_type == vllm::kS8 || b_type == vllm::kU4B8 || b_type == vllm::kU8B128; - // see comments of dequant.h for more details - constexpr bool dequant_skip_flop = is_a_8bit || b_type == vllm::kFE4M3fn || b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn || has_zp && !is_zp_float && !std::is_same::value || has_zp && !is_zp_float && !(b_type == vllm::kU8); - - float global_scale_f32 = 1.0f; - - if constexpr (b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn) { - global_scale_f32 = global_scale_ptr[0]; - } - - constexpr bool has_act_order = group_blocks == 0; - constexpr int m_block_size = m_block_size_8 ? 8 : (16 * thread_m_blocks); - - extern __shared__ int4 sh[]; - float *sh_a_s = reinterpret_cast(sh); - int4 *sh_new = sh + (is_a_8bit ? (4 * thread_m_blocks) : 0); - constexpr int pack_factor = 32 / b_type.size_bits(); - static_assert(thread_m_blocks == 1 || !m_block_size_8); - - // For larger GEMMs we run multiple batchsize 64 versions in parallel for a - // better partitioning with less reductions - int parallel = 1; - if (prob_m > m_block_size) { - parallel = prob_m / m_block_size; - prob_m = m_block_size; - } - - int k_tiles = prob_k / 16 / thread_k_blocks; - int n_tiles = prob_n / 16 / thread_n_blocks; - - int global_mn_tiles = parallel * n_tiles; - int part2_mn_tiles = global_mn_tiles; - int part1_mn_iters = 0; - bool in_part2 = false; - - if (global_mn_tiles > gridDim.x) { - part2_mn_tiles = global_mn_tiles % gridDim.x; - if (part2_mn_tiles * 3 <= gridDim.x) { - part2_mn_tiles += gridDim.x; - } - part1_mn_iters = (global_mn_tiles - part2_mn_tiles) / gridDim.x; - } - - int iters = div_ceil(k_tiles * part2_mn_tiles, gridDim.x); - - if constexpr (!has_act_order && group_blocks != -1) { - if (group_blocks >= thread_k_blocks) { - // Ensure that the number of tiles in each stripe is a multiple of the - // groupsize; this avoids an annoying special case where a stripe starts - // in the middle of group. - iters = (group_blocks / thread_k_blocks) * div_ceil(iters, (group_blocks / thread_k_blocks)); - } - } - - int slice_row = 0; - int slice_col_par = blockIdx.x; - int slice_col; - int slice_iters = k_tiles; // number of threadblock tiles in the current slice - // total number of active threadblocks in the current slice - int slice_count = 1; - // index of threadblock in current slice; numbered bottom to top - int slice_idx = 0; - - int par_id = 0; - int locks_off = 0; - - if (part2_mn_tiles >= gridDim.x) { - // when part2_mn_tiles >= sms - // then there are at most $sms$ conflict tile blocks - locks_off = blockIdx.x; - } else { - locks_off = (iters * blockIdx.x) / k_tiles - 1; - } - - // Compute all information about the current slice which is required for - // synchronization. - bool first_init = true; - auto init_part2_slice = [&]() { - slice_iters = iters * (blockIdx.x + 1) - (k_tiles * slice_col_par + slice_row); - if (slice_iters < 0 || slice_col_par >= part2_mn_tiles) { - slice_iters = 0; - } - if (slice_iters == 0) { - return; - } - if (slice_row + slice_iters > k_tiles) { - slice_iters = k_tiles - slice_row; - } - slice_count = 1; - slice_idx = 0; - int col_first = iters * div_ceil(k_tiles * slice_col_par, iters); - if (col_first <= k_tiles * (slice_col_par + 1)) { - int col_off = col_first - k_tiles * slice_col_par; - slice_count = div_ceil(k_tiles - col_off, iters); - if (col_off > 0) { - slice_count++; - } - int delta_first = iters * blockIdx.x - col_first; - if (delta_first < 0 || (col_off == 0 && delta_first == 0)) { - slice_idx = slice_count - 1; - } else { - slice_idx = slice_count - 1 - delta_first / iters; - if (col_off > 0) { - slice_idx--; - } - } - } - if (part2_mn_tiles >= gridDim.x) { - if (slice_count > 1 && slice_idx == slice_count - 1) { - locks_off++; - } - } else { - locks_off++; - } - - if (first_init && use_atomic_add && slice_count > 1 && slice_idx == 0) { - constexpr int threads_per_m = 16 * thread_n_blocks / 8; - int m_per_thread = div_ceil(thread_m_blocks * 16, threads / threads_per_m); - if (m_block_size_8) { - m_per_thread = div_ceil(8, threads / threads_per_m); - } - for (int i = 0; i < m_per_thread; i++) { - int row = threads / threads_per_m * i + threadIdx.x / threads_per_m; - if (row < prob_m) { - int col = slice_col * 16 * thread_n_blocks / 8 + threadIdx.x % threads_per_m; - C[row * prob_n / 8 + col] = {0, 0, 0, 0}; - } - } - // After write zero to output, write a negative value to lock. - // Every SM that processes the same slice would wait for - // the negative value, and then atomicAdd 1 to it. - // After all SMs are processed, the lock value would back to 0 again. - __syncthreads(); - if (threadIdx.x == 0) { - locks[locks_off] = 1 - slice_count; - } - } - - if (slice_col == n_tiles) { - A += 16 * thread_m_blocks * lda / (is_a_8bit ? 16 : 8); - C += 16 * thread_m_blocks * prob_n / 8; - slice_col = 0; - par_id++; - } - if (is_a_8bit && (first_init || slice_col == 0)) { - __syncthreads(); - int a_s_gl_rd = par_id * 16 * thread_m_blocks + threadIdx.x; - cp_async1_ca_pred(&sh_a_s[threadIdx.x], &a_scales_ptr[a_s_gl_rd], - threadIdx.x < prob_m); - } - }; - - auto init_part1_slice = [&]() { - if (part1_mn_iters) { - part1_mn_iters--; - par_id = slice_col_par / n_tiles; - slice_col = slice_col_par % n_tiles; - slice_iters = k_tiles; - A = A0 + 16 * thread_m_blocks / (is_a_8bit ? 16 : 8) * par_id * lda; - C = C0 + 16 * thread_m_blocks / 8 * par_id * prob_n; - if (is_a_8bit) { - __syncthreads(); - int a_s_gl_rd = par_id * 16 * thread_m_blocks + threadIdx.x; - cp_async1_ca_pred(&sh_a_s[threadIdx.x], &a_scales_ptr[a_s_gl_rd], - threadIdx.x < prob_m); - } - } - }; - - auto init_slice = [&]() { - if (!in_part2 && !part1_mn_iters) { - in_part2 = true; - slice_col_par = (iters * blockIdx.x) / k_tiles; - slice_row = (iters * blockIdx.x) % k_tiles; - slice_col = (slice_col_par + global_mn_tiles - part2_mn_tiles) % n_tiles; - par_id = (slice_col_par + global_mn_tiles - part2_mn_tiles) / n_tiles; - A = A0 + 16 * thread_m_blocks / (is_a_8bit ? 16 : 8) * par_id * lda; - C = C0 + 16 * thread_m_blocks / 8 * par_id * prob_n; - } - if (!in_part2) { - init_part1_slice(); - } else { - init_part2_slice(); - first_init = false; - } - }; - - init_slice(); - - // A sizes/strides - - // stride of the A matrix in global memory - int a_gl_stride = lda / (is_a_8bit ? 16 : 8); - // stride of an A matrix tile in shared memory - constexpr int a_sh_stride = 16 * thread_k_blocks / (is_a_8bit ? 16 : 8); - // delta between subsequent A tiles in global memory - constexpr int a_gl_rd_delta_o = 16 * thread_k_blocks / (is_a_8bit ? 16 : 8); - // between subsequent accesses within a tile - int a_gl_rd_delta_i = a_gl_stride * (threads / a_gl_rd_delta_o); - // between shared memory writes - constexpr int a_sh_wr_delta = a_sh_stride * (threads / a_gl_rd_delta_o); - // within a shared memory tile - constexpr int a_sh_rd_delta_i = a_sh_stride * 16; - // overall size of a tile - constexpr int a_sh_stage = a_sh_stride * m_block_size; - // number of shared write iterations for a tile - constexpr int a_sh_wr_iters = div_ceil(a_sh_stage, a_sh_wr_delta); - - // B sizes/strides - int b_gl_stride = 16 * prob_n / (pack_factor * (is_a_8bit ? 2 : 4)); - constexpr int b_sh_stride = ((thread_n_blocks * 16) * 16 / pack_factor) / (is_a_8bit ? 2 : 4); - constexpr int b_thread_vecs = b_type.size_bits() == 4 ? 1 : 2; - constexpr int b_sh_stride_threads = b_sh_stride / b_thread_vecs; - - int b_gl_rd_delta_o = b_gl_stride * thread_k_blocks / (is_a_8bit ? 2 : 1); - constexpr int b_sh_wr_delta = threads * b_thread_vecs; - constexpr int b_sh_stage = b_sh_stride * thread_k_blocks / (is_a_8bit ? 2 : 1); - constexpr int b_sh_wr_iters = b_sh_stage / b_sh_wr_delta; - - // Scale sizes/strides without act_order - int s_gl_stride = prob_n / (b_type == vllm::kFE2M1f ? 16 : 8); - constexpr int s_sh_stride = 16 * thread_n_blocks / (b_type == vllm::kFE2M1f ? 16 : 8); - constexpr int s_tb_groups = !has_act_order && group_blocks != -1 && group_blocks < thread_k_blocks - ? thread_k_blocks / group_blocks - : 1; - constexpr int s_sh_stage = s_tb_groups * s_sh_stride; - int s_gl_rd_delta = s_gl_stride; - - // Scale size/strides with act_order - constexpr int tb_k = 16 * thread_k_blocks; - constexpr int g_idx_stage = has_act_order ? (tb_k * sizeof(int)) / 16 : 0; - // constexpr int act_s_row_stride = 1; - // int act_s_col_stride = act_s_row_stride * num_groups; - constexpr int act_s_max_num_groups = 32; - int act_s_col_stride = 1; - int act_s_col_warp_stride = act_s_col_stride * 8; - - constexpr int tb_n_warps = thread_n_blocks / (is_a_8bit ? 2 : 4); - int act_s_col_tb_stride = act_s_col_warp_stride * tb_n_warps; - - // Zero-points sizes/strides - int zp_gl_stride = is_zp_float ? prob_n / 8 : (prob_n / pack_factor) / 4; - constexpr int zp_sh_stride = is_zp_float - ? 16 * thread_n_blocks / 8 - : ((16 * thread_n_blocks) / pack_factor) / 4; - constexpr int zp_tb_groups = s_tb_groups; - constexpr int zp_sh_stage = has_zp ? zp_tb_groups * zp_sh_stride : 0; - int zp_gl_rd_delta = zp_gl_stride; - - // Global A read index of current thread. - int a_gl_rd = a_gl_stride * (threadIdx.x / a_gl_rd_delta_o) + (threadIdx.x % a_gl_rd_delta_o); - a_gl_rd += a_gl_rd_delta_o * slice_row; - // Shared write index of current thread. - int a_sh_wr = a_sh_stride * (threadIdx.x / a_gl_rd_delta_o) + (threadIdx.x % a_gl_rd_delta_o); - // Shared read index. - int a_sh_rd = a_sh_stride * ((threadIdx.x % 32) % (16 / (m_block_size_8 ? 2 : 1))) + (threadIdx.x % 32) / (16 / (m_block_size_8 ? 2 : 1)); - a_sh_rd += 2 * ((threadIdx.x / 32) / tb_n_warps) * b_sh_wr_iters; - - int b_gl_rd; - if (threads <= b_sh_stride) { - b_gl_rd = threadIdx.x; - } else { - b_gl_rd = b_gl_stride * (threadIdx.x / b_sh_stride) + (threadIdx.x % b_sh_stride); - } - - b_gl_rd += b_sh_stride * slice_col; - b_gl_rd += b_gl_rd_delta_o * slice_row; - auto b_sh_rd = threadIdx.x * b_thread_vecs; - b_sh_rd += b_sh_rd / b_sh_stride * (b_sh_stride * (b_sh_wr_iters - 1)); - - // For act_order - int slice_k_start = tb_k * slice_row; - int slice_k_finish = slice_k_start + tb_k * slice_iters; - int slice_k_start_shared_fetch = slice_k_start; - int slice_n_offset = act_s_col_tb_stride * slice_col; - - // No act_order - int s_gl_rd; - if constexpr (!has_act_order) { - if constexpr (group_blocks == -1) { - s_gl_rd = s_sh_stride * slice_col + threadIdx.x; - } else if constexpr (group_blocks >= thread_k_blocks) { - s_gl_rd = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) + s_sh_stride * slice_col + threadIdx.x; - } else { - s_gl_rd = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks + threadIdx.x / s_sh_stride) + s_sh_stride * slice_col + threadIdx.x % s_sh_stride; - } - } - auto s_sh_wr = threadIdx.x; - bool s_sh_wr_pred = threadIdx.x < s_sh_stage; - - // Zero-points - int zp_gl_rd; - if constexpr (has_zp) { - if constexpr (group_blocks == -1) { - zp_gl_rd = zp_sh_stride * slice_col + threadIdx.x; - } else if constexpr (group_blocks >= thread_k_blocks) { - zp_gl_rd = zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) + zp_sh_stride * slice_col + threadIdx.x; - } else { - zp_gl_rd = zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks + threadIdx.x / zp_sh_stride) + zp_sh_stride * slice_col + threadIdx.x % zp_sh_stride; - } - } - auto zp_sh_wr = threadIdx.x; - bool zp_sh_wr_pred = zp_sh_stage > 0 && threadIdx.x < zp_sh_stage; - - // We use a different scale layout for grouped and column-wise quantization as - // we scale a `half2` tile in column-major layout in the former and in - // row-major in the latter case. - int s_sh_rd; - if constexpr (is_a_8bit) { - s_sh_rd = 4 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 4); - } else if constexpr (group_blocks != -1) { - s_sh_rd = 8 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 32) / 4; - } else if constexpr (group_blocks == -1 && (m_block_size_8 || (has_zp && !dequant_skip_flop))) { - s_sh_rd = 8 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 32) / 8; - } else { - s_sh_rd = 8 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 32) % 4; - } - - int bias_sh_rd; - if constexpr (m_block_size_8) { - bias_sh_rd = 8 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 32) / 8; - } else { - bias_sh_rd = (is_a_8bit ? 4 : 8) * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 32) % 4; - } - - int bias_sh_wr = threadIdx.x; - int bias_gl_rd = (thread_n_blocks * 16 / 8) * slice_col + threadIdx.x; - - // Zero-points have the same read layout as the scales - // (without column-wise case) - constexpr int num_col_threads = 8; - constexpr int num_row_threads = 4; - constexpr int num_ints_per_thread = 8 / pack_factor; - int zp_sh_rd; - if constexpr (has_zp) { - if constexpr (is_zp_float) { - if constexpr (group_blocks != -1) { - zp_sh_rd = 8 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 32) / 4; - } - } else if (is_a_8bit) { - zp_sh_rd = num_ints_per_thread * num_col_threads * ((threadIdx.x / 32) % tb_n_warps / 2) + num_ints_per_thread * ((threadIdx.x % 32) / num_row_threads); - } else { - zp_sh_rd = num_ints_per_thread * num_col_threads * ((threadIdx.x / 32) % tb_n_warps) + num_ints_per_thread * ((threadIdx.x % 32) / num_row_threads); - } - } - - // Precompute which thread should not read memory in which iterations; this is - // needed if there are more threads than required for a certain tilesize or - // when the batchsize is not a multiple of 16. - bool a_sh_wr_pred[a_sh_wr_iters]; -#pragma unroll - for (int i = 0; i < a_sh_wr_iters; i++) { - a_sh_wr_pred[i] = a_sh_wr_delta * i + a_sh_wr < a_sh_stride * prob_m; - } - - // To ensure that writing and reading A tiles to/from shared memory, the - // latter in fragment format, is fully bank conflict free, we need to use a - // rather fancy XOR-based layout. The key here is that neither reads nor - // writes of the 16-byte `int4` blocks of 8 consecutive threads involve the - // same shared memory banks. Further, it seems (based on NSight-Compute) that - // each warp must also write a consecutive memory segment? - auto transform_a = [&](int i) { - int row = i / a_gl_rd_delta_o; - return a_gl_rd_delta_o * row + (i % a_gl_rd_delta_o) ^ (row % 8); - }; - // Since the computation of this remapping is non-trivial and, due to our main - // loop unrolls, all shared memory accesses are static, we simply precompute - // both transformed reads and writes. - int a_sh_wr_trans[a_sh_wr_iters]; -#pragma unroll - for (int i = 0; i < a_sh_wr_iters; i++) { - a_sh_wr_trans[i] = transform_a(a_sh_wr_delta * i + a_sh_wr); - } - int a_sh_rd_trans[b_sh_wr_iters][thread_m_blocks]; -#pragma unroll - for (int i = 0; i < b_sh_wr_iters; i++) { -#pragma unroll - for (int j = 0; j < thread_m_blocks; j++) { - a_sh_rd_trans[i][j] = transform_a(2 * i + a_sh_rd_delta_i * j + a_sh_rd); - } - } - - // Since B-accesses have non-constant stride they have to be computed at - // runtime; we break dependencies between subsequent accesses with a tile by - // maintining multiple pointers (we have enough registers), a tiny - // optimization. - - // Shared memory storage for global fetch pipelines. - constexpr int sh_red_size = (2 * thread_n_blocks + 1) * 16 * thread_m_blocks; - constexpr int sh_b_size = stages * b_sh_stage; - int4 *sh_b = sh_new; - int4 *sh_red = sh_new; - constexpr int sh_size_b_red_min = (sh_red_size < sh_b_size ? sh_red_size : sh_b_size); - constexpr int sh_size_b_red_max = (sh_red_size > sh_b_size ? sh_red_size : sh_b_size); - constexpr int sh_bias_size = (thread_n_blocks * 16 / 8); - constexpr int sh_b_red_bias_size = sh_size_b_red_max > (sh_size_b_red_min + sh_bias_size) - ? sh_size_b_red_max - : (sh_size_b_red_min + sh_bias_size); - - int4 *sh_bias = sh_new + sh_size_b_red_min; - int4 *sh_g_idx = sh_new + sh_b_red_bias_size; - int4 *sh_zp = sh_g_idx + (stages * g_idx_stage); - constexpr int sh_s_size = has_act_order ? (act_s_max_num_groups * s_sh_stride) - : (stages * s_sh_stage); - int4 *sh_s = sh_zp + (stages * zp_sh_stage); - int4 *sh_a = sh_s + sh_s_size; - - // Register storage for double buffer of shared memory reads. - FragA frag_a[2][thread_m_blocks]; - I4 frag_b_quant[2][b_thread_vecs]; - FragC frag_c[thread_m_blocks][is_a_8bit ? 2 : 4][2]; - FragC frag_c_tmp[thread_m_blocks][is_a_8bit ? 2 : 4][2]; - FragS frag_s[2][4]; // No act-order - FragS frag_bias[2][4]; - FragS act_frag_s[2][4][4]; // For act-order - int frag_qzp[2][num_ints_per_thread]; // Zero-points - FragZP frag_zp; // Zero-points in fp16 - FragZP frag_zpf[2]; // Zero-points in fp16 in HQQ - - if constexpr (is_a_8bit) { -#pragma unroll - for (int j = 0; j < 2; j++) { -#pragma unroll - for (int i = 0; i < thread_m_blocks; i++) { -#pragma unroll - for (int g = 0; g < 4; g++) { - frag_c_tmp[i][j][0][g] = 0.0f; - } - -#pragma unroll - for (int g = 0; g < 4; g++) { - frag_c_tmp[i][j][1][g] = 0.0f; - } - } - } - } - - // Zero accumulators. - auto zero_accums = [&]() { -#pragma unroll - for (int i = 0; i < thread_m_blocks * 4 * 2 * 4; i++) - reinterpret_cast(frag_c)[i] = 0; - }; - - int sh_first_group_id = -1; - int sh_num_groups = -1; - - auto fetch_act_order_scales_to_shared = [&](bool is_async, int first_group_id, - int last_group_id) { - sh_first_group_id = first_group_id; - sh_num_groups = last_group_id - first_group_id + 1; - - if (sh_num_groups > act_s_max_num_groups) { - sh_num_groups = act_s_max_num_groups; - } - - if (sh_first_group_id + sh_num_groups > num_groups) { - sh_num_groups = num_groups - sh_first_group_id; - } - - int row_offset = first_group_id * s_gl_stride; - - if (is_async) { - for (int i = 0; i < sh_num_groups; i++) { - if (threadIdx.x < s_sh_stride) { - cp_async4_pred(&sh_s[(i * s_sh_stride) + threadIdx.x], - &scales_ptr[row_offset + (i * s_gl_stride) + slice_n_offset + threadIdx.x]); - } - } - } else { - for (int i = 0; i < sh_num_groups; i++) { - if (threadIdx.x < s_sh_stride) { - sh_s[(i * s_sh_stride) + threadIdx.x] = scales_ptr[row_offset + (i * s_gl_stride) + slice_n_offset + threadIdx.x]; - } - } - } - }; - // Asynchronously fetch the next A, B and s tile from global to the next - // shared memory pipeline location. - auto fetch_to_shared = [&](int pipe, int a_off, bool pred = true) { - if (pred) { - int4 *sh_a_stage = sh_a + a_sh_stage * pipe; -#pragma unroll - for (int i = 0; i < a_sh_wr_iters; i++) { - cp_async4_pred( - &sh_a_stage[a_sh_wr_trans[i]], - &A[a_gl_rd_delta_i * i + a_gl_rd + a_gl_rd_delta_o * a_off], - a_sh_wr_pred[i]); - } - int4 *sh_b_stage = sh_b + b_sh_stage * pipe; -#pragma unroll - for (int i = 0; i < (b_sh_wr_iters * b_thread_vecs); i++) { - constexpr int count = div_ceil(b_sh_stride, threads); - int b_gl_idx = b_gl_rd + (i % count) * threads + b_gl_stride * (i / count) * div_ceil(threads, b_sh_stride); - - cp_async4(&sh_b_stage[threads * i + threadIdx.x], &B[b_gl_idx]); - } - - b_gl_rd += b_gl_rd_delta_o; - - if constexpr (has_act_order) { - // Fetch g_idx thread-block portion - int full_pipe = a_off; - int cur_k = slice_k_start_shared_fetch + tb_k * full_pipe; - if (cur_k < prob_k && cur_k < slice_k_finish) { - int4 *sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe; - - int4 const *cur_g_idx_stage_ptr = reinterpret_cast(&g_idx[cur_k]); - - if (threadIdx.x < g_idx_stage) { - cp_async4_pred(&sh_g_idx_stage[threadIdx.x], - &cur_g_idx_stage_ptr[threadIdx.x]); - } - } - } else { - if constexpr (group_blocks != -1) { - int4 *sh_s_stage = sh_s + s_sh_stage * pipe; - - // Only fetch scales if this tile starts a new group - if (pipe % div_ceil(group_blocks, thread_k_blocks) == 0) { - if (s_sh_wr_pred) { - cp_async4(&sh_s_stage[s_sh_wr], &scales_ptr[s_gl_rd]); - } - s_gl_rd += s_gl_rd_delta * s_tb_groups; - } - } - - if constexpr (has_zp && group_blocks != -1) { - int4 *sh_zp_stage = sh_zp + zp_sh_stage * pipe; - - // Only fetch zero points if this tile starts a new group - if (pipe % div_ceil(group_blocks, thread_k_blocks) == 0) { - if (zp_sh_wr_pred) { - cp_async4(&sh_zp_stage[zp_sh_wr], &zp_ptr[zp_gl_rd]); - } - zp_gl_rd += zp_gl_rd_delta * zp_tb_groups; - } - } - } - } - // Insert a fence even when we are winding down the pipeline to ensure that - // waiting is also correct at this point. - cp_async_fence(); - }; - - auto fetch_col_zp_to_shared = [&]() { - if (zp_sh_wr_pred) { - cp_async4(&sh_zp[zp_sh_wr], &zp_ptr[zp_gl_rd]); - } - }; - - auto fetch_col_scale_to_shared = [&]() { - if (s_sh_wr_pred) { - cp_async4(&sh_s[s_sh_wr], &scales_ptr[s_gl_rd]); - } - }; - - // Wait until the next thread tile has been loaded to shared memory. - auto wait_for_stage = [&]() { - // We only have `stages - 2` active fetches since we are double buffering - // and can only issue the next fetch when it is guaranteed that the previous - // shared memory load is fully complete (as it may otherwise be - // overwritten). - cp_async_wait(); - __syncthreads(); - }; - - // Load the next sub-tile from the current location in the shared memory pipe - // into the current register buffer. - auto fetch_to_registers = [&](int k, int pipe) { - int4 *sh_a_stage = sh_a + a_sh_stage * pipe; -#pragma unroll - for (int i = 0; i < thread_m_blocks; i++) - ldsm( - frag_a[k % 2][i], &sh_a_stage[a_sh_rd_trans[k % b_sh_wr_iters][i]]); - int4 *sh_b_stage = sh_b + b_sh_stage * pipe; - -#pragma unroll - for (int i = 0; i < b_thread_vecs; i++) { - frag_b_quant[k % 2][i] = *reinterpret_cast( - &sh_b_stage[b_sh_stride * (k % b_sh_wr_iters) + b_sh_rd + i]); - } - }; - - bool is_same_group[stages]; - int same_group_id[stages]; - - auto init_same_group = [&](int pipe) { - if constexpr (!has_act_order) { - return; - } - - int4 *sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe; - int *sh_g_idx_int_ptr = reinterpret_cast(sh_g_idx_stage); - - int group_id_1 = sh_g_idx_int_ptr[0]; - int group_id_2 = sh_g_idx_int_ptr[tb_k - 1]; - - is_same_group[pipe] = group_id_1 == group_id_2; - same_group_id[pipe] = group_id_1; - }; - - auto fetch_scales_to_registers = [&](int k, int full_pipe) { - int pipe = full_pipe % stages; - using IT1 = typename std::conditional_t; - using IT0 = typename std::conditional_t; - constexpr int group_blocks2 = div_ceil(group_blocks, is_a_8bit ? 2 : 1); - - if constexpr (!has_act_order) { - // No act-order case - if constexpr (group_blocks == -1) { - // load only when starting a new slice - if (k == 0 && full_pipe == 0 && dequant_skip_flop) { - reinterpret_cast(&frag_s)[0] = sh_s[s_sh_rd]; - reinterpret_cast(&frag_s)[1] = sh_s[s_sh_rd + 4]; - } - } else if constexpr (group_blocks != -1) { - if constexpr (group_blocks >= thread_k_blocks) { - constexpr int g = group_blocks / thread_k_blocks; - if (pipe % g == 0) { - if (k % b_sh_wr_iters == 0) { - int4 *sh_s_stage = sh_s + s_sh_stage * (g * (pipe / g)); - reinterpret_cast(&frag_s[k % 2])[0] = sh_s_stage[s_sh_rd]; - } else { - reinterpret_cast(&frag_s[1])[0] = reinterpret_cast(&frag_s[0])[0]; - } - } - } else if (group_blocks2 < b_sh_wr_iters || k % b_sh_wr_iters == 0) { - auto warp_id = threadIdx.x / 32; - int warp_row = warp_id / tb_n_warps; - - int k_blocks = b_sh_wr_iters * warp_row + k % b_sh_wr_iters; - int cur_group_id = k_blocks / group_blocks2; - - int4 *sh_s_stage = sh_s + s_sh_stage * pipe; - - if constexpr (b_type_id != vllm::kFE2M1f.id()) { - reinterpret_cast(&frag_s[k % 2])[0] = sh_s_stage[s_sh_rd + cur_group_id * s_sh_stride]; - } else { - reinterpret_cast(&frag_s[k % 2])[0] = reinterpret_cast( - sh_s_stage)[s_sh_rd + cur_group_id * (2 * s_sh_stride)]; - } - } else if (group_blocks >= b_sh_wr_iters) { - if constexpr (b_type_id != vllm::kFE2M1f.id()) { - reinterpret_cast(&frag_s[1])[0] = reinterpret_cast(&frag_s[0])[0]; - } else { - reinterpret_cast(&frag_s[1])[0] = reinterpret_cast(&frag_s[0])[0]; - } - } - } - - return; - } - - // Act-order case - - // Determine K of the "current" thread-block - int cur_k = slice_k_start + tb_k * full_pipe; - if (cur_k >= prob_k || cur_k >= slice_k_finish) { - return; - } - - // Reset (to current thread-block) since we read g_idx portion from the - // shared memory - cur_k = 0; - - // Progress to current iteration - cur_k += 16 * (k % b_sh_wr_iters); - - // Determine "position" inside the thread-block (based on warp and - // thread-id) - auto warp_id = threadIdx.x / 32; - int warp_row = warp_id / tb_n_warps; - int warp_col = warp_id % tb_n_warps; - - cur_k += warp_row * 16 * b_sh_wr_iters; - - auto th_id = threadIdx.x % 32; - cur_k += (th_id % 4) * 2; // Due to tensor-core layout for fp16 B matrix - - int s_col_shift = - /*slice_n_offset +*/ (act_s_col_warp_stride * warp_col) + (th_id / 4) * act_s_col_stride; - - if (is_same_group[pipe]) { - if (k % 2 == 0) { - *(reinterpret_cast(&(act_frag_s[k % 2][0][0]))) = sh_s[(same_group_id[pipe] - sh_first_group_id) * s_sh_stride + s_col_shift]; - } else { - *(reinterpret_cast(&(act_frag_s[k % 2][0][0]))) = *(reinterpret_cast(&(act_frag_s[(k - 1) % 2][0][0]))); - } - - for (int i = 1; i < 4; i++) { - *(reinterpret_cast(&(act_frag_s[k % 2][i][0]))) = *(reinterpret_cast(&(act_frag_s[k % 2][0][0]))); - } - return; - } - - int4 *sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe; - int *sh_g_idx_int_ptr = reinterpret_cast(sh_g_idx_stage); - - constexpr int k_frag_offsets[4] = {0, 1, 8, - 9}; // Tensor core offsets per thread - -#pragma unroll - for (int i = 0; i < 4; i++) { - int actual_k = cur_k + k_frag_offsets[i]; - - int group_id = sh_g_idx_int_ptr[actual_k]; - int rel_group_id = group_id - sh_first_group_id; - - *(reinterpret_cast(&(act_frag_s[k % 2][i][0]))) = sh_s[rel_group_id * s_sh_stride + s_col_shift]; - } - }; - - auto fetch_zp_to_registers = [&](int k, int full_pipe) { - // This code does not handle group_blocks == 0, - // which signifies act_order. - // has_zp implies AWQ, which doesn't have act_order, - static_assert(!has_zp || group_blocks != 0); - - if constexpr (has_zp && !is_zp_float) { - int pipe = full_pipe % stages; - - if constexpr (group_blocks == -1) { - // load only when starting a new slice - if (k == 0 && full_pipe == 0 || is_a_8bit) { -#pragma unroll - for (int i = 0; i < num_ints_per_thread; i++) { - frag_qzp[k % 2][i] = (reinterpret_cast(sh_zp))[zp_sh_rd + i]; - } - } - } else if constexpr (group_blocks >= thread_k_blocks) { - constexpr int g = group_blocks / thread_k_blocks; - if (pipe % g == 0 && k % b_sh_wr_iters == 0 || is_a_8bit) { - int4 *sh_zp_stage = sh_zp + zp_sh_stage * (g * (pipe / g)); -#pragma unroll - for (int i = 0; i < num_ints_per_thread; i++) { - frag_qzp[k % 2][i] = (reinterpret_cast(sh_zp_stage))[zp_sh_rd + i]; - } - } - } else { - auto warp_id = threadIdx.x / 32; - - int warp_row = warp_id / tb_n_warps; - - int k_blocks = b_sh_wr_iters * warp_row + k % b_sh_wr_iters; - int cur_group_id = k_blocks / div_ceil(group_blocks, is_a_8bit ? 2 : 1); - - int4 *sh_zp_stage = sh_zp + zp_sh_stage * pipe; - - sh_zp_stage += cur_group_id * zp_sh_stride; - -#pragma unroll - for (int i = 0; i < num_ints_per_thread; i++) { - frag_qzp[k % 2][i] = (reinterpret_cast(sh_zp_stage))[zp_sh_rd + i]; - } - } - } - - else if constexpr (has_zp && is_zp_float) { - int pipe = full_pipe % stages; - - if constexpr (group_blocks != -1) { - if constexpr (group_blocks >= thread_k_blocks) { - constexpr int g = group_blocks / thread_k_blocks; - if (pipe % g == 0 && k % b_sh_wr_iters == 0) { - int4 *sh_zp_stage = sh_zp + zp_sh_stage * (g * (pipe / g)); - reinterpret_cast(&frag_zpf[k % 2])[0] = sh_zp_stage[zp_sh_rd]; - } - } else if (group_blocks < b_sh_wr_iters || k % b_sh_wr_iters == 0) { - auto warp_id = threadIdx.x / 32; - - int warp_row = warp_id / tb_n_warps; - int k_blocks = b_sh_wr_iters * warp_row + k % b_sh_wr_iters; - int cur_group_id = k_blocks / group_blocks; - - int4 *sh_zp_stage = sh_zp + zp_sh_stage * pipe; - - reinterpret_cast(&frag_zpf[k % 2])[0] = sh_zp_stage[zp_sh_rd + cur_group_id * zp_sh_stride]; - } - } - } - }; - - auto dequant_data = [&](int q, scalar_32bit_t *frag_b_ptr, int zp = 0) { - if constexpr (a_type.size_bits() != b_type.size_bits()) { - if constexpr (is_a_8bit && has_zp) { - sub_zp_and_dequant( - q, frag_b_ptr, zp); - } else { - dequant(q, frag_b_ptr); - } - } - }; - - // Execute the actual tensor core matmul of a sub-tile. - bool is_first_matmul_in_slice = true; - auto matmul = [&](int k, int pipe) { - if (is_a_8bit) { - return; - } - int k2 = k % 2; - constexpr int g = group_blocks > 0 ? div_ceil(group_blocks, thread_k_blocks) : 1; - const bool is_new_zp = (group_blocks == 0) || ((group_blocks > 0) && (group_blocks < b_sh_wr_iters || k == 0)) && (pipe % g == 0) || (group_blocks == -1 && is_first_matmul_in_slice); - if constexpr (has_zp && !is_zp_float) { - if (is_new_zp) { - if constexpr (group_blocks == -1) { - is_first_matmul_in_slice = false; - } - int zp_quant_0, zp_quant_1; - - if constexpr (b_type.size_bits() == 4) { - zp_quant_0 = frag_qzp[k2][0]; - zp_quant_1 = zp_quant_0 >> 8; - } else { - static_assert(b_type.size_bits() == 8); - zp_quant_0 = frag_qzp[k2][0]; - zp_quant_1 = frag_qzp[k2][1]; - } - - dequant_data(zp_quant_0, reinterpret_cast(&frag_zp)); - dequant_data(zp_quant_1, - reinterpret_cast(&frag_zp) + 2); - } - } - if constexpr (!dequant_skip_flop && has_zp && is_zp_float) { - if (is_new_zp) { - reinterpret_cast(&frag_zp)[0] = reinterpret_cast(&frag_zpf[k2])[0]; - } - } - - if constexpr (b_type == vllm::kFE2M1f) { - int s_quant_0 = reinterpret_cast(frag_s[k2])[0]; - int s_quant_1 = reinterpret_cast(frag_s[k2])[1]; - - dequant_fp8_scales( - s_quant_0, reinterpret_cast(&frag_s[k2])); - dequant_fp8_scales( - s_quant_1, reinterpret_cast(&frag_s[k2]) + 2); - } - -// We have the m dimension as the inner loop in order to encourage overlapping -// dequantization and matmul operations. -#pragma unroll - for (int j = 0; j < 4; j++) { - FragB frag_b0; - FragB frag_b1; - int b_quant_0, b_quant_1; - - if constexpr (b_type_id == vllm::kFE2M1f.id()) { - b_quant_1 = frag_b_quant[k2][0][j]; - b_quant_0 = b_quant_1 << 8; - } else if constexpr (b_type.size_bits() == 4) { - b_quant_0 = frag_b_quant[k2][0][j]; - b_quant_1 = b_quant_0 >> 8; - } else { - static_assert(b_type.size_bits() == 8); - int *frag_b_quant_ptr = reinterpret_cast(frag_b_quant[k2]); - b_quant_0 = frag_b_quant_ptr[j * 2 + 0]; - b_quant_1 = frag_b_quant_ptr[j * 2 + 1]; - } - - dequant_data(b_quant_0, reinterpret_cast(&frag_b0)); - dequant_data(b_quant_1, reinterpret_cast(&frag_b1)); - - if constexpr (dequant_skip_flop && has_zp && !is_zp_float && !is_a_8bit) { - sub_zp(frag_b0, frag_zp[j], 0); - sub_zp(frag_b1, frag_zp[j], 1); - } - - // Apply scale to frag_b0 - if constexpr (has_act_order && !is_a_8bit) { - static_assert(group_blocks != -1); - scale4(frag_b0, act_frag_s[k2][0][j], act_frag_s[k2][1][j], - act_frag_s[k2][2][j], act_frag_s[k2][3][j], 0); - scale4(frag_b1, act_frag_s[k2][0][j], act_frag_s[k2][1][j], - act_frag_s[k2][2][j], act_frag_s[k2][3][j], 1); - } else if constexpr (!dequant_skip_flop && has_zp && !is_zp_float && group_blocks == -1 && !is_a_8bit) { - int idx = (threadIdx.x / 4) % 2; - scalar_t2 s2 = Adtype::nums2num2( - reinterpret_cast(&frag_s[j / 2][j % 2 * 2 + 0])[idx], - reinterpret_cast(&frag_s[j / 2][j % 2 * 2 + 1])[idx]); - if (is_new_zp) - frag_zp[j] = __hmul2(frag_zp[j], s2); - scale_and_sub(frag_b0, s2.x, frag_zp[j].x); - scale_and_sub(frag_b1, s2.y, frag_zp[j].y); - } else if constexpr (!dequant_skip_flop && has_zp && group_blocks != -1 && !is_a_8bit) { - if (is_new_zp) - frag_zp[j] = __hmul2(frag_zp[j], - *reinterpret_cast(&frag_s[k2][j])); - scale_and_sub(frag_b0, frag_s[k2][j][0].x, frag_zp[j].x); - scale_and_sub(frag_b1, frag_s[k2][j][0].y, frag_zp[j].y); - } else if constexpr (group_blocks != -1 && !is_a_8bit) { - scale(frag_b0, frag_s[k2][j], 0); - scale(frag_b1, frag_s[k2][j], 1); - } - -#pragma unroll - for (int i = 0; i < thread_m_blocks; i++) { - if constexpr (m_block_size_8) { - mma_trans(frag_a[k2][i], frag_b0, frag_b1, - frag_c[i][j][0]); - } else { - mma(frag_a[k2][i], frag_b0, - frag_c[i][j][0]); - mma(frag_a[k2][i], frag_b1, - frag_c[i][j][1]); - } - } - } - }; - - auto matmul_a8 = [&](int k) { - int k2 = k % 2; -#pragma unroll - for (int j = 0; j < 2; j++) { - FragB frag_b[2]; - - if (is_a_8bit && b_type.size_bits() == 4 && !has_zp) { - dequant_data(frag_b_quant[k2][0][j * 2], - reinterpret_cast(&frag_b)); - dequant_data(frag_b_quant[k2][0][j * 2 + 1], - reinterpret_cast(&frag_b) + 2); - } else if (is_a_8bit && b_type.size_bits() == 4 && has_zp) { - int off = (threadIdx.x / 32) % 2 * 2 + j; - int zp = (frag_qzp[k2][0] >> (off * 8)) & 0xF; - dequant_data(frag_b_quant[k2][0][j * 2], - reinterpret_cast(&frag_b), zp); - zp = (frag_qzp[k2][0] >> (off * 8 + 4)) & 0xF; - dequant_data(frag_b_quant[k2][0][j * 2 + 1], - reinterpret_cast(&frag_b) + 2, zp); - } else { - reinterpret_cast(&frag_b)[0] = reinterpret_cast(&frag_b_quant[k2][j])[0]; - reinterpret_cast(&frag_b)[1] = reinterpret_cast(&frag_b_quant[k2][j])[1]; - } - -#pragma unroll - for (int i = 0; i < thread_m_blocks; i++) { - mma( - frag_a[k2][i], frag_b[0], - (group_blocks == -1 ? frag_c : frag_c_tmp)[i][j][0]); - mma( - frag_a[k2][i], frag_b[1], - (group_blocks == -1 ? frag_c : frag_c_tmp)[i][j][1]); - } - - if constexpr (group_blocks != -1) { - if (group_blocks == 2 || k == 1) { - if constexpr (a_type == vllm::kS8) { - int2 s_vals[2]; - s_vals[0] = { - (int)reinterpret_cast(&frag_s[k2][j * 2][0])[0], - (int)reinterpret_cast(&frag_s[k2][j * 2][0])[1]}; - s_vals[1] = { - (int)reinterpret_cast(&frag_s[k2][j * 2 + 1][0])[0], - (int)reinterpret_cast(&frag_s[k2][j * 2 + 1][0])[1]}; - -#pragma unroll - for (int i = 0; i < thread_m_blocks; i++) { -#pragma unroll - for (int g = 0; g < 4; g++) { - int scale = reinterpret_cast(&s_vals[0])[g % 2]; - *reinterpret_cast(&frag_c[i][j][0][g]) += *reinterpret_cast(&frag_c_tmp[i][j][0][g]) * scale; - frag_c_tmp[i][j][0][g] = 0.0f; - } - -#pragma unroll - for (int g = 0; g < 4; g++) { - int scale = reinterpret_cast(&s_vals[1])[g % 2]; - *reinterpret_cast(&frag_c[i][j][1][g]) += *reinterpret_cast(&frag_c_tmp[i][j][1][g]) * scale; - frag_c_tmp[i][j][1][g] = 0.0f; - } - } - } else { - float2 s_vals[2]; - if constexpr (s_type_id != vllm::kFE8M0fnu.id()) { - static_assert(a_type.size_bits() == 16 || s_type.size_bits() == 16); - s_vals[0] = Cdtype::num22float2(frag_s[k2][j * 2][0]); - s_vals[1] = Cdtype::num22float2(frag_s[k2][j * 2 + 1][0]); - } else { - int32_t *s_vals_int = reinterpret_cast(&s_vals[0]); - int32_t s_vals_e8m0 = *reinterpret_cast(&frag_s[k2][j][0]); - - s_vals_int[0] = (s_vals_e8m0 & 0xFF) << 23; - s_vals_int[1] = (s_vals_e8m0 & 0xFF00) << 15; - s_vals_int[2] = (s_vals_e8m0 & 0xFF0000) << 7; - s_vals_int[3] = (s_vals_e8m0 & 0xFF000000) >> 1; - } - -#pragma unroll - for (int i = 0; i < thread_m_blocks; i++) { -#pragma unroll - for (int g = 0; g < 4; g++) { - float scale = reinterpret_cast(&s_vals[0])[g % 2]; - frag_c[i][j][0][g] += frag_c_tmp[i][j][0][g] * scale; - frag_c_tmp[i][j][0][g] = 0.0f; - } - -#pragma unroll - for (int g = 0; g < 4; g++) { - float scale = reinterpret_cast(&s_vals[1])[g % 2]; - frag_c[i][j][1][g] += frag_c_tmp[i][j][1][g] * scale; - frag_c_tmp[i][j][1][g] = 0.0f; - } - } - } - } - } - } - }; - - // Since we slice across the k dimension of a tile in order to increase the - // number of warps while keeping the n dimension of a tile reasonable, we have - // multiple warps that accumulate their partial sums of the same output - // location; which we have to reduce over in the end. We do in shared memory. - auto thread_block_reduce = [&]() { - constexpr int red_off = threads / b_sh_stride_threads / 2; - if (red_off >= 1) { - auto red_idx = threadIdx.x / b_sh_stride_threads; - constexpr int red_sh_stride = b_sh_stride_threads * (is_a_8bit ? 2 : 4) * 2; - constexpr int red_sh_delta = b_sh_stride_threads; - int red_sh_rd = red_sh_stride * (threadIdx.x / b_sh_stride_threads) + (threadIdx.x % b_sh_stride_threads); - - // Parallel logarithmic shared memory reduction. We make sure to avoid any - // unnecessary read or write iterations, e.g., for two warps we write only - // once by warp 1 and read only once by warp 0. - -#pragma unroll - for (int m_block = 0; m_block < thread_m_blocks; m_block++) { -#pragma unroll - for (int i = red_off; i > 0; i /= 2) { - if (i <= red_idx && red_idx < 2 * i) { -#pragma unroll - for (int j = 0; j < (is_a_8bit ? 2 : 4) * 2; - j += (m_block_size_8 ? 2 : 1)) { - int red_sh_wr = red_sh_delta * j + (red_sh_rd - red_sh_stride * i); - if (i < red_off) { - float *c_rd = reinterpret_cast( - &sh_red[red_sh_delta * j + red_sh_rd]); - float *c_wr = reinterpret_cast(&sh_red[red_sh_wr]); -#pragma unroll - for (int k = 0; k < 4; k++) - reinterpret_cast( - frag_c)[(is_a_8bit ? 2 : 4) * 2 * m_block + j][k] - += c_rd[k] + c_wr[k]; - } - sh_red[red_sh_wr] = reinterpret_cast( - &frag_c)[(is_a_8bit ? 2 : 4) * 2 * m_block + j]; - } - } - __syncthreads(); - } - if (red_idx == 0) { -#pragma unroll - for (int i = 0; i < (is_a_8bit ? 2 : 4) * 2; - i += (m_block_size_8 ? 2 : 1)) { - float *c_rd = reinterpret_cast(&sh_red[red_sh_delta * i + red_sh_rd]); -#pragma unroll - for (int j = 0; j < 4; j++) - reinterpret_cast( - frag_c)[(is_a_8bit ? 2 : 4) * 2 * m_block + i][j] - += c_rd[j]; - } - } - __syncthreads(); - } - } - }; - - // Since multiple threadblocks may process parts of the same column slice, we - // finally have to globally reduce over the results. As the striped - // partitioning minimizes the number of such reductions and our outputs are - // usually rather small, we perform this reduction serially in L2 cache. - auto global_reduce_fp16 = [&](bool first = false, bool last = false) { - // We are very careful here to reduce directly in the output buffer to - // maximize L2 cache utilization in this step. To do this, we write out - // results in FP16 (but still reduce with FP32 compute). - constexpr int active_threads = 32 * tb_n_warps; - if (threadIdx.x < active_threads) { - int c_gl_stride = prob_n / 8; - int c_gl_wr_delta_o = 8 * c_gl_stride * (is_a_8bit ? 2 : 1); - int c_gl_wr_delta_i = 4 * (active_threads / 32); - int c_gl_wr; - if constexpr (m_block_size_8) { - c_gl_wr = c_gl_stride * ((threadIdx.x % 4) * 2) + 4 * (threadIdx.x / 32) + (threadIdx.x % 32) / 8; - c_gl_wr += (2 * thread_n_blocks) * slice_col; - } else { - c_gl_wr = c_gl_stride * ((threadIdx.x % 32) / 4) * (is_a_8bit ? 2 : 1) + 4 * (threadIdx.x / 32) + threadIdx.x % 4; - c_gl_wr += (2 * thread_n_blocks) * slice_col * (is_a_8bit ? 2 : 1); - } - constexpr int c_sh_wr_delta = active_threads; - auto c_sh_wr = threadIdx.x; - - int row = (threadIdx.x % 32) / 4; - - if (!first) { -// Interestingly, doing direct global accesses here really seems to mess up -// the compiler and lead to slowdowns, hence we also use async-copies even -// though these fetches are not actually asynchronous. -#pragma unroll - for (int i = 0; i < (m_block_size_8 ? 2 : thread_m_blocks * 4); i++) { - if constexpr (m_block_size_8) { - cp_async4_pred(&sh_red[c_sh_wr + c_sh_wr_delta * i], - &C[c_gl_wr + i * c_gl_stride + (threadIdx.x % 8) / 4 * c_gl_wr_delta_i], - (threadIdx.x % 4) * 2 + i < prob_m); - } else if constexpr (is_a_8bit) { - int2 *sh_red_int2 = reinterpret_cast(sh_red); - int2 *c_int2 = reinterpret_cast(C); - cp_async2_ca_pred( - &sh_red_int2[c_sh_wr + c_sh_wr_delta * i], - &c_int2[c_gl_wr + c_gl_wr_delta_o * (i / 2) + c_gl_wr_delta_i * (i % 2)], - i < (thread_m_blocks - 1) * 4 || 8 * (i / 2) + row < prob_m); - } else { - cp_async4_pred( - &sh_red[c_sh_wr + c_sh_wr_delta * i], - &C[c_gl_wr + c_gl_wr_delta_o * (i / 2) + c_gl_wr_delta_i * (i % 2)], - i < (thread_m_blocks - 1) * 4 || 8 * (i / 2) + row < prob_m); - } - } - cp_async_fence(); - cp_async_wait<0>(); - } - -#pragma unroll - for (int i = 0; i < (m_block_size_8 ? 2 : thread_m_blocks * 4); i++) { - bool mask = (!m_block_size_8) && (i < (thread_m_blocks - 1) * 4 || 8 * (i / 2) + row < prob_m) || (m_block_size_8) && ((threadIdx.x % 4) * 2 + i < prob_m); - if (mask) { - if (!first) { - c_scalar_t *c_red_f16; - if constexpr (is_a_8bit) { - int2 tmp = reinterpret_cast(sh_red)[c_sh_wr + i * c_sh_wr_delta]; - c_red_f16 = reinterpret_cast(&tmp); - } else { - int4 tmp = sh_red[c_sh_wr + i * c_sh_wr_delta]; - c_red_f16 = reinterpret_cast(&tmp); - } -#pragma unroll - for (int j = 0; j < 2 * (is_a_8bit ? 2 : 4); j++) { - int delta = 0; - if constexpr (m_block_size_8) { - delta = j % 2 == 1 ? -2 : 0; - } - reinterpret_cast( - &frag_c)[(is_a_8bit ? 2 : 4) * 2 * 4 * (i / 4) + 4 * j + (i % 4) + delta] - += Cdtype::num2float(c_red_f16[j]); - } - } - if (!last) { - c_scalar_t c_f16[is_a_8bit ? 4 : 8]; -#pragma unroll - for (int j = 0; j < 2 * (is_a_8bit ? 2 : 4); j++) { - int delta = 0; - if constexpr (m_block_size_8) { - delta = j % 2 == 1 ? -2 : 0; - } - c_f16[j] = Cdtype::float2num(reinterpret_cast( - &frag_c)[(is_a_8bit ? 2 : 4) * 2 * 4 * (i / 4) + 4 * j + (i % 4) + delta]); - } - if constexpr (m_block_size_8) { - C[c_gl_wr + i * c_gl_stride + (threadIdx.x % 8) / 4 * c_gl_wr_delta_i] = *reinterpret_cast(c_f16); - } else if constexpr (is_a_8bit) { - int2 *c_int2 = reinterpret_cast(C); - c_int2[c_gl_wr + c_gl_wr_delta_o * (i / 2) + c_gl_wr_delta_i * (i % 2)] = *reinterpret_cast(c_f16); - } else { - C[c_gl_wr + c_gl_wr_delta_o * (i / 2) + c_gl_wr_delta_i * (i % 2)] = *reinterpret_cast(c_f16); - } - } - } - } - } - }; - - // Globally reduce over threadblocks that compute the same column block. - // We use a tmp C buffer to reduce in full fp32 precision. - auto global_reduce_fp32 = [&](bool first = false, bool last = false) { - constexpr int tb_m = thread_m_blocks * 16; - constexpr int tb_n = thread_n_blocks * 16; - - constexpr int c_size = tb_m * tb_n * sizeof(float) / 16; - - constexpr int active_threads = 32 * tb_n_warps; - bool is_th_active = threadIdx.x < active_threads; - - constexpr int num_floats = thread_m_blocks * (is_a_8bit ? 2 : 4) * 2 * 4; - constexpr int th_size = num_floats * sizeof(float) / 16; - - int c_cur_offset = locks_off * c_size; - - if (!is_th_active) { - return; - } - - if (!first) { - float *frag_c_ptr = reinterpret_cast(&frag_c); -#pragma unroll - for (int k = 0; k < th_size; k += (m_block_size_8 ? 2 : 1)) { - sh_red[threadIdx.x] = C_tmp[c_cur_offset + active_threads * k + threadIdx.x]; - - float *sh_c_ptr = reinterpret_cast(&sh_red[threadIdx.x]); -#pragma unroll - for (int f = 0; f < 4; f++) { - frag_c_ptr[k * 4 + f] += sh_c_ptr[f]; - } - } - } - - if (!last) { - int4 *frag_c_ptr = reinterpret_cast(&frag_c); -#pragma unroll - for (int k = 0; k < th_size; k += (m_block_size_8 ? 2 : 1)) { - C_tmp[c_cur_offset + active_threads * k + threadIdx.x] = frag_c_ptr[k]; - } - } - }; - - // Write out the reduce final result in the correct layout. We only actually - // reshuffle matrix fragments in this step, the reduction above is performed - // in fragment layout. - auto write_result = [&](bool last) { - int c_gl_stride = prob_n / 8; - constexpr int c_sh_stride = 2 * thread_n_blocks + 1; - int c_gl_wr_delta = c_gl_stride * (threads / (2 * thread_n_blocks)); - constexpr int c_sh_rd_delta = c_sh_stride * (threads / (2 * thread_n_blocks)); - - int c_gl_wr = c_gl_stride * (threadIdx.x / (2 * thread_n_blocks)) + (threadIdx.x % (2 * thread_n_blocks)); - c_gl_wr += (2 * thread_n_blocks) * slice_col; - int c_sh_wr; - if constexpr (m_block_size_8) { - c_sh_wr = (8 * c_sh_stride) * ((threadIdx.x % 32) % 4 * 2) + (threadIdx.x % 32) / 4; - c_sh_wr += 64 * (threadIdx.x / 32); - } else { - c_sh_wr = (4 * c_sh_stride) * ((threadIdx.x % 32) / 4) + (threadIdx.x % 32) % 4; - c_sh_wr += (is_a_8bit ? 16 : 32) * (threadIdx.x / 32); - } - - int c_sh_rd = c_sh_stride * (threadIdx.x / (2 * thread_n_blocks)) + (threadIdx.x % (2 * thread_n_blocks)); - - int c_gl_wr_end = c_gl_stride * prob_m; - // We first reorder in shared memory to guarantee the most efficient final - // global write patterns - auto write = [&](int idx, float c0, float c1, FragS &s, FragS &b_bias) { - if constexpr (b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn) { - c0 *= global_scale_f32; - c1 *= global_scale_f32; - } - c_scalar_t2 res = Cdtype::nums2num2(Cdtype::float2num(c0), Cdtype::float2num(c1)); - - // For per-column quantization we finally apply the scale here (only for - // 4-bit) - if constexpr (!has_act_order && group_blocks == -1 && !is_a_8bit && b_type.size_bits() == 4 && (has_zp && dequant_skip_flop || !has_zp)) { - c_scalar_t2 tmp_scale = s[0]; - if constexpr (m_block_size_8) { - tmp_scale = Cdtype::num2num2( - reinterpret_cast(&s[0])[(threadIdx.x % 8) / 4]); - } - res = __hmul2(res, tmp_scale); - } - if (has_bias && last) { - c_scalar_t2 tmp_bias = b_bias[0]; - if constexpr (m_block_size_8) { - tmp_bias = Cdtype::num2num2( - reinterpret_cast(&b_bias[0])[(threadIdx.x % 8) / 4]); - } - res = __hadd2(res, tmp_bias); - } - - if constexpr (m_block_size_8) { - ((c_scalar_t *)sh_red)[idx] = res.x; - ((c_scalar_t *)sh_red)[idx + 8 * c_sh_stride] = res.y; - } else { - ((c_scalar_t2 *)sh_red)[idx] = res; - } - }; - - if (threadIdx.x / 32 < tb_n_warps) { -#pragma unroll - for (int i = 0; i < thread_m_blocks; i++) { -#pragma unroll - for (int j = 0; j < (is_a_8bit ? 2 : 4); j++) { - if constexpr (m_block_size_8) { - int wr = c_sh_wr + 16 * j; - write(wr, frag_c[i][j][0][0], frag_c[i][j][0][1], - frag_s[j / 2][2 * (j % 2) + 0], - frag_bias[j / 2][2 * (j % 2) + 0]); - write(wr + 8, frag_c[i][j][0][2], frag_c[i][j][0][3], - frag_s[j / 2][2 * (j % 2) + 1], - frag_bias[j / 2][2 * (j % 2) + 1]); - } else { - int wr = c_sh_wr + 8 * j; - write(wr + (4 * c_sh_stride) * 0 + 0, frag_c[i][j][0][0], - frag_c[i][j][0][1], frag_s[j / 2][2 * (j % 2) + 0], - frag_bias[j / 2][2 * (j % 2) + 0]); - write(wr + (4 * c_sh_stride) * 8 + 0, frag_c[i][j][0][2], - frag_c[i][j][0][3], frag_s[j / 2][2 * (j % 2) + 0], - frag_bias[j / 2][2 * (j % 2) + 0]); - write(wr + (4 * c_sh_stride) * 0 + 4, frag_c[i][j][1][0], - frag_c[i][j][1][1], frag_s[j / 2][2 * (j % 2) + 1], - frag_bias[j / 2][2 * (j % 2) + 1]); - write(wr + (4 * c_sh_stride) * 8 + 4, frag_c[i][j][1][2], - frag_c[i][j][1][3], frag_s[j / 2][2 * (j % 2) + 1], - frag_bias[j / 2][2 * (j % 2) + 1]); - } - } - c_sh_wr += 16 * (4 * c_sh_stride); - } - } - __syncthreads(); - -#pragma unroll - for (int i = 0; - i < div_ceil(16 * thread_m_blocks, threads / (2 * thread_n_blocks)); - i++) { - if (c_gl_wr < c_gl_wr_end) { - if (use_atomic_add && slice_count > 1) { - c_scalar_t2 *C_half2 = reinterpret_cast(&C[c_gl_wr]); - c_scalar_t2 *sh_red_half2 = reinterpret_cast(&sh_red[c_sh_rd]); -#pragma unroll - for (int a = 0; a < 4; a++) { - atomicAdd(&C_half2[a], sh_red_half2[a]); - } - } else { - C[c_gl_wr] = sh_red[c_sh_rd]; - } - c_gl_wr += c_gl_wr_delta; - c_sh_rd += c_sh_rd_delta; - } - } - __syncthreads(); - }; - - // Start global fetch and register load pipelines. - auto start_pipes = [&]() { - -#pragma unroll - for (int i = 0; i < stages - 1; i++) { - if (has_act_order && i == 0) { - int last_g_idx = slice_k_start + stages * tb_k * 2; - if (last_g_idx >= prob_k) { - last_g_idx = prob_k - 1; - } - fetch_act_order_scales_to_shared(true, g_idx[slice_k_start], - g_idx[last_g_idx]); - } - - if constexpr (has_zp && !is_zp_float && group_blocks == -1) { - if (i == 0) { - fetch_col_zp_to_shared(); - if constexpr (!dequant_skip_flop) { - fetch_col_scale_to_shared(); - } - } - } - fetch_to_shared(i, i, i < slice_iters); - } - - zero_accums(); - wait_for_stage(); - init_same_group(0); - fetch_to_registers(0, 0); - fetch_scales_to_registers(0, 0); - fetch_zp_to_registers(0, 0); - a_gl_rd += a_gl_rd_delta_o * (stages - 1); - if constexpr (has_act_order) { - slice_k_start_shared_fetch += tb_k * (stages - 1); - } - }; - if (slice_iters) { - start_pipes(); - } - - // Main loop. - while (slice_iters) { - // We unroll over both the global fetch and the register load pipeline to - // ensure all shared memory accesses are static. Note that both pipelines - // have even length meaning that the next iteration will always start at - // index 0. - -#pragma unroll - for (int pipe = 0; pipe < stages;) { -#pragma unroll - for (int k = 0; k < b_sh_wr_iters; k++) { - fetch_to_registers(k + 1, pipe % stages); - fetch_scales_to_registers(k + 1, pipe); - fetch_zp_to_registers(k + 1, pipe); - if (k == b_sh_wr_iters - 2) { - fetch_to_shared((pipe + stages - 1) % stages, pipe, - slice_iters >= stages); - pipe++; - wait_for_stage(); - init_same_group(pipe % stages); - } - - if constexpr (!is_a_8bit) { - matmul(k, pipe - (k >= b_sh_wr_iters - 2 ? 1 : 0)); - } else { - static_assert(group_blocks != 0 && group_blocks != 1); - matmul_a8(k); - } - } - slice_iters--; - if (slice_iters == 0) { - break; - } - } - - a_gl_rd += a_gl_rd_delta_o * stages; - - if constexpr (has_act_order) { - slice_k_start += tb_k * stages; - - if (slice_k_start < prob_k) { - slice_k_start_shared_fetch += tb_k * stages; - int first_group_id = g_idx[slice_k_start]; - int last_g_idx = slice_k_start + stages * tb_k * 2; - if (last_g_idx >= prob_k) { - last_g_idx = prob_k - 1; - } - int last_group_id = g_idx[last_g_idx]; - if (last_group_id >= sh_first_group_id + sh_num_groups) { - fetch_act_order_scales_to_shared(false, first_group_id, - last_group_id); - __syncthreads(); - } - } - } - - // Process results and, if necessary, proceed to the next column slice. - // While this pattern may not be the most readable, other ways of writing - // the loop seemed to noticeably worse performance after compilation. - if (slice_iters == 0) { - // convert fp16 accum to fp32 for reduction - if constexpr (use_fp16_accum) { -#pragma unroll - for (int i = 0; i < (thread_m_blocks * (is_a_8bit ? 2 : 4) * 2); i++) { - float *frag_c_part_float = reinterpret_cast(frag_c) + i * 4; - scalar_t *frag_c_part_half = reinterpret_cast(frag_c_part_float); - -#pragma unroll - for (int i = 3; i >= 0; i--) { - frag_c_part_float[i] = Cdtype::num2float(frag_c_part_half[i]); - } - } - } - - if constexpr (is_a_8bit) { - float frag_a_s[2 * thread_m_blocks]; - - for (int i = 0; i < 2 * thread_m_blocks; i++) { - frag_a_s[i] = sh_a_s[i * 8 + (threadIdx.x % 32) / 4]; - } - -#pragma unroll - for (int j = 0; j < 2; j++) { -#pragma unroll - for (int i = 0; i < thread_m_blocks; i++) { -#pragma unroll - for (int g = 0; g < 4; g++) { - float c_val = frag_c[i][j][0][g]; - - if constexpr (a_type == vllm::kS8) { - c_val = __int2float_rn(*reinterpret_cast(&c_val)); - } - float s_val = frag_a_s[i * 2 + g / 2]; - frag_c[i][j][0][g] = c_val * s_val; - } -#pragma unroll - for (int g = 0; g < 4; g++) { - float c_val = frag_c[i][j][1][g]; - - if constexpr (a_type == vllm::kS8) { - c_val = __int2float_rn(*reinterpret_cast(&c_val)); - } - float s_val = frag_a_s[i * 2 + g / 2]; - frag_c[i][j][1][g] = c_val * s_val; - } - } - } - } - - cp_async_wait<0>(); - bool last = slice_idx == slice_count - 1; - // For per-column scales, we only fetch them here in the final step before - // write-out - if constexpr (!has_act_order && group_blocks == -1 && (has_zp && dequant_skip_flop || !has_zp)) { - if (b_type.size_bits() == 8 || (last || use_atomic_add) || is_a_8bit) { - if (s_sh_wr_pred) { - cp_async4(&sh_s[s_sh_wr], &scales_ptr[s_gl_rd]); - } - cp_async_fence(); - } - } - - thread_block_reduce(); - - if (has_bias && last) { - __syncthreads(); - cp_async4_pred(&sh_bias[bias_sh_wr], &b_bias_ptr[bias_gl_rd], - threadIdx.x < 16 * thread_n_blocks / 8); - cp_async_fence(); - } - - if constexpr (!has_act_order && group_blocks == -1 && (has_zp && dequant_skip_flop || !has_zp || is_a_8bit)) { - if constexpr (is_a_8bit) { - cp_async_wait<0>(); - __syncthreads(); - if (threadIdx.x / 32 < tb_n_warps) { - reinterpret_cast(&frag_s)[0] = sh_s[s_sh_rd + 0]; - } - } else if (b_type.size_bits() == 8 || (last || use_atomic_add)) { - cp_async_wait<0>(); - __syncthreads(); - if (threadIdx.x / 32 < tb_n_warps) { - reinterpret_cast(&frag_s)[0] = sh_s[s_sh_rd + 0]; - reinterpret_cast(&frag_s)[1] = sh_s[s_sh_rd + 4]; - if constexpr (m_block_size_8) { - int idx = (threadIdx.x / 4) % 2; - c_scalar_t2 *frag_s_half2 = reinterpret_cast(frag_s); -#pragma unroll - for (int i = 0; i < 8; i++) { - frag_s_half2[i] = Cdtype::num2num2( - reinterpret_cast(&frag_s_half2[i])[idx]); - } - } - } - } - } - - // For 8-bit channelwise, we apply the scale before the global reduction - // that converts the fp32 results to fp16 (so that we avoid possible - // overflow in fp16) - if constexpr (!has_act_order && group_blocks == -1 && is_a_8bit) { -#pragma unroll - for (int j = 0; j < 2; j++) { - float2 aa[2]; - aa[0] = Cdtype::num22float2(frag_s[0][j * 2][0]); - aa[1] = Cdtype::num22float2(frag_s[0][j * 2 + 1][0]); - -#pragma unroll - for (int i = 0; i < thread_m_blocks; i++) { -#pragma unroll - for (int g = 0; g < 4; g++) { - float scale = reinterpret_cast(&aa[0])[g % 2]; - frag_c[i][j][0][g] *= scale; - } - -#pragma unroll - for (int g = 0; g < 4; g++) { - float scale = reinterpret_cast(&aa[1])[g % 2]; - frag_c[i][j][1][g] *= scale; - } - } - } - } else if (!has_act_order && group_blocks == -1 && b_type.size_bits() == 8 && (has_zp && dequant_skip_flop || !has_zp)) { - if (threadIdx.x / 32 < tb_n_warps) { -#pragma unroll - for (int i = 0; i < thread_m_blocks; i++) { -#pragma unroll - for (int j = 0; j < 4; j++) { - scale_float( - reinterpret_cast(&frag_c[i][j][0][0]), - frag_s[j / 2][2 * (j % 2) + 0]); - scale_float( - reinterpret_cast(&frag_c[i][j][0][2]), - frag_s[j / 2][2 * (j % 2) + (m_block_size_8 ? 1 : 0)]); - - if constexpr (!m_block_size_8) { - scale_float( - reinterpret_cast(&frag_c[i][j][1][0]), - frag_s[j / 2][2 * (j % 2) + 1]); - scale_float( - reinterpret_cast(&frag_c[i][j][1][2]), - frag_s[j / 2][2 * (j % 2) + 1]); - } - } - } - } - } - - if (slice_count > 1 && !use_atomic_add) { - // only globally reduce if there is more than one block in a slice - barrier_acquire(&locks[locks_off], slice_idx); - if (use_fp32_reduce) { - global_reduce_fp32(slice_idx == 0, last); - } else { - global_reduce_fp16(slice_idx == 0, last); - } - barrier_release(&locks[locks_off], last); - } - - if (has_bias && last) { - cp_async_wait<0>(); - __syncthreads(); - reinterpret_cast(&frag_bias)[0] = sh_bias[bias_sh_rd]; - if constexpr (!is_a_8bit) { - reinterpret_cast(&frag_bias)[1] = sh_bias[bias_sh_rd + 4]; - } - __syncthreads(); - } - - if (use_atomic_add && slice_count > 1 && slice_idx != 0) { - wait_negative_and_add(&locks[locks_off]); - } - if (last || use_atomic_add) { - // only the last block in a slice actually writes the result - write_result(last); - } - slice_row = 0; - if (!in_part2) { - slice_col_par += gridDim.x; - } else { - slice_col_par++; - slice_col++; - } - is_first_matmul_in_slice = true; - init_slice(); - - if (slice_iters) { - a_gl_rd = a_gl_stride * (threadIdx.x / a_gl_rd_delta_o) + (threadIdx.x % a_gl_rd_delta_o); - a_gl_rd += a_gl_rd_delta_o * slice_row; - b_gl_rd = b_gl_stride * (threadIdx.x / b_sh_stride) + (threadIdx.x % b_sh_stride); - b_gl_rd += b_sh_stride * slice_col + b_gl_rd_delta_o * slice_row; - - bias_gl_rd = (thread_n_blocks * 16 / 8) * slice_col + threadIdx.x; - // Update slice k/n for scales loading - if constexpr (has_act_order) { - slice_k_start = tb_k * slice_row; - slice_k_finish = slice_k_start + tb_k * slice_iters; - slice_k_start_shared_fetch = slice_k_start; - slice_n_offset = act_s_col_tb_stride * slice_col; - } else { - if constexpr (group_blocks == -1) { - s_gl_rd = s_sh_stride * slice_col + threadIdx.x; - zp_gl_rd = zp_sh_stride * slice_col + threadIdx.x; - } else if constexpr (group_blocks >= thread_k_blocks) { - s_gl_rd = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) + s_sh_stride * slice_col + threadIdx.x; - zp_gl_rd = zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) + zp_sh_stride * slice_col + threadIdx.x; - } else { - s_gl_rd = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks + threadIdx.x / s_sh_stride) + s_sh_stride * slice_col + threadIdx.x % s_sh_stride; - zp_gl_rd = zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks + threadIdx.x / zp_sh_stride) + zp_sh_stride * slice_col + threadIdx.x % zp_sh_stride; - } - } - start_pipes(); - } - } - } -} - -} // namespace MARLIN_NAMESPACE_NAME - -#endif diff --git a/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/nvidia/kernel.cuh b/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/nvidia/kernel.cuh index 21e5367b8..161faa64b 100644 --- a/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/nvidia/kernel.cuh +++ b/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/nvidia/kernel.cuh @@ -28,7 +28,7 @@ #endif #include "../core/utils.h" -#include "../marlin/kernel.h" +#include "marlin/kernel.h" #define STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t) \ static_assert(std::is_same::value || std::is_same::value, \ diff --git a/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/nvidia/kernel_selector.h b/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/nvidia/kernel_selector.h deleted file mode 100644 index 5a07f4219..000000000 --- a/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/nvidia/kernel_selector.h +++ /dev/null @@ -1,545 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// Copyright contributors to the vLLM project. -// Generated from InfiniCore 1ab85ef1d3c2bce49edf1fc73a5b48becaf7552e, -// src/infiniop/ops/awq_marlin_gemm/nvidia/generate_kernels.py. -// Modified for InfiniOps: W4A16-only dispatch. See THIRD_PARTY.txt and LICENSE. -if (a_type == vllm::kFloat16 && b_type == vllm::kU4 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == -1 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == true && stages == 4 && group_blocks == -1 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == -1 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 2 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 3 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 4 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == 2 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == true && stages == 4 && group_blocks == 2 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == 2 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 2 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 3 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 4 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == 4 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == true && stages == 4 && group_blocks == 4 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == 4 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 2 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 3 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 4 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == 8 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == true && stages == 4 && group_blocks == 8 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == 8 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 2 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 3 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 4 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == -1 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == true && stages == 4 && group_blocks == -1 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == -1 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 2 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 3 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 4 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == 0 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == true && stages == 4 && group_blocks == 0 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == 0 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 2 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 3 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 4 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == 2 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == true && stages == 4 && group_blocks == 2 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == 2 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 2 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 3 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 4 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == 4 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == true && stages == 4 && group_blocks == 4 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == 4 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 2 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 3 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 4 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == 8 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == true && stages == 4 && group_blocks == 8 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == 8 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 2 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 3 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 4 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == -1 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == true && stages == 4 && group_blocks == -1 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == -1 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 2 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 3 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 4 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == 2 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == true && stages == 4 && group_blocks == 2 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == 2 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 2 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 3 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 4 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == 4 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == true && stages == 4 && group_blocks == 4 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == 4 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 2 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 3 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 4 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == 8 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == true && stages == 4 && group_blocks == 8 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == 8 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 2 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 3 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 4 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == -1 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == true && stages == 4 && group_blocks == -1 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == -1 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 2 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 3 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 4 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == 0 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == true && stages == 4 && group_blocks == 0 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == 0 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 2 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 3 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 4 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == 2 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == true && stages == 4 && group_blocks == 2 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == 2 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 2 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 3 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 4 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == 4 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == true && stages == 4 && group_blocks == 4 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == 4 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 2 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 3 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 4 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == 8 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == true && stages == 4 && group_blocks == 8 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == 8 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 2 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 3 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 4 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) - kernel = Marlin; -else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) - kernel = Marlin; diff --git a/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/nvidia/sm80_kernel_bfloat16_u4_bfloat16.cu b/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/nvidia/sm80_kernel_bfloat16_u4_bfloat16.cu deleted file mode 100644 index c9f344cc8..000000000 --- a/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/nvidia/sm80_kernel_bfloat16_u4_bfloat16.cu +++ /dev/null @@ -1,70 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// Copyright contributors to the vLLM project. -// Generated from InfiniCore 1ab85ef1d3c2bce49edf1fc73a5b48becaf7552e, -// src/infiniop/ops/awq_marlin_gemm/nvidia/generate_kernels.py. -// Modified for InfiniOps: W4A16-only instantiations and relative includes. -// See repository THIRD_PARTY.txt and LICENSE. -#include "../marlin/kernel.h" -#include "../marlin/marlin_template.h" -namespace MARLIN_NAMESPACE_NAME { -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -} diff --git a/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/nvidia/sm80_kernel_bfloat16_u4b8_bfloat16.cu b/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/nvidia/sm80_kernel_bfloat16_u4b8_bfloat16.cu deleted file mode 100644 index f3af34444..000000000 --- a/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/nvidia/sm80_kernel_bfloat16_u4b8_bfloat16.cu +++ /dev/null @@ -1,85 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// Copyright contributors to the vLLM project. -// Generated from InfiniCore 1ab85ef1d3c2bce49edf1fc73a5b48becaf7552e, -// src/infiniop/ops/awq_marlin_gemm/nvidia/generate_kernels.py. -// Modified for InfiniOps: W4A16-only instantiations and relative includes. -// See repository THIRD_PARTY.txt and LICENSE. -#include "../marlin/kernel.h" -#include "../marlin/marlin_template.h" -namespace MARLIN_NAMESPACE_NAME { -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -} diff --git a/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/nvidia/sm80_kernel_float16_u4_float16.cu b/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/nvidia/sm80_kernel_float16_u4_float16.cu deleted file mode 100644 index bf51abb2c..000000000 --- a/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/nvidia/sm80_kernel_float16_u4_float16.cu +++ /dev/null @@ -1,70 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// Copyright contributors to the vLLM project. -// Generated from InfiniCore 1ab85ef1d3c2bce49edf1fc73a5b48becaf7552e, -// src/infiniop/ops/awq_marlin_gemm/nvidia/generate_kernels.py. -// Modified for InfiniOps: W4A16-only instantiations and relative includes. -// See repository THIRD_PARTY.txt and LICENSE. -#include "../marlin/kernel.h" -#include "../marlin/marlin_template.h" -namespace MARLIN_NAMESPACE_NAME { -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -} diff --git a/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/nvidia/sm80_kernel_float16_u4b8_float16.cu b/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/nvidia/sm80_kernel_float16_u4b8_float16.cu deleted file mode 100644 index d35826b7c..000000000 --- a/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/nvidia/sm80_kernel_float16_u4b8_float16.cu +++ /dev/null @@ -1,85 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// Copyright contributors to the vLLM project. -// Generated from InfiniCore 1ab85ef1d3c2bce49edf1fc73a5b48becaf7552e, -// src/infiniop/ops/awq_marlin_gemm/nvidia/generate_kernels.py. -// Modified for InfiniOps: W4A16-only instantiations and relative includes. -// See repository THIRD_PARTY.txt and LICENSE. -#include "../marlin/kernel.h" -#include "../marlin/marlin_template.h" -namespace MARLIN_NAMESPACE_NAME { -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -template __global__ void Marlin(MARLIN_KERNEL_PARAMS); -} diff --git a/src/native/cuda/nvidia/ops/w4_marlin/vendor/gptq_repack/cuda/kernel.cuh b/src/native/cuda/nvidia/ops/w4_marlin/vendor/gptq_repack/cuda/kernel.cuh deleted file mode 100644 index 8b9bed21f..000000000 --- a/src/native/cuda/nvidia/ops/w4_marlin/vendor/gptq_repack/cuda/kernel.cuh +++ /dev/null @@ -1,256 +0,0 @@ -// Adapted from InfiniCore 1ab85ef1d3c2bce49edf1fc73a5b48becaf7552e: -// src/infiniop/ops/gptq_marlin_repack/cuda/kernel.cuh (Marlin/vLLM). -// Modified for InfiniOps: relative include paths; kernel mathematics retained. -// Original licenses and copyright notices: repository THIRD_PARTY.txt and LICENSE. -#include "../../gemm/marlin/marlin.cuh" - -namespace marlin { - -template -__device__ void gptq_marlin_repack_kernel( - uint32_t const *__restrict__ b_q_weight_ptr, - uint32_t const *__restrict__ perm_ptr, uint32_t *__restrict__ out_ptr, - int size_k, int size_n) { - constexpr int pack_factor = 32 / num_bits; - - constexpr int target_tile_n_size = tile_n_size / (is_a_8bit ? 2 : 1); - constexpr int target_tile_k_size = tile_k_size * (is_a_8bit ? 2 : 1); - int k_tiles = size_k / target_tile_k_size; - int n_tiles = size_n / target_tile_n_size; - int block_k_tiles = div_ceil(k_tiles, gridDim.x); - - auto start_k_tile = blockIdx.x * block_k_tiles; - if (start_k_tile >= k_tiles) { - return; - } - - int finish_k_tile = min(start_k_tile + block_k_tiles, k_tiles); - - // Wait until the next thread tile has been loaded to shared memory. - auto wait_for_stage = [&]() { - // We only have `stages - 2` active fetches since we are double buffering - // and can only issue the next fetch when it is guaranteed that the previous - // shared memory load is fully complete (as it may otherwise be - // overwritten). - cp_async_wait(); - __syncthreads(); - }; - - extern __shared__ int4 sh[]; - - constexpr int perm_size = target_tile_k_size / 4; - - int4 *sh_perm_ptr = sh; - int4 *sh_pipe_ptr = sh_perm_ptr; - if constexpr (has_perm) { - sh_pipe_ptr += perm_size; - } - - constexpr int tile_ints = target_tile_k_size / pack_factor; - - constexpr int stage_n_threads = target_tile_n_size / 4; - constexpr int stage_k_threads = has_perm ? target_tile_k_size : tile_ints; - constexpr int stage_size = stage_k_threads * stage_n_threads; - - auto load_perm_to_shared = [&](int k_tile_id) { - int first_k_int4 = (k_tile_id * target_tile_k_size) / 4; - - int4 const *perm_int4_ptr = reinterpret_cast(perm_ptr); - - if (threadIdx.x < perm_size) { - sh_perm_ptr[threadIdx.x] = perm_int4_ptr[first_k_int4 + threadIdx.x]; - } - __syncthreads(); - }; - - auto fetch_to_shared = [&](int pipe, int k_tile_id, int n_tile_id) { - if (n_tile_id >= n_tiles) { - cp_async_fence(); - return; - } - - int first_n = n_tile_id * target_tile_n_size; - - int4 *sh_ptr = sh_pipe_ptr + stage_size * pipe; - - if constexpr (has_perm) { - if (threadIdx.x < stage_size) { - auto k_id = threadIdx.x / stage_n_threads; - auto n_id = threadIdx.x % stage_n_threads; - - uint32_t const *sh_perm_int_ptr = reinterpret_cast(sh_perm_ptr); - - int src_k = sh_perm_int_ptr[k_id]; - int src_k_packed = src_k / pack_factor; - - cp_async4( - &sh_ptr[k_id * stage_n_threads + n_id], - reinterpret_cast(&( - b_q_weight_ptr[src_k_packed * size_n + first_n + (n_id * 4)]))); - } - } else { - if (threadIdx.x < stage_size) { - auto k_id = threadIdx.x / stage_n_threads; - auto n_id = threadIdx.x % stage_n_threads; - - int first_k = k_tile_id * target_tile_k_size; - int first_k_packed = first_k / pack_factor; - - cp_async4(&sh_ptr[k_id * stage_n_threads + n_id], - reinterpret_cast( - &(b_q_weight_ptr[(first_k_packed + k_id) * size_n + first_n + (n_id * 4)]))); - } - } - - cp_async_fence(); - }; - - auto repack_tile = [&](int pipe, int k_tile_id, int n_tile_id) { - if (n_tile_id >= n_tiles) { - return; - } - - auto warp_id = threadIdx.x / 32; - auto th_id = threadIdx.x % 32; - - if (warp_id >= 4) { - return; - } - - int tc_col = th_id / 4; - int tc_row = (th_id % 4) * (is_a_8bit ? 4 : 2); - - constexpr int tc_offsets[4] = {0, 1, 8, 9}; - - int cur_n = (warp_id / (is_a_8bit ? 2 : 1)) * 16 + tc_col; - - constexpr int sh_stride = target_tile_n_size; - constexpr uint32_t mask = (1 << num_bits) - 1; - - int4 *sh_stage_ptr = sh_pipe_ptr + stage_size * pipe; - uint32_t *sh_stage_int_ptr = reinterpret_cast(sh_stage_ptr); - - uint32_t *sh_perm_int_ptr = reinterpret_cast(sh_perm_ptr); - - uint32_t vals[8]; - - if constexpr (has_perm) { - static_assert(!is_a_8bit); - for (int i = 0; i < 4; i++) { - int k_idx = tc_row + tc_offsets[i]; - - uint32_t src_k = sh_perm_int_ptr[k_idx]; - uint32_t src_k_pos = src_k % pack_factor; - - uint32_t b1_val = sh_stage_int_ptr[k_idx * sh_stride + cur_n]; - uint32_t b1_cur_val = (b1_val >> (src_k_pos * num_bits)) & mask; - - uint32_t b2_val = sh_stage_int_ptr[k_idx * sh_stride + cur_n + 8]; - uint32_t b2_cur_val = (b2_val >> (src_k_pos * num_bits)) & mask; - - vals[i] = b1_cur_val; - vals[4 + i] = b2_cur_val; - } - } else { - uint32_t b1_vals[tile_ints]; - uint32_t b2_vals[tile_ints]; - -#pragma unroll - for (int i = 0; i < tile_ints; i++) { - if constexpr (is_a_8bit) { - b1_vals[i] = sh_stage_int_ptr[cur_n + sh_stride * i + (warp_id % 2) * 8]; - } else { - b1_vals[i] = sh_stage_int_ptr[cur_n + sh_stride * i]; - b2_vals[i] = sh_stage_int_ptr[cur_n + 8 + sh_stride * i]; - } - } - -#pragma unroll - for (int i = 0; i < 4; i++) { - int cur_elem = tc_row + (is_a_8bit ? i : tc_offsets[i]); - int cur_int = cur_elem / pack_factor; - int cur_pos = cur_elem % pack_factor; - - vals[i] = (b1_vals[cur_int] >> (cur_pos * num_bits)) & mask; - if constexpr (is_a_8bit) - vals[4 + i] = (b1_vals[cur_int + tile_ints / 2] >> (cur_pos * num_bits)) & mask; - else - vals[4 + i] = (b2_vals[cur_int] >> (cur_pos * num_bits)) & mask; - } - } - - constexpr int tile_size = target_tile_k_size * target_tile_n_size / pack_factor; - int out_offset = (k_tile_id * n_tiles + n_tile_id) * tile_size; - - // Result of: - // https://github.com/NVIDIA/FasterTransformer/blob/main/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h - if constexpr (!is_a_8bit && num_bits == 4) { - int pack_idx[8] = {0, 2, 4, 6, 1, 3, 5, 7}; - - uint32_t res = 0; -#pragma unroll - for (int i = 0; i < 8; i++) { - res |= vals[pack_idx[i]] << (i * 4); - } - - out_ptr[out_offset + th_id * 4 + warp_id] = res; - } else if constexpr (is_a_8bit && num_bits == 4) { - int pack_idx[8] = {0, 4, 1, 5, 2, 6, 3, 7}; - - uint32_t res = 0; -#pragma unroll - for (int i = 0; i < 8; i++) { - res |= vals[pack_idx[i]] << (i * 4); - } - - out_ptr[out_offset + th_id * 4 + warp_id] = res; - } else { - constexpr int pack_idx[4] = {0, 2, 1, 3}; - - uint32_t res1 = 0; - uint32_t res2 = 0; -#pragma unroll - for (int i = 0; i < 4; i++) { - const int ii = is_a_8bit ? i : pack_idx[i]; - res1 |= vals[ii] << (i * 8); - res2 |= vals[4 + ii] << (i * 8); - } - - out_ptr[out_offset + th_id * 8 + (warp_id * 2) + 0] = res1; - out_ptr[out_offset + th_id * 8 + (warp_id * 2) + 1] = res2; - } - }; - - auto start_pipes = [&](int k_tile_id, int n_tile_id) { -#pragma unroll - for (int pipe = 0; pipe < repack_stages - 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 (has_perm) { - 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 < repack_stages; pipe++) { - fetch_to_shared((pipe + repack_stages - 1) % repack_stages, k_tile_id, - n_tile_id + pipe + repack_stages - 1); - repack_tile(pipe, k_tile_id, n_tile_id + pipe); - wait_for_stage(); - } - n_tile_id += repack_stages; - } - } -} - -} // namespace marlin diff --git a/src/native/cuda/ops/add/kernel.h b/src/native/cuda/ops/add/kernel.h index eadda2918..82941cca9 100644 --- a/src/native/cuda/ops/add/kernel.h +++ b/src/native/cuda/ops/add/kernel.h @@ -4,14 +4,13 @@ #include #include #include -#include -#include #include #include #include "base/add.h" #include "common/generic_utils.h" #include "host_range_profiler.h" +#include "native/cuda/device_buffer.h" #include "native/cuda/kernel_commons.cuh" #include "native/cuda/ops/add/kernel.cuh" #include "native/cuda/runtime_utils.h" @@ -27,75 +26,43 @@ class CudaAdd : public Add { const size_t shape_size = ndim_ * sizeof(*d_input_shape_); const size_t strides_size = ndim_ * sizeof(*d_input_strides_); const size_t metadata_size = 3 * (shape_size + strides_size); - if (Backend::GetDevice(&owner_device_) != Backend::kSuccess) - throw std::runtime_error("CudaAdd cannot query metadata owner device"); if (!metadata_size) return; std::vector metadata(metadata_size); - void* allocation = nullptr; - const auto allocation_status = Backend::Malloc(&allocation, metadata_size); - if (allocation_status != Backend::kSuccess || !allocation) { - if (allocation && Backend::Free(allocation) != Backend::kSuccess) - std::terminate(); - throw std::bad_alloc(); - } - d_metadata_ = static_cast(allocation); - - try { - size_t offset = 0; - d_input_shape_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, input_shape_.data(), shape_size); - offset += shape_size; - - d_other_shape_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, other_shape_.data(), shape_size); - offset += shape_size; - - d_out_shape_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, out_shape_.data(), shape_size); - offset += shape_size; - - d_input_strides_ = - reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, input_strides_.data(), - strides_size); - offset += strides_size; - - d_other_strides_ = - reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, other_strides_.data(), - strides_size); - offset += strides_size; - - d_out_strides_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, out_strides_.data(), strides_size); - - if (Backend::Memcpy(d_metadata_, metadata.data(), metadata_size, - Backend::kMemcpyHostToDevice) != Backend::kSuccess) - throw std::runtime_error("CudaAdd metadata copy failed"); - } catch (...) { - if (Backend::Free(d_metadata_) != Backend::kSuccess) std::terminate(); - d_metadata_ = nullptr; - throw; - } + auto* d_metadata_ = metadata_.Allocate(metadata_size); + + size_t offset = 0; + d_input_shape_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, input_shape_.data(), shape_size); + offset += shape_size; + + d_other_shape_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, other_shape_.data(), shape_size); + offset += shape_size; + + d_out_shape_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, out_shape_.data(), shape_size); + offset += shape_size; + + d_input_strides_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, input_strides_.data(), strides_size); + offset += strides_size; + + d_other_strides_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, other_strides_.data(), strides_size); + offset += strides_size; + + d_out_strides_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, out_strides_.data(), strides_size); + + if (Backend::Memcpy(d_metadata_, metadata.data(), metadata_size, + Backend::kMemcpyHostToDevice) != Backend::kSuccess) + throw std::runtime_error("CudaAdd metadata copy failed"); } CudaAdd(const Tensor input, const Tensor other, Tensor out) : CudaAdd{input, other, 1.0, out} {} - ~CudaAdd() { - if (!d_metadata_) return; - int previous = -1; - if (Backend::GetDevice(&previous) != Backend::kSuccess) std::terminate(); - const bool switched = previous != owner_device_; - if (switched && Backend::SetDevice(owner_device_) != Backend::kSuccess) - std::terminate(); - const auto freed = Backend::Free(d_metadata_); - if (switched && Backend::SetDevice(previous) != Backend::kSuccess) - std::terminate(); - if (freed != Backend::kSuccess) std::terminate(); - } - void operator()(const Tensor input, const Tensor other, const double alpha, Tensor out) const override { if (output_size_ == 0) { @@ -105,9 +72,8 @@ class CudaAdd : public Add { [[maybe_unused]] HostRangeScope host_range_backend_submit{ HostRangeLayer::kBackendSubmit}; - int block_size = RuntimeUtils::GetOptimalBlockSize(); DispatchFunc( - {static_cast(out_type_), block_size}, + {static_cast(out_type_), block_size_}, [&](auto list_tag) { using T = TypeMapType(list_tag)>; constexpr int kBlockSize = ListGet<1>(list_tag); @@ -115,7 +81,7 @@ class CudaAdd : public Add { auto cuda_stream = static_cast(stream_ ? stream_ : 0); dim3 blockDims( - std::min(static_cast(block_size), output_size_)); + std::min(static_cast(block_size_), output_size_)); dim3 gridDims(utils::CeilDiv(output_size_, blockDims.x)); T* d_out = reinterpret_cast(out.data()); @@ -133,8 +99,9 @@ class CudaAdd : public Add { } private: - std::byte* d_metadata_{nullptr}; - int owner_device_{-1}; + const int block_size_ = + RuntimeUtils::GetOptimalBlockSize(); + CudaDeviceBuffer metadata_; Tensor::Size* d_input_shape_{nullptr}; diff --git a/src/native/cuda/ops/copy/kernel.h b/src/native/cuda/ops/copy/kernel.h index 004e7c878..c91e70f82 100644 --- a/src/native/cuda/ops/copy/kernel.h +++ b/src/native/cuda/ops/copy/kernel.h @@ -4,8 +4,6 @@ #include #include #include -#include -#include #include #include @@ -13,6 +11,7 @@ #include "common/generic_utils.h" #include "data_type.h" #include "dispatcher.h" +#include "native/cuda/device_buffer.h" #include "native/cuda/kernel_commons.cuh" #include "native/cuda/ops/copy/kernel.cuh" #include "native/cuda/runtime_utils.h" @@ -24,64 +23,34 @@ class CudaCopy : public Copy { public: CudaCopy(const Tensor input, const bool non_blocking, Tensor out) : Copy{input, non_blocking, out} { - if (Backend::GetDevice(&owner_device_) != Backend::kSuccess) - throw std::runtime_error("CudaCopy cannot query metadata owner device"); if (!output_size_ || (is_input_contiguous_ && is_out_contiguous_)) return; size_t shape_size = ndim_ * sizeof(*d_input_shape_); size_t strides_size = ndim_ * sizeof(*d_input_strides_); const size_t metadata_size = 2 * (shape_size + strides_size); std::vector metadata(metadata_size); - void* allocation = nullptr; - const auto allocation_status = Backend::Malloc(&allocation, metadata_size); - if (allocation_status != Backend::kSuccess || !allocation) { - if (allocation && Backend::Free(allocation) != Backend::kSuccess) - std::terminate(); - throw std::bad_alloc(); - } - d_metadata_ = static_cast(allocation); - - try { - size_t offset = 0; - d_input_shape_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, input_shape_.data(), shape_size); - offset += shape_size; - - d_out_shape_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, out_shape_.data(), shape_size); - offset += shape_size; - - d_input_strides_ = - reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, input_strides_.data(), - strides_size); - offset += strides_size; - - d_out_strides_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, out_strides_.data(), strides_size); - - const bool copied = - Backend::Memcpy(d_metadata_, metadata.data(), metadata_size, - Backend::kMemcpyHostToDevice) == Backend::kSuccess; - if (!copied) throw std::runtime_error("CudaCopy metadata copy failed"); - } catch (...) { - if (Backend::Free(d_metadata_) != Backend::kSuccess) std::terminate(); - d_metadata_ = nullptr; - throw; - } - } + auto* d_metadata_ = metadata_.Allocate(metadata_size); + + size_t offset = 0; + d_input_shape_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, input_shape_.data(), shape_size); + offset += shape_size; + + d_out_shape_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, out_shape_.data(), shape_size); + offset += shape_size; + + d_input_strides_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, input_strides_.data(), strides_size); + offset += strides_size; + + d_out_strides_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, out_strides_.data(), strides_size); - ~CudaCopy() { - if (!d_metadata_) return; - int previous = -1; - if (Backend::GetDevice(&previous) != Backend::kSuccess) std::terminate(); - const bool switched = previous != owner_device_; - if (switched && Backend::SetDevice(owner_device_) != Backend::kSuccess) - std::terminate(); - const auto freed = Backend::Free(d_metadata_); - if (switched && Backend::SetDevice(previous) != Backend::kSuccess) - std::terminate(); - if (freed != Backend::kSuccess) std::terminate(); + const bool copied = + Backend::Memcpy(d_metadata_, metadata.data(), metadata_size, + Backend::kMemcpyHostToDevice) == Backend::kSuccess; + if (!copied) throw std::runtime_error("CudaCopy metadata copy failed"); } void operator()(const Tensor input, const bool /*non_blocking*/, @@ -90,13 +59,11 @@ class CudaCopy : public Copy { auto cuda_stream = static_cast(stream_ ? stream_ : 0); - int block_size = std::min( - RuntimeUtils::GetOptimalBlockSize(), 1024); - dim3 block(std::min(static_cast(block_size), output_size_)); + dim3 block(std::min(static_cast(block_size_), output_size_)); dim3 grid(utils::CeilDiv(output_size_, block.x)); DispatchFunc>( - {static_cast(out_type_), block_size}, + {static_cast(out_type_), block_size_}, [&](auto list_tag) { using T = TypeMapType(list_tag)>; constexpr int kBlockSize = ListGet<1>(list_tag); @@ -111,8 +78,9 @@ class CudaCopy : public Copy { } private: - std::byte* d_metadata_{nullptr}; - int owner_device_{-1}; + const int block_size_ = + std::min(RuntimeUtils::GetOptimalBlockSize(), 1024); + CudaDeviceBuffer metadata_; Tensor::Size* d_input_shape_{nullptr}; diff --git a/src/native/cuda/ops/embedding/kernel.h b/src/native/cuda/ops/embedding/kernel.h index 837559d78..a0ce60977 100644 --- a/src/native/cuda/ops/embedding/kernel.h +++ b/src/native/cuda/ops/embedding/kernel.h @@ -4,8 +4,6 @@ #include #include #include -#include -#include #include #include @@ -13,6 +11,7 @@ #include "common/generic_utils.h" #include "data_type.h" #include "dispatcher.h" +#include "native/cuda/device_buffer.h" #include "native/cuda/kernel_commons.cuh" #include "native/cuda/ops/embedding/kernel.cuh" @@ -34,65 +33,38 @@ class CudaEmbedding : public Embedding { is_out_contiguous_{out.IsContiguous()}, weight_row_stride_{weight.stride(0)}, weight_col_stride_{weight.stride(1)} { - if (Backend::GetDevice(&owner_device_) != Backend::kSuccess) - throw std::runtime_error( - "CudaEmbedding cannot query metadata owner device"); const size_t input_shape_size = input_ndim_ * sizeof(*d_input_shape_); const size_t input_strides_size = input_ndim_ * sizeof(*d_input_strides_); const size_t out_shape_size = out_ndim_ * sizeof(*d_out_shape_); const size_t out_strides_size = out_ndim_ * sizeof(*d_out_strides_); const size_t metadata_size = input_shape_size + input_strides_size + out_shape_size + out_strides_size; - try { - if (num_indices_ != 0 && !(is_input_contiguous_ && is_out_contiguous_)) { - std::vector metadata(metadata_size); - void* allocation = nullptr; - const auto status = Backend::Malloc(&allocation, metadata_size); - if (status != Backend::kSuccess || !allocation) { - if (allocation && Backend::Free(allocation) != Backend::kSuccess) - std::terminate(); - throw std::bad_alloc(); - } - d_metadata_ = static_cast(allocation); - size_t offset = 0; - d_input_shape_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, input_shape_.data(), - input_shape_size); - offset += input_shape_size; - d_input_strides_ = - reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, input_strides_.data(), - input_strides_size); - offset += input_strides_size; - d_out_shape_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, out_shape_.data(), - out_shape_size); - offset += out_shape_size; - d_out_strides_ = - reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, out_strides_.data(), - out_strides_size); - if (Backend::Memcpy(d_metadata_, metadata.data(), metadata_size, - Backend::kMemcpyHostToDevice) != Backend::kSuccess) - throw std::runtime_error("CudaEmbedding metadata copy failed"); - } - if (max_norm.has_value() && vocab_size_ > 0) { - void* allocation = nullptr; - const auto status = - Backend::Malloc(&allocation, vocab_size_ * sizeof(*d_visited_)); - if (status != Backend::kSuccess || !allocation) { - if (allocation && Backend::Free(allocation) != Backend::kSuccess) - std::terminate(); - throw std::bad_alloc(); - } - d_visited_ = static_cast(allocation); - } - } catch (...) { - if (d_visited_ && Backend::Free(d_visited_) != Backend::kSuccess) - std::terminate(); - if (d_metadata_ && Backend::Free(d_metadata_) != Backend::kSuccess) - std::terminate(); - throw; + if (num_indices_ != 0 && !(is_input_contiguous_ && is_out_contiguous_)) { + std::vector metadata(metadata_size); + auto* d_metadata_ = metadata_.Allocate(metadata_size); + size_t offset = 0; + d_input_shape_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, input_shape_.data(), + input_shape_size); + offset += input_shape_size; + d_input_strides_ = + reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, input_strides_.data(), + input_strides_size); + offset += input_strides_size; + d_out_shape_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, out_shape_.data(), out_shape_size); + offset += out_shape_size; + d_out_strides_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, out_strides_.data(), + out_strides_size); + if (Backend::Memcpy(d_metadata_, metadata.data(), metadata_size, + Backend::kMemcpyHostToDevice) != Backend::kSuccess) + throw std::runtime_error("CudaEmbedding metadata copy failed"); + } + if (max_norm.has_value() && vocab_size_ > 0) { + d_visited_ = reinterpret_cast( + visited_.Allocate(vocab_size_ * sizeof(*d_visited_))); } } @@ -109,22 +81,6 @@ class CudaEmbedding : public Embedding { : CudaEmbedding(input, weight, padding_idx, std::nullopt, 2.0, scale_grad_by_freq, sparse, out) {} - ~CudaEmbedding() override { - if (!d_metadata_ && !d_visited_) return; - int previous = -1; - if (Backend::GetDevice(&previous) != Backend::kSuccess) std::terminate(); - const bool switched = previous != owner_device_; - if (switched && Backend::SetDevice(owner_device_) != Backend::kSuccess) - std::terminate(); - const auto visited_freed = - !d_visited_ || Backend::Free(d_visited_) == Backend::kSuccess; - const auto metadata_freed = - !d_metadata_ || Backend::Free(d_metadata_) == Backend::kSuccess; - if (switched && Backend::SetDevice(previous) != Backend::kSuccess) - std::terminate(); - if (!visited_freed || !metadata_freed) std::terminate(); - } - void operator()(const Tensor input, const Tensor weight, const std::optional /*padding_idx*/, const std::optional max_norm, const double norm_type, @@ -147,14 +103,7 @@ class CudaEmbedding : public Embedding { vocab_size_); } - size_t block_size = 256; - if (embedding_dim_ <= 64) { - block_size = 512; - } else if (embedding_dim_ >= 1024) { - block_size = 128; - } - - size_t grid_size = utils::CeilDiv(num_indices_, block_size); + size_t grid_size = utils::CeilDiv(num_indices_, block_size_); DispatchFunc, ConcatType, ReducedFloatTypes>>( @@ -177,7 +126,7 @@ class CudaEmbedding : public Embedding { } EmbeddingKernel - <<>>( + <<>>( reinterpret_cast(out.data()), reinterpret_cast(input.data()), reinterpret_cast(weight.data()), num_indices_, @@ -190,6 +139,8 @@ class CudaEmbedding : public Embedding { } private: + const size_t block_size_ = + embedding_dim_ <= 64 ? 512 : (embedding_dim_ >= 1024 ? 128 : 256); Tensor::Size input_ndim_{0}; Tensor::Size out_ndim_{0}; @@ -204,8 +155,8 @@ class CudaEmbedding : public Embedding { int* d_visited_{nullptr}; - std::byte* d_metadata_{nullptr}; - int owner_device_{-1}; + CudaDeviceBuffer metadata_; + CudaDeviceBuffer visited_; Tensor::Size* d_input_shape_{nullptr}; diff --git a/src/native/cuda/ops/fill/kernel.h b/src/native/cuda/ops/fill/kernel.h index 1bea676a0..b1a60249b 100644 --- a/src/native/cuda/ops/fill/kernel.h +++ b/src/native/cuda/ops/fill/kernel.h @@ -5,8 +5,6 @@ #include #include #include -#include -#include #include #include @@ -14,6 +12,7 @@ #include "common/generic_utils.h" #include "data_type.h" #include "dispatcher.h" +#include "native/cuda/device_buffer.h" #include "native/cuda/kernel_commons.cuh" #include "native/cuda/ops/fill/kernel.cuh" #include "native/cuda/runtime_utils.h" @@ -36,8 +35,6 @@ class CudaFill : public Fill { "`CudaFill` requires `input` and `out` on the same device."); assert(!out.HasBroadcastDim() && "`CudaFill` output must not have broadcast dimensions."); - if (Backend::GetDevice(&owner_device_) != Backend::kSuccess) - throw std::runtime_error("CudaFill cannot query metadata owner device"); if (output_size_ == 0 || is_out_contiguous_) return; size_t shape_size = ndim_ * sizeof(*d_out_shape_); @@ -45,45 +42,19 @@ class CudaFill : public Fill { const size_t metadata_size = shape_size + strides_size; std::vector metadata(metadata_size); - void* allocation = nullptr; - const auto allocation_status = Backend::Malloc(&allocation, metadata_size); - if (allocation_status != Backend::kSuccess || !allocation) { - if (allocation && Backend::Free(allocation) != Backend::kSuccess) - std::terminate(); - throw std::bad_alloc(); - } - d_metadata_ = static_cast(allocation); - - try { - size_t offset = 0; - d_out_shape_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, out_shape_.data(), shape_size); - offset += shape_size; - - d_out_strides_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, out_strides_.data(), strides_size); - - if (Backend::Memcpy(d_metadata_, metadata.data(), metadata_size, - Backend::kMemcpyHostToDevice) != Backend::kSuccess) - throw std::runtime_error("CudaFill metadata copy failed"); - } catch (...) { - if (Backend::Free(d_metadata_) != Backend::kSuccess) std::terminate(); - d_metadata_ = nullptr; - throw; - } - } + auto* d_metadata_ = metadata_.Allocate(metadata_size); - ~CudaFill() override { - if (!d_metadata_) return; - int previous = -1; - if (Backend::GetDevice(&previous) != Backend::kSuccess) std::terminate(); - const bool switched = previous != owner_device_; - if (switched && Backend::SetDevice(owner_device_) != Backend::kSuccess) - std::terminate(); - const auto freed = Backend::Free(d_metadata_); - if (switched && Backend::SetDevice(previous) != Backend::kSuccess) - std::terminate(); - if (freed != Backend::kSuccess) std::terminate(); + size_t offset = 0; + d_out_shape_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, out_shape_.data(), shape_size); + offset += shape_size; + + d_out_strides_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, out_strides_.data(), strides_size); + + if (Backend::Memcpy(d_metadata_, metadata.data(), metadata_size, + Backend::kMemcpyHostToDevice) != Backend::kSuccess) + throw std::runtime_error("CudaFill metadata copy failed"); } void operator()(const Tensor input, const double value, @@ -95,13 +66,11 @@ class CudaFill : public Fill { auto cuda_stream = static_cast(stream_ ? stream_ : 0); - int block_size = std::min( - RuntimeUtils::GetOptimalBlockSize(), 1024); - dim3 block(std::min(static_cast(block_size), output_size_)); + dim3 block(std::min(static_cast(block_size_), output_size_)); dim3 grid(utils::CeilDiv(output_size_, block.x)); DispatchFunc>( - {static_cast(out_type_), block_size}, + {static_cast(out_type_), block_size_}, [&](auto list_tag) { using T = TypeMapType(list_tag)>; constexpr int kBlockSize = ListGet<1>(list_tag); @@ -124,15 +93,15 @@ class CudaFill : public Fill { } private: + const int block_size_ = + std::min(RuntimeUtils::GetOptimalBlockSize(), 1024); Tensor::Size output_size_{0}; Tensor::Size ndim_{0}; bool is_out_contiguous_{false}; - std::byte* d_metadata_{nullptr}; - - int owner_device_{-1}; + CudaDeviceBuffer metadata_; Tensor::Size* d_out_shape_{nullptr}; diff --git a/src/native/cuda/ops/gelu/kernel.h b/src/native/cuda/ops/gelu/kernel.h index 3fbe5200e..f5b8ccd89 100644 --- a/src/native/cuda/ops/gelu/kernel.h +++ b/src/native/cuda/ops/gelu/kernel.h @@ -5,8 +5,6 @@ #include #include #include -#include -#include #include #include @@ -14,6 +12,7 @@ #include "common/generic_utils.h" #include "data_type.h" #include "dispatcher.h" +#include "native/cuda/device_buffer.h" #include "native/cuda/kernel_commons.cuh" #include "native/cuda/ops/gelu_infinilm/kernel.cuh" #include "native/cuda/ops/gelutanh_infinilm/kernel.cuh" @@ -34,8 +33,6 @@ class CudaGelu : public Gelu { "`Gelu` input and output devices must match"); assert(!out.HasBroadcastDim() && "`Gelu` output must not have broadcasted dimensions"); - if (Backend::GetDevice(&owner_device_) != Backend::kSuccess) - throw std::runtime_error("CudaGelu cannot query metadata owner device"); if (output_size_ == 0 || (is_input_contiguous_ && is_out_contiguous_)) return; @@ -44,55 +41,27 @@ class CudaGelu : public Gelu { const size_t metadata_size = 2 * (shape_size + strides_size); std::vector metadata(metadata_size); - void* allocation = nullptr; - const auto allocation_status = Backend::Malloc(&allocation, metadata_size); - if (allocation_status != Backend::kSuccess || !allocation) { - if (allocation && Backend::Free(allocation) != Backend::kSuccess) - std::terminate(); - throw std::bad_alloc(); - } - d_metadata_ = static_cast(allocation); - - try { - size_t offset = 0; - d_input_shape_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, input_shape_.data(), shape_size); - offset += shape_size; - - d_out_shape_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, out_shape_.data(), shape_size); - offset += shape_size; - - d_input_strides_ = - reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, input_strides_.data(), - strides_size); - offset += strides_size; - - d_out_strides_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, out_strides_.data(), strides_size); - - if (Backend::Memcpy(d_metadata_, metadata.data(), metadata_size, - Backend::kMemcpyHostToDevice) != Backend::kSuccess) - throw std::runtime_error("CudaGelu metadata copy failed"); - } catch (...) { - if (Backend::Free(d_metadata_) != Backend::kSuccess) std::terminate(); - d_metadata_ = nullptr; - throw; - } - } + auto* d_metadata_ = metadata_.Allocate(metadata_size); + + size_t offset = 0; + d_input_shape_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, input_shape_.data(), shape_size); + offset += shape_size; + + d_out_shape_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, out_shape_.data(), shape_size); + offset += shape_size; + + d_input_strides_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, input_strides_.data(), strides_size); + offset += strides_size; + + d_out_strides_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, out_strides_.data(), strides_size); - ~CudaGelu() override { - if (!d_metadata_) return; - int previous = -1; - if (Backend::GetDevice(&previous) != Backend::kSuccess) std::terminate(); - const bool switched = previous != owner_device_; - if (switched && Backend::SetDevice(owner_device_) != Backend::kSuccess) - std::terminate(); - const auto freed = Backend::Free(d_metadata_); - if (switched && Backend::SetDevice(previous) != Backend::kSuccess) - std::terminate(); - if (freed != Backend::kSuccess) std::terminate(); + if (Backend::Memcpy(d_metadata_, metadata.data(), metadata_size, + Backend::kMemcpyHostToDevice) != Backend::kSuccess) + throw std::runtime_error("CudaGelu metadata copy failed"); } void operator()(const Tensor input, const std::string approximate, @@ -105,13 +74,11 @@ class CudaGelu : public Gelu { auto cuda_stream = static_cast(stream_ ? stream_ : 0); - int block_size = std::min( - RuntimeUtils::GetOptimalBlockSize(), 1024); - dim3 block(std::min(static_cast(block_size), output_size_)); + dim3 block(std::min(static_cast(block_size_), output_size_)); dim3 grid(utils::CeilDiv(output_size_, block.x)); DispatchFunc>( - {static_cast(out_type_), block_size}, + {static_cast(out_type_), block_size_}, [&](auto list_tag) { using T = TypeMapType(list_tag)>; constexpr int kBlockSize = ListGet<1>(list_tag); @@ -138,8 +105,9 @@ class CudaGelu : public Gelu { } private: - std::byte* d_metadata_{nullptr}; - int owner_device_{-1}; + const int block_size_ = + std::min(RuntimeUtils::GetOptimalBlockSize(), 1024); + CudaDeviceBuffer metadata_; Tensor::Size* d_input_shape_{nullptr}; diff --git a/src/native/cuda/ops/mamba_selective_scan/kernel.h b/src/native/cuda/ops/mamba_selective_scan/kernel.h index ed6bdbbda..dd4e3cf9d 100644 --- a/src/native/cuda/ops/mamba_selective_scan/kernel.h +++ b/src/native/cuda/ops/mamba_selective_scan/kernel.h @@ -78,18 +78,41 @@ class CudaMambaSelectiveScan : public MambaSelectiveScan { Tensor state) { try { ValidateMetadata({out, x, dt, b, c, a_log, d, gate, dt_bias, state}); - return true; + return a_log.dtype() == x.dtype() || + Backend::kDeviceType == Device::Type::kNvidia; } catch (const std::invalid_argument&) { return false; } catch (const std::overflow_error&) { return false; } } - using MambaSelectiveScan::MambaSelectiveScan; + CudaMambaSelectiveScan(Tensor out, const Tensor x, const Tensor dt, + const Tensor b, const Tensor c, const Tensor a_log, + const Tensor d, const Tensor gate, + const Tensor dt_bias, Tensor state) + : MambaSelectiveScan(out, x, dt, b, c, a_log, d, gate, dt_bias, state), + prepared_metadata_{out, x, dt, b, c, a_log, d, gate, dt_bias, state} { + if (a_log.dtype() != x.dtype() && + Backend::kDeviceType != Device::Type::kNvidia) + throw std::invalid_argument( + "CUDA MambaSelectiveScan backend does not support mixed F32 " + "parameters"); + } void operator()(Tensor out, const Tensor x, const Tensor dt, const Tensor b, const Tensor c, const Tensor a_log, const Tensor d, const Tensor gate, const Tensor dt_bias, Tensor state) const override { + const std::array tensors{out, x, dt, b, c, + a_log, d, gate, dt_bias, state}; + for (size_t i = 0; i < tensors.size(); ++i) { + const auto& t = tensors[i]; + const auto& cached = prepared_metadata_[i]; + if (t.shape() != cached.shape() || t.strides() != cached.strides() || + t.dtype() != cached.dtype() || t.device() != cached.device() || + (t.numel() && !t.data())) + throw std::invalid_argument( + "MambaSelectiveScan metadata changed or nonempty storage is null"); + } if (!layout_.batch || !layout_.length || !layout_.channels) return; DispatchFunc< List>( @@ -141,6 +164,10 @@ class CudaMambaSelectiveScan : public MambaSelectiveScan { }, "CudaMambaSelectiveScan"); } + + private: + // TensorView owns metadata only; snapshots do not retain tensor allocations. + std::array prepared_metadata_; }; } // namespace infini::ops #endif diff --git a/src/native/cuda/ops/mul/kernel.h b/src/native/cuda/ops/mul/kernel.h index 6f5e5a78f..032d44022 100644 --- a/src/native/cuda/ops/mul/kernel.h +++ b/src/native/cuda/ops/mul/kernel.h @@ -4,13 +4,12 @@ #include #include #include -#include -#include #include #include #include "base/mul.h" #include "common/generic_utils.h" +#include "native/cuda/device_buffer.h" #include "native/cuda/kernel_commons.cuh" #include "native/cuda/ops/mul/kernel.cuh" #include "native/cuda/runtime_utils.h" @@ -25,73 +24,41 @@ class CudaMul : public Mul { size_t shape_size = ndim_ * sizeof(*d_input_shape_); size_t strides_size = ndim_ * sizeof(*d_input_strides_); const size_t metadata_size = 3 * (shape_size + strides_size); - if (Backend::GetDevice(&owner_device_) != Backend::kSuccess) - throw std::runtime_error("CudaMul cannot query metadata owner device"); if (metadata_size == 0) { return; } std::vector metadata(metadata_size); - void* allocation = nullptr; - const auto allocation_status = Backend::Malloc(&allocation, metadata_size); - if (allocation_status != Backend::kSuccess || !allocation) { - if (allocation && Backend::Free(allocation) != Backend::kSuccess) - std::terminate(); - throw std::bad_alloc(); - } - d_metadata_ = static_cast(allocation); - - try { - size_t offset = 0; - d_input_shape_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, input_shape_.data(), shape_size); - offset += shape_size; - - d_other_shape_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, other_shape_.data(), shape_size); - offset += shape_size; - - d_out_shape_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, out_shape_.data(), shape_size); - offset += shape_size; - - d_input_strides_ = - reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, input_strides_.data(), - strides_size); - offset += strides_size; - - d_other_strides_ = - reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, other_strides_.data(), - strides_size); - offset += strides_size; - - d_out_strides_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, out_strides_.data(), strides_size); - - if (Backend::Memcpy(d_metadata_, metadata.data(), metadata_size, - Backend::kMemcpyHostToDevice) != Backend::kSuccess) - throw std::runtime_error("CudaMul metadata copy failed"); - } catch (...) { - if (Backend::Free(d_metadata_) != Backend::kSuccess) std::terminate(); - d_metadata_ = nullptr; - throw; - } - } + auto* d_metadata_ = metadata_.Allocate(metadata_size); + + size_t offset = 0; + d_input_shape_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, input_shape_.data(), shape_size); + offset += shape_size; + + d_other_shape_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, other_shape_.data(), shape_size); + offset += shape_size; + + d_out_shape_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, out_shape_.data(), shape_size); + offset += shape_size; + + d_input_strides_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, input_strides_.data(), strides_size); + offset += strides_size; + + d_other_strides_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, other_strides_.data(), strides_size); + offset += strides_size; + + d_out_strides_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, out_strides_.data(), strides_size); - ~CudaMul() { - if (!d_metadata_) return; - int previous = -1; - if (Backend::GetDevice(&previous) != Backend::kSuccess) std::terminate(); - const bool switched = previous != owner_device_; - if (switched && Backend::SetDevice(owner_device_) != Backend::kSuccess) - std::terminate(); - const auto freed = Backend::Free(d_metadata_); - if (switched && Backend::SetDevice(previous) != Backend::kSuccess) - std::terminate(); - if (freed != Backend::kSuccess) std::terminate(); + if (Backend::Memcpy(d_metadata_, metadata.data(), metadata_size, + Backend::kMemcpyHostToDevice) != Backend::kSuccess) + throw std::runtime_error("CudaMul metadata copy failed"); } void operator()(const Tensor input, const Tensor other, @@ -100,9 +67,8 @@ class CudaMul : public Mul { return; } - int block_size = RuntimeUtils::GetOptimalBlockSize(); DispatchFunc( - {static_cast(out_type_), block_size}, + {static_cast(out_type_), block_size_}, [&](auto list_tag) { using T = TypeMapType(list_tag)>; constexpr int kBlockSize = ListGet<1>(list_tag); @@ -110,7 +76,7 @@ class CudaMul : public Mul { auto cuda_stream = static_cast(stream_ ? stream_ : 0); dim3 block_dims( - std::min(static_cast(block_size), output_size_)); + std::min(static_cast(block_size_), output_size_)); dim3 grid_dims(utils::CeilDiv(output_size_, block_dims.x)); T* d_out = reinterpret_cast(out.data()); @@ -128,8 +94,9 @@ class CudaMul : public Mul { } private: - std::byte* d_metadata_{nullptr}; - int owner_device_{-1}; + const int block_size_ = + RuntimeUtils::GetOptimalBlockSize(); + CudaDeviceBuffer metadata_; Tensor::Size* d_input_shape_{nullptr}; diff --git a/src/native/cuda/ops/sigmoid/kernel.h b/src/native/cuda/ops/sigmoid/kernel.h index f333cd595..bdea727f0 100644 --- a/src/native/cuda/ops/sigmoid/kernel.h +++ b/src/native/cuda/ops/sigmoid/kernel.h @@ -5,8 +5,6 @@ #include #include #include -#include -#include #include #include @@ -14,6 +12,7 @@ #include "common/generic_utils.h" #include "data_type.h" #include "dispatcher.h" +#include "native/cuda/device_buffer.h" #include "native/cuda/kernel_commons.cuh" #include "native/cuda/ops/sigmoid_infinilm/kernel.cuh" #include "native/cuda/runtime_utils.h" @@ -33,9 +32,6 @@ class CudaSigmoid : public Sigmoid { assert(!out.HasBroadcastDim() && "`Sigmoid` output must not have broadcasted dimensions"); - if (Backend::GetDevice(&owner_device_) != Backend::kSuccess) - throw std::runtime_error( - "CudaSigmoid cannot query metadata owner device"); if (output_size_ == 0 || (is_input_contiguous_ && is_out_contiguous_)) return; @@ -47,55 +43,27 @@ class CudaSigmoid : public Sigmoid { } std::vector metadata(metadata_size); - void* allocation = nullptr; - const auto allocation_status = Backend::Malloc(&allocation, metadata_size); - if (allocation_status != Backend::kSuccess || !allocation) { - if (allocation && Backend::Free(allocation) != Backend::kSuccess) - std::terminate(); - throw std::bad_alloc(); - } - d_metadata_ = static_cast(allocation); - - try { - size_t offset = 0; - d_input_shape_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, input_shape_.data(), shape_size); - offset += shape_size; - - d_out_shape_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, out_shape_.data(), shape_size); - offset += shape_size; - - d_input_strides_ = - reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, input_strides_.data(), - strides_size); - offset += strides_size; - - d_out_strides_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, out_strides_.data(), strides_size); - - if (Backend::Memcpy(d_metadata_, metadata.data(), metadata_size, - Backend::kMemcpyHostToDevice) != Backend::kSuccess) - throw std::runtime_error("CudaSigmoid metadata copy failed"); - } catch (...) { - if (Backend::Free(d_metadata_) != Backend::kSuccess) std::terminate(); - d_metadata_ = nullptr; - throw; - } - } + auto* d_metadata_ = metadata_.Allocate(metadata_size); + + size_t offset = 0; + d_input_shape_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, input_shape_.data(), shape_size); + offset += shape_size; + + d_out_shape_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, out_shape_.data(), shape_size); + offset += shape_size; + + d_input_strides_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, input_strides_.data(), strides_size); + offset += strides_size; + + d_out_strides_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, out_strides_.data(), strides_size); - ~CudaSigmoid() { - if (!d_metadata_) return; - int previous = -1; - if (Backend::GetDevice(&previous) != Backend::kSuccess) std::terminate(); - const bool switched = previous != owner_device_; - if (switched && Backend::SetDevice(owner_device_) != Backend::kSuccess) - std::terminate(); - const auto freed = Backend::Free(d_metadata_); - if (switched && Backend::SetDevice(previous) != Backend::kSuccess) - std::terminate(); - if (freed != Backend::kSuccess) std::terminate(); + if (Backend::Memcpy(d_metadata_, metadata.data(), metadata_size, + Backend::kMemcpyHostToDevice) != Backend::kSuccess) + throw std::runtime_error("CudaSigmoid metadata copy failed"); } void operator()(const Tensor input, Tensor out) const override { @@ -105,13 +73,11 @@ class CudaSigmoid : public Sigmoid { auto cuda_stream = static_cast(stream_ ? stream_ : 0); - int block_size = std::min( - RuntimeUtils::GetOptimalBlockSize(), 1024); - dim3 block(std::min(static_cast(block_size), output_size_)); + dim3 block(std::min(static_cast(block_size_), output_size_)); dim3 grid(utils::CeilDiv(output_size_, block.x)); DispatchFunc>( - {static_cast(out_type_), block_size}, + {static_cast(out_type_), block_size_}, [&](auto list_tag) { using T = TypeMapType(list_tag)>; constexpr int kBlockSize = ListGet<1>(list_tag); @@ -128,8 +94,9 @@ class CudaSigmoid : public Sigmoid { } private: - std::byte* d_metadata_{nullptr}; - int owner_device_{-1}; + const int block_size_ = + std::min(RuntimeUtils::GetOptimalBlockSize(), 1024); + CudaDeviceBuffer metadata_; Tensor::Size* d_input_shape_{nullptr}; diff --git a/src/native/cuda/ops/silu/kernel.h b/src/native/cuda/ops/silu/kernel.h index b4ec1d9d9..1fa37e69f 100644 --- a/src/native/cuda/ops/silu/kernel.h +++ b/src/native/cuda/ops/silu/kernel.h @@ -4,8 +4,6 @@ #include #include #include -#include -#include #include #include @@ -13,6 +11,7 @@ #include "common/generic_utils.h" #include "data_type.h" #include "dispatcher.h" +#include "native/cuda/device_buffer.h" #include "native/cuda/kernel_commons.cuh" #include "native/cuda/ops/silu/kernel.cuh" #include "native/cuda/runtime_utils.h" @@ -23,8 +22,6 @@ template class CudaSilu : public Silu { public: CudaSilu(const Tensor input, Tensor out) : Silu{input, out} { - if (Backend::GetDevice(&owner_device_) != Backend::kSuccess) - throw std::runtime_error("CudaSilu cannot query metadata owner device"); if (output_size_ == 0 || (is_input_contiguous_ && is_out_contiguous_)) return; @@ -33,55 +30,27 @@ class CudaSilu : public Silu { const size_t metadata_size = 2 * (shape_size + strides_size); std::vector metadata(metadata_size); - void* allocation = nullptr; - const auto allocation_status = Backend::Malloc(&allocation, metadata_size); - if (allocation_status != Backend::kSuccess || !allocation) { - if (allocation && Backend::Free(allocation) != Backend::kSuccess) - std::terminate(); - throw std::bad_alloc(); - } - d_metadata_ = static_cast(allocation); - - try { - size_t offset = 0; - d_input_shape_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, input_shape_.data(), shape_size); - offset += shape_size; - - d_out_shape_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, out_shape_.data(), shape_size); - offset += shape_size; - - d_input_strides_ = - reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, input_strides_.data(), - strides_size); - offset += strides_size; - - d_out_strides_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, out_strides_.data(), strides_size); - - if (Backend::Memcpy(d_metadata_, metadata.data(), metadata_size, - Backend::kMemcpyHostToDevice) != Backend::kSuccess) - throw std::runtime_error("CudaSilu metadata copy failed"); - } catch (...) { - if (Backend::Free(d_metadata_) != Backend::kSuccess) std::terminate(); - d_metadata_ = nullptr; - throw; - } - } + auto* d_metadata_ = metadata_.Allocate(metadata_size); + + size_t offset = 0; + d_input_shape_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, input_shape_.data(), shape_size); + offset += shape_size; + + d_out_shape_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, out_shape_.data(), shape_size); + offset += shape_size; + + d_input_strides_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, input_strides_.data(), strides_size); + offset += strides_size; + + d_out_strides_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, out_strides_.data(), strides_size); - ~CudaSilu() { - if (!d_metadata_) return; - int previous = -1; - if (Backend::GetDevice(&previous) != Backend::kSuccess) std::terminate(); - const bool switched = previous != owner_device_; - if (switched && Backend::SetDevice(owner_device_) != Backend::kSuccess) - std::terminate(); - const auto freed = Backend::Free(d_metadata_); - if (switched && Backend::SetDevice(previous) != Backend::kSuccess) - std::terminate(); - if (freed != Backend::kSuccess) std::terminate(); + if (Backend::Memcpy(d_metadata_, metadata.data(), metadata_size, + Backend::kMemcpyHostToDevice) != Backend::kSuccess) + throw std::runtime_error("CudaSilu metadata copy failed"); } void operator()(const Tensor input, Tensor out) const override { @@ -89,9 +58,8 @@ class CudaSilu : public Silu { return; } - int block_size = RuntimeUtils::GetOptimalBlockSize(); DispatchFunc( - {static_cast(out_type_), block_size}, + {static_cast(out_type_), block_size_}, [&](auto list_tag) { using T = TypeMapType(list_tag)>; constexpr int kBlockSize = ListGet<1>(list_tag); @@ -99,7 +67,7 @@ class CudaSilu : public Silu { auto cuda_stream = static_cast(stream_ ? stream_ : 0); dim3 blockDims( - std::min(static_cast(block_size), output_size_)); + std::min(static_cast(block_size_), output_size_)); dim3 gridDims(utils::CeilDiv(output_size_, blockDims.x)); T* d_out = reinterpret_cast(out.data()); @@ -115,8 +83,9 @@ class CudaSilu : public Silu { } private: - std::byte* d_metadata_{nullptr}; - int owner_device_{-1}; + const int block_size_ = + RuntimeUtils::GetOptimalBlockSize(); + CudaDeviceBuffer metadata_; Tensor::Size* d_input_shape_{nullptr}; diff --git a/src/native/cuda/ops/silu_and_mul/kernel.h b/src/native/cuda/ops/silu_and_mul/kernel.h index 83d50c0af..2174c799c 100644 --- a/src/native/cuda/ops/silu_and_mul/kernel.h +++ b/src/native/cuda/ops/silu_and_mul/kernel.h @@ -4,13 +4,12 @@ #include #include #include -#include -#include #include #include #include "base/silu_and_mul.h" #include "common/generic_utils.h" +#include "native/cuda/device_buffer.h" #include "native/cuda/ops/silu_and_mul/kernel.cuh" #include "native/cuda/runtime_utils.h" @@ -20,9 +19,6 @@ template class CudaSiluAndMul : public SiluAndMul { public: CudaSiluAndMul(const Tensor input, Tensor out) : SiluAndMul{input, out} { - if (Backend::GetDevice(&owner_device_) != Backend::kSuccess) - throw std::runtime_error( - "CudaSiluAndMul cannot query metadata owner device"); if (output_size_ == 0 || (is_input_contiguous_ && is_out_contiguous_)) return; @@ -32,55 +28,27 @@ class CudaSiluAndMul : public SiluAndMul { const size_t metadata_size = 2 * (shape_size + strides_size); std::vector metadata(metadata_size); - void* allocation = nullptr; - const auto allocation_status = Backend::Malloc(&allocation, metadata_size); - if (allocation_status != Backend::kSuccess || !allocation) { - if (allocation && Backend::Free(allocation) != Backend::kSuccess) - std::terminate(); - throw std::bad_alloc(); - } - d_metadata_ = static_cast(allocation); - - try { - size_t offset = 0; - d_input_shape_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, input_shape_.data(), shape_size); - offset += shape_size; - - d_out_shape_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, out_shape_.data(), shape_size); - offset += shape_size; - - d_input_strides_ = - reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, input_strides_.data(), - strides_size); - offset += strides_size; - - d_out_strides_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, out_strides_.data(), strides_size); - - if (Backend::Memcpy(d_metadata_, metadata.data(), metadata_size, - Backend::kMemcpyHostToDevice) != Backend::kSuccess) - throw std::runtime_error("CudaSiluAndMul metadata copy failed"); - } catch (...) { - if (Backend::Free(d_metadata_) != Backend::kSuccess) std::terminate(); - d_metadata_ = nullptr; - throw; - } - } + auto* d_metadata_ = metadata_.Allocate(metadata_size); + + size_t offset = 0; + d_input_shape_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, input_shape_.data(), shape_size); + offset += shape_size; + + d_out_shape_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, out_shape_.data(), shape_size); + offset += shape_size; + + d_input_strides_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, input_strides_.data(), strides_size); + offset += strides_size; + + d_out_strides_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, out_strides_.data(), strides_size); - ~CudaSiluAndMul() { - if (!d_metadata_) return; - int previous = -1; - if (Backend::GetDevice(&previous) != Backend::kSuccess) std::terminate(); - const bool switched = previous != owner_device_; - if (switched && Backend::SetDevice(owner_device_) != Backend::kSuccess) - std::terminate(); - const auto freed = Backend::Free(d_metadata_); - if (switched && Backend::SetDevice(previous) != Backend::kSuccess) - std::terminate(); - if (freed != Backend::kSuccess) std::terminate(); + if (Backend::Memcpy(d_metadata_, metadata.data(), metadata_size, + Backend::kMemcpyHostToDevice) != Backend::kSuccess) + throw std::runtime_error("CudaSiluAndMul metadata copy failed"); } void operator()(const Tensor input, Tensor out) const override { @@ -88,9 +56,8 @@ class CudaSiluAndMul : public SiluAndMul { return; } - int block_size = RuntimeUtils::GetOptimalBlockSize(); DispatchFunc( - {static_cast(out_type_), block_size}, + {static_cast(out_type_), block_size_}, [&](auto list_tag) { using T = TypeMapType(list_tag)>; constexpr int kBlockSize = ListGet<1>(list_tag); @@ -98,7 +65,7 @@ class CudaSiluAndMul : public SiluAndMul { auto cuda_stream = static_cast(stream_ ? stream_ : 0); dim3 blockDims( - std::min(static_cast(block_size), output_size_)); + std::min(static_cast(block_size_), output_size_)); dim3 gridDims(utils::CeilDiv(output_size_, blockDims.x)); T* d_out = reinterpret_cast(out.data()); @@ -113,8 +80,9 @@ class CudaSiluAndMul : public SiluAndMul { } private: - std::byte* d_metadata_{nullptr}; - int owner_device_{-1}; + const int block_size_ = + RuntimeUtils::GetOptimalBlockSize(); + CudaDeviceBuffer metadata_; Tensor::Size* d_input_shape_{nullptr}; diff --git a/src/native/cuda/ops/softmax/kernel.h b/src/native/cuda/ops/softmax/kernel.h index 8bba2c0f7..9f909ae12 100644 --- a/src/native/cuda/ops/softmax/kernel.h +++ b/src/native/cuda/ops/softmax/kernel.h @@ -5,8 +5,6 @@ #include #include #include -#include -#include #include #include @@ -14,6 +12,7 @@ #include "common/generic_utils.h" #include "data_type.h" #include "dispatcher.h" +#include "native/cuda/device_buffer.h" #include "native/cuda/kernel_commons.cuh" #include "native/cuda/ops/softmax_infinilm/kernel.cuh" #include "native/cuda/runtime_utils.h" @@ -39,9 +38,6 @@ class CudaSoftmax : public Softmax { assert(!out.HasBroadcastDim() && "`Softmax` output must not have broadcasted dimensions"); - if (Backend::GetDevice(&owner_device_) != Backend::kSuccess) - throw std::runtime_error( - "CudaSoftmax cannot query metadata owner device"); block_size_ = std::min( RuntimeUtils::GetOptimalBlockSize(), 1024); if (row_count_ == 0 || dim_size_ == 0) return; @@ -54,51 +50,23 @@ class CudaSoftmax : public Softmax { } std::vector metadata(metadata_size); - void* allocation = nullptr; - const auto status = Backend::Malloc(&allocation, metadata_size); - if (status != Backend::kSuccess || !allocation) { - if (allocation && Backend::Free(allocation) != Backend::kSuccess) - std::terminate(); - throw std::bad_alloc(); - } - d_metadata_ = static_cast(allocation); - - try { - size_t offset = 0; - d_shape_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, out_shape_.data(), shape_size); - offset += shape_size; - - d_input_strides_ = - reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, input_strides_.data(), - strides_size); - offset += strides_size; - - d_out_strides_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, out_strides_.data(), strides_size); - - if (Backend::Memcpy(d_metadata_, metadata.data(), metadata_size, - Backend::kMemcpyHostToDevice) != Backend::kSuccess) - throw std::runtime_error("CudaSoftmax metadata copy failed"); - } catch (...) { - if (Backend::Free(d_metadata_) != Backend::kSuccess) std::terminate(); - d_metadata_ = nullptr; - throw; - } - } + auto* d_metadata_ = metadata_.Allocate(metadata_size); + + size_t offset = 0; + d_shape_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, out_shape_.data(), shape_size); + offset += shape_size; + + d_input_strides_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, input_strides_.data(), strides_size); + offset += strides_size; + + d_out_strides_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, out_strides_.data(), strides_size); - ~CudaSoftmax() { - if (!d_metadata_) return; - int previous = -1; - if (Backend::GetDevice(&previous) != Backend::kSuccess) std::terminate(); - const bool switched = previous != owner_device_; - if (switched && Backend::SetDevice(owner_device_) != Backend::kSuccess) - std::terminate(); - const auto freed = Backend::Free(d_metadata_); - if (switched && Backend::SetDevice(previous) != Backend::kSuccess) - std::terminate(); - if (freed != Backend::kSuccess) std::terminate(); + if (Backend::Memcpy(d_metadata_, metadata.data(), metadata_size, + Backend::kMemcpyHostToDevice) != Backend::kSuccess) + throw std::runtime_error("CudaSoftmax metadata copy failed"); } void operator()(const Tensor input, const int64_t, @@ -132,8 +100,7 @@ class CudaSoftmax : public Softmax { } private: - std::byte* d_metadata_{nullptr}; - int owner_device_{-1}; + CudaDeviceBuffer metadata_; int block_size_{0}; Tensor::Size* d_shape_{nullptr}; diff --git a/src/operator.h b/src/operator.h index 64175aab5..7edfa78f9 100644 --- a/src/operator.h +++ b/src/operator.h @@ -401,6 +401,15 @@ class Operator : public OperatorBase { [](auto, auto, auto, auto&...) {}, tensor, args...); } + template + static bool CanMake(const Config& config, const std::vector tensors, + const Args&... args) { + return !tensors.empty() && + SelectWithDevice( + config, tensors.front().device().type(), false, + [](auto, auto, auto, auto&...) {}, tensors, args...); + } + template static void Call(const Handle& handle, const Config& config, const Args&... args) { diff --git a/tests/test_generate_wrappers.py b/tests/test_generate_wrappers.py index 9ff7148a4..82298650a 100644 --- a/tests/test_generate_wrappers.py +++ b/tests/test_generate_wrappers.py @@ -111,8 +111,12 @@ def test_operator_call_instantiations_externalize_default_implementation_lookup( ) assert declarations.count(f"extern template {default_config_signature}") == 1 assert definitions.count(f"template {default_config_signature}") == 1 - assert not any("::CanMake<" in entry for entry in declarations) - assert not any("::CanMake<" in entry for entry in definitions) + eligibility = ( + "bool Operator<::infini::ops::Abs>::CanMake" + "(const Config& config, const Tensor input, const Tensor& out);" + ) + assert f"extern template {eligibility}" in declarations + assert f"template {eligibility}" in definitions def test_operator_call_instantiations_externalize_active_implementation_query(): @@ -441,6 +445,12 @@ class Cat { assert "converted_first_tensor.at(0).device().type()" in text assert "std::move(converted_first_tensor)" not in text assert "DeviceFromPybind11Handle(inputs.at(0))" not in text + declarations, _ = module._generate_operator_call_instantiation_entries(operator) + assert any( + "::CanMake(const Config& config, " + "const std::vector inputs, const Tensor& out)" in declaration + for declaration in declarations + ) def test_optional_tensor_vector_is_preserved_across_generated_wrappers( @@ -807,7 +817,17 @@ def test_select_ops_from_config_preserves_path_and_backend_descriptors(): }, } - assert module._select_ops_from_config({}, config, "ops.json") == { + assert module._select_ops_from_config( + { + "add": [ + module._Implementation( + pathlib.Path("src/native/cpu/ops/add/add.h"), "native" + ) + ] + }, + config, + "ops.json", + ) == { "add": [ module._Implementation( pathlib.Path("src/native/cpu/ops/add/add.h"), "native" @@ -1180,3 +1200,27 @@ def test_cmake_configures_cambricon_half_compatibility(): assert "#include " in compatibility_header assert "#include " in compatibility_header assert "#define __mlu_host__ inline" in compatibility_header + + +def test_builtin_header_requires_enabled_provider(): + module = _load_generator_module() + config = { + "flash_attn_with_kvcache": { + "headers": ["src/native/cuda/nvidia/ops/flash_attn_with_kvcache/kernel.h"], + "implementations": None, + } + } + with pytest.raises(ValueError, match="not enabled by this build"): + module._select_ops_from_config({}, config, "ops.json") + + +def test_optional_provider_flags_with_absolute_source_root(monkeypatch): + module = _load_generator_module() + source = pathlib.Path(__file__).resolve().parents[1] / "src" + monkeypatch.setattr(module, "_SRC_DIR", source) + monkeypatch.setattr(module, "_BASE_DIR", source / "base") + off = module._get_all_ops(["nvidia"]) + on = module._get_all_ops(["nvidia"], with_native_flash=True, with_native_w4=True) + for name in ("flash_attn_with_kvcache", "flash_attn_varlen_func", "w4_marlin_gemm"): + assert not off.get(name) + assert on[name] diff --git a/third_party/flash-attention b/third_party/flash-attention new file mode 160000 index 000000000..060c9188b --- /dev/null +++ b/third_party/flash-attention @@ -0,0 +1 @@ +Subproject commit 060c9188beec3a8b62b33a3bfa6d5d2d44975fab diff --git a/third_party/vllm b/third_party/vllm new file mode 160000 index 000000000..ffc4f08c8 --- /dev/null +++ b/third_party/vllm @@ -0,0 +1 @@ +Subproject commit ffc4f08c8ee130d4ea6347c1bf31ffd4f8af28ab From a989fcbd2dfa12905e8d41b0db445b8c5d71050b Mon Sep 17 00:00:00 2001 From: whjthu Date: Tue, 20 Oct 2026 00:49:44 +0800 Subject: [PATCH 19/25] fix: correct native provider and source preparation contracts --- cmake/NativeFlashAttention.cmake | 4 +- cmake/NativeW4Marlin.cmake | 4 +- docs/build-nvidia-inference.md | 7 +- patches/marlin/w4a16.patch | 67 +------------------ ...a_dense_sampling_torch_free_no_marlin.json | 56 ++++++++++++++++ scripts/generate_wrappers.py | 11 +-- scripts/prepare_flash_attention_inference.py | 29 ++++++-- scripts/prepare_marlin.py | 53 +++++++++++---- .../chunk_gated_delta_rule_legacy_kernel.cuh | 5 +- .../recurrent_delta_rule_common.cuh | 5 +- src/native/cuda/nvidia/ops/w4_marlin/gemm.cu | 2 + .../cuda/nvidia/ops/w4_marlin/repack.cc | 3 +- tests/test_generate_wrappers.py | 30 ++++++++- 13 files changed, 166 insertions(+), 110 deletions(-) create mode 100644 scripts/configs/infiniops_ops_nvidia_dense_sampling_torch_free_no_marlin.json diff --git a/cmake/NativeFlashAttention.cmake b/cmake/NativeFlashAttention.cmake index c0db8ad8f..99148b920 100644 --- a/cmake/NativeFlashAttention.cmake +++ b/cmake/NativeFlashAttention.cmake @@ -1,7 +1,7 @@ # Build the pinned inference kernels separately. No Torch include/link flags or # Python runtime enter this target; Python is used only to prepare pinned source. function(infiniops_add_native_flash target) - find_package(Python3 REQUIRED COMPONENTS Interpreter) + find_package(Python REQUIRED COMPONENTS Interpreter) find_package(CUDAToolkit REQUIRED) find_program(_flash_patch NAMES patch REQUIRED) set(_provider "${CMAKE_CURRENT_FUNCTION_LIST_DIR}/../src/native/cuda/nvidia/ops/flash_attn_with_kvcache") @@ -30,7 +30,7 @@ function(infiniops_add_native_flash target) endif() endforeach() set_property(DIRECTORY APPEND PROPERTY CMAKE_CONFIGURE_DEPENDS ${_inputs}) - execute_process(COMMAND "${Python3_EXECUTABLE}" + execute_process(COMMAND "${Python_EXECUTABLE}" "${_source}/scripts/prepare_flash_attention_inference.py" --source-dir "${_upstream}" --output "${_prepared}" --patch-executable "${_flash_patch}" diff --git a/cmake/NativeW4Marlin.cmake b/cmake/NativeW4Marlin.cmake index ebc27c1c7..d84aaad41 100644 --- a/cmake/NativeW4Marlin.cmake +++ b/cmake/NativeW4Marlin.cmake @@ -2,7 +2,7 @@ if(CMAKE_CUDA_COMPILER_VERSION VERSION_LESS 12.3) message(FATAL_ERROR "NVIDIA Marlin requires CUDA >=12.3 and C++20") endif() set(_w4_root "${PROJECT_SOURCE_DIR}/src/native/cuda/nvidia/ops/w4_marlin") -find_package(Python3 REQUIRED COMPONENTS Interpreter) +find_package(Python REQUIRED COMPONENTS Interpreter) find_program(_w4_patch NAMES patch REQUIRED) set(_w4_upstream "${PROJECT_SOURCE_DIR}/third_party/vllm") set(_w4_upstream_marlin "${_w4_upstream}/csrc/libtorch_stable/quantization/marlin") @@ -23,7 +23,7 @@ foreach(_w4_input IN LISTS _w4_inputs) endforeach() set_property(DIRECTORY APPEND PROPERTY CMAKE_CONFIGURE_DEPENDS ${_w4_inputs}) execute_process( - COMMAND "${Python3_EXECUTABLE}" "${PROJECT_SOURCE_DIR}/scripts/prepare_marlin.py" + COMMAND "${Python_EXECUTABLE}" "${PROJECT_SOURCE_DIR}/scripts/prepare_marlin.py" --source-dir "${_w4_upstream}" --output "${_w4_prepared}" --patch-executable "${_w4_patch}" RESULT_VARIABLE _w4_status ERROR_VARIABLE _w4_prepare_error) diff --git a/docs/build-nvidia-inference.md b/docs/build-nvidia-inference.md index e979a0c82..6efd9bdde 100644 --- a/docs/build-nvidia-inference.md +++ b/docs/build-nvidia-inference.md @@ -123,8 +123,11 @@ cmake --install build/ops ``` Marlin is enabled here because the LM profile explicitly contains W4 operators. -A dense-only SDK can set it `OFF` and use an allowlist without `w4_marlin_*`; -then its CUDA >=12.3/C++20 check and provider registration are both omitted. +A dense-only SDK can set it `OFF`. When building LM with its checkpoint +adapters, select `infiniops_ops_nvidia_dense_sampling_torch_free_no_marlin.json`: +its explicit `public_api_only` entries preserve the three W4 API symbols without +providers. Other consumers can omit those entries entirely. Disabled Marlin +has no CUDA >=12.3/C++20 check or provider registration and cannot execute W4. Both native components require the standard `patch` command when enabled. Their preparation scripts patch build-directory copies, leaving submodules unchanged. Jinja2 is used only to generate Marlin instances at build time. diff --git a/patches/marlin/w4a16.patch b/patches/marlin/w4a16.patch index 5d4f18805..2af0f7260 100644 --- a/patches/marlin/w4a16.patch +++ b/patches/marlin/w4a16.patch @@ -61,63 +61,7 @@ --- a/gemm/marlin/marlin_template.h +++ b/gemm/marlin/marlin_template.h -@@ -327,9 +327,6 @@ - if constexpr (b_type == vllm::kFE2M1f) { - static_assert(s_type == vllm::kFE4M3fn && group_blocks == 1 || - s_type == vllm::kFE8M0fnu && group_blocks == 2); -- } else if constexpr (s_type == vllm::kFE8M0fnu) { -- // MXFP8: FP8 weights with e8m0 microscaling block scales -- static_assert(b_type == vllm::kFE4M3fn && group_blocks == 2); - } else if constexpr (std::is_same::value) { - static_assert(s_type == vllm::kBFloat16); - } else if constexpr (std::is_same::value) { -@@ -337,7 +334,6 @@ - } - - constexpr bool is_a_8bit = a_type.size_bits() == 8; -- constexpr bool is_8bit_scale = s_type.size_bits() == 8; - if constexpr (!is_a_8bit) { - static_assert(std::is_same::value); - } -@@ -347,7 +343,7 @@ - b_type == vllm::kU4B8 || b_type == vllm::kU8B128; - // see comments of dequant.h for more details - constexpr bool dequant_skip_flop = -- is_a_8bit || (b_type == vllm::kFE4M3fn && !(s_type == vllm::kFE8M0fnu)) || -+ is_a_8bit || b_type == vllm::kFE4M3fn || - b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn || - has_zp && !is_zp_float && !std::is_same::value || - has_zp && !is_zp_float && !(b_type == vllm::kU8); -@@ -559,8 +555,8 @@ - constexpr int b_sh_wr_iters = b_sh_stage / b_sh_wr_delta; - - // Scale sizes/strides without act_order -- int s_gl_stride = prob_n / (is_8bit_scale ? 16 : 8); -- constexpr int s_sh_stride = 16 * thread_n_blocks / (is_8bit_scale ? 16 : 8); -+ int s_gl_stride = prob_n / (b_type == vllm::kFE2M1f ? 16 : 8); -+ constexpr int s_sh_stride = 16 * thread_n_blocks / (b_type == vllm::kFE2M1f ? 16 : 8); - constexpr int s_tb_groups = - !has_act_order && group_blocks != -1 && group_blocks < thread_k_blocks - ? thread_k_blocks / group_blocks -@@ -1000,7 +996,7 @@ - - int4* sh_s_stage = sh_s + s_sh_stage * pipe; - -- if constexpr (!is_8bit_scale) { -+ if constexpr (b_type_id != vllm::kFE2M1f.id()) { - reinterpret_cast(&frag_s[k % 2])[0] = - sh_s_stage[s_sh_rd + cur_group_id * s_sh_stride]; - } else { -@@ -1009,7 +1005,7 @@ - sh_s_stage)[s_sh_rd + cur_group_id * (2 * s_sh_stride)]; - } - } else if (group_blocks >= b_sh_wr_iters) { -- if constexpr (!is_8bit_scale) { -+ if constexpr (b_type_id != vllm::kFE2M1f.id()) { - reinterpret_cast(&frag_s[1])[0] = - reinterpret_cast(&frag_s[0])[0]; - } else { -@@ -1035,7 +1031,7 @@ +@@ -1035,7 +1035,7 @@ cur_k = 0; // Progress to current iteration @@ -126,15 +70,6 @@ // Determine "position" inside the thread-block (based on warp and // thread-id) -@@ -1210,7 +1206,7 @@ - } - } - -- if constexpr (s_type == vllm::kFE4M3fn || s_type == vllm::kFE8M0fnu) { -+ if constexpr (b_type == vllm::kFE2M1f) { - int s_quant_0 = reinterpret_cast(frag_s[k2])[0]; - int s_quant_1 = reinterpret_cast(frag_s[k2])[1]; - --- a/awq_repack/cuda/kernel.cuh +++ b/awq_repack/cuda/kernel.cuh @@ -1,18 +1,9 @@ diff --git a/scripts/configs/infiniops_ops_nvidia_dense_sampling_torch_free_no_marlin.json b/scripts/configs/infiniops_ops_nvidia_dense_sampling_torch_free_no_marlin.json new file mode 100644 index 000000000..0b76aa23d --- /dev/null +++ b/scripts/configs/infiniops_ops_nvidia_dense_sampling_torch_free_no_marlin.json @@ -0,0 +1,56 @@ +{ + "add": {"implementations": [0]}, + "argmax": {"implementations": [0]}, + "cast": {"implementations": [0]}, + "cat": {"implementations": [0]}, + "causal_conv1d": {"implementations": [0]}, + "chunk_gated_delta_rule": {"implementations": [0]}, + "convolution": {"implementations": [0]}, + "copy": {"implementations": [0]}, + "deepseek_moe": {"implementations": [0]}, + "embedding": {"implementations": [0]}, + "fill": {"implementations": [0]}, + "flash_attn_varlen_func": {"implementations": [0]}, + "flash_attn_with_kvcache": {"implementations": [0]}, + "fused_add_rms_norm": {"implementations": [0]}, + "fused_gated_delta_net_gating": {"implementations": [0]}, + "gelu": {"implementations": [0]}, + "gemm": {"implementations": [0]}, + "i8_gemm": {"implementations": [0]}, + "layer_norm": {"implementations": [0]}, + "mamba_selective_scan": {"implementations": [0]}, + "moe_align_block_size": {"implementations": [0]}, + "moe_fused_dense": {"implementations": [0]}, + "moe_fused_gate": {"implementations": [0]}, + "mrope": {"implementations": [0]}, + "mul": {"implementations": [0]}, + "mul_scalar": {"implementations": [0]}, + "pad_constant": {"implementations": [0]}, + "per_channel_quant_i8": {"implementations": [0]}, + "per_tensor_dequant_i8": {"implementations": [0]}, + "per_tensor_quant_i8": {"implementations": [0]}, + "quick_gelu": {"implementations": [0]}, + "recurrent_gated_delta_rule": {"implementations": [0]}, + "relu": {"implementations": [0]}, + "reshape_and_cache_flash": {"implementations": [0]}, + "rms_norm": {"implementations": [0]}, + "rotary_embedding": {"implementations": [0]}, + "row_l2_normalize": {"implementations": [0]}, + "separate_table_rope": {"implementations": [0]}, + "sigmoid": {"implementations": [0]}, + "silu": {"implementations": [0]}, + "silu_and_mul": {"implementations": [0]}, + "softmax": {"implementations": [0]}, + "sum": {"implementations": [0]}, + "take": {"implementations": [0]}, + "top_k_top_p_sampling_from_logits": {"implementations": [16]}, + "topk_sigmoid": {"implementations": [0]}, + "topk_softmax": {"implementations": [0]}, + "tril": {"implementations": [0]}, + "triu": {"implementations": [0]}, + "upsample_bicubic2d": {"public_api_only": true}, + "upsample_bilinear2d": {"implementations": [0]}, + "w4_marlin_gemm": {"public_api_only": true}, + "w4_marlin_parameters": {"public_api_only": true}, + "w4_marlin_repack": {"public_api_only": true} +} diff --git a/scripts/generate_wrappers.py b/scripts/generate_wrappers.py index 31b272525..1d73dcc68 100644 --- a/scripts/generate_wrappers.py +++ b/scripts/generate_wrappers.py @@ -2028,11 +2028,7 @@ def _get_all_ops( impl_headers_by_operator.get(_op_relative_type(op_name), ()) ) - if not implementations and op_name not in { - "w4_marlin_gemm", - "w4_marlin_repack", - "w4_marlin_parameters", - }: + if not implementations: continue ops[op_name] = implementations @@ -2262,11 +2258,6 @@ def _dispatch_gen_batch_size(): if config_path is not None: config = ops_config.load_ops_config(config_path) - # Optional Marlin providers do not remove the public checkpoint API. - # Reuse the existing empty-provider generation path when none is built. - for name in ("w4_marlin_gemm", "w4_marlin_repack", "w4_marlin_parameters"): - if name not in config or (name in ops and not ops[name]): - config[name] = {"public_api_only": True} ops = _select_ops_from_config(ops, config, config_path) ops = _filter_ops( diff --git a/scripts/prepare_flash_attention_inference.py b/scripts/prepare_flash_attention_inference.py index 395e4f325..c42867692 100644 --- a/scripts/prepare_flash_attention_inference.py +++ b/scripts/prepare_flash_attention_inference.py @@ -28,7 +28,10 @@ "static_switch.h", "utils.h", ) -PATCH = pathlib.Path(__file__).resolve().parents[1] / "patches/flash-attention/inference.patch" +PATCH = ( + pathlib.Path(__file__).resolve().parents[1] + / "patches/flash-attention/inference.patch" +) def write_if_changed(path, content): @@ -51,7 +54,9 @@ def read_source(path): def prepare(source, output, patch_executable="patch"): source, output = pathlib.Path(source).resolve(), pathlib.Path(output).resolve() if output == source or source in output.parents: - raise ValueError("FlashAttention preparation output must be outside its source tree") + raise ValueError( + "FlashAttention preparation output must be outside its source tree" + ) kernel_source = source / "csrc/flash_attn/src" headers = {name: read_source(kernel_source / name) for name in HEADERS} api = read_source(source / "csrc/flash_attn/flash_api.cpp") @@ -61,15 +66,22 @@ def prepare(source, output, patch_executable="patch"): raise ValueError(f"missing FlashAttention inference adaptation: {PATCH}") output.parent.mkdir(parents=True, exist_ok=True) - with tempfile.TemporaryDirectory(prefix="flash-inference-", dir=output.parent) as temp: + with tempfile.TemporaryDirectory( + prefix="flash-inference-", dir=output.parent + ) as temp: staging = pathlib.Path(temp) for name, text in headers.items(): write_if_changed(staging / name, text) try: applied = subprocess.run( [ - str(patch_executable), "--batch", "--forward", "--fuzz=0", - "-p1", "--input", str(PATCH), + str(patch_executable), + "--batch", + "--forward", + "--fuzz=0", + "-p1", + "--input", + str(PATCH), ], cwd=staging, text=True, @@ -86,9 +98,12 @@ def prepare(source, output, patch_executable="patch"): if applied.returncode: raise ValueError( "pinned FlashAttention inference patch failed:\n" - + applied.stdout + applied.stderr + + applied.stdout + + applied.stderr ) - headers = {name: (staging / name).read_text(encoding="utf-8") for name in HEADERS} + headers = { + name: (staging / name).read_text(encoding="utf-8") for name in HEADERS + } # Keep the upstream scheduling heuristic, without introducing another policy. begin = api.index("inline int num_splits_heuristic(") diff --git a/scripts/prepare_marlin.py b/scripts/prepare_marlin.py index e121decd6..90978e448 100644 --- a/scripts/prepare_marlin.py +++ b/scripts/prepare_marlin.py @@ -13,15 +13,22 @@ SOURCES = { pathlib.Path("gemm/marlin") / name: UPSTREAM / name for name in ( - "dequant.h", "kernel.h", "marlin.cuh", "marlin_dtypes.cuh", - "marlin_mma.h", "marlin_template.h", + "dequant.h", + "kernel.h", + "marlin.cuh", + "marlin_dtypes.cuh", + "marlin_mma.h", + "marlin_template.h", ) } -SOURCES.update({ - pathlib.Path("awq_repack/cuda/kernel.cuh"): UPSTREAM / "awq_marlin_repack.cu", - pathlib.Path("gptq_repack/cuda/kernel.cuh"): UPSTREAM / "gptq_marlin_repack.cu", - pathlib.Path("gemm/nvidia/generate_kernels.py"): UPSTREAM / "generate_kernels.py", -}) +SOURCES.update( + { + pathlib.Path("awq_repack/cuda/kernel.cuh"): UPSTREAM / "awq_marlin_repack.cu", + pathlib.Path("gptq_repack/cuda/kernel.cuh"): UPSTREAM / "gptq_marlin_repack.cu", + pathlib.Path("gemm/nvidia/generate_kernels.py"): UPSTREAM + / "generate_kernels.py", + } +) GENERATED = ( "kernel_selector.h", "sm80_kernel_float16_u4_float16.cu", @@ -52,24 +59,42 @@ def prepare(source, output, patch_executable): ) target = staging / destination target.parent.mkdir(parents=True, exist_ok=True) - target.write_text(upstream.read_text(encoding="utf-8"), encoding="utf-8", newline="\n") + target.write_text( + upstream.read_text(encoding="utf-8"), encoding="utf-8", newline="\n" + ) applied = subprocess.run( - [str(patch_executable), "--batch", "--forward", "--fuzz=0", "-p1", - "--input", str(REPO / "patches/marlin/w4a16.patch")], - cwd=staging, text=True, capture_output=True, check=False, timeout=60, + [ + str(patch_executable), + "--batch", + "--forward", + "--fuzz=0", + "-p1", + "--input", + str(REPO / "patches/marlin/w4a16.patch"), + ], + cwd=staging, + text=True, + capture_output=True, + check=False, + timeout=60, ) if applied.returncode: - raise ValueError("Pinned Marlin adaptation failed:\n" + applied.stdout + applied.stderr) + raise ValueError( + "Pinned Marlin adaptation failed:\n" + applied.stdout + applied.stderr + ) # The upstream generator's single result/config list emits both CUDA # instantiations and the selector. The patch retains only its W4A16 set. subprocess.run( [sys.executable, str(staging / "gemm/nvidia/generate_kernels.py"), "8.0"], - cwd=staging, check=True, + cwd=staging, + check=True, ) prepared = [path for path in SOURCES if path.suffix != ".py"] prepared.extend(pathlib.Path("gemm/nvidia") / name for name in GENERATED) for destination in prepared: - write_if_changed(output / "vendor" / destination, (staging / destination).read_bytes()) + write_if_changed( + output / "vendor" / destination, (staging / destination).read_bytes() + ) def main(): diff --git a/src/native/cuda/nvidia/ops/chunk_gated_delta_rule/chunk_gated_delta_rule_legacy_kernel.cuh b/src/native/cuda/nvidia/ops/chunk_gated_delta_rule/chunk_gated_delta_rule_legacy_kernel.cuh index 733b71c1c..5f5e513dd 100644 --- a/src/native/cuda/nvidia/ops/chunk_gated_delta_rule/chunk_gated_delta_rule_legacy_kernel.cuh +++ b/src/native/cuda/nvidia/ops/chunk_gated_delta_rule/chunk_gated_delta_rule_legacy_kernel.cuh @@ -103,9 +103,10 @@ __device__ void chunkGatedDeltaRuleRecurrentKernel( : loadOptionalIndex(final_state_indices, final_state_indices_i64, batch_idx, batch_idx); - if (read_slot < 0 || write_slot < 0 || + if (read_slot < 0 || read_slot >= static_cast(pool_size) || - write_slot >= static_cast(pool_size)) { + (final_state_indices != nullptr && + (write_slot < 0 || write_slot >= static_cast(pool_size)))) { return; } } diff --git a/src/native/cuda/nvidia/ops/recurrent_gated_delta_rule/recurrent_delta_rule_common.cuh b/src/native/cuda/nvidia/ops/recurrent_gated_delta_rule/recurrent_delta_rule_common.cuh index e5855f8ab..37db9ef83 100644 --- a/src/native/cuda/nvidia/ops/recurrent_gated_delta_rule/recurrent_delta_rule_common.cuh +++ b/src/native/cuda/nvidia/ops/recurrent_gated_delta_rule/recurrent_delta_rule_common.cuh @@ -133,9 +133,10 @@ __device__ void recurrentDeltaRuleWarpSequence( ? static_cast(batch_idx) : loadOptionalIndex(final_state_indices, final_state_indices_i64, batch_idx, batch_idx); - if (read_slot < 0 || write_slot < 0 || + if (read_slot < 0 || read_slot >= static_cast(pool_size) || - write_slot >= static_cast(pool_size)) { + (final_state_indices != nullptr && + (write_slot < 0 || write_slot >= static_cast(pool_size)))) { if (valid_value_dim && lane_idx == 0) { const int token_batch = has_cu_seqlens ? 0 : batch_idx; for (int64_t token_idx = token_begin; token_idx < token_end; diff --git a/src/native/cuda/nvidia/ops/w4_marlin/gemm.cu b/src/native/cuda/nvidia/ops/w4_marlin/gemm.cu index 8db4f1d00..00ec09c9a 100644 --- a/src/native/cuda/nvidia/ops/w4_marlin/gemm.cu +++ b/src/native/cuda/nvidia/ops/w4_marlin/gemm.cu @@ -33,6 +33,7 @@ Operator::Operator( const float alpha, Tensor out) : W4MarlinGemm(a,b,scales,zeros,g_idx,perm,bias,scheme,group_size,is_k_full,alpha,out), device_{a.device().index()} { w4::require(a.device().type() == Device::Type::kNvidia, "W4 requires one NVIDIA device"); + if (!is_k_full && perm.numel() && g_idx.data()) w4::address(g_idx, 16); int shared; native_w4::info(device_,sms_,shared); // Exactly the original FP32 reduction requirement plus A permutation temp. const auto max_m = std::min(((m_+15)/16)*16,64); @@ -49,6 +50,7 @@ void Operator::operator()( w4::require(scheme == scheme_ && group_size == group_ && is_k_full == full_ && alpha == alpha_, "W4 GEMM attributes changed"); for (const auto& t : {a,b,scales,zeros,out}) w4::address(t,16); for (const auto& t : {g_idx,perm,bias}) w4::address(t,4); + if (!full_ && perm.numel()) w4::address(g_idx,16); for (const auto& t : {a,b,scales,zeros,g_idx,perm,bias}) w4::disjoint(out,t); w4::require(workspace_ && workspace_size_in_bytes_ >= bytes_ && reinterpret_cast(workspace_)%256 == 0, "W4 GEMM requires aligned external workspace"); diff --git a/src/native/cuda/nvidia/ops/w4_marlin/repack.cc b/src/native/cuda/nvidia/ops/w4_marlin/repack.cc index 686e06dd1..0786e7da5 100644 --- a/src/native/cuda/nvidia/ops/w4_marlin/repack.cc +++ b/src/native/cuda/nvidia/ops/w4_marlin/repack.cc @@ -9,6 +9,7 @@ Operator::Operator( device_{qweight.device().index()} { w4::require(qweight.device().type() == Device::Type::kNvidia, "W4 requires one NVIDIA device"); + if (scheme == 1 && perm.numel() && perm.data()) w4::address(perm, 16); native_w4::info(device_, sms_, shared_); native_w4::DeviceGuard guard(device_); if (scheme == 0) @@ -23,7 +24,7 @@ void Operator::operator()( "W4 repack attributes changed"); w4::address(qweight, 16); w4::address(out, 16); - w4::address(perm, 4); + w4::address(perm, scheme == 1 && perm.numel() ? 16 : 4); w4::disjoint(out, qweight); w4::disjoint(out, perm); native_w4::DeviceGuard guard(device_); diff --git a/tests/test_generate_wrappers.py b/tests/test_generate_wrappers.py index 82298650a..b8672e852 100644 --- a/tests/test_generate_wrappers.py +++ b/tests/test_generate_wrappers.py @@ -1221,6 +1221,32 @@ def test_optional_provider_flags_with_absolute_source_root(monkeypatch): monkeypatch.setattr(module, "_BASE_DIR", source / "base") off = module._get_all_ops(["nvidia"]) on = module._get_all_ops(["nvidia"], with_native_flash=True, with_native_w4=True) - for name in ("flash_attn_with_kvcache", "flash_attn_varlen_func", "w4_marlin_gemm"): - assert not off.get(name) + for name in ( + "flash_attn_with_kvcache", + "flash_attn_varlen_func", + "w4_marlin_gemm", + "w4_marlin_repack", + "w4_marlin_parameters", + ): + assert name not in off assert on[name] + + +def test_w4_public_api_selection_is_explicit(tmp_path): + module = _load_generator_module() + name = "w4_marlin_gemm" + external = module._Implementation(tmp_path / "w4.h", "custom") + config = { + name: { + "headers": [{"path": str(external.path), "backend": "custom"}], + "implementations": None, + } + } + assert module._select_ops_from_config({}, config, "ops.json") == {name: [external]} + assert module._select_ops_from_config( + {}, {name: {"public_api_only": True}}, "ops.json" + ) == {name: []} + with pytest.raises(ValueError, match="not available"): + module._select_ops_from_config( + {}, {name: {"headers": None, "implementations": (0,)}}, "ops.json" + ) From 7ecc523a0118ee519b481c4e5a53f57973e6e9c7 Mon Sep 17 00:00:00 2001 From: whjthu Date: Tue, 20 Oct 2026 01:41:32 +0800 Subject: [PATCH 20/25] refactor: simplify native inference source integration --- .github/workflows/clang-format.yml | 1 - .gitmodules | 3 - CMakeLists.txt | 10 + LICENSES/Apache-2.0-components.txt | 96 +-- LICENSES/FlashInfer-0.6.7/LICENSE.cutlass.txt | 27 - .../LICENSE.flashattention3.txt | 29 - LICENSES/FlashInfer-0.6.7/LICENSE.fmt.txt | 27 - LICENSES/FlashInfer-0.6.7/LICENSE.spdlog.txt | 25 - LICENSES/TVM-FFI-0.1.10/LICENSE.pytorch.txt | 84 --- LICENSES/TVM-FFI-0.1.10/NOTICE.pytorch.txt | 456 ------------- THIRD_PARTY.txt | 53 +- cmake/MarlinSource.cmake | 31 + cmake/NativeW4Marlin.cmake | 27 +- docs/build-nvidia-inference.md | 48 +- patches/marlin/w4a16.patch | 597 ++++++++++-------- pyproject.toml | 15 +- scripts/prepare_marlin.py | 25 +- src/CMakeLists.txt | 19 +- src/base/flash_attn_with_kvcache.h | 4 - src/base/separate_table_rope.h | 2 +- .../ops/rotary_embedding/cann82_compat.hpp | 111 ---- .../ascend/ops/rotary_embedding/kernel.h | 6 - src/native/cambricon/ops/add/kernel.h | 132 ---- src/native/cpu/ops/fill/fill.h | 3 - src/native/cpu/ops/mul_scalar/mul_scalar.h | 2 - .../ops/row_l2_normalize/row_l2_normalize.h | 2 - src/native/cuda/nvidia/ops/cast/kernel.h | 2 - src/native/cuda/nvidia/ops/cat/kernel.h | 3 - .../cuda/nvidia/ops/causal_conv1d/kernel.h | 1 - .../ops/chunk_gated_delta_rule/kernel.h | 1 - .../ops/recurrent_gated_delta_rule/kernel.h | 1 - src/native/cuda/nvidia/ops/rms_norm/kernel.h | 7 +- src/native/cuda/nvidia/ops/sum/kernel.h | 2 +- src/native/cuda/nvidia/ops/take/kernel.h | 4 +- .../nvidia/ops/upsample_bilinear2d/kernel.h | 1 - src/native/cuda/nvidia/ops/w4_marlin/gemm.cu | 5 +- src/native/cuda/nvidia/ops/w4_marlin/kernel.h | 3 - src/native/cuda/nvidia/ops/w4_marlin/launch.h | 17 + .../cuda/nvidia/ops/w4_marlin/marlin_check.h | 26 + .../cuda/nvidia/ops/w4_marlin/repack_awq.cu | 9 +- .../cuda/nvidia/ops/w4_marlin/repack_gptq.cu | 14 +- .../vendor/gemm/core/scalar_type.hpp | 240 ------- .../vendor/gemm/core/source_location.h | 44 -- .../ops/w4_marlin/vendor/gemm/core/utils.h | 142 ----- .../w4_marlin/vendor/gemm/nvidia/kernel.cuh | 553 ---------------- .../cuda/ops/separate_table_rope/kernel.h | 4 +- third_party/vllm | 1 - 47 files changed, 572 insertions(+), 2343 deletions(-) delete mode 100644 LICENSES/FlashInfer-0.6.7/LICENSE.cutlass.txt delete mode 100644 LICENSES/FlashInfer-0.6.7/LICENSE.flashattention3.txt delete mode 100644 LICENSES/FlashInfer-0.6.7/LICENSE.fmt.txt delete mode 100644 LICENSES/FlashInfer-0.6.7/LICENSE.spdlog.txt delete mode 100644 LICENSES/TVM-FFI-0.1.10/LICENSE.pytorch.txt delete mode 100644 LICENSES/TVM-FFI-0.1.10/NOTICE.pytorch.txt create mode 100644 cmake/MarlinSource.cmake delete mode 100644 src/native/ascend/ops/rotary_embedding/cann82_compat.hpp create mode 100644 src/native/cuda/nvidia/ops/w4_marlin/marlin_check.h delete mode 100644 src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/core/scalar_type.hpp delete mode 100644 src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/core/source_location.h delete mode 100644 src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/core/utils.h delete mode 100644 src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/nvidia/kernel.cuh delete mode 160000 third_party/vllm diff --git a/.github/workflows/clang-format.yml b/.github/workflows/clang-format.yml index 2ba0a570a..fcebc327c 100644 --- a/.github/workflows/clang-format.yml +++ b/.github/workflows/clang-format.yml @@ -9,4 +9,3 @@ jobs: with: clang-format-version: "21" include-regex: '^.*\.(h|cc|cuh|mlu)$' - exclude-regex: '^\./src/native/cuda/nvidia/ops/(native_flash|w4_marlin)/vendor/' diff --git a/.gitmodules b/.gitmodules index 77094cdd7..c313414c1 100644 --- a/.gitmodules +++ b/.gitmodules @@ -4,6 +4,3 @@ [submodule "third_party/flash-attention"] path = third_party/flash-attention url = https://github.com/Dao-AILab/flash-attention.git -[submodule "third_party/vllm"] - path = third_party/vllm - url = https://github.com/vllm-project/vllm.git diff --git a/CMakeLists.txt b/CMakeLists.txt index b4cdcb9f1..198127e79 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1,4 +1,14 @@ cmake_minimum_required(VERSION 3.18) +if(SKBUILD_STATE STREQUAL "sdist") + # Source packaging must not probe compilers, GPUs or installed SDKs. + project(InfiniOps LANGUAGES NONE) + include(cmake/MarlinSource.cmake) + foreach(_marlin_file IN LISTS _marlin_source_files) + configure_file("${vllm_marlin_SOURCE_DIR}/${_marlin_file}" + "${CMAKE_BINARY_DIR}/marlin-source/${_marlin_file}" COPYONLY) + endforeach() + return() +endif() project(InfiniOps LANGUAGES CXX) if(POLICY CMP0116) diff --git a/LICENSES/Apache-2.0-components.txt b/LICENSES/Apache-2.0-components.txt index 9b7be775a..72477a894 100644 --- a/LICENSES/Apache-2.0-components.txt +++ b/LICENSES/Apache-2.0-components.txt @@ -1,82 +1,40 @@ -Apache-2.0 components +Apache-2.0 components in the fixed native sampling/runtime recipe The complete common Apache License, Version 2.0, including sections 1-9 and -the application appendix, is in the repository LICENSE, installed beside -this file as share/licenses/InfiniOps/LICENSE. Distribute that file together -with this component list and the accompanying notices. - -This list consolidates the common Apache terms from the fixed releases below. -Their distinct copyrights, NOTICE files and secondary licenses remain -applicable. It is not a license inventory for arbitrary resolved versions. +the application appendix, is in repository LICENSE, installed beside this file +as share/licenses/InfiniOps/LICENSE. Distribute it with these copyrights, +the original project NOTICE files and the applicable secondary terms. +This is not a notice inventory for complete Python wheels or arbitrary binaries. FlashInfer 0.6.7 https://github.com/flashinfer-ai/flashinfer/tree/v0.6.7 +Commit: 56ed54098411946525da17eca2bcfd7ec72a515e Copyright 2025-2026 NVIDIA Copyright 2023-2026 FlashInfer community (https://flashinfer.ai/) -License: Apache-2.0, complete terms in LICENSE. -Original NOTICE retained in FlashInfer-0.6.7/NOTICE.txt. - -The following component list is retained from the release LICENSE: -------------------------------------------------------------------------------------------------- -Some of the code in this project are adapted from other open-source projects with different -licenses. This product also bundles some third-party components under other open source licenses. -This section summarizes those components and their licenses. -See licenses/ for text of these licenses. - -BSD 3-Clause License --------------------- - -include/flashinfer/attention/hopper/epilogue.cuh -include/flashinfer/attention/hopper/mainloop.cuh -include/flashinfer/attention/hopper/kernel_traits.cuh -include/flashinfer/attention/hopper/named_barrier.cuh -include/flashinfer/attention/hopper/tile_scheduler.cuh -include/flashinfer/attention/hopper/utils.cuh - -BSD 3-Clause "New" License --------------------------- - -3rdparty/cutlass -include/flashinfer/attention/hopper/block_sparse_gather.cuh - -MIT License ------------ - -3rdparty/spdlog -3rdparty/spdlog/include/spdlog/fmt/bundled (fmt library) - -The release's licenses/ materials are retained here in FlashInfer-0.6.7/: -LICENSE.flashattention3.txt, LICENSE.cutlass.txt, LICENSE.spdlog.txt, -LICENSE.fmt.txt. These keep each original copyright and complete terms. +Native source headers also retain Copyright (c) 2023, 2024, 2025 and 2026 by +FlashInfer team. Shared terms: LICENSE. Original NOTICE: FlashInfer-0.6.7/NOTICE.txt. +Sampling's three TUs are csrc/sampling.cu, csrc/renorm.cu and +csrc/flashinfer_sampling_binding.cu (flashinfer/jit/sampling.py). +The fixed source include closure does not use the release's separate +FA3/CUTLASS/fmt/spdlog components. Complete FlashInfer distributions retain their +own additional applicable materials, which this native-only list does not replace. Apache TVM FFI 0.1.10 https://github.com/apache/tvm-ffi/tree/v0.1.10 +Commit: 1fed0ae0421e614d45662e8ee6bcae353d3ab2ea Copyright 2024-present The Apache Software Foundation -License: Apache-2.0, complete terms in LICENSE. -Original NOTICE retained in TVM-FFI-0.1.10/NOTICE.txt. - -The following component list is retained from the release LICENSE: ------------------------------------------------------------------------------------- -This product bundles various third-party components under other open source licenses. -This section summarizes those components and their licenses. See licenses/ -for text of these licenses. - - -Apache Software Foundation License 2.0 --------------------------------------- - -3rdparty/dlpack - -BSD 3-Clause "New" or "Revised" License ---------------------------------------- - -3rdparty/libbacktrace +Shared terms: LICENSE. Original NOTICE: TVM-FFI-0.1.10/NOTICE.txt. +Native object source lists: CMakeLists.txt (_tvm_ffi_objs_sources and extras). +The separate python/tvm_ffi/_optional_torch_c_dlpack.py contains PyTorch-derived +code, but is not in the native runtime source list. Its notices still apply to +distributions including that addon; this list does not describe those distributions. DLPack: Copyright 2017 by Contributors -License: Apache-2.0, complete terms in LICENSE. - -The release's other licenses/ materials are retained in TVM-FFI-0.1.10/: -LICENSE.libbacktrace.txt, LICENSE.pytorch.txt, NOTICE.pytorch.txt. -Their complete original terms and copyrights are retained, including the -PyTorch secondary notices. This consolidation does not assert that every -component in the full releases is compiled into sampling.so or libtvm_ffi.so. +Shared Apache-2.0 terms: LICENSE. Its interfaces are used by the native runtime. +libbacktrace: original complete copyright and BSD-3-Clause terms retained in +TVM-FFI-0.1.10/LICENSE.libbacktrace.txt; retain them for the optional compiled support. +CUDA toolkit supplied headers/libraries remain subject to their own applicable +redistribution terms; this list does not replace those materials. + +Actual resolved-package license inventories are collected separately and can +include conservative full-package materials outside the two native DSOs. diff --git a/LICENSES/FlashInfer-0.6.7/LICENSE.cutlass.txt b/LICENSES/FlashInfer-0.6.7/LICENSE.cutlass.txt deleted file mode 100644 index 525500841..000000000 --- a/LICENSES/FlashInfer-0.6.7/LICENSE.cutlass.txt +++ /dev/null @@ -1,27 +0,0 @@ -Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -SPDX-License-Identifier: BSD-3-Clause - -Redistribution and use in source and binary forms, with or without -modification, are permitted provided that the following conditions are met: - -1. Redistributions of source code must retain the above copyright notice, this -list of conditions and the following disclaimer. - -2. Redistributions in binary form must reproduce the above copyright notice, -this list of conditions and the following disclaimer in the documentation -and/or other materials provided with the distribution. - -3. Neither the name of the copyright holder nor the names of its -contributors may be used to endorse or promote products derived from -this software without specific prior written permission. - -THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" -AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE -IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE -DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE -FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL -DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR -SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER -CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, -OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE -OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/LICENSES/FlashInfer-0.6.7/LICENSE.flashattention3.txt b/LICENSES/FlashInfer-0.6.7/LICENSE.flashattention3.txt deleted file mode 100644 index 5860e4b33..000000000 --- a/LICENSES/FlashInfer-0.6.7/LICENSE.flashattention3.txt +++ /dev/null @@ -1,29 +0,0 @@ -BSD 3-Clause License - -Copyright (c) 2022, the respective contributors, as shown by the AUTHORS file. -All rights reserved. - -Redistribution and use in source and binary forms, with or without -modification, are permitted provided that the following conditions are met: - -* Redistributions of source code must retain the above copyright notice, this - list of conditions and the following disclaimer. - -* Redistributions in binary form must reproduce the above copyright notice, - this list of conditions and the following disclaimer in the documentation - and/or other materials provided with the distribution. - -* Neither the name of the copyright holder nor the names of its - contributors may be used to endorse or promote products derived from - this software without specific prior written permission. - -THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" -AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE -IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE -DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE -FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL -DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR -SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER -CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, -OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE -OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/LICENSES/FlashInfer-0.6.7/LICENSE.fmt.txt b/LICENSES/FlashInfer-0.6.7/LICENSE.fmt.txt deleted file mode 100644 index f0ec3db4d..000000000 --- a/LICENSES/FlashInfer-0.6.7/LICENSE.fmt.txt +++ /dev/null @@ -1,27 +0,0 @@ -Copyright (c) 2012 - present, Victor Zverovich - -Permission is hereby granted, free of charge, to any person obtaining -a copy of this software and associated documentation files (the -"Software"), to deal in the Software without restriction, including -without limitation the rights to use, copy, modify, merge, publish, -distribute, sublicense, and/or sell copies of the Software, and to -permit persons to whom the Software is furnished to do so, subject to -the following conditions: - -The above copyright notice and this permission notice shall be -included in all copies or substantial portions of the Software. - -THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, -EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF -MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND -NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE -LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION -OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION -WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. - ---- Optional exception to the license --- - -As an exception, if, as a result of your compiling your source code, portions -of this Software are embedded into a machine-executable object form of such -source code, you may redistribute such embedded portions in such object form -without including the above copyright and permission notices. diff --git a/LICENSES/FlashInfer-0.6.7/LICENSE.spdlog.txt b/LICENSES/FlashInfer-0.6.7/LICENSE.spdlog.txt deleted file mode 100644 index c46826e51..000000000 --- a/LICENSES/FlashInfer-0.6.7/LICENSE.spdlog.txt +++ /dev/null @@ -1,25 +0,0 @@ -The MIT License (MIT) - -Copyright (c) 2016 Gabi Melman. - -Permission is hereby granted, free of charge, to any person obtaining a copy -of this software and associated documentation files (the "Software"), to deal -in the Software without restriction, including without limitation the rights -to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -copies of the Software, and to permit persons to whom the Software is -furnished to do so, subject to the following conditions: - -The above copyright notice and this permission notice shall be included in -all copies or substantial portions of the Software. - -THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN -THE SOFTWARE. - --- NOTE: Third party dependency used by this software -- -This software depends on the fmt lib (MIT License), -and users must comply to its license: https://raw.githubusercontent.com/fmtlib/fmt/master/LICENSE diff --git a/LICENSES/TVM-FFI-0.1.10/LICENSE.pytorch.txt b/LICENSES/TVM-FFI-0.1.10/LICENSE.pytorch.txt deleted file mode 100644 index 966a609b6..000000000 --- a/LICENSES/TVM-FFI-0.1.10/LICENSE.pytorch.txt +++ /dev/null @@ -1,84 +0,0 @@ -From PyTorch: - -Copyright (c) 2016- Facebook, Inc (Adam Paszke) -Copyright (c) 2014- Facebook, Inc (Soumith Chintala) -Copyright (c) 2011-2014 Idiap Research Institute (Ronan Collobert) -Copyright (c) 2012-2014 Deepmind Technologies (Koray Kavukcuoglu) -Copyright (c) 2011-2012 NEC Laboratories America (Koray Kavukcuoglu) -Copyright (c) 2011-2013 NYU (Clement Farabet) -Copyright (c) 2006-2010 NEC Laboratories America (Ronan Collobert, Leon Bottou, Iain Melvin, Jason Weston) -Copyright (c) 2006 Idiap Research Institute (Samy Bengio) -Copyright (c) 2001-2004 Idiap Research Institute (Ronan Collobert, Samy Bengio, Johnny Mariethoz) - -From Caffe2: - -Copyright (c) 2016-present, Facebook Inc. All rights reserved. - -All contributions by Facebook: -Copyright (c) 2016 Facebook Inc. - -All contributions by Google: -Copyright (c) 2015 Google Inc. -All rights reserved. - -All contributions by Yangqing Jia: -Copyright (c) 2015 Yangqing Jia -All rights reserved. - -All contributions by Kakao Brain: -Copyright 2019-2020 Kakao Brain - -All contributions by Cruise LLC: -Copyright (c) 2022 Cruise LLC. -All rights reserved. - -All contributions by Tri Dao: -Copyright (c) 2024 Tri Dao. -All rights reserved. - -All contributions by Arm: -Copyright (c) 2021, 2023-2024 Arm Limited and/or its affiliates - -All contributions from Caffe: -Copyright(c) 2013, 2014, 2015, the respective contributors -All rights reserved. - -All other contributions: -Copyright(c) 2015, 2016 the respective contributors -All rights reserved. - -Caffe2 uses a copyright model similar to Caffe: each contributor holds -copyright over their contributions to Caffe2. The project versioning records -all such contribution and copyright details. If a contributor wants to further -mark their specific copyright on a particular contribution, they should -indicate their copyright solely in the commit message of the change when it is -committed. - -All rights reserved. - -Redistribution and use in source and binary forms, with or without -modification, are permitted provided that the following conditions are met: - -1. Redistributions of source code must retain the above copyright - notice, this list of conditions and the following disclaimer. - -2. Redistributions in binary form must reproduce the above copyright - notice, this list of conditions and the following disclaimer in the - documentation and/or other materials provided with the distribution. - -3. Neither the names of Facebook, Deepmind Technologies, NYU, NEC Laboratories America - and IDIAP Research Institute nor the names of its contributors may be - used to endorse or promote products derived from this software without - specific prior written permission. - -THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" -AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE -IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE -ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE -LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR -CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF -SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS -INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN -CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) -ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE -POSSIBILITY OF SUCH DAMAGE. diff --git a/LICENSES/TVM-FFI-0.1.10/NOTICE.pytorch.txt b/LICENSES/TVM-FFI-0.1.10/NOTICE.pytorch.txt deleted file mode 100644 index 6effb8b5d..000000000 --- a/LICENSES/TVM-FFI-0.1.10/NOTICE.pytorch.txt +++ /dev/null @@ -1,456 +0,0 @@ -======================================================================= -Software under third_party -======================================================================= -Software libraries under third_party are provided as github submodule -links, and their content is not part of the Caffe2 codebase. Their -licences can be found under the respective software repositories. - -======================================================================= -Earlier BSD License -======================================================================= -Early development of Caffe2 in 2015 and early 2016 is licensed under the -BSD license. The license is attached below: - -All contributions by Facebook: -Copyright (c) 2016 Facebook Inc. - -All contributions by Google: -Copyright (c) 2015 Google Inc. -All rights reserved. - -All contributions by Yangqing Jia: -Copyright (c) 2015 Yangqing Jia -All rights reserved. - -All contributions by Kakao Brain: -Copyright 2019-2020 Kakao Brain - -All other contributions: -Copyright(c) 2015, 2016 the respective contributors -All rights reserved. - -Redistribution and use in source and binary forms, with or without -modification, are permitted provided that the following conditions are met: - -1. Redistributions of source code must retain the above copyright notice, this - list of conditions and the following disclaimer. -2. Redistributions in binary form must reproduce the above copyright notice, - this list of conditions and the following disclaimer in the documentation - and/or other materials provided with the distribution. - -THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND -ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED -WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE -DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR -ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES -(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; -LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND -ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT -(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS -SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. - - -======================================================================= -Caffe's BSD License -======================================================================= -Some parts of the caffe2 code is derived from the original Caffe code, which is -created by Yangqing Jia and is now a BSD-licensed open-source project. The Caffe -license is as follows: - -COPYRIGHT - -All contributions by the University of California: -Copyright (c) 2014, The Regents of the University of California (Regents) -All rights reserved. - -All other contributions: -Copyright (c) 2014, the respective contributors -All rights reserved. - -Caffe uses a shared copyright model: each contributor holds copyright over -their contributions to Caffe. The project versioning records all such -contribution and copyright details. If a contributor wants to further mark -their specific copyright on a particular contribution, they should indicate -their copyright solely in the commit message of the change when it is -committed. - -LICENSE - -Redistribution and use in source and binary forms, with or without -modification, are permitted provided that the following conditions are met: - -1. Redistributions of source code must retain the above copyright notice, this - list of conditions and the following disclaimer. -2. Redistributions in binary form must reproduce the above copyright notice, - this list of conditions and the following disclaimer in the documentation - and/or other materials provided with the distribution. - -THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND -ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED -WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE -DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR -ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES -(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; -LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND -ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT -(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS -SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. - -CONTRIBUTION AGREEMENT - -By contributing to the BVLC/caffe repository through pull-request, comment, -or otherwise, the contributor releases their content to the -license and copyright terms herein. - -======================================================================= -Caffe2's Apache License -======================================================================= - -This repo contains Caffe2 code, which was previously licensed under -Apache License Version 2.0: - - Apache License - Version 2.0, January 2004 - http://www.apache.org/licenses/ - - TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION - - 1. Definitions. - - "License" shall mean the terms and conditions for use, reproduction, - and distribution as defined by Sections 1 through 9 of this document. - - "Licensor" shall mean the copyright owner or entity authorized by - the copyright owner that is granting the License. - - "Legal Entity" shall mean the union of the acting entity and all - other entities that control, are controlled by, or are under common - control with that entity. For the purposes of this definition, - "control" means (i) the power, direct or indirect, to cause the - direction or management of such entity, whether by contract or - otherwise, or (ii) ownership of fifty percent (50%) or more of the - outstanding shares, or (iii) beneficial ownership of such entity. - - "You" (or "Your") shall mean an individual or Legal Entity - exercising permissions granted by this License. - - "Source" form shall mean the preferred form for making modifications, - including but not limited to software source code, documentation - source, and configuration files. - - "Object" form shall mean any form resulting from mechanical - transformation or translation of a Source form, including but - not limited to compiled object code, generated documentation, - and conversions to other media types. - - "Work" shall mean the work of authorship, whether in Source or - Object form, made available under the License, as indicated by a - copyright notice that is included in or attached to the work - (an example is provided in the Appendix below). - - "Derivative Works" shall mean any work, whether in Source or Object - form, that is based on (or derived from) the Work and for which the - editorial revisions, annotations, elaborations, or other modifications - represent, as a whole, an original work of authorship. For the purposes - of this License, Derivative Works shall not include works that remain - separable from, or merely link (or bind by name) to the interfaces of, - the Work and Derivative Works thereof. - - "Contribution" shall mean any work of authorship, including - the original version of the Work and any modifications or additions - to that Work or Derivative Works thereof, that is intentionally - submitted to Licensor for inclusion in the Work by the copyright owner - or by an individual or Legal Entity authorized to submit on behalf of - the copyright owner. For the purposes of this definition, "submitted" - means any form of electronic, verbal, or written communication sent - to the Licensor or its representatives, including but not limited to - communication on electronic mailing lists, source code control systems, - and issue tracking systems that are managed by, or on behalf of, the - Licensor for the purpose of discussing and improving the Work, but - excluding communication that is conspicuously marked or otherwise - designated in writing by the copyright owner as "Not a Contribution." - - "Contributor" shall mean Licensor and any individual or Legal Entity - on behalf of whom a Contribution has been received by Licensor and - subsequently incorporated within the Work. - - 2. Grant of Copyright License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - copyright license to reproduce, prepare Derivative Works of, - publicly display, publicly perform, sublicense, and distribute the - Work and such Derivative Works in Source or Object form. - - 3. Grant of Patent License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - (except as stated in this section) patent license to make, have made, - use, offer to sell, sell, import, and otherwise transfer the Work, - where such license applies only to those patent claims licensable - by such Contributor that are necessarily infringed by their - Contribution(s) alone or by combination of their Contribution(s) - with the Work to which such Contribution(s) was submitted. If You - institute patent litigation against any entity (including a - cross-claim or counterclaim in a lawsuit) alleging that the Work - or a Contribution incorporated within the Work constitutes direct - or contributory patent infringement, then any patent licenses - granted to You under this License for that Work shall terminate - as of the date such litigation is filed. - - 4. Redistribution. You may reproduce and distribute copies of the - Work or Derivative Works thereof in any medium, with or without - modifications, and in Source or Object form, provided that You - meet the following conditions: - - (a) You must give any other recipients of the Work or - Derivative Works a copy of this License; and - - (b) You must cause any modified files to carry prominent notices - stating that You changed the files; and - - (c) You must retain, in the Source form of any Derivative Works - that You distribute, all copyright, patent, trademark, and - attribution notices from the Source form of the Work, - excluding those notices that do not pertain to any part of - the Derivative Works; and - - (d) If the Work includes a "NOTICE" text file as part of its - distribution, then any Derivative Works that You distribute must - include a readable copy of the attribution notices contained - within such NOTICE file, excluding those notices that do not - pertain to any part of the Derivative Works, in at least one - of the following places: within a NOTICE text file distributed - as part of the Derivative Works; within the Source form or - documentation, if provided along with the Derivative Works; or, - within a display generated by the Derivative Works, if and - wherever such third-party notices normally appear. The contents - of the NOTICE file are for informational purposes only and - do not modify the License. You may add Your own attribution - notices within Derivative Works that You distribute, alongside - or as an addendum to the NOTICE text from the Work, provided - that such additional attribution notices cannot be construed - as modifying the License. - - You may add Your own copyright statement to Your modifications and - may provide additional or different license terms and conditions - for use, reproduction, or distribution of Your modifications, or - for any such Derivative Works as a whole, provided Your use, - reproduction, and distribution of the Work otherwise complies with - the conditions stated in this License. - - 5. Submission of Contributions. Unless You explicitly state otherwise, - any Contribution intentionally submitted for inclusion in the Work - by You to the Licensor shall be under the terms and conditions of - this License, without any additional terms or conditions. - Notwithstanding the above, nothing herein shall supersede or modify - the terms of any separate license agreement you may have executed - with Licensor regarding such Contributions. - - 6. Trademarks. This License does not grant permission to use the trade - names, trademarks, service marks, or product names of the Licensor, - except as required for reasonable and customary use in describing the - origin of the Work and reproducing the content of the NOTICE file. - - 7. Disclaimer of Warranty. Unless required by applicable law or - agreed to in writing, Licensor provides the Work (and each - Contributor provides its Contributions) on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or - implied, including, without limitation, any warranties or conditions - of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A - PARTICULAR PURPOSE. You are solely responsible for determining the - appropriateness of using or redistributing the Work and assume any - risks associated with Your exercise of permissions under this License. - - 8. Limitation of Liability. In no event and under no legal theory, - whether in tort (including negligence), contract, or otherwise, - unless required by applicable law (such as deliberate and grossly - negligent acts) or agreed to in writing, shall any Contributor be - liable to You for damages, including any direct, indirect, special, - incidental, or consequential damages of any character arising as a - result of this License or out of the use or inability to use the - Work (including but not limited to damages for loss of goodwill, - work stoppage, computer failure or malfunction, or any and all - other commercial damages or losses), even if such Contributor - has been advised of the possibility of such damages. - - 9. Accepting Warranty or Additional Liability. While redistributing - the Work or Derivative Works thereof, You may choose to offer, - and charge a fee for, acceptance of support, warranty, indemnity, - or other liability obligations and/or rights consistent with this - License. However, in accepting such obligations, You may act only - on Your own behalf and on Your sole responsibility, not on behalf - of any other Contributor, and only if You agree to indemnify, - defend, and hold each Contributor harmless for any liability - incurred by, or claims asserted against, such Contributor by reason - of your accepting any such warranty or additional liability. - -======================================================================= -Cephes's 3-Clause BSD License -======================================================================= - -Code derived from implementations in the Cephes Math Library should mention -its derivation and reference the following license: - - 3-Clause BSD License for the Cephes Math Library - Copyright (c) 2018, Steven Moshier - All rights reserved. - - Redistribution and use in source and binary forms, with or without - modification, are permitted provided that the following conditions are met: - - * Redistributions of source code must retain the above copyright - notice, this list of conditions and the following disclaimer. - - * Redistributions in binary form must reproduce the above copyright - notice, this list of conditions and the following disclaimer in the - documentation and/or other materials provided with the distribution. - - * Neither the name of the nor the - names of its contributors may be used to endorse or promote products - derived from this software without specific prior written permission. - - THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND - ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED - WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE - DISCLAIMED. IN NO EVENT SHALL Steven Moshier BE LIABLE FOR ANY - DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES - (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; - LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND - ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT - (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS - SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. - - -======================================================================= -SciPy's 3-Clause BSD License -======================================================================= - -Code derived from implementations in SciPy should mention its derivation -and reference the following license: - - Copyright (c) 2001-2002 Enthought, Inc. 2003-2019, SciPy Developers. - All rights reserved. - - Redistribution and use in source and binary forms, with or without - modification, are permitted provided that the following conditions - are met: - - 1. Redistributions of source code must retain the above copyright - notice, this list of conditions and the following disclaimer. - - 2. Redistributions in binary form must reproduce the above - copyright notice, this list of conditions and the following - disclaimer in the documentation and/or other materials provided - with the distribution. - - 3. Neither the name of the copyright holder nor the names of its - contributors may be used to endorse or promote products derived - from this software without specific prior written permission. - - THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS - "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT - LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR - A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT - OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, - SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT - LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, - DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY - THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT - (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE - OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. - -======================================================================= -Boost's 1.0 Software License -======================================================================= - -Code derived from implementations in Boost 1.0 should mention its -derivation and reference the following license: - - Boost Software License - Version 1.0 - August 17th, 2003 - - Permission is hereby granted, free of charge, to any person or organization - obtaining a copy of the software and accompanying documentation covered by - this license (the "Software") to use, reproduce, display, distribute, - execute, and transmit the Software, and to prepare derivative works of the - Software, and to permit third-parties to whom the Software is furnished to - do so, all subject to the following: - - The copyright notices in the Software and this entire statement, including - the above license grant, this restriction and the following disclaimer, - must be included in all copies of the Software, in whole or in part, and - all derivative works of the Software, unless such copies or derivative - works are solely in the form of machine-executable object code generated by - a source language processor. - - THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR - IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, - FITNESS FOR A PARTICULAR PURPOSE, TITLE AND NON-INFRINGEMENT. IN NO EVENT - SHALL THE COPYRIGHT HOLDERS OR ANYONE DISTRIBUTING THE SOFTWARE BE LIABLE - FOR ANY DAMAGES OR OTHER LIABILITY, WHETHER IN CONTRACT, TORT OR OTHERWISE, - ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER - DEALINGS IN THE SOFTWARE. - - END OF TERMS AND CONDITIONS - - APPENDIX: How to apply the Apache License to your work. - - To apply the Apache License to your work, attach the following - boilerplate notice, with the fields enclosed by brackets "[]" - replaced with your own identifying information. (Don't include - the brackets!) The text should be enclosed in the appropriate - comment syntax for the file format. We also recommend that a - file or class name and description of purpose be included on the - same "printed page" as the copyright notice for easier - identification within third-party archives. - - Copyright [yyyy] [name of copyright owner] - - Licensed under the Apache License, Version 2.0 (the "License"); - you may not use this file except in compliance with the License. - You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - - Unless required by applicable law or agreed to in writing, software - distributed under the License is distributed on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - See the License for the specific language governing permissions and - limitations under the License. - -======================================================================= -PILLOW-SIMD Software License -======================================================================= - -Code derived from implementations in PILLOW-SIMD should mention its derivation -and reference the following license: - - The Python Imaging Library (PIL) is - - Copyright © 1997-2011 by Secret Labs AB - Copyright © 1995-2011 by Fredrik Lundh - - Pillow is the friendly PIL fork. It is - - Copyright © 2010-2022 by Alex Clark and contributors - - Like PIL, Pillow is licensed under the open source HPND License: - - By obtaining, using, and/or copying this software and/or its associated - documentation, you agree that you have read, understood, and will comply - with the following terms and conditions: - - Permission to use, copy, modify, and distribute this software and its - associated documentation for any purpose and without fee is hereby granted, - provided that the above copyright notice appears in all copies, and that - both that copyright notice and this permission notice appear in supporting - documentation, and that the name of Secret Labs AB or the author not be - used in advertising or publicity pertaining to distribution of the software - without specific, written prior permission. - - SECRET LABS AB AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH REGARD TO THIS - SOFTWARE, INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS. - IN NO EVENT SHALL SECRET LABS AB OR THE AUTHOR BE LIABLE FOR ANY SPECIAL, - INDIRECT OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM - LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE - OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR - PERFORMANCE OF THIS SOFTWARE. diff --git a/THIRD_PARTY.txt b/THIRD_PARTY.txt index f9965ac8a..5afd54862 100644 --- a/THIRD_PARTY.txt +++ b/THIRD_PARTY.txt @@ -10,29 +10,32 @@ Earlier upstream revisions are not asserted where only the fixed direct Core don Path prefixes: P = src/native/cuda/nvidia/ops/ C = src/infiniop/ops/ in the pinned direct Core donor -W = P/w4_marlin/vendor/ W4 / Marlin Upstream projects: https://github.com/IST-DASLab/marlin and https://github.com/vllm-project/vllm -Current kernel source: third_party/vllm, pinned gitlink +Current kernel source: the fixed vLLM source archive at commit ffc4f08c8ee130d4ea6347c1bf31ffd4f8af28ab, -csrc/libtorch_stable/quantization/marlin. Only its source kernels and generator +SHA256 d8f3d663347dff35ea8123d3eb570f67d2dd6c4258ee8b8a8ec2996c42ada678, +csrc/libtorch_stable/quantization/marlin and csrc/core/scalar_type.hpp. +Its source kernels, native pointer host entry, scalar definition and generator are consumed; the vLLM/Torch runtime is not a dependency of this native target. +cmake/MarlinSource.cmake acquires that archive through standard FetchContent. +Complete source distributions contain LICENSE plus the same eleven inputs under +third_party/marlin-source, and use those files without downloading. The original Ops port came through the fixed Core donor above; this selected vLLM revision is not asserted to be the unique historical origin of that port. Preparation applies patches/marlin to build-directory copies, retaining the existing Torch-free includes, repack bindings and activation-order index behavior. CUDA template instances and their selector are generated together in the build -directory. The four host/type adaptations below remain in Ops. +directory. The upstream marlin.cu native pointer entry and scalar_type.hpp are +prepared alongside the kernels; no copied host/type/helper subtree is retained +in Ops. A private declaration and small local check helper adapt the native +entry without Torch types, while existing CUDA error propagation is retained. Marlin notices retain Copyright 2024 Elias Frantar, modifications by Neural Magic, Apache-2.0. The original generator identifies Apache-2.0 / Copyright contributors to the vLLM project. -Marlin original Apache headers remain in the prepared source and nvidia/kernel.cuh; +Marlin original Apache headers remain in the prepared source, including marlin.cu; Apache text: LICENSE. Original source notices remain in prepared files. -Remaining adaptations <- direct Core donor: -COPY W/gemm/core/scalar_type.hpp <- C/awq_marlin_gemm/core/scalar_type.hpp -COPY W/gemm/core/source_location.h <- C/awq_marlin_gemm/core/source_location.h -ADAPTED W/gemm/core/utils.h <- C/awq_marlin_gemm/core/utils.h (unused pointer/range/dtype helpers removed; runtime checks unchanged) -The prepared marlin/dequant.h INT4/INT8 conversion portions also derive from NVIDIA FasterTransformer: +The prepared dequant.h INT4/INT8 conversion portions also derive from NVIDIA FasterTransformer: https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h The dequant.h references map INT4 FP16/BF16 to upstream lines 215-287/327-385 and INT8 FP16/BF16 to upstream lines 53-85/125-175; packing, bias and scalar @@ -41,7 +44,6 @@ Original Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES and the complete BSD-3-Clause terms are retained in the prepared dequant.h and LICENSES/FasterTransformer-BSD-3-Clause.txt. The separate CUTLASS dependency's 2017-2026 notice does not replace this original copyright notice. -COPY W/gemm/nvidia/kernel.cuh <- C/awq_marlin_gemm/nvidia/kernel.cuh ADAPTED P/w4_marlin/parameters.cu implements the scale/zero permutations from vLLM's vllm/model_executor/layers/quantization/utils/marlin_utils.py; the fixed reference revision and loading-only CUDA adaptation are noted in that file. @@ -91,14 +93,21 @@ CUTLASS-BSD-3-Clause.txt records the separate CMake dependency at Linked libraries are not relicensed by Ops: distributions bundling sampling.so or libtvm_ffi.so must retain the resolved FlashInfer/TVM-FFI licenses, NOTICE files and applicable third-party notices. See docs/build-nvidia-inference.md for the pinned recipe. -LICENSES/Apache-2.0-components.txt maps the corresponding release components -and copyrights to the complete common Apache terms in LICENSE, and preserves -their original secondary-component lists. LICENSES/FlashInfer-0.6.7 and -LICENSES/TVM-FFI-0.1.10 retain the upstream NOTICE and secondary-license texts -verbatim; the repeated common Apache terms are consolidated, not omitted. -They are not a substitute for the notices of a different resolved binary version. -These release bundles conservatively retain upstream secondary notices, including -PyTorch-derived TVM-FFI code; not importing Torch does not remove source-derived -obligations. The full Python wheels contain more components than sampling.so and -libtvm_ffi.so alone. Do not infer a native-only notice subset from Python imports -or from an absent wheel license directory; check the compiled source provenance. +LICENSES/Apache-2.0-components.txt records the fixed native recipe's copyrights +and maps the common Apache terms to LICENSE. The original FlashInfer and TVM-FFI +NOTICE files and libbacktrace terms are retained verbatim. DLPack's Apache terms +are shared, with its distinct copyright recorded in that component list. +FlashInfer v0.6.7 (56ed54098411946525da17eca2bcfd7ec72a515e) generates sampling +from csrc/{sampling.cu,renorm.cu,flashinfer_sampling_binding.cu}; their source +include closure does not contain FA3, CUTLASS, fmt or spdlog. TVM-FFI v0.1.10 +(1fed0ae0421e614d45662e8ee6bcae353d3ab2ea) builds libtvm_ffi from the native +object lists in CMakeLists.txt, including DLPack interfaces and optional +libbacktrace. Its PyTorch-derived optional Python DLPack addon is not in those +object lists. These source-based distinctions, not Python import behavior, +scope the checked-in fallback materials to the two native artifacts. +They do not cover arbitrary replacement binaries or the complete Python wheels. +The resolver additionally collects existing license/NOTICE/COPYING materials +from actual package metadata. That conservative package inventory can include +materials for components outside these two DSOs; it is intentionally not filtered +by fixed hashes or used to restrict build capability. A complete wheel or SDK +redistribution must retain its own applicable notices, including secondary ones. diff --git a/cmake/MarlinSource.cmake b/cmake/MarlinSource.cmake new file mode 100644 index 000000000..f2e285e90 --- /dev/null +++ b/cmake/MarlinSource.cmake @@ -0,0 +1,31 @@ +# Source acquisition only; never configure or link the upstream vLLM project. +include(FetchContent) +set(_marlin_bundled "${CMAKE_CURRENT_LIST_DIR}/../third_party/marlin-source") +if(NOT FETCHCONTENT_SOURCE_DIR_VLLM_MARLIN AND IS_DIRECTORY "${_marlin_bundled}") + # Complete sdists carry the same eleven upstream inputs, without Git metadata. + set(FETCHCONTENT_SOURCE_DIR_VLLM_MARLIN "${_marlin_bundled}") +endif() +cmake_policy(PUSH) +if(POLICY CMP0135) + cmake_policy(SET CMP0135 NEW) +endif() +FetchContent_Declare(vllm_marlin + URL "https://github.com/vllm-project/vllm/archive/ffc4f08c8ee130d4ea6347c1bf31ffd4f8af28ab.tar.gz" + URL_HASH SHA256=d8f3d663347dff35ea8123d3eb570f67d2dd6c4258ee8b8a8ec2996c42ada678 + SOURCE_SUBDIR csrc/libtorch_stable/quantization/marlin) +FetchContent_MakeAvailable(vllm_marlin) +cmake_policy(POP) + +set(_marlin_source_files LICENSE) +foreach(_marlin_file dequant.h kernel.h marlin.cuh marlin_dtypes.cuh + marlin_mma.h marlin_template.h awq_marlin_repack.cu + gptq_marlin_repack.cu generate_kernels.py marlin.cu) + list(APPEND _marlin_source_files + "csrc/libtorch_stable/quantization/marlin/${_marlin_file}") +endforeach() +list(APPEND _marlin_source_files "csrc/core/scalar_type.hpp") +foreach(_marlin_file IN LISTS _marlin_source_files) + if(NOT EXISTS "${vllm_marlin_SOURCE_DIR}/${_marlin_file}") + message(FATAL_ERROR "Pinned Marlin source is missing ${_marlin_file}") + endif() +endforeach() diff --git a/cmake/NativeW4Marlin.cmake b/cmake/NativeW4Marlin.cmake index d84aaad41..e7cc9ed57 100644 --- a/cmake/NativeW4Marlin.cmake +++ b/cmake/NativeW4Marlin.cmake @@ -4,21 +4,19 @@ endif() set(_w4_root "${PROJECT_SOURCE_DIR}/src/native/cuda/nvidia/ops/w4_marlin") find_package(Python REQUIRED COMPONENTS Interpreter) find_program(_w4_patch NAMES patch REQUIRED) -set(_w4_upstream "${PROJECT_SOURCE_DIR}/third_party/vllm") -set(_w4_upstream_marlin "${_w4_upstream}/csrc/libtorch_stable/quantization/marlin") +include("${CMAKE_CURRENT_LIST_DIR}/MarlinSource.cmake") +set(_w4_upstream "${vllm_marlin_SOURCE_DIR}") set(_w4_prepared "${CMAKE_CURRENT_BINARY_DIR}/native-w4-marlin") set(_w4_inputs "${PROJECT_SOURCE_DIR}/scripts/prepare_marlin.py" "${PROJECT_SOURCE_DIR}/patches/marlin/w4a16.patch") -foreach(_w4_input dequant.h kernel.h marlin.cuh marlin_dtypes.cuh - marlin_mma.h marlin_template.h awq_marlin_repack.cu - gptq_marlin_repack.cu generate_kernels.py) - list(APPEND _w4_inputs "${_w4_upstream_marlin}/${_w4_input}") +foreach(_w4_input IN LISTS _marlin_source_files) + list(APPEND _w4_inputs "${_w4_upstream}/${_w4_input}") endforeach() foreach(_w4_input IN LISTS _w4_inputs) if(NOT EXISTS "${_w4_input}") message(FATAL_ERROR - "Native Marlin is missing ${_w4_input}. Initialize third_party/vllm.") + "Native Marlin is missing ${_w4_input}. Check the pinned source archive or FETCHCONTENT_SOURCE_DIR_VLLM_MARLIN.") endif() endforeach() set_property(DIRECTORY APPEND PROPERTY CMAKE_CONFIGURE_DEPENDS ${_w4_inputs}) @@ -36,22 +34,23 @@ add_library(infiniops_w4_marlin OBJECT "${_w4_root}/repack_gptq.cu" "${_w4_root}/parameters.cu" "${_w4_root}/gemm.cu" - "${_w4_prepared}/vendor/gemm/nvidia/sm80_kernel_bfloat16_u4_bfloat16.cu" - "${_w4_prepared}/vendor/gemm/nvidia/sm80_kernel_bfloat16_u4b8_bfloat16.cu" - "${_w4_prepared}/vendor/gemm/nvidia/sm80_kernel_float16_u4_float16.cu" - "${_w4_prepared}/vendor/gemm/nvidia/sm80_kernel_float16_u4b8_float16.cu") + "${_w4_prepared}/marlin.cu" + "${_w4_prepared}/sm80_kernel_bfloat16_u4_bfloat16.cu" + "${_w4_prepared}/sm80_kernel_bfloat16_u4b8_bfloat16.cu" + "${_w4_prepared}/sm80_kernel_float16_u4_float16.cu" + "${_w4_prepared}/sm80_kernel_float16_u4b8_float16.cu") set_target_properties(infiniops_w4_marlin PROPERTIES POSITION_INDEPENDENT_CODE ON CXX_STANDARD 20 CXX_STANDARD_REQUIRED ON CUDA_STANDARD 20 CUDA_STANDARD_REQUIRED ON CUDA_ARCHITECTURES 80) target_include_directories(infiniops_w4_marlin PRIVATE - "${_w4_prepared}" "${_w4_prepared}/vendor/gemm" - "${_w4_prepared}/vendor/gemm/nvidia" "${_w4_root}/vendor/gemm" + "${_w4_prepared}" "${_w4_root}" "${PROJECT_SOURCE_DIR}/src" "${PROJECT_SOURCE_DIR}/include" ${INFINI_RT_INCLUDE_DIRS}) target_link_libraries(infiniops_w4_marlin PRIVATE infinirt CUDA::cudart) -target_compile_definitions(infiniops_w4_marlin PRIVATE WITH_NVIDIA=1) +target_compile_definitions(infiniops_w4_marlin PRIVATE + WITH_NVIDIA=1 INFINI_OPS_MARLIN_KERNEL_ONLY=1) target_compile_options(infiniops_w4_marlin PRIVATE $<$:--expt-relaxed-constexpr;--expt-extended-lambda>) target_sources(infiniops PRIVATE $) diff --git a/docs/build-nvidia-inference.md b/docs/build-nvidia-inference.md index 6efd9bdde..c537c0682 100644 --- a/docs/build-nvidia-inference.md +++ b/docs/build-nvidia-inference.md @@ -15,7 +15,8 @@ execute their Python inference APIs. No `flashinfer-python` or Torch is required | TVM FFI | `apache_tvm_ffi-0.1.10-cp312-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl` | `73d478395a8625dd92fde7b7fd92b4719f18f480b78336e422cb66cc7985213d` | | CUTLASS | commit `087c84df83d254b5fb295a7a408f1a1d554085cf` archive | `99340f7a8c87f5295b91aecf8295df6693ad954807672c3869a7fb3d0084f709` | -Native Attention and W4 kernels use the fixed submodule revisions below. They +Native Attention uses the fixed submodule revisions below. W4 kernels use the +fixed vLLM source archive through CMake FetchContent. They are source inputs only; neither the Torch extension nor the vLLM Python package is built or installed. @@ -23,7 +24,22 @@ is built or installed. |---|---| | `third_party/flash-attention` | `060c9188beec3a8b62b33a3bfa6d5d2d44975fab` | | Its `csrc/cutlass` submodule | `dc4817921edda44a549197ff3a9dcf5df0636e7b` | -| `third_party/vllm` (Marlin only) | `ffc4f08c8ee130d4ea6347c1bf31ffd4f8af28ab` | +| vLLM source archive (Marlin only) | `ffc4f08c8ee130d4ea6347c1bf31ffd4f8af28ab` | + +The vLLM archive SHA256 is +`d8f3d663347dff35ea8123d3eb570f67d2dd6c4258ee8b8a8ec2996c42ada678`. +It is a build source input, not a Git submodule or an installed package. +The upstream project is never configured. For an offline checkout build, pass +`-DFETCHCONTENT_SOURCE_DIR_VLLM_MARLIN=/absolute/path/to/pinned/extracted/vllm`; +the standard `FETCHCONTENT_FULLY_DISCONNECTED=ON` can then disable downloading. +Complete sdists bundle LICENSE and the eleven original Marlin inputs under +`third_party/marlin-source`, which CMake selects automatically without network. +This includes the native `marlin.cu` host entry and `csrc/core/scalar_type.hpp`; +Ops keeps only its private declaration and small error adaptation, not copies. +Building an sdist uses the standard scikit-build `sdist.cmake` flow and +`build/sdist` staging; an offline packager can pass the same source override. +Keep that sdist build directory unless also overriding its force-include path. +No generated Marlin CUDA files are stored in the source package. The other CUTLASS input above remains the existing dependency of ordinary Ops providers. Its version is not substituted for FlashAttention's private headers. @@ -41,7 +57,7 @@ version/hash and run the corresponding direct operator tests again. `LICENSES/Apache-2.0-components.txt` records the release component copyrights and maps their shared Apache terms to `LICENSE`. `LICENSES/FlashInfer-0.6.7` and `LICENSES/TVM-FFI-0.1.10` retain the original -release NOTICE and secondary-license texts for these exact inputs. +NOTICE files and native-runtime libbacktrace terms for these exact inputs. If replacing either wheel, update those redistribution materials from the actual release as well. SDK installation puts these materials, the native FlashAttention notices and the Ops attribution list in `share/licenses/InfiniOps`; @@ -59,7 +75,7 @@ export INFINI_ROOT="$PWD/sdk" export CUDA_HOME=/usr/local/cuda export INFINI_RT_ROOT="$INFINI_ROOT" git -C "$CORE/submodules/InfiniOps" submodule update --init \ - third_party/flash-attention third_party/vllm + third_party/flash-attention git -C "$CORE/submodules/InfiniOps/third_party/flash-attention" \ submodule update --init csrc/cutlass mkdir artifacts @@ -146,15 +162,21 @@ FlashInfer 0.6.7 and TVM-FFI 0.1.10 inputs into SDK license directory when bundling the native libraries. No additional download or manual copy into the LM source tree is needed for these inputs. -The checked-in release notices cover a conservative superset of the two native -libraries. `sampling.so` is one JIT-cache artifact, not the entire FlashInfer -wheel; `libtvm_ffi.so` is the TVM-FFI runtime, not its full Python package. -In particular, TVM-FFI contains code adapted from PyTorch, and can compile -libbacktrace and DLPack interfaces into its runtime. A Torch-free execution path -does not establish that these secondary notices can be removed. Retain the -release bundles unless a source-to-binary audit establishes a smaller set for -the exact library hashes. Distributing the complete wheels requires their -complete applicable notices, independently of the Ops SDK. +The fixed fallback is installed only for a linked TVM-FFI transport, not for +every linked backend. `sampling.so` is one JIT-cache artifact, not the entire +FlashInfer wheel; `libtvm_ffi.so` is the native runtime, not its Python package. +The fixed source closures described in `THIRD_PARTY.txt` exclude FlashInfer's +FA3/CUTLASS/fmt/spdlog components and TVM-FFI's separate PyTorch-derived Python +addon. Their six secondary license/notice copies are therefore not part of the +checked-in native fallback. Common Apache terms, original project NOTICEs, +DLPack attribution and libbacktrace terms remain. + +The resolver also preserves real license inventories from the distributions +actually selected. They can conservatively include full-package materials +beyond these DSOs; the final installed SDK is not claimed to be reduced to the +fixed fallback alone. For different binaries, validate source correspondence +and preserve the actual applicable materials. Complete wheel/SDK distributions +retain their own applicable notices independently of this native-only recipe. Preserve the separate FlashAttention/CUTLASS and Marlin notices described in `THIRD_PARTY.txt` when distributing those compiled providers. Do not treat this diff --git a/patches/marlin/w4a16.patch b/patches/marlin/w4a16.patch index 2af0f7260..e3bfcd0b3 100644 --- a/patches/marlin/w4a16.patch +++ b/patches/marlin/w4a16.patch @@ -1,5 +1,5 @@ ---- a/gemm/marlin/dequant.h -+++ b/gemm/marlin/dequant.h +--- a/dequant.h ++++ b/dequant.h @@ -1,3 +1,39 @@ +// Source: InfiniCore 1ab85ef1d3c2bce49edf1fc73a5b48becaf7552e, +// src/infiniop/ops/awq_marlin_gemm/marlin/dequant.h. @@ -40,27 +40,8 @@ /* Fast Dequantization (Converting INT4/INT8/FP4/FP8 to FP16/BF16) ---- a/gemm/marlin/marlin.cuh -+++ b/gemm/marlin/marlin.cuh -@@ -10,6 +10,16 @@ - #ifndef MARLIN_NAMESPACE_NAME - #define MARLIN_NAMESPACE_NAME marlin - #endif -+ -+template -+__device__ __forceinline__ uint32_t __cvta_generic_to_shared(T *ptr) { -+ size_t smem_addr; -+ asm volatile( -+ "cvta.to.shared.u64 %0, %1;" -+ : "=l"(smem_addr) -+ : "l"(ptr)); -+ return static_cast(smem_addr); -+} - - namespace MARLIN_NAMESPACE_NAME { - ---- a/gemm/marlin/marlin_template.h -+++ b/gemm/marlin/marlin_template.h +--- a/marlin_template.h ++++ b/marlin_template.h @@ -1035,7 +1035,7 @@ cur_k = 0; @@ -70,240 +51,78 @@ // Determine "position" inside the thread-block (based on warp and // thread-id) ---- a/awq_repack/cuda/kernel.cuh -+++ b/awq_repack/cuda/kernel.cuh -@@ -1,18 +1,9 @@ --#include "marlin.cuh" -- --#include --#include --#include --#include --#include --#include -- --#include "libtorch_stable/torch_utils.h" -+#include "../../gemm/marlin/marlin.cuh" +--- a/awq_marlin_repack.cu ++++ b/awq_marlin_repack.cu +@@ -1,5 +1,6 @@ + #include "marlin.cuh" + ++#ifndef INFINI_OPS_MARLIN_KERNEL_ONLY + #include + #include + #include +@@ -8,6 +9,7 @@ + #include + + #include "libtorch_stable/torch_utils.h" ++#endif namespace marlin { - template --__global__ void awq_marlin_repack_kernel( -+__device__ void awq_marlin_repack_kernel( - uint32_t const* __restrict__ b_q_weight_ptr, uint32_t* __restrict__ out_ptr, - int size_k, int size_n) { - constexpr int pack_factor = 32 / num_bits; -@@ -212,83 +203,3 @@ - } +@@ -213,6 +215,7 @@ } // namespace marlin -- --#define CALL_IF(NUM_BITS, IS_A_8BIT) \ -- else if (num_bits == NUM_BITS && is_a_8bit == IS_A_8BIT) { \ -- cudaFuncSetAttribute( \ -- marlin::awq_marlin_repack_kernel, \ -- cudaFuncAttributeMaxDynamicSharedMemorySize, max_shared_mem); \ -- marlin::awq_marlin_repack_kernel \ -- <<>>( \ -- b_q_weight_ptr, out_ptr, size_k, size_n); \ -- } -- --torch::stable::Tensor awq_marlin_repack(torch::stable::Tensor& b_q_weight, -- int64_t size_k, int64_t size_n, -- int64_t num_bits, bool is_a_8bit) { -- // Verify compatibility with marlin tile of 16x64 -- STD_TORCH_CHECK(size_k % marlin::tile_k_size == 0, "size_k = ", size_k, -- " is not divisible by tile_k_size = ", marlin::tile_k_size); -- STD_TORCH_CHECK(size_n % marlin::tile_n_size == 0, "size_n = ", size_n, -- " is not divisible by tile_n_size = ", marlin::tile_n_size); -- -- STD_TORCH_CHECK(num_bits == 4 || num_bits == 8, -- "num_bits must be 4 or 8. Got = ", num_bits); -- int const pack_factor = 32 / num_bits; -- -- // Verify B -- STD_TORCH_CHECK(b_q_weight.size(0) == size_k, -- "b_q_weight.size(0) = ", b_q_weight.size(0), -- " is not size_k = ", size_k); -- STD_TORCH_CHECK((size_n / pack_factor) == b_q_weight.size(1), -- "Shape mismatch: b_q_weight.size(1) = ", b_q_weight.size(1), -- ", size_n = ", size_n, ", pack_factor = ", pack_factor); -- -- // Verify device and strides -- STD_TORCH_CHECK(b_q_weight.is_cuda(), "b_q_weight is not on GPU"); -- STD_TORCH_CHECK(b_q_weight.is_contiguous(), "b_q_weight is not contiguous"); -- STD_TORCH_CHECK( -- b_q_weight.scalar_type() == torch::headeronly::ScalarType::Int, -- "b_q_weight type is not kInt"); -- -- const int32_t device_index = b_q_weight.get_device_index(); -- torch::stable::accelerator::DeviceGuard device_guard(device_index); -- const cudaStream_t stream = get_current_cuda_stream(device_index); -- -- // Alloc buffers -- torch::stable::Tensor out = torch::stable::empty( -- {size_k / marlin::tile_size, size_n * marlin::tile_size / pack_factor}, -- b_q_weight.scalar_type(), std::nullopt, b_q_weight.device()); -- -- // Get ptrs -- uint32_t const* b_q_weight_ptr = -- reinterpret_cast(b_q_weight.const_data_ptr()); -- uint32_t* out_ptr = reinterpret_cast(out.mutable_data_ptr()); -- -- int blocks; -- cudaDeviceGetAttribute(&blocks, cudaDevAttrMultiProcessorCount, device_index); -- -- int max_shared_mem = 0; -- cudaDeviceGetAttribute(&max_shared_mem, -- cudaDevAttrMaxSharedMemoryPerBlockOptin, device_index); -- STD_TORCH_CHECK(max_shared_mem > 0); -- -- if (false) { -- } -- CALL_IF(4, false) -- CALL_IF(8, false) -- CALL_IF(4, true) -- CALL_IF(8, true) -- else { -- STD_TORCH_CHECK(false, "Unsupported repack config: num_bits = ", num_bits, -- ", is_a_8bit = ", is_a_8bit); -- } -- -- return out; --} -- --STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) { -- m.impl("awq_marlin_repack", TORCH_BOX(&awq_marlin_repack)); --} ---- a/gptq_repack/cuda/kernel.cuh -+++ b/gptq_repack/cuda/kernel.cuh -@@ -1,19 +1,10 @@ --#include "marlin.cuh" -- --#include --#include --#include --#include --#include --#include -- --#include "libtorch_stable/torch_utils.h" -+#include "../../gemm/marlin/marlin.cuh" + ++#ifndef INFINI_OPS_MARLIN_KERNEL_ONLY + #define CALL_IF(NUM_BITS, IS_A_8BIT) \ + else if (num_bits == NUM_BITS && is_a_8bit == IS_A_8BIT) { \ + cudaFuncSetAttribute( \ +@@ -292,3 +295,4 @@ + STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) { + m.impl("awq_marlin_repack", TORCH_BOX(&awq_marlin_repack)); + } ++#endif +--- a/gptq_marlin_repack.cu ++++ b/gptq_marlin_repack.cu +@@ -1,5 +1,6 @@ + #include "marlin.cuh" + ++#ifndef INFINI_OPS_MARLIN_KERNEL_ONLY + #include + #include + #include +@@ -8,6 +9,7 @@ + #include + + #include "libtorch_stable/torch_utils.h" ++#endif namespace marlin { - template --__global__ void gptq_marlin_repack_kernel( -+__device__ void gptq_marlin_repack_kernel( - uint32_t const* __restrict__ b_q_weight_ptr, - uint32_t const* __restrict__ perm_ptr, uint32_t* __restrict__ out_ptr, - int size_k, int size_n) { -@@ -268,99 +259,3 @@ - } +@@ -269,6 +271,7 @@ } // namespace marlin -- --#define CALL_IF(NUM_BITS, HAS_PERM, IS_A_8BIT) \ -- else if (num_bits == NUM_BITS && has_perm == HAS_PERM && \ -- is_a_8bit == IS_A_8BIT) { \ -- cudaFuncSetAttribute( \ -- marlin::gptq_marlin_repack_kernel, \ -- cudaFuncAttributeMaxDynamicSharedMemorySize, max_shared_mem); \ -- marlin::gptq_marlin_repack_kernel \ -- <<>>( \ -- b_q_weight_ptr, perm_ptr, out_ptr, size_k, size_n); \ -- } -- --torch::stable::Tensor gptq_marlin_repack(torch::stable::Tensor& b_q_weight, -- torch::stable::Tensor& perm, -- int64_t size_k, int64_t size_n, -- int64_t num_bits, bool is_a_8bit) { -- // Verify compatibility with marlin tile of 16x64 -- STD_TORCH_CHECK(size_k % marlin::tile_k_size == 0, "size_k = ", size_k, -- " is not divisible by tile_k_size = ", marlin::tile_k_size); -- STD_TORCH_CHECK(size_n % marlin::tile_n_size == 0, "size_n = ", size_n, -- " is not divisible by tile_n_size = ", marlin::tile_n_size); -- -- STD_TORCH_CHECK(num_bits == 4 || num_bits == 8, -- "num_bits must be 4 or 8. Got = ", num_bits); -- int const pack_factor = 32 / num_bits; -- -- // Verify B -- STD_TORCH_CHECK((size_k / pack_factor) == b_q_weight.size(0), -- "Shape mismatch: b_q_weight.size(0) = ", b_q_weight.size(0), -- ", size_k = ", size_k, ", pack_factor = ", pack_factor); -- STD_TORCH_CHECK(b_q_weight.size(1) == size_n, -- "b_q_weight.size(1) = ", b_q_weight.size(1), -- " is not size_n = ", size_n); -- -- // Verify device and strides -- STD_TORCH_CHECK(b_q_weight.is_cuda(), "b_q_weight is not on GPU"); -- STD_TORCH_CHECK(b_q_weight.is_contiguous(), "b_q_weight is not contiguous"); -- STD_TORCH_CHECK( -- b_q_weight.scalar_type() == torch::headeronly::ScalarType::Int, -- "b_q_weight type is not kInt"); -- -- STD_TORCH_CHECK(perm.is_cuda(), "perm is not on GPU"); -- STD_TORCH_CHECK(perm.is_contiguous(), "perm is not contiguous"); -- STD_TORCH_CHECK(perm.scalar_type() == torch::headeronly::ScalarType::Int, -- "perm type is not at::kInt"); -- -- const int32_t device_index = b_q_weight.get_device_index(); -- torch::stable::accelerator::DeviceGuard device_guard(device_index); -- const cudaStream_t stream = get_current_cuda_stream(device_index); -- -- // Alloc buffers -- torch::stable::Tensor out = torch::stable::empty( -- {size_k / marlin::tile_size, size_n * marlin::tile_size / pack_factor}, -- b_q_weight.scalar_type(), std::nullopt, b_q_weight.device()); -- -- // Detect if there is act_order -- bool has_perm = perm.size(0) != 0; -- -- // Get ptrs -- uint32_t const* b_q_weight_ptr = -- reinterpret_cast(b_q_weight.const_data_ptr()); -- uint32_t const* perm_ptr = -- reinterpret_cast(perm.const_data_ptr()); -- uint32_t* out_ptr = reinterpret_cast(out.mutable_data_ptr()); -- -- int blocks; -- cudaDeviceGetAttribute(&blocks, cudaDevAttrMultiProcessorCount, device_index); -- -- int max_shared_mem = 0; -- cudaDeviceGetAttribute(&max_shared_mem, -- cudaDevAttrMaxSharedMemoryPerBlockOptin, device_index); -- STD_TORCH_CHECK(max_shared_mem > 0); -- -- if (false) { -- } -- CALL_IF(4, false, false) -- CALL_IF(4, true, false) -- CALL_IF(8, false, false) -- CALL_IF(8, true, false) -- -- CALL_IF(4, false, true) -- CALL_IF(8, false, true) -- -- else { -- STD_TORCH_CHECK(false, "Unsupported repack config: num_bits = ", num_bits, -- ", has_perm = ", has_perm, ", is_a_8bit = ", is_a_8bit); -- } -- -- return out; --} -- --STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) { -- m.impl("gptq_marlin_repack", TORCH_BOX(&gptq_marlin_repack)); --} ---- a/gemm/nvidia/generate_kernels.py -+++ b/gemm/nvidia/generate_kernels.py -@@ -27,15 +27,17 @@ + ++#ifndef INFINI_OPS_MARLIN_KERNEL_ONLY + #define CALL_IF(NUM_BITS, HAS_PERM, IS_A_8BIT) \ + else if (num_bits == NUM_BITS && has_perm == HAS_PERM && \ + is_a_8bit == IS_A_8BIT) { \ +@@ -364,3 +367,4 @@ + STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) { + m.impl("gptq_marlin_repack", TORCH_BOX(&gptq_marlin_repack)); + } ++#endif +--- a/generate_kernels.py ++++ b/generate_kernels.py +@@ -3,7 +3,7 @@ + import glob + import itertools + import os +-import subprocess ++import pathlib + import sys + + import jinja2 +@@ -27,7 +27,9 @@ SUPPORT_SM75 = True FILE_HEAD_COMMENT = """ @@ -314,17 +133,7 @@ // clang-format off """.lstrip() - FILE_HEAD = ( - FILE_HEAD_COMMENT - + """ --#include "kernel.h" --#include "marlin_template.h" -+#include "../marlin/kernel.h" -+#include "../marlin/marlin_template.h" - - namespace MARLIN_NAMESPACE_NAME { - """ -@@ -162,6 +164,10 @@ +@@ -162,12 +164,16 @@ ] @@ -334,8 +143,16 @@ + def remove_old_kernels(): for filename in glob.glob(os.path.dirname(__file__) + "/*kernel_*.cu"): - subprocess.call(["rm", "-f", filename]) -@@ -300,17 +306,9 @@ +- subprocess.call(["rm", "-f", filename]) ++ pathlib.Path(filename).unlink(missing_ok=True) + + filename = os.path.dirname(__file__) + "/kernel_selector.h" +- subprocess.call(["rm", "-f", filename]) ++ pathlib.Path(filename).unlink(missing_ok=True) + + + def generate_new_kernels(): +@@ -300,13 +306,6 @@ with open(os.path.join(os.path.dirname(__file__), filename), "w") as f: f.write(file_content) @@ -349,7 +166,261 @@ with open(os.path.join(os.path.dirname(__file__), "kernel_selector.h"), "w") as f: f.write(kernel_selector_str) +--- a/marlin.cu ++++ b/marlin.cu +@@ -1,3 +1,4 @@ ++// Modified by InfiniOps: isolate the native host entry and preserve CUDA error propagation. + /* + * Modified by Neural Magic + * Copyright (C) Marlin.2024 Elias Frantar +@@ -24,7 +25,9 @@ + #endif + + #include "kernel.h" +- ++#include "launch.h" ++ ++#ifndef INFINI_OPS_MARLIN_KERNEL_ONLY + #include + #include + #include +@@ -33,6 +36,7 @@ + #include + + #include "libtorch_stable/torch_utils.h" ++#endif + + #define STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t) \ + static_assert(std::is_same::value || \ +@@ -54,6 +58,7 @@ + + } // namespace marlin + ++#ifndef INFINI_OPS_MARLIN_KERNEL_ONLY + torch::stable::Tensor marlin_gemm( + torch::stable::Tensor& a, std::optional c_or_none, + torch::stable::Tensor& b_q_weight, +@@ -72,6 +77,8 @@ + return torch::stable::empty({1, 1}); + } + ++#endif // INFINI_OPS_MARLIN_KERNEL_ONLY ++ + #else + + // For a given "a" of size [M,K] performs a permutation of the K columns based +@@ -334,18 +341,18 @@ + int thread_n_init, int sms, bool use_atomic_add, + bool use_fp32_reduce, bool is_zp_float) { + bool is_a_8bit = a_type.size_bits() == 8; +- STD_TORCH_CHECK(prob_m > 0 && prob_n > 0 && prob_k > 0, "Invalid MNK = [", ++ INFINI_MARLIN_CHECK(prob_m > 0 && prob_n > 0 && prob_k > 0, "Invalid MNK = [", + prob_m, ", ", prob_n, ", ", prob_k, "]"); + + int group_blocks = 0; + if (has_act_order) { + if (is_k_full) { +- STD_TORCH_CHECK(group_size != -1); ++ INFINI_MARLIN_CHECK(group_size != -1); + group_blocks = group_size / 16; +- STD_TORCH_CHECK(prob_k % group_blocks == 0, "prob_k = ", prob_k, ++ INFINI_MARLIN_CHECK(prob_k % group_blocks == 0, "prob_k = ", prob_k, + " is not divisible by group_blocks = ", group_blocks); + } else { +- STD_TORCH_CHECK(group_size == 0); ++ INFINI_MARLIN_CHECK(group_size == 0); + group_blocks = 0; + } + } else { +@@ -353,7 +360,7 @@ + group_blocks = -1; + } else { + group_blocks = group_size / 16; +- STD_TORCH_CHECK(prob_k % group_blocks == 0, "prob_k = ", prob_k, ++ INFINI_MARLIN_CHECK(prob_k % group_blocks == 0, "prob_k = ", prob_k, + " is not divisible by group_blocks = ", group_blocks); + } + } +@@ -395,25 +402,25 @@ + int max_shared_mem = 0; + cudaDeviceGetAttribute(&max_shared_mem, + cudaDevAttrMaxSharedMemoryPerBlockOptin, dev); +- STD_TORCH_CHECK(max_shared_mem > 0); ++ INFINI_MARLIN_CHECK(max_shared_mem > 0); + + int major_capability, minor_capability; + cudaDeviceGetAttribute(&major_capability, cudaDevAttrComputeCapabilityMajor, + dev); + cudaDeviceGetAttribute(&minor_capability, cudaDevAttrComputeCapabilityMinor, + dev); +- STD_TORCH_CHECK(major_capability * 10 + minor_capability >= 75, ++ INFINI_MARLIN_CHECK(major_capability * 10 + minor_capability >= 75, + "marlin kernel only support Turing or newer GPUs."); + int stages = 4; + if (major_capability == 7 && minor_capability == 5) { + stages = 2; +- STD_TORCH_CHECK(a_type == vllm::kFloat16 || a_type == vllm::kS8, ++ INFINI_MARLIN_CHECK(a_type == vllm::kFloat16 || a_type == vllm::kS8, + "Turing only support FP16 or INT8 activation."); + } + if (a_type == vllm::kFE4M3fn) { +- STD_TORCH_CHECK(major_capability * 10 + minor_capability >= 89, ++ INFINI_MARLIN_CHECK(major_capability * 10 + minor_capability >= 89, + "FP8 only support Ada Lovelace or newer GPUs."); +- STD_TORCH_CHECK( ++ INFINI_MARLIN_CHECK( + major_capability * 10 + minor_capability == 89 || + major_capability == 12, + "Marlin W4A8-FP8 only support SM89 or SM12x device (It is slower than " +@@ -443,9 +450,9 @@ + if (thread_k != -1 && thread_n != -1) { + thread_tfg = thread_config_t{thread_k, thread_n, default_threads}; + exec_cfg = exec_config_t{1, thread_tfg}; +- STD_TORCH_CHECK(prob_n % thread_n == 0, "prob_n = ", prob_n, ++ INFINI_MARLIN_CHECK(prob_n % thread_n == 0, "prob_n = ", prob_n, + " is not divisible by thread_n = ", thread_n); +- STD_TORCH_CHECK(prob_k % thread_k == 0, "prob_k = ", prob_k, ++ INFINI_MARLIN_CHECK(prob_k % thread_k == 0, "prob_k = ", prob_k, + " is not divisible by thread_k = ", thread_k); + } else { + // Auto config +@@ -485,7 +492,7 @@ + int thread_k_blocks = thread_k / 16; + int thread_n_blocks = thread_n / 16; + +- STD_TORCH_CHECK( ++ INFINI_MARLIN_CHECK( + is_valid_config(thread_tfg, thread_m_blocks, prob_m_split, prob_n, + prob_k, num_bits, group_size, has_act_order, is_k_full, + has_zp, is_zp_float, is_a_8bit, stages, +@@ -506,7 +513,7 @@ + num_threads, is_zp_float, stages); + + if (kernel == MarlinDefault) { +- STD_TORCH_CHECK( ++ INFINI_MARLIN_CHECK( + false, "Unsupported shapes: MNK = [", prob_m, ", ", prob_n, ", ", + prob_k, "]", ", has_act_order = ", has_act_order, + ", num_groups = ", num_groups, ", group_size = ", group_size, +@@ -517,8 +524,14 @@ + ", num_threads = ", num_threads, ", num_bits = ", num_bits); + } + +- cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, +- max_shared_mem_new); ++ cudaStreamCaptureStatus capture_status = cudaStreamCaptureStatusNone; ++ ::infini::ops::native_w4::check( ++ cudaStreamIsCapturing(stream, &capture_status), "cudaStreamIsCapturing"); ++ if (capture_status == cudaStreamCaptureStatusNone) { ++ ::infini::ops::native_w4::check( ++ cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, ++ max_shared_mem_new), "cudaFuncSetAttribute"); ++ } + + bool part_use_atomic_add = + use_atomic_add && div_ceil(prob_m_split, 64) * prob_n <= 2048; +@@ -542,6 +555,7 @@ + + } // namespace marlin + ++#ifndef INFINI_OPS_MARLIN_KERNEL_ONLY + torch::stable::Tensor marlin_gemm( + torch::stable::Tensor& a, std::optional c_or_none, + torch::stable::Tensor& b_q_weight, +@@ -893,8 +907,12 @@ + return c; + } + ++#endif // INFINI_OPS_MARLIN_KERNEL_ONLY ++ + #endif + ++#ifndef INFINI_OPS_MARLIN_KERNEL_ONLY + STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) { + m.impl("marlin_gemm", TORCH_BOX(&marlin_gemm)); + } ++#endif // INFINI_OPS_MARLIN_KERNEL_ONLY +--- a/core/scalar_type.hpp ++++ b/core/scalar_type.hpp +@@ -1,3 +1,4 @@ ++// Modified by InfiniOps: replace Torch-only error checks with a private standard-C++ helper. + #pragma once + + #include +@@ -6,8 +7,7 @@ + #include + #include + +-// For STD_TORCH_CHECK +-#include ++#include "marlin_check.h" + + namespace vllm { + +@@ -51,7 +51,7 @@ + // IEEE 754 compliant floating point type + static constexpr ScalarType float_IEEE754(uint8_t exponent, + uint8_t mantissa) { +- STD_TORCH_CHECK(mantissa > 0 && exponent > 0); ++ INFINI_MARLIN_CHECK(mantissa > 0 && exponent > 0); + return ScalarType(exponent, mantissa, true, 0, false, NAN_IEEE_754); + } + +@@ -59,9 +59,9 @@ + static constexpr ScalarType float_(uint8_t exponent, uint8_t mantissa, + bool finite_values_only, + NanRepr nan_repr) { +- STD_TORCH_CHECK(nan_repr < NAN_REPR_ID_MAX, "Invalid NanRepr"); +- STD_TORCH_CHECK(mantissa > 0 && exponent > 0); +- STD_TORCH_CHECK( ++ INFINI_MARLIN_CHECK(nan_repr < NAN_REPR_ID_MAX, "Invalid NanRepr"); ++ INFINI_MARLIN_CHECK(mantissa > 0 && exponent > 0); ++ INFINI_MARLIN_CHECK( + nan_repr != NAN_IEEE_754, + "use `float_IEEE754` constructor for floating point types that " + "follow IEEE 754 conventions"); +@@ -183,7 +183,7 @@ + + private: + double _floating_point_max() const { +- STD_TORCH_CHECK(mantissa <= 52 && exponent <= 11, ++ INFINI_MARLIN_CHECK(mantissa <= 52 && exponent <= 11, + "Cannot represent max/min as a double for type ", str()); + + uint64_t max_mantissa = (uint64_t(1) << mantissa) - 1; +@@ -193,7 +193,7 @@ + + uint64_t max_exponent = (uint64_t(1) << exponent) - 2; + if (nan_repr == NAN_EXTD_RANGE_MAX_MIN || nan_repr == NAN_NONE) { +- STD_TORCH_CHECK(exponent < 11, ++ INFINI_MARLIN_CHECK(exponent < 11, + "Cannot represent max/min as a double for type ", str()); + max_exponent += 1; + } +@@ -223,7 +223,7 @@ + if (is_floating_point()) { + return {_floating_point_max()}; + } else { +- STD_TORCH_CHECK(size_bits() < 64 || size_bits() == 64 && is_signed(), ++ INFINI_MARLIN_CHECK(size_bits() < 64 || size_bits() == 64 && is_signed(), + "Cannot represent max as a int64_t"); + return {(int64_t(1) << mantissa) - 1}; + } +@@ -231,7 +231,7 @@ - if __name__ == "__main__": -- remove_old_kernels() - generate_new_kernels() + constexpr std::variant _raw_min() const { + if (is_floating_point()) { +- STD_TORCH_CHECK( ++ INFINI_MARLIN_CHECK( + is_signed(), + "We currently assume all floating point types are signed"); + constexpr uint64_t sign_bit_double = (uint64_t(1) << 63); +@@ -241,7 +241,7 @@ + uint64_t min_raw = max_raw | sign_bit_double; + return {*reinterpret_cast(&min_raw)}; + } else { +- STD_TORCH_CHECK(!is_signed() || size_bits() <= 64, ++ INFINI_MARLIN_CHECK(!is_signed() || size_bits() <= 64, + "Cannot represent min as a int64_t"); + if (is_signed()) { + // set the top bit to 1 (i.e. INT64_MIN) and the rest to 0 diff --git a/pyproject.toml b/pyproject.toml index 68a883585..ae7a52600 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,5 +1,5 @@ [build-system] -requires = ["scikit-build-core", "pybind11", "libclang", "packaging", "pyyaml", "jinja2"] +requires = ["scikit-build-core>=1.1.1", "pybind11", "libclang", "packaging", "pyyaml", "jinja2"] build-backend = "scikit_build_core.build" [project] @@ -13,7 +13,9 @@ dev = ["pytest", "pytest-cov", "pytest-xdist", "ruff==0.15.22", "torch", "pyyaml install-dir = "infini" [tool.scikit-build.sdist] -exclude = ["third_party/flash-attention/**", "third_party/vllm/**"] +cmake = true +exclude = ["third_party/flash-attention/**"] +force-include = { "build/sdist/marlin-source" = "third_party/marlin-source" } include = [ "third_party/flash-attention/LICENSE", "third_party/flash-attention/AUTHORS", @@ -22,13 +24,12 @@ include = [ "third_party/flash-attention/csrc/flash_attn/flash_api.cpp", "third_party/flash-attention/csrc/cutlass/include/**", "third_party/flash-attention/csrc/cutlass/LICENSE.txt", - "third_party/vllm/LICENSE", - "third_party/vllm/csrc/libtorch_stable/quantization/marlin/*.h", - "third_party/vllm/csrc/libtorch_stable/quantization/marlin/*.cuh", - "third_party/vllm/csrc/libtorch_stable/quantization/marlin/*repack.cu", - "third_party/vllm/csrc/libtorch_stable/quantization/marlin/generate_kernels.py", ] +[[tool.scikit-build.overrides]] +if.state = "sdist" +build-dir = "build/sdist" + [tool.scikit-build.cmake.define] AUTO_DETECT_DEVICES = "ON" AUTO_DETECT_BACKENDS = "ON" diff --git a/scripts/prepare_marlin.py b/scripts/prepare_marlin.py index 90978e448..c831d771c 100644 --- a/scripts/prepare_marlin.py +++ b/scripts/prepare_marlin.py @@ -11,7 +11,7 @@ REPO = pathlib.Path(__file__).resolve().parents[1] UPSTREAM = pathlib.Path("csrc/libtorch_stable/quantization/marlin") SOURCES = { - pathlib.Path("gemm/marlin") / name: UPSTREAM / name + pathlib.Path(name): UPSTREAM / name for name in ( "dequant.h", "kernel.h", @@ -19,15 +19,14 @@ "marlin_dtypes.cuh", "marlin_mma.h", "marlin_template.h", + "awq_marlin_repack.cu", + "gptq_marlin_repack.cu", + "generate_kernels.py", + "marlin.cu", ) } -SOURCES.update( - { - pathlib.Path("awq_repack/cuda/kernel.cuh"): UPSTREAM / "awq_marlin_repack.cu", - pathlib.Path("gptq_repack/cuda/kernel.cuh"): UPSTREAM / "gptq_marlin_repack.cu", - pathlib.Path("gemm/nvidia/generate_kernels.py"): UPSTREAM - / "generate_kernels.py", - } +SOURCES[pathlib.Path("core/scalar_type.hpp")] = pathlib.Path( + "csrc/core/scalar_type.hpp" ) GENERATED = ( "kernel_selector.h", @@ -55,7 +54,7 @@ def prepare(source, output, patch_executable): if not upstream.is_file(): raise ValueError( f"Missing pinned vLLM Marlin source: {upstream}; " - "initialize third_party/vllm (or include it in the source archive)" + "check the pinned source archive or offline FetchContent override" ) target = staging / destination target.parent.mkdir(parents=True, exist_ok=True) @@ -85,16 +84,14 @@ def prepare(source, output, patch_executable): # The upstream generator's single result/config list emits both CUDA # instantiations and the selector. The patch retains only its W4A16 set. subprocess.run( - [sys.executable, str(staging / "gemm/nvidia/generate_kernels.py"), "8.0"], + [sys.executable, str(staging / "generate_kernels.py"), "8.0"], cwd=staging, check=True, ) prepared = [path for path in SOURCES if path.suffix != ".py"] - prepared.extend(pathlib.Path("gemm/nvidia") / name for name in GENERATED) + prepared.extend(pathlib.Path(name) for name in GENERATED) for destination in prepared: - write_if_changed( - output / "vendor" / destination, (staging / destination).read_bytes() - ) + write_if_changed(output / destination, (staging / destination).read_bytes()) def main(): diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 258c8e375..e6147ea01 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -1704,14 +1704,17 @@ if(WITH_NVIDIA) DESTINATION "${CMAKE_INSTALL_DATADIR}/licenses/InfiniOps") endif() if(WITH_LINKED) - # Keep the fixed recipe's secondary notices as a conservative fallback; - # additionally preserve the inventory of the distributions actually used. - install(FILES "${PROJECT_SOURCE_DIR}/LICENSES/Apache-2.0-components.txt" - DESTINATION "${CMAKE_INSTALL_DATADIR}/licenses/InfiniOps") - install(DIRECTORY - "${PROJECT_SOURCE_DIR}/LICENSES/FlashInfer-0.6.7" - "${PROJECT_SOURCE_DIR}/LICENSES/TVM-FFI-0.1.10" - DESTINATION "${CMAKE_INSTALL_DATADIR}/licenses/InfiniOps") + if(_infini_ops_linked_uses_tvm_ffi) + # Fixed native sampling/runtime recipe; not the complete Python wheels. + install(FILES "${PROJECT_SOURCE_DIR}/LICENSES/Apache-2.0-components.txt" + DESTINATION "${CMAKE_INSTALL_DATADIR}/licenses/InfiniOps") + install(FILES "${PROJECT_SOURCE_DIR}/LICENSES/FlashInfer-0.6.7/NOTICE.txt" + DESTINATION "${CMAKE_INSTALL_DATADIR}/licenses/InfiniOps/FlashInfer-0.6.7") + install(FILES "${PROJECT_SOURCE_DIR}/LICENSES/TVM-FFI-0.1.10/NOTICE.txt" + "${PROJECT_SOURCE_DIR}/LICENSES/TVM-FFI-0.1.10/LICENSE.libbacktrace.txt" + DESTINATION "${CMAKE_INSTALL_DATADIR}/licenses/InfiniOps/TVM-FFI-0.1.10") + endif() + # Keep the real inventory for resolved packages, including unknown versions. foreach(_notice IN LISTS INFINI_OPS_LINKED_NOTICE_FILES) get_filename_component(_notice_dir "${_notice}" DIRECTORY) string(SHA256 _notice_group "${_notice_dir}") diff --git a/src/base/flash_attn_with_kvcache.h b/src/base/flash_attn_with_kvcache.h index ef4abddb2..6b26fceba 100644 --- a/src/base/flash_attn_with_kvcache.h +++ b/src/base/flash_attn_with_kvcache.h @@ -25,10 +25,6 @@ class FlashAttnWithKvcacheConfig : public Config { return std::make_unique(*this); } - bool requires_config_aware_provider() const override { - return logical_kv_upper_bound_ != 0; - } - int64_t logical_kv_upper_bound() const { return logical_kv_upper_bound_; } static int64_t BoundFrom(const Config& config) { diff --git a/src/base/separate_table_rope.h b/src/base/separate_table_rope.h index 0378e7301..02e4f861a 100644 --- a/src/base/separate_table_rope.h +++ b/src/base/separate_table_rope.h @@ -20,7 +20,7 @@ class SeparateTableRope : public Operator { const Tensor sin_table, const Tensor cos_table, bool /*is_neox*/, const Tensor out) { const auto rank = input.ndim(); - // Preserve the installed Rs4 CPU domain. NVIDIA additionally retains LM's + // Preserve the CPU table dtype domain. NVIDIA additionally retains LM's // independently typed tables and padded per-batch position rows. const bool nvidia = input.device().type() == Device::Type::kNvidia; if ((rank != 3 && rank != 4) || out.shape() != input.shape() || diff --git a/src/native/ascend/ops/rotary_embedding/cann82_compat.hpp b/src/native/ascend/ops/rotary_embedding/cann82_compat.hpp deleted file mode 100644 index 5dc7a92c5..000000000 --- a/src/native/ascend/ops/rotary_embedding/cann82_compat.hpp +++ /dev/null @@ -1,111 +0,0 @@ -#ifndef INFINI_OPS_ASCEND_ROTARY_EMBEDDING_CANN82_COMPAT_HPP_ -#define INFINI_OPS_ASCEND_ROTARY_EMBEDDING_CANN82_COMPAT_HPP_ - -#include -#include -#include - -#include "aclnn_copy.h" -#include "aclnnop/level2/aclnn_rope_with_sin_cos_cache.h" -#include "base/rotary_embedding.h" -#include "native/ascend/common.h" - -namespace infini::ops { - -// CANN 8.2's packed-cache RoPE API accepts 2D token/head views. Its two -// outputs are staged separately so single-query calls and fused QKV views -// never alias an input while the kernel reads it. All math stays on device. -template <> -class Operator : public RotaryEmbedding { - public: - Operator(const Tensor positions, Tensor query, std::optional key, - const Tensor cos_sin_cache, int64_t head_size, bool is_neox, - int64_t rope_dim_offset = 0, bool inverse = false) - : RotaryEmbedding(positions, query, key, cos_sin_cache, head_size, - is_neox, rope_dim_offset, inverse) { - if (rope_dim_offset != 0 || inverse) - throw std::invalid_argument("CANN 8.2 packed-cache RoPE has no offset/inverse mode"); - } - - void operator()(const Tensor positions, Tensor query, - std::optional key, const Tensor cos_sin_cache, - int64_t head_size, bool is_neox, int64_t rope_dim_offset = 0, - bool inverse = false) const override { - if (rope_dim_offset != 0 || inverse) - throw std::invalid_argument("CANN 8.2 packed-cache RoPE has no offset/inverse mode"); - auto stream = static_cast(stream_); - const auto tokens = num_tokens_; - const auto q_width = num_heads_ * head_size_; - const auto k_width = key ? num_kv_heads_ * head_size_ : q_width; - const Tensor q_input{query.data(), {tokens, q_width}, query.dtype(), query.device(), {query_token_stride_, 1}}; - const Tensor k_input{key ? key->data() : query.data(), {tokens, k_width}, query.dtype(), query.device(), {key ? key_token_stride_ : query_token_stride_, 1}}; - const Tensor pos{const_cast(positions.data()), {tokens}, DataType::kInt64, positions.device(), {1}}; - void* q_data = nullptr; - void* k_data = nullptr; - const auto q_bytes = tokens * q_width * query.element_size(); - const auto k_bytes = tokens * k_width * query.element_size(); - Check(aclrtMalloc(&q_data, q_bytes, ACL_MEM_MALLOC_NORMAL_ONLY), "query staging allocation"); - Check(aclrtMalloc(&k_data, k_bytes, ACL_MEM_MALLOC_NORMAL_ONLY), "key staging allocation"); - if (!q_data || !k_data) FailStop("null staging allocation", -1); - const Tensor q_output{q_data, {tokens, q_width}, query.dtype(), query.device()}; - const Tensor k_output{k_data, {tokens, k_width}, query.dtype(), query.device()}; - auto t_pos = ascend::BuildAclTensor(pos); - auto t_q = ascend::BuildAclTensor(q_input); - auto t_k = ascend::BuildAclTensor(k_input); - auto t_cache = ascend::BuildAclTensor(cos_sin_cache); - auto t_q_out = ascend::BuildAclTensor(q_output); - auto t_k_out = ascend::BuildAclTensor(k_output); - for (auto tensor : {t_pos, t_q, t_k, t_cache, t_q_out, t_k_out}) - if (!tensor) FailStop("descriptor creation", -1); - uint64_t workspace_size = 0; - aclOpExecutor* executor = nullptr; - Check(aclnnRopeWithSinCosCacheGetWorkspaceSize(t_pos, t_q, t_k, t_cache, - nullptr, head_size, is_neox, t_q_out, t_k_out, &workspace_size, &executor), "RoPE workspace query"); - if (!executor) FailStop("null RoPE executor", -1); - void* workspace = AllocateWorkspace(workspace_size); - Check(aclnnRopeWithSinCosCache(workspace, workspace_size, executor, stream), "RoPE execution"); - Check(aclrtSynchronizeStream(stream), "RoPE completion"); - if (workspace) Check(aclrtFree(workspace), "RoPE workspace release"); - CopyBack(t_q, t_q_out, stream); - if (key) CopyBack(t_k, t_k_out, stream); - for (auto tensor : {t_pos, t_q, t_k, t_cache, t_q_out, t_k_out}) - Check(aclDestroyTensor(tensor), "descriptor release"); - Check(aclrtFree(q_data), "query staging release"); - Check(aclrtFree(k_data), "key staging release"); - } - - private: - static void* AllocateWorkspace(uint64_t size) { - void* data = nullptr; - if (size) { - Check(aclrtMalloc(&data, size, ACL_MEM_MALLOC_NORMAL_ONLY), "workspace allocation"); - if (!data) FailStop("null workspace allocation", -1); - } - return data; - } - - static void CopyBack(aclTensor* destination, aclTensor* source, aclrtStream stream) { - uint64_t size = 0; - aclOpExecutor* executor = nullptr; - Check(aclnnInplaceCopyGetWorkspaceSize(destination, source, &size, &executor), "copy workspace query"); - if (!executor) FailStop("null copy executor", -1); - void* workspace = AllocateWorkspace(size); - Check(aclnnInplaceCopy(workspace, size, executor, stream), "copy execution"); - Check(aclrtSynchronizeStream(stream), "copy completion"); - if (workspace) Check(aclrtFree(workspace), "copy workspace release"); - } - - static void Check(int64_t status, const char* phase) { - if (status != ACL_SUCCESS) FailStop(phase, status); - } - - [[noreturn]] static void FailStop(const char* phase, int64_t status) { - std::fprintf(stderr, "Ascend RotaryEmbedding %s failed (%lld)\n", phase, - static_cast(status)); - std::abort(); - } -}; - -} // namespace infini::ops - -#endif diff --git a/src/native/ascend/ops/rotary_embedding/kernel.h b/src/native/ascend/ops/rotary_embedding/kernel.h index e9f51fa41..b8e922e5b 100644 --- a/src/native/ascend/ops/rotary_embedding/kernel.h +++ b/src/native/ascend/ops/rotary_embedding/kernel.h @@ -1,8 +1,6 @@ #ifndef INFINI_OPS_ASCEND_ROTARY_EMBEDDING_KERNEL_H_ #define INFINI_OPS_ASCEND_ROTARY_EMBEDDING_KERNEL_H_ -#if __has_include("aclnnop/aclnn_rotary_position_embedding.h") - #include #include #include @@ -339,8 +337,4 @@ class Operator } // namespace infini::ops -#else -#include "native/ascend/ops/rotary_embedding/cann82_compat.hpp" -#endif - #endif diff --git a/src/native/cambricon/ops/add/kernel.h b/src/native/cambricon/ops/add/kernel.h index 6dee73975..8a701d39b 100644 --- a/src/native/cambricon/ops/add/kernel.h +++ b/src/native/cambricon/ops/add/kernel.h @@ -1,14 +1,7 @@ #ifndef INFINI_OPS_CAMBRICON_ADD_H_ #define INFINI_OPS_CAMBRICON_ADD_H_ -#include -#include -#include -#include - #include "base/add.h" -#include "native/cambricon/cnnl_utils.h" -#include "native/cambricon/cnrt_utils.h" #include "native/cambricon/common.h" #include "native/cambricon/data_type_.h" @@ -32,7 +25,6 @@ class Operator : public Add { cnrt_utils::GetLaunchConfig(input.device(), &core_per_cluster, &cluster_count); cnrtMalloc(&default_workspace_, workspace_size_in_bytes()); - PrepareMaskAdd(input, other, alpha, out); } Operator(const Tensor input, const Tensor other, Tensor out) @@ -43,12 +35,6 @@ class Operator : public Add { auto queue = static_cast(stream_ ? stream_ : 0); auto workspace{workspace_ ? workspace_ : default_workspace_}; - if (mask_op_ && out.shape() == out_shape_ && - CanUseMaskAdd(input, other, alpha, out)) { - RunMaskAdd(queue, input, other, out); - return; - } - bool fast_path = is_input_contiguous_ && is_other_contiguous_ && is_out_contiguous_ && input_shape_ == out_shape_ && other_shape_ == out_shape_; @@ -78,124 +64,6 @@ class Operator : public Add { void* default_workspace_{nullptr}; int core_per_cluster = 0; int cluster_count = 0; - - private: - // This leaf uses the same prepared-owner lifetime as the upstream Softmax: - // caller completion must precede owner/cache release and stream destruction. - // Capture remains opt-in false through OperatorBase. - struct MaskOpDeleter { - void operator()(cnnlOpTensorDescriptor_t p) const noexcept { - if (p) (void)cnnlDestroyOpTensorDescriptor(p); - } - }; - using MaskOp = - std::unique_ptr, - MaskOpDeleter>; - static void CheckMaskStatus(cnnlStatus_t status, const char* operation) { - if (status != CNNL_STATUS_SUCCESS) { - std::fprintf(stderr, "MASK_ADD_CNNL_FAIL %s %d\n", operation, - int(status)); - std::abort(); - } - } - static bool CanUseMaskAdd(const Tensor input, const Tensor other, - double alpha, const Tensor out) { - if (input.dtype() != DataType::kFloat32 || - other.dtype() != DataType::kFloat32 || - out.dtype() != DataType::kFloat32 || alpha != 1.0 || - input.device() != out.device() || other.device() != out.device() || - out.ndim() != 3 || input.shape() != out.shape() || other.ndim() != 2 || - !input.IsContiguous() || !other.IsContiguous() || !out.IsContiguous() || - !input.data() || !other.data() || input.data() != out.data()) - return false; - const auto h = out.size(0), r = out.size(1), c = out.size(2); - if (h < 2 || h > INT_MAX || r < 2 || r > INT_MAX || c < 2 || c > INT_MAX || - other.size(0) != r || other.size(1) != c || r > size_t(INT_MAX) / c) - return false; - const auto plane = r * c; - if (h > size_t(INT_MAX) / plane) return false; - const auto a = reinterpret_cast(input.data()); - const auto b = reinterpret_cast(other.data()); - if (a % sizeof(float) || b % sizeof(float)) return false; - return b > a ? b - a >= h * plane * sizeof(float) - : a - b >= plane * sizeof(float); - } - void PrepareMaskAdd(const Tensor input, const Tensor other, double alpha, - const Tensor out) { - if (!CanUseMaskAdd(input, other, alpha, out)) return; - cnnlHandle_t raw_handle = nullptr; - auto status = cnnlCreate(&raw_handle); - mask_handle_.reset(raw_handle); - CheckMaskStatus(status, "create handle"); - auto create_desc = [](cnnl_utils::TensorDescriptor& owner) { - cnnlTensorDescriptor_t raw = nullptr; - const auto result = cnnlCreateTensorDescriptor(&raw); - owner.reset(raw); - CheckMaskStatus(result, "create tensor descriptor"); - }; - create_desc(mask_input_desc_); - create_desc(mask_plane_desc_); - create_desc(mask_out_desc_); - const int dims[3] = {int(out.size(0)), int(out.size(1)), int(out.size(2))}; - const int plane[3] = {1, dims[1], dims[2]}; - CheckMaskStatus( - cnnlSetTensorDescriptor(mask_input_desc_.get(), CNNL_LAYOUT_ARRAY, - CNNL_DTYPE_FLOAT, 3, dims), - "input descriptor"); - CheckMaskStatus( - cnnlSetTensorDescriptor(mask_plane_desc_.get(), CNNL_LAYOUT_ARRAY, - CNNL_DTYPE_FLOAT, 3, plane), - "plane descriptor"); - CheckMaskStatus( - cnnlSetTensorDescriptor(mask_out_desc_.get(), CNNL_LAYOUT_ARRAY, - CNNL_DTYPE_FLOAT, 3, dims), - "output descriptor"); - cnnlOpTensorDescriptor_t raw_op = nullptr; - status = cnnlCreateOpTensorDescriptor(&raw_op); - mask_op_.reset(raw_op); - CheckMaskStatus(status, "create operation descriptor"); - CheckMaskStatus( - cnnlSetOpTensorDescriptor(mask_op_.get(), CNNL_OP_TENSOR_ADD, - CNNL_DTYPE_FLOAT, CNNL_PROPAGATE_NAN), - "operation descriptor"); - CheckMaskStatus( - cnnlGetOpTensorWorkspaceSize( - mask_handle_.get(), mask_input_desc_.get(), mask_plane_desc_.get(), - mask_out_desc_.get(), &mask_workspace_bytes_), - "workspace query"); - if (mask_workspace_bytes_) { - void* raw = nullptr; - const auto result = cnrtMalloc(&raw, mask_workspace_bytes_); - mask_workspace_.reset(raw); - if (result != cnrtSuccess) { - std::fprintf(stderr, "MASK_ADD_CNRT_FAIL allocate %d\n", int(result)); - std::abort(); - } - } - } - void RunMaskAdd(cnrtQueue_t queue, const Tensor input, const Tensor other, - Tensor out) const { - CheckMaskStatus(cnnlSetQueue(mask_handle_.get(), queue), "set queue"); - static const float one = 1.0F, zero = 0.0F; - const auto status = cnnlOpTensor( - mask_handle_.get(), mask_op_.get(), &one, mask_input_desc_.get(), - input.data(), &one, mask_plane_desc_.get(), other.data(), - mask_workspace_.get(), mask_workspace_bytes_, &zero, - mask_out_desc_.get(), out.data()); - if (status != CNNL_STATUS_SUCCESS) { - // A failed call can have queued partial work. Never unwind its resources. - const auto completion = cnrtQueueSync(queue); - std::fprintf(stderr, "MASK_ADD_ENQUEUE_FAIL cnnl=%d drain=%d\n", - int(status), int(completion)); - std::abort(); - } - } - cnnl_utils::Handle mask_handle_{}; - cnnl_utils::TensorDescriptor mask_input_desc_{}, mask_plane_desc_{}, - mask_out_desc_{}; - MaskOp mask_op_{}; - cnrt_utils::DeviceBuffer mask_workspace_{}; - size_t mask_workspace_bytes_{0}; }; } // namespace infini::ops diff --git a/src/native/cpu/ops/fill/fill.h b/src/native/cpu/ops/fill/fill.h index 3e5e67bda..4ae669cf2 100644 --- a/src/native/cpu/ops/fill/fill.h +++ b/src/native/cpu/ops/fill/fill.h @@ -30,9 +30,6 @@ class Operator : public Fill, "CPU Fill does not supply a Tensor-valued provider"); } - // The former CPU fill boundary qualified host replay and reported true here. - bool capture_safe() const override { return true; } - void operator()(const Tensor input, const double value, Tensor out) const override { if (input.shape() != input_shape_ || input.strides() != input_strides_ || diff --git a/src/native/cpu/ops/mul_scalar/mul_scalar.h b/src/native/cpu/ops/mul_scalar/mul_scalar.h index 182a90d1d..49cb092a5 100644 --- a/src/native/cpu/ops/mul_scalar/mul_scalar.h +++ b/src/native/cpu/ops/mul_scalar/mul_scalar.h @@ -41,8 +41,6 @@ class Operator : public MulScalar, count_ = bytes / Width(dtype_); } - bool capture_safe() const override { return false; } - void operator()(const Tensor input, const double alpha, Tensor out) const override { if (input.shape() != input_shape_ || out.shape() != out_shape_ || diff --git a/src/native/cpu/ops/row_l2_normalize/row_l2_normalize.h b/src/native/cpu/ops/row_l2_normalize/row_l2_normalize.h index 72e741d33..e9facbc26 100644 --- a/src/native/cpu/ops/row_l2_normalize/row_l2_normalize.h +++ b/src/native/cpu/ops/row_l2_normalize/row_l2_normalize.h @@ -38,8 +38,6 @@ class Operator Operator(const Tensor input, const double eps, Tensor out) : RowL2Normalize{Checked(input, eps, out), eps, out} {} - bool capture_safe() const override { return false; } - void operator()(const Tensor input, const double eps, Tensor out) const override { if (input.shape() != input_shape_ || out.shape() != out_shape_ || diff --git a/src/native/cuda/nvidia/ops/cast/kernel.h b/src/native/cuda/nvidia/ops/cast/kernel.h index df0afc5ce..7bd6a4aa5 100644 --- a/src/native/cuda/nvidia/ops/cast/kernel.h +++ b/src/native/cuda/nvidia/ops/cast/kernel.h @@ -52,8 +52,6 @@ class Operator : public Cast { Operator(const Tensor input, Tensor out) : Cast{Checked(input, out), out}, device_{input.device()} {} - bool capture_safe() const override { return false; } - void operator()(const Tensor input, Tensor out) const override { if (input.dtype() != input_dtype_ || out.dtype() != out_dtype_ || input.shape() != input_shape_ || out.shape() != out_shape_ || diff --git a/src/native/cuda/nvidia/ops/cat/kernel.h b/src/native/cuda/nvidia/ops/cat/kernel.h index 46c4e7bc4..229bd2ec1 100644 --- a/src/native/cuda/nvidia/ops/cat/kernel.h +++ b/src/native/cuda/nvidia/ops/cat/kernel.h @@ -29,9 +29,6 @@ class Operator : public Cat { if (out.device().type() != Device::Type::kNvidia || !Cat::supports(tensors, dim, out)) return false; - const auto normalized = - dim < 0 ? dim + static_cast(out.ndim()) : dim; - if (FastAxis0(tensors, normalized, out)) return true; return true; } diff --git a/src/native/cuda/nvidia/ops/causal_conv1d/kernel.h b/src/native/cuda/nvidia/ops/causal_conv1d/kernel.h index 9353b11c6..c1023f653 100644 --- a/src/native/cuda/nvidia/ops/causal_conv1d/kernel.h +++ b/src/native/cuda/nvidia/ops/causal_conv1d/kernel.h @@ -132,7 +132,6 @@ class Operator : public CausalConv1d { } // Device capture/state-pool lifetime is not supported for this stateful path. // Zero device allocations does not alone establish capture eligibility. - bool capture_safe() const override { return false; } void operator()(const Tensor qkv, const Tensor weight, Tensor conv_state, const std::optional bias, const std::optional cu_seqlens, diff --git a/src/native/cuda/nvidia/ops/chunk_gated_delta_rule/kernel.h b/src/native/cuda/nvidia/ops/chunk_gated_delta_rule/kernel.h index 451443059..267ddf400 100644 --- a/src/native/cuda/nvidia/ops/chunk_gated_delta_rule/kernel.h +++ b/src/native/cuda/nvidia/ops/chunk_gated_delta_rule/kernel.h @@ -84,7 +84,6 @@ class Operator p_.s[i][j] = metadata_[i].strides[j]; } size_t workspace_size_in_bytes() const override { return bytes_; } - bool capture_safe() const override { return false; } void operator()(const Tensor q, const Tensor k, const Tensor v, const Tensor g, const Tensor beta, Tensor initial_state, const std::optional cu_seqlens, diff --git a/src/native/cuda/nvidia/ops/recurrent_gated_delta_rule/kernel.h b/src/native/cuda/nvidia/ops/recurrent_gated_delta_rule/kernel.h index c687268a1..78b11cf74 100644 --- a/src/native/cuda/nvidia/ops/recurrent_gated_delta_rule/kernel.h +++ b/src/native/cuda/nvidia/ops/recurrent_gated_delta_rule/kernel.h @@ -69,7 +69,6 @@ class Operator for (size_t j = 0; j < metadata_[i].strides.size(); ++j) p_.s[i][j] = metadata_[i].strides[j]; } - bool capture_safe() const override { return false; } void operator()(const Tensor q, const Tensor k, const Tensor v, const Tensor g, const Tensor beta, Tensor initial_state, const std::optional initial_state_indices, diff --git a/src/native/cuda/nvidia/ops/rms_norm/kernel.h b/src/native/cuda/nvidia/ops/rms_norm/kernel.h index 87266bc50..85769299c 100644 --- a/src/native/cuda/nvidia/ops/rms_norm/kernel.h +++ b/src/native/cuda/nvidia/ops/rms_norm/kernel.h @@ -12,12 +12,7 @@ class Operator : public CudaRmsNorm> { public: using Base = CudaRmsNorm>; - - Operator(const Tensor input, const Tensor weight, float eps, Tensor out) - : Base(input, weight, eps, out) {} - - Operator(const Tensor input, const Tensor weight, Tensor out) - : Operator(input, weight, 1e-6f, out) {} + using Base::Base; // Mixed weight types dispatch to kernels too; replay allocates no metadata. bool capture_safe() const override { return true; } diff --git a/src/native/cuda/nvidia/ops/sum/kernel.h b/src/native/cuda/nvidia/ops/sum/kernel.h index c5a3d1eb4..34f1685de 100644 --- a/src/native/cuda/nvidia/ops/sum/kernel.h +++ b/src/native/cuda/nvidia/ops/sum/kernel.h @@ -71,7 +71,7 @@ class Operator : public Sum { for (std::size_t d = 0; d < rank; ++d) { const auto extent = input.shape()[d]; const auto stride = input.strides()[d]; - if (!extent || stride < 0 || (extent > 1 && stride == 0) || + if (!extent || stride < 0 || extent > static_cast( std::numeric_limits::max()) || static_cast(stride) > diff --git a/src/native/cuda/nvidia/ops/take/kernel.h b/src/native/cuda/nvidia/ops/take/kernel.h index 4a6c8b995..c44bb9bba 100644 --- a/src/native/cuda/nvidia/ops/take/kernel.h +++ b/src/native/cuda/nvidia/ops/take/kernel.h @@ -83,9 +83,7 @@ class Operator : public Take { alignof(std::int64_t)))) { return false; } - return !Overlaps(input.data(), input_count * width, index.data(), - output_count * sizeof(std::int64_t)) && - !Overlaps(input.data(), input_count * width, out.data(), + return !Overlaps(input.data(), input_count * width, out.data(), output_count * width) && !Overlaps(index.data(), output_count * sizeof(std::int64_t), out.data(), output_count * width); diff --git a/src/native/cuda/nvidia/ops/upsample_bilinear2d/kernel.h b/src/native/cuda/nvidia/ops/upsample_bilinear2d/kernel.h index 4b3eadc51..3590ba063 100644 --- a/src/native/cuda/nvidia/ops/upsample_bilinear2d/kernel.h +++ b/src/native/cuda/nvidia/ops/upsample_bilinear2d/kernel.h @@ -85,7 +85,6 @@ class Operator const auto extent = input.shape()[d]; const auto stride = input.strides()[d]; if (!extent || !out.shape()[d] || stride < 0 || - (extent > 1 && stride == 0) || static_cast(stride) > static_cast( std::numeric_limits::max()) || diff --git a/src/native/cuda/nvidia/ops/w4_marlin/gemm.cu b/src/native/cuda/nvidia/ops/w4_marlin/gemm.cu index 00ec09c9a..e9d8b0e12 100644 --- a/src/native/cuda/nvidia/ops/w4_marlin/gemm.cu +++ b/src/native/cuda/nvidia/ops/w4_marlin/gemm.cu @@ -1,6 +1,9 @@ +#include + #include "kernel.h" #include "launch.h" -#include "vendor/gemm/nvidia/kernel.cuh" +#include "core/scalar_type.hpp" +#include "marlin.cuh" namespace infini::ops::native_w4 { inline std::size_t align256(std::size_t bytes) { diff --git a/src/native/cuda/nvidia/ops/w4_marlin/kernel.h b/src/native/cuda/nvidia/ops/w4_marlin/kernel.h index b3e8a3c63..7fd2b47a3 100644 --- a/src/native/cuda/nvidia/ops/w4_marlin/kernel.h +++ b/src/native/cuda/nvidia/ops/w4_marlin/kernel.h @@ -14,7 +14,6 @@ class Operator void operator()(const Tensor qweight, const Tensor perm, const int64_t scheme, const int64_t size_k, const int64_t size_n, Tensor out) const override; - bool capture_safe() const override { return false; } private: int sms_{0}, shared_{0}, device_{0}; @@ -31,7 +30,6 @@ class Operator const int64_t scheme, const int64_t size_k, const int64_t size_n, const int64_t group_size, Tensor scales_out, Tensor zeros_out) const override; - bool capture_safe() const override { return false; } private: int device_{0}; @@ -53,7 +51,6 @@ class Operator : public W4MarlinGemm { bool requires_workspace() const override { return true; } // Original host selector still sets CUDA function attributes and queries the // device at invocation. No capture promise without a qualified frozen launch. - bool capture_safe() const override { return false; } private: int device_{0}, sms_{0}; diff --git a/src/native/cuda/nvidia/ops/w4_marlin/launch.h b/src/native/cuda/nvidia/ops/w4_marlin/launch.h index da32faf17..d0a9f1dfc 100644 --- a/src/native/cuda/nvidia/ops/w4_marlin/launch.h +++ b/src/native/cuda/nvidia/ops/w4_marlin/launch.h @@ -57,3 +57,20 @@ void launch_gptq_repack(const uint32_t* input, const uint32_t* perm, uint32_t* output, int k, int n, int sms, int shared, cudaStream_t stream); } // namespace infini::ops::native_w4 + +// Native pointer entry from the fixed upstream marlin.cu; no Tensor wrapper. +namespace vllm { +class ScalarType; +} +namespace marlin { +void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias, + void* a_s, void* b_s, void* g_s, void* zp, void* g_idx, + void* perm, void* a_tmp, int prob_m, int prob_n, int prob_k, + int lda, void* workspace, vllm::ScalarType const& a_type, + vllm::ScalarType const& b_type, vllm::ScalarType const& c_type, + vllm::ScalarType const& s_type, bool has_bias, + bool has_act_order, bool is_k_full, bool has_zp, int num_groups, + int group_size, int dev, cudaStream_t stream, int thread_k_init, + int thread_n_init, int sms, bool use_atomic_add, + bool use_fp32_reduce, bool is_zp_float); +} // namespace marlin diff --git a/src/native/cuda/nvidia/ops/w4_marlin/marlin_check.h b/src/native/cuda/nvidia/ops/w4_marlin/marlin_check.h new file mode 100644 index 000000000..085902de1 --- /dev/null +++ b/src/native/cuda/nvidia/ops/w4_marlin/marlin_check.h @@ -0,0 +1,26 @@ +#pragma once + +#include +#include +#include + +namespace infini::ops::native_w4 { +template +[[noreturn]] inline void marlin_fail(Args&&... args) { + std::ostringstream message; + message << "Marlin check failed"; + if constexpr (sizeof...(args) > 0) { + message << ": "; + (message << ... << std::forward(args)); + } + throw std::runtime_error(message.str()); +} +} // namespace infini::ops::native_w4 + +// Local replacement for the upstream checks, including constexpr factories. +#define INFINI_MARLIN_CHECK(condition, ...) \ + do { \ + if (!(condition)) { \ + ::infini::ops::native_w4::marlin_fail(__VA_ARGS__); \ + } \ + } while (false) diff --git a/src/native/cuda/nvidia/ops/w4_marlin/repack_awq.cu b/src/native/cuda/nvidia/ops/w4_marlin/repack_awq.cu index 28830d524..6bc12c71c 100644 --- a/src/native/cuda/nvidia/ops/w4_marlin/repack_awq.cu +++ b/src/native/cuda/nvidia/ops/w4_marlin/repack_awq.cu @@ -1,16 +1,13 @@ #include "launch.h" -#include "vendor/awq_repack/cuda/kernel.cuh" +#include "awq_marlin_repack.cu" namespace infini::ops::native_w4 { -__global__ void awq_repack(const uint32_t* input, uint32_t* output, int k, int n) { - marlin::awq_marlin_repack_kernel(input, output, k, n); -} void prepare_awq_repack(int shared) { - check(cudaFuncSetAttribute(awq_repack, cudaFuncAttributeMaxDynamicSharedMemorySize, shared), "W4 AWQ repack attributes"); + check(cudaFuncSetAttribute(marlin::awq_marlin_repack_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shared), "W4 AWQ repack attributes"); } void launch_awq_repack(const uint32_t* input, uint32_t* output, int k, int n, int sms, int shared, cudaStream_t stream) { - awq_repack<<>>(input, output, k, n); + marlin::awq_marlin_repack_kernel<<>>(input, output, k, n); check(cudaGetLastError(), "W4 AWQ repack launch"); } } // namespace infini::ops::native_w4 diff --git a/src/native/cuda/nvidia/ops/w4_marlin/repack_gptq.cu b/src/native/cuda/nvidia/ops/w4_marlin/repack_gptq.cu index 8da720fe6..bdfdf99fd 100644 --- a/src/native/cuda/nvidia/ops/w4_marlin/repack_gptq.cu +++ b/src/native/cuda/nvidia/ops/w4_marlin/repack_gptq.cu @@ -1,19 +1,15 @@ #include "launch.h" -#include "vendor/gptq_repack/cuda/kernel.cuh" +#include "gptq_marlin_repack.cu" namespace infini::ops::native_w4 { -template -__global__ void gptq_repack(const uint32_t* input, const uint32_t* perm, uint32_t* output, int k, int n) { - marlin::gptq_marlin_repack_kernel(input, perm, output, k, n); -} void prepare_gptq_repack(int shared, bool permutation) { - if (permutation) check(cudaFuncSetAttribute(gptq_repack, cudaFuncAttributeMaxDynamicSharedMemorySize, shared), "W4 GPTQ repack attributes"); - else check(cudaFuncSetAttribute(gptq_repack, cudaFuncAttributeMaxDynamicSharedMemorySize, shared), "W4 GPTQ repack attributes"); + if (permutation) check(cudaFuncSetAttribute(marlin::gptq_marlin_repack_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shared), "W4 GPTQ repack attributes"); + else check(cudaFuncSetAttribute(marlin::gptq_marlin_repack_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shared), "W4 GPTQ repack attributes"); } void launch_gptq_repack(const uint32_t* input, const uint32_t* perm, uint32_t* output, int k, int n, int sms, int shared, cudaStream_t stream) { - if (perm) gptq_repack<<>>(input, perm, output, k, n); - else gptq_repack<<>>(input, nullptr, output, k, n); + if (perm) marlin::gptq_marlin_repack_kernel<<>>(input, perm, output, k, n); + else marlin::gptq_marlin_repack_kernel<<>>(input, nullptr, output, k, n); check(cudaGetLastError(), "W4 GPTQ repack launch"); } } // namespace infini::ops::native_w4 diff --git a/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/core/scalar_type.hpp b/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/core/scalar_type.hpp deleted file mode 100644 index ba45d7245..000000000 --- a/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/core/scalar_type.hpp +++ /dev/null @@ -1,240 +0,0 @@ -// Source: InfiniCore 1ab85ef1d3c2bce49edf1fc73a5b48becaf7552e, -// src/infiniop/ops/awq_marlin_gemm/core/scalar_type.hpp. -// Donor implementation retained; this attribution notice was added by InfiniOps. -// Marlin/vLLM and direct-donor licenses: repository THIRD_PARTY.txt. -#pragma once - -#include -#include -#include -#include -#include -#include -#include - -namespace vllm { - -class ScalarType { -public: - enum NanRepr : uint8_t { - NAN_NONE = 0, - NAN_IEEE_754 = 1, - NAN_EXTD_RANGE_MAX_MIN = 2, - - NAN_REPR_ID_MAX - }; - - constexpr ScalarType(uint8_t exponent, uint8_t mantissa, bool signed_, - int32_t bias, bool finite_values_only = false, - NanRepr nan_repr = NAN_IEEE_754) - : exponent(exponent), - mantissa(mantissa), - signed_(signed_), - bias(bias), - finite_values_only(finite_values_only), - nan_repr(nan_repr) {} - - static constexpr ScalarType int_(uint8_t size_bits, int32_t bias = 0) { - return ScalarType(0, size_bits - 1, true, bias); - } - - static constexpr ScalarType uint(uint8_t size_bits, int32_t bias = 0) { - return ScalarType(0, size_bits, false, bias); - } - - static constexpr ScalarType float_IEEE754(uint8_t exponent, - uint8_t mantissa) { - return ScalarType(exponent, mantissa, true, 0, false, NAN_IEEE_754); - } - - static constexpr ScalarType float_(uint8_t exponent, uint8_t mantissa, - bool finite_values_only, - NanRepr nan_repr) { - return ScalarType(exponent, mantissa, true, 0, - finite_values_only, nan_repr); - } - - static inline ScalarType float_checked(uint8_t exponent, - uint8_t mantissa, - bool finite_values_only, - NanRepr nan_repr) { - if (!(nan_repr < NAN_REPR_ID_MAX)) { - throw std::runtime_error("Invalid NanRepr"); - } - - if (!(mantissa > 0 && exponent > 0)) { - throw std::runtime_error("mantissa/exponent must > 0"); - } - - if (nan_repr == NAN_IEEE_754) { - throw std::runtime_error("use float_IEEE754"); - } - - return float_(exponent, mantissa, finite_values_only, nan_repr); - } - - uint8_t const exponent; - uint8_t const mantissa; - bool const signed_; - int32_t const bias; - - bool const finite_values_only; - NanRepr const nan_repr; - - using Id = int64_t; - -private: - template - static constexpr size_t member_id_field_width() { - using T = std::decay_t; - return std::is_same::value ? 1 : sizeof(T) * 8; - } - - template - static constexpr auto reduce_members_helper(Fn f, Init val, Member member, - Rest... rest) { - auto new_val = f(val, member); - if constexpr (sizeof...(rest) > 0) { - return reduce_members_helper(f, new_val, rest...); - } else { - return new_val; - } - } - - template - constexpr auto reduce_members(Fn f, Init init) const { - return reduce_members_helper(f, init, exponent, mantissa, signed_, bias, - finite_values_only, nan_repr); - } - - template - static constexpr auto reduce_member_types(Fn f, Init init) { - constexpr auto dummy = ScalarType(0, 0, false, 0, false, NAN_NONE); - return dummy.reduce_members(f, init); - } - - static constexpr auto id_size_bits() { - return reduce_member_types( - [](int acc, auto member) -> int { - return acc + member_id_field_width(); - }, - 0); - } - -public: - constexpr Id id() const { - static_assert(id_size_bits() <= sizeof(Id) * 8, - "ScalarType id too large"); - - auto fn = [](std::pair result, auto member) { - auto [id, offset] = result; - constexpr auto bits = member_id_field_width(); - return std::pair{ - id | (int64_t(member) & ((uint64_t(1) << bits) - 1)) << offset, - offset + bits}; - }; - - return reduce_members(fn, std::pair{}).first; - } - - static constexpr ScalarType from_id(Id id) { - auto fn = [id](auto result, auto member) { - using T = decltype(member); - auto [tuple, offset] = result; - constexpr auto bits = member_id_field_width(); - auto val = static_cast((id >> offset) & ((uint64_t(1) << bits) - 1)); - return std::pair{std::tuple_cat(tuple, std::make_tuple(val)), offset + bits}; - }; - - auto [args, _] = reduce_member_types(fn, std::pair, int>{}); - - return std::apply([](auto... xs) { return ScalarType(xs...); }, args); - } - - constexpr int64_t size_bits() const { - return mantissa + exponent + (signed_ ? 1 : 0); - } - - constexpr bool is_signed() const { return signed_; } - constexpr bool is_integer() const { return exponent == 0; } - constexpr bool is_floating_point() const { return exponent > 0; } - - constexpr bool is_ieee_754() const { - return is_floating_point() && !finite_values_only && nan_repr == NAN_IEEE_754; - } - - constexpr bool has_nans() const { - return is_floating_point() && nan_repr != NAN_NONE; - } - - constexpr bool has_infs() const { - return is_floating_point() && !finite_values_only; - } - - constexpr bool has_bias() const { return bias != 0; } - - std::string str() const { - if (is_floating_point()) { - auto ret = "float" + std::to_string(size_bits()) + "_e" + std::to_string(exponent) + "m" + std::to_string(mantissa); - - if (!is_ieee_754()) { - if (finite_values_only) { - ret += "f"; - } - if (nan_repr != NAN_NONE) { - ret += "n"; - } - } - return ret; - } else { - auto ret = (signed_ ? "int" : "uint") + std::to_string(size_bits()); - if (has_bias()) { - ret += "b" + std::to_string(bias); - } - return ret; - } - } - - constexpr bool operator==(ScalarType const &other) const { - return mantissa == other.mantissa && exponent == other.exponent && bias == other.bias && signed_ == other.signed_ && finite_values_only == other.finite_values_only && nan_repr == other.nan_repr; - } -}; - -using ScalarTypeId = ScalarType::Id; - -static inline constexpr auto kS4 = ScalarType::int_(4); -static inline constexpr auto kU4 = ScalarType::uint(4); -static inline constexpr auto kU4B8 = ScalarType::uint(4, 8); -static inline constexpr auto kS8 = ScalarType::int_(8); -static inline constexpr auto kU8 = ScalarType::uint(8); -static inline constexpr auto kU8B128 = ScalarType::uint(8, 128); - -static inline constexpr auto kFE2M1f = ScalarType::float_(2, 1, true, ScalarType::NAN_NONE); -static inline constexpr auto kFE3M2f = ScalarType::float_(3, 2, true, ScalarType::NAN_NONE); -static inline constexpr auto kFE4M3fn = ScalarType::float_(4, 3, true, ScalarType::NAN_EXTD_RANGE_MAX_MIN); -static inline constexpr auto kFE8M0fnu = ScalarType(8, 0, false, 0, true, ScalarType::NAN_EXTD_RANGE_MAX_MIN); -static inline constexpr auto kFE5M2 = ScalarType::float_IEEE754(5, 2); -static inline constexpr auto kFE8M7 = ScalarType::float_IEEE754(8, 7); -static inline constexpr auto kFE5M10 = ScalarType::float_IEEE754(5, 10); - -static inline constexpr auto kInt4 = kS4; -static inline constexpr auto kUint4 = kU4; -static inline constexpr auto kUint4b8 = kU4B8; -static inline constexpr auto kInt8 = kS8; -static inline constexpr auto kUint8 = kU8; -static inline constexpr auto kUint8b128 = kU8B128; - -static inline constexpr auto kFloat4_e2m1f = kFE2M1f; -static inline constexpr auto kFloat6_e3m2f = kFE3M2f; -static inline constexpr auto kFloat8_e4m3fn = kFE4M3fn; -static inline constexpr auto kFloat8_e5m2 = kFE5M2; -static inline constexpr auto kFloat16_e8m7 = kFE8M7; -static inline constexpr auto kFloat16_e5m10 = kFE5M10; - -static inline constexpr auto kHalf = kFE5M10; -static inline constexpr auto kFloat16 = kHalf; -static inline constexpr auto kBFloat16 = kFE8M7; - -static inline constexpr auto kFloat16Id = kFloat16.id(); - -} // namespace vllm diff --git a/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/core/source_location.h b/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/core/source_location.h deleted file mode 100644 index 74ef9c422..000000000 --- a/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/core/source_location.h +++ /dev/null @@ -1,44 +0,0 @@ -// Source: InfiniCore 1ab85ef1d3c2bce49edf1fc73a5b48becaf7552e, -// src/infiniop/ops/awq_marlin_gemm/core/source_location.h. -// Donor implementation retained; this attribution notice was added by InfiniOps. -// Marlin/vLLM and direct-donor licenses: repository THIRD_PARTY.txt. -/// \file source_location.h -/// \brief Portable `source_location` wrapper. -/// -/// Uses `std::source_location` when available (C++20), otherwise falls -/// back to a minimal stub that returns empty/zero values. - -#pragma once -#include - -/// NOTE: fallback to a minimal source_location implementation -#if defined(__cpp_lib_source_location) -#include - -using source_location_t = std::source_location; - -#else - -struct source_location_fallback { -public: - static constexpr source_location_fallback current() noexcept { - return source_location_fallback{}; - } - constexpr source_location_fallback() noexcept = default; - constexpr unsigned line() const noexcept { - return 0; - } - constexpr unsigned column() const noexcept { - return 0; - } - constexpr const char *file_name() const noexcept { - return ""; - } - constexpr const char *function_name() const noexcept { - return ""; - } -}; - -using source_location_t = source_location_fallback; - -#endif diff --git a/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/core/utils.h b/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/core/utils.h deleted file mode 100644 index f1497cd60..000000000 --- a/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/core/utils.h +++ /dev/null @@ -1,142 +0,0 @@ -// Adapted from InfiniCore 1ab85ef1d3c2bce49edf1fc73a5b48becaf7552e: -// src/infiniop/ops/awq_marlin_gemm/core/utils.h. -// Modified for InfiniOps: unused pointer, range and dtype helpers removed. -// See repository THIRD_PARTY.txt and LICENSES/InfiniCore-MIT.txt. -/// \file utils.h -/// \brief Host-side C++ utilities used by JIT kernel wrappers. -/// -/// Provides: -/// - `DebugInfo` - wraps `std::source_location` for error reporting. -/// - `RuntimeCheck` - runtime assertion with formatted error messages. -/// - `Panic` - unconditional abort with formatted error messages. - -#pragma once - -// ref: https://forums.developer.nvidia.com/t/c-20s-source-location-compilation-error-when-using-nvcc-12-1/258026/3 -#ifdef __CUDACC__ -#include -#if CUDA_VERSION <= 12010 - -#pragma push_macro("__cpp_consteval") -#pragma push_macro("_NODISCARD") -#pragma push_macro("__builtin_LINE") - -#pragma clang diagnostic push -#pragma clang diagnostic ignored "-Wbuiltin-macro-redefined" -#define __cpp_consteval 201811L -#pragma clang diagnostic pop - -#ifdef _NODISCARD -#undef _NODISCARD -#define _NODISCARD -#endif - -#define consteval constexpr - -#include "source_location.h" - -#undef consteval -#pragma pop_macro("__cpp_consteval") -#pragma pop_macro("_NODISCARD") -#else // __CUDACC__ && CUDA_VERSION > 12010 -#include "source_location.h" -#endif -#else // no __CUDACC__ -#include "source_location.h" -#endif - -// #include -#include -#include -// #include -#include -#include - -namespace host { - -/// \brief Source-location wrapper for debug/error messages. -struct DebugInfo : public source_location_t { - DebugInfo(source_location_t loc = source_location_t::current()) : source_location_t(loc) {} -}; - -/// \brief Exception type thrown by `RuntimeCheck` and `Panic`. -struct PanicError : public std::runtime_error { -public: - explicit PanicError(std::string msg) : runtime_error(msg), m_message(std::move(msg)) {} - auto root_cause() const -> std::string_view { - const auto str = std::string_view{m_message}; - const auto pos = str.find(": "); - return pos == std::string_view::npos ? str : str.substr(pos + 2); - } - -private: - std::string m_message; -}; - -/// \brief Unconditionally abort with a formatted error message. -template -[[noreturn]] inline auto panic(DebugInfo location, Args &&...args) -> void { - std::ostringstream os; - os << "Runtime check failed at " << location.file_name() << ":" << location.line(); - if constexpr (sizeof...(args) > 0) { - os << ": "; - (os << ... << std::forward(args)); - } else { - os << " in " << location.function_name(); - } - throw PanicError(std::move(os).str()); -} - -/** - * \brief Runtime assertion: panics with a formatted message when `condition` - * is false. Extra `args` are streamed to the error message. - * - * Example: - * \code - * RuntimeCheck(n > 0, "n must be positive, got ", n); - * \endcode - */ -template -struct RuntimeCheck { - template - explicit RuntimeCheck(Cond &&condition, Args &&...args, DebugInfo location = {}) { - if (condition) { - return; - } - [[unlikely]] ::host::panic(location, std::forward(args)...); - } - template - explicit RuntimeCheck(DebugInfo location, Cond &&condition, Args &&...args) { - if (condition) { - return; - } - [[unlikely]] ::host::panic(location, std::forward(args)...); - } -}; - -template -struct Panic { - explicit Panic(Args &&...args, DebugInfo location = {}) { - ::host::panic(location, std::forward(args)...); - } - explicit Panic(DebugInfo location, Args &&...args) { - ::host::panic(location, std::forward(args)...); - } - [[noreturn]] ~Panic() { - std::terminate(); - } -}; - -template -explicit RuntimeCheck(Cond &&, Args &&...) -> RuntimeCheck; - -template -explicit RuntimeCheck(DebugInfo, Cond &&, Args &&...) -> RuntimeCheck; - -template -explicit Panic(Args &&...) -> Panic; - -template -explicit Panic(DebugInfo, Args &&...) -> Panic; - -} // namespace host diff --git a/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/nvidia/kernel.cuh b/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/nvidia/kernel.cuh deleted file mode 100644 index 161faa64b..000000000 --- a/src/native/cuda/nvidia/ops/w4_marlin/vendor/gemm/nvidia/kernel.cuh +++ /dev/null @@ -1,553 +0,0 @@ -// Source: InfiniCore 1ab85ef1d3c2bce49edf1fc73a5b48becaf7552e, -// src/infiniop/ops/awq_marlin_gemm/nvidia/kernel.cuh. -// Donor implementation retained; this attribution notice was added by InfiniOps. -// Marlin/vLLM and direct-donor licenses: repository THIRD_PARTY.txt. -/* - * Modified by Neural Magic - * Copyright (C) Marlin.2024 Elias Frantar - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -/* - * Adapted from https://github.com/IST-DASLab/marlin - */ - -#ifndef MARLIN_NAMESPACE_NAME -#define MARLIN_NAMESPACE_NAME marlin -#endif - -#include "../core/utils.h" -#include "marlin/kernel.h" - -#define STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t) \ - static_assert(std::is_same::value || std::is_same::value, \ - "only float16 and bfloat16 is supported"); - -namespace marlin { - -__global__ void MarlinDefault(MARLIN_KERNEL_PARAMS){}; - -using MarlinFuncPtr = void (*)(MARLIN_KERNEL_PARAMS); - -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 750 - -__global__ void permute_cols_kernel(int4 const *__restrict__ a_int4_ptr, - int const *__restrict__ perm_int_ptr, - int4 *__restrict__ out_int4_ptr, int size_m, - int size_k, int lda, int block_rows) {} - -} // namespace marlin - -template -void awq_marlin_gemm_kernel( - const void *a, - void *c, - const void *b_q_weight, - void *b_bias, - void *b_scales, - void *a_scales, - void *global_scale, - void *b_zeros, - void *g_idx, - void *perm, - int64_t b_type_id, - bool is_k_full, - bool use_atomic_add, - bool use_fp32_reduce, - bool is_zp_float, - int size_m, - int size_k, - int size_n, - int b_q_size_0, - int b_q_size_1, - int a_stride_0, - int b_zeros_size_1, - int num_groups, - cudaStream_t stream) { -} - -#else - -// For a given "a" of size [M,K] performs a permutation of the K columns based -// on the given "perm" indices. -__global__ void permute_cols_kernel(int4 const *__restrict__ a_int4_ptr, - int const *__restrict__ perm_int_ptr, - int4 *__restrict__ out_int4_ptr, int size_m, - int size_k, int lda, int block_rows) { - auto start_row = block_rows * blockIdx.x; - int finish_row = start_row + block_rows; - if (finish_row > size_m) { - finish_row = size_m; - } - int cur_block_rows = finish_row - start_row; - - int input_row_stride = lda * sizeof(half) / 16; - int output_row_stride = size_k * sizeof(half) / 16; - - auto permute_row = [&](int row) { - int iters = size_k / default_threads; - int rest = size_k % default_threads; - - int input_offset = row * input_row_stride; - int output_offset = row * output_row_stride; - - half const *a_row_half = reinterpret_cast(a_int4_ptr + input_offset); - half *out_half = reinterpret_cast(out_int4_ptr + output_offset); - - int base_k = 0; - - for (int i = 0; i < iters; i++) { - auto cur_k = base_k + threadIdx.x; - int src_pos = perm_int_ptr[cur_k]; - - out_half[cur_k] = a_row_half[src_pos]; - - base_k += default_threads; - } - - if (rest) { - if (threadIdx.x < rest) { - auto cur_k = base_k + threadIdx.x; - int src_pos = perm_int_ptr[cur_k]; - - out_half[cur_k] = a_row_half[src_pos]; - } - } - }; - - for (int i = 0; i < cur_block_rows; i++) { - int cur_row = start_row + i; - if (cur_row < size_m) { - permute_row(cur_row); - } - } -} - -typedef struct -{ - int thread_k; - int thread_n; - int num_threads; -} thread_config_t; - -thread_config_t small_batch_thread_configs[] = { - // Ordered by priority - - // thread_k, thread_n, num_threads - {128, 128, 256}, - {64, 128, 128}, - {128, 64, 128}}; - -thread_config_t large_batch_thread_configs[] = { - // Ordered by priority - - // thread_k, thread_n, num_threads - {64, 256, 256}, - {64, 128, 128}, - {128, 64, 128}}; - -typedef struct -{ - int blocks_per_sm; - thread_config_t tb_cfg; -} exec_config_t; - -int get_scales_cache_size(thread_config_t const &th_config, int prob_m, - int prob_n, int prob_k, int num_bits, int group_size, - bool has_act_order, bool is_k_full, int stages) { - bool cache_scales_chunk = has_act_order && !is_k_full; - - int tb_n = th_config.thread_n; - int tb_k = th_config.thread_k; - - // Get max scale groups per thread-block - int tb_groups; - if (group_size == -1) { - tb_groups = 1; - } else if (group_size == 0) { - tb_groups = div_ceil(tb_k, 32); // Worst case is 32 group size - } else { - tb_groups = div_ceil(tb_k, group_size); - } - - if (cache_scales_chunk) { - int load_groups = tb_groups * stages * 2; // Chunk size is 2x pipeline over dim K - load_groups = max(load_groups, 32); // We load at least 32 scale groups - return load_groups * tb_n * 2; - } else { - int tb_scales = tb_groups * tb_n * 2; - - return tb_scales * stages; - } -} - -int get_kernel_cache_size(thread_config_t const &th_config, int thread_m_blocks, - int prob_m, int prob_n, int prob_k, int num_bits, - int group_size, bool has_act_order, bool is_k_full, - int has_zp, bool is_zp_float, bool is_a_8bit, - int stages) { - int pack_factor = 32 / num_bits; - - // Get B size - int tb_k = th_config.thread_k; - int tb_n = th_config.thread_n; - int tb_m = thread_m_blocks * 16; - int sh_a_size = stages * (tb_m * tb_k) * (is_a_8bit ? 1 : 2); - int sh_b_size = stages * (tb_k * tb_n / pack_factor) * 4; - int sh_red_size = tb_m * (tb_n + 8) * 2; - int sh_bias_size = tb_n * 2; - int tmp_size = (sh_b_size > sh_red_size ? sh_red_size : sh_b_size) + sh_bias_size; - tmp_size = max(max(sh_b_size, sh_red_size), tmp_size); - - int sh_s_size = get_scales_cache_size(th_config, prob_m, prob_n, prob_k, num_bits, - group_size, has_act_order, is_k_full, stages); - int sh_g_idx_size = has_act_order && !is_k_full ? stages * tb_k / 4 : 0; - int sh_zp_size = 0; - if (has_zp) { - if (is_zp_float) { - sh_zp_size = sh_s_size; - } else if (num_bits == 4) { - sh_zp_size = sh_s_size / 4; - } else if (num_bits == 8) { - sh_zp_size = sh_s_size / 2; - } - } - - int total_size = tmp_size + sh_a_size + sh_s_size + sh_zp_size + sh_g_idx_size; - - return total_size; -} - -bool is_valid_config(thread_config_t const &th_config, int thread_m_blocks, - int prob_m, int prob_n, int prob_k, int num_bits, - int group_size, bool has_act_order, bool is_k_full, - int has_zp, bool is_zp_float, bool is_a_8bit, int stages, - int max_shared_mem) { - // Sanity - if (th_config.thread_k == -1 || th_config.thread_n == -1 || th_config.num_threads == -1) { - return false; - } - - // Verify K/N are divisible by thread K/N - if (prob_k % th_config.thread_k != 0 || prob_n % th_config.thread_n != 0) { - return false; - } - - // Verify min for thread K/N - if (th_config.thread_n < min_thread_n || th_config.thread_k < min_thread_k) { - return false; - } - - // num_threads must be at least 128 (= 4 warps) - if (th_config.num_threads < 128) { - return false; - } - - // Check that pipeline fits into cache - int cache_size = get_kernel_cache_size( - th_config, thread_m_blocks, prob_m, prob_n, prob_k, num_bits, group_size, - has_act_order, is_k_full, has_zp, is_zp_float, is_a_8bit, stages); - return cache_size <= max_shared_mem; -} - -MarlinFuncPtr get_marlin_kernel( - const vllm::ScalarType a_type, const vllm::ScalarType b_type, - const vllm::ScalarType c_type, const vllm::ScalarType s_type, - int thread_m_blocks, int thread_n_blocks, int thread_k_blocks, - bool m_block_size_8, bool has_act_order, bool has_zp, int group_blocks, - int threads, bool is_zp_float, int stages) { - int num_bits = b_type.size_bits(); - auto kernel = MarlinDefault; - -#include "kernel_selector.h" - - return kernel; -} - -exec_config_t determine_exec_config( - const vllm::ScalarType &a_type, const vllm::ScalarType &b_type, - const vllm::ScalarType &c_type, const vllm::ScalarType &s_type, int prob_m, - int prob_n, int prob_k, int thread_m_blocks, bool m_block_size_8, - int num_bits, int group_size, bool has_act_order, bool is_k_full, - bool has_zp, bool is_zp_float, int is_a_8bit, int stages, - int max_shared_mem, int sms) { - exec_config_t exec_cfg = exec_config_t{1, thread_config_t{-1, -1, -1}}; - thread_config_t *thread_configs = thread_m_blocks > 1 - ? large_batch_thread_configs - : small_batch_thread_configs; - int thread_configs_size = thread_m_blocks > 1 - ? sizeof(large_batch_thread_configs) / sizeof(thread_config_t) - : sizeof(small_batch_thread_configs) / sizeof(thread_config_t); - - for (int i = 0; i < thread_configs_size; i++) { - thread_config_t th_config = thread_configs[i]; - - if (!is_valid_config(th_config, thread_m_blocks, prob_m, prob_n, prob_k, - num_bits, group_size, has_act_order, is_k_full, has_zp, - is_zp_float, is_a_8bit, stages, - max_shared_mem - 512)) { - continue; - } - - int cache_size = get_kernel_cache_size(th_config, thread_m_blocks, prob_m, - prob_n, prob_k, num_bits, group_size, - has_act_order, is_k_full, has_zp, - is_zp_float, is_a_8bit, stages); - - int group_blocks = 0; - if (!has_act_order) { - group_blocks = group_size == -1 ? -1 : group_size / 16; - } - - auto kernel = get_marlin_kernel(a_type, b_type, c_type, s_type, thread_m_blocks, - th_config.thread_n / 16, th_config.thread_k / 16, - m_block_size_8, has_act_order, has_zp, group_blocks, - th_config.num_threads, is_zp_float, stages); - - if (kernel == MarlinDefault) { - continue; - } - - return {1, th_config}; - } - - return exec_cfg; -} - -void marlin_mm(const void *A, const void *B, void *C, void *C_tmp, void *b_bias, - void *a_s, void *b_s, void *g_s, void *zp, void *g_idx, - void *perm, void *a_tmp, int prob_m, int prob_n, int prob_k, - int lda, void *workspace, vllm::ScalarType const &a_type, - vllm::ScalarType const &b_type, vllm::ScalarType const &c_type, - vllm::ScalarType const &s_type, bool has_bias, - bool has_act_order, bool is_k_full, bool has_zp, int num_groups, - int group_size, int dev, cudaStream_t stream, int thread_k_init, - int thread_n_init, int sms, bool use_atomic_add, - bool use_fp32_reduce, bool is_zp_float) { - bool is_a_8bit = a_type.size_bits() == 8; - host::RuntimeCheck(prob_m > 0 && prob_n > 0 && prob_k > 0, "Invalid MNK = [", prob_m, - ", ", prob_n, ", ", prob_k, "]"); - - int group_blocks = 0; - if (has_act_order) { - if (is_k_full) { - host::RuntimeCheck(group_size != -1); - group_blocks = group_size / 16; - host::RuntimeCheck(prob_k % group_blocks == 0, "prob_k = ", prob_k, - " is not divisible by group_blocks = ", group_blocks); - } else { - host::RuntimeCheck(group_size == 0); - group_blocks = 0; - } - } else { - if (group_size == -1) { - group_blocks = -1; - } else { - group_blocks = group_size / 16; - host::RuntimeCheck(prob_k % group_blocks == 0, "prob_k = ", prob_k, - " is not divisible by group_blocks = ", group_blocks); - } - } - - int num_bits = b_type.size_bits(); - const int4 *A_ptr = (const int4 *)A; - const int4 *B_ptr = (const int4 *)B; - int4 *C_ptr = (int4 *)C; - int4 *C_tmp_ptr = (int4 *)C_tmp; - - const int4 *bias_ptr = (const int4 *)b_bias; - const float *a_s_ptr = (const float *)a_s; - const int4 *b_s_ptr = (const int4 *)b_s; - const float *g_s_ptr = (const float *)g_s; - - const int4 *zp_ptr = (const int4 *)zp; - const int *g_idx_ptr = (const int *)g_idx; - const int *perm_ptr = (const int *)perm; - int4 *a_tmp_ptr = (int4 *)a_tmp; - int *locks = (int *)workspace; - - if (has_act_order) { - // Permute A columns - int block_rows = div_ceil(prob_m, sms); - // avoid ">>>" being formatted to "> > >" - // clang-format off - permute_cols_kernel<<>>( - A_ptr, perm_ptr, a_tmp_ptr, prob_m, prob_k, lda, block_rows); - // clang-format on - A_ptr = a_tmp_ptr; - lda = prob_k; - - // If we have a full K, then we can run the non-act-order version of Marlin - // (since the weight rows are reordered by increasing group ids, and by - // having a full K, we have full original groups) - if (is_k_full) { - has_act_order = false; - } - } - - int max_shared_mem = 0; - cudaDeviceGetAttribute(&max_shared_mem, - cudaDevAttrMaxSharedMemoryPerBlockOptin, dev); - host::RuntimeCheck(max_shared_mem > 0); - - int major_capability, minor_capability; - cudaDeviceGetAttribute(&major_capability, cudaDevAttrComputeCapabilityMajor, - dev); - cudaDeviceGetAttribute(&minor_capability, cudaDevAttrComputeCapabilityMinor, - dev); - host::RuntimeCheck(major_capability * 10 + minor_capability >= 75, - "marlin kernel only support Turing or newer GPUs."); - int stages = 4; - if (major_capability == 7 && minor_capability == 5) { - stages = 2; - host::RuntimeCheck(a_type == vllm::kFloat16 || a_type == vllm::kS8, - "Turing only support FP16 or INT8 activation."); - } - if (a_type == vllm::kFE4M3fn) { - host::RuntimeCheck( - major_capability * 10 + minor_capability == 89 || major_capability * 10 + minor_capability == 120, - "Marlin W4A8-FP8 only support SM89 or SM120 device (It is slower than " - "Marlin W4A16 on other devices)."); - } - - int max_par = 16; - if (prob_n <= 4096) { - max_par = 16 * 8; - } - int max_shared_mem_new = max_shared_mem; - int rest_m = prob_m; - int max_thread_m_blocks = 4; - while (rest_m) { - int par_count = rest_m / (max_thread_m_blocks * 16); - if (par_count > max_par) { - par_count = max_par; - } - int prob_m_split = par_count > 0 ? (par_count * (max_thread_m_blocks * 16)) : rest_m; - - int thread_k = thread_k_init; - int thread_n = thread_n_init; - - int thread_m_blocks = min(div_ceil(prob_m_split, 16), max_thread_m_blocks); - int m_block_size_8 = prob_m_split <= 8 && a_type.size_bits() == 16; - - // Set thread config - exec_config_t exec_cfg; - thread_config_t thread_tfg; - if (thread_k != -1 && thread_n != -1) { - thread_tfg = thread_config_t{thread_k, thread_n, default_threads}; - exec_cfg = exec_config_t{1, thread_tfg}; - host::RuntimeCheck(prob_n % thread_n == 0, "prob_n = ", prob_n, - " is not divisible by thread_n = ", thread_n); - host::RuntimeCheck(prob_k % thread_k == 0, "prob_k = ", prob_k, - " is not divisible by thread_k = ", thread_k); - } else { - // Auto config - exec_cfg = determine_exec_config( - a_type, b_type, c_type, s_type, prob_m_split, prob_n, prob_k, - thread_m_blocks, m_block_size_8, num_bits, group_size, has_act_order, - is_k_full, has_zp, is_zp_float, is_a_8bit, stages, max_shared_mem, - sms); - thread_tfg = exec_cfg.tb_cfg; - if (thread_tfg.thread_n != -1) { - if (prob_n / thread_tfg.thread_n * div_ceil(prob_m_split, thread_m_blocks * 16) * 4 <= sms) { - if (is_valid_config({128, 64, 128}, thread_m_blocks, prob_m_split, - prob_n, prob_k, num_bits, group_size, - has_act_order, is_k_full, has_zp, is_zp_float, - is_a_8bit, stages, max_shared_mem_new)) { - thread_tfg = {128, 64, 128}; - exec_cfg = {1, thread_tfg}; - } - } - } - - if (thread_tfg.thread_k == -1 && max_thread_m_blocks > 1) { - max_thread_m_blocks--; - continue; - } - } - - int num_threads = thread_tfg.num_threads; - thread_k = thread_tfg.thread_k; - thread_n = thread_tfg.thread_n; - int blocks = sms * exec_cfg.blocks_per_sm; - if (exec_cfg.blocks_per_sm > 1) { - max_shared_mem_new = max_shared_mem / exec_cfg.blocks_per_sm - 1024; - } - - int thread_k_blocks = thread_k / 16; - int thread_n_blocks = thread_n / 16; - - host::RuntimeCheck( - is_valid_config(thread_tfg, thread_m_blocks, prob_m_split, prob_n, - prob_k, num_bits, group_size, has_act_order, is_k_full, - has_zp, is_zp_float, is_a_8bit, stages, - max_shared_mem_new), - "Invalid thread config: thread_m_blocks = ", thread_m_blocks, - ", thread_k = ", thread_tfg.thread_k, - ", thread_n = ", thread_tfg.thread_n, - ", num_threads = ", thread_tfg.num_threads, " for MKN = [", prob_m, - ", ", prob_k, ", ", prob_n, "] and num_bits = ", num_bits, - ", prob_m_split = ", prob_m_split, ", group_size = ", group_size, - ", has_act_order = ", has_act_order, ", is_k_full = ", is_k_full, - ", has_zp = ", has_zp, ", is_zp_float = ", is_zp_float, - ", stages = ", stages, ", max_shared_mem_new = ", max_shared_mem_new); - - auto kernel = get_marlin_kernel( - a_type, b_type, c_type, s_type, thread_m_blocks, thread_n_blocks, - thread_k_blocks, m_block_size_8, has_act_order, has_zp, group_blocks, - num_threads, is_zp_float, stages); - - if (kernel == MarlinDefault) { - host::RuntimeCheck(false, "Unsupported shapes: MNK = [", prob_m, ", ", prob_n, - ", ", prob_k, "]", ", has_act_order = ", has_act_order, - ", num_groups = ", num_groups, ", group_size = ", group_size, - ", prob_m_split = ", prob_m_split, - ", thread_m_blocks = ", thread_m_blocks, - ", thread_n_blocks = ", thread_n_blocks, - ", thread_k_blocks = ", thread_k_blocks, - ", num_threads = ", num_threads, ", num_bits = ", num_bits); - } - - cudaStreamCaptureStatus capture_status = cudaStreamCaptureStatusNone; - host::RuntimeCheck( - cudaStreamIsCapturing(stream, &capture_status) == cudaSuccess, - "cudaStreamIsCapturing failed"); - if (capture_status == cudaStreamCaptureStatusNone) { - host::RuntimeCheck( - cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, max_shared_mem_new) == cudaSuccess, - "cudaFuncSetAttribute failed"); - } - - bool part_use_atomic_add = use_atomic_add && div_ceil(prob_m_split, 64) * prob_n <= 2048; - - // avoid ">>>" being formatted to "> > >" - // clang-format off - kernel<<>>( - A_ptr, B_ptr, C_ptr, C_tmp_ptr, bias_ptr, a_s_ptr, b_s_ptr, g_s_ptr, zp_ptr, - g_idx_ptr, num_groups, - prob_m_split, prob_n, prob_k, lda, locks, has_bias, part_use_atomic_add, - use_fp32_reduce, max_shared_mem_new); - // clang-format on - - bool is_a_8bit = a_type.size_bits() == 8; - A_ptr += prob_m_split * (lda / (is_a_8bit ? 16 : 8)); - a_s_ptr += prob_m_split; - C_ptr += prob_m_split * (prob_n / 8); - rest_m -= prob_m_split; - } -} - -} // namespace marlin - -#endif diff --git a/src/native/cuda/ops/separate_table_rope/kernel.h b/src/native/cuda/ops/separate_table_rope/kernel.h index cbbcfd8b7..5d0a44e6a 100644 --- a/src/native/cuda/ops/separate_table_rope/kernel.h +++ b/src/native/cuda/ops/separate_table_rope/kernel.h @@ -45,8 +45,8 @@ class CudaSeparateTableRope : public SeparateTableRope { using Angle = TypeMapType(tag)>; constexpr int kBlock = ListGet<2>(tag); // Existing kernel true = interleaved pair (GPT-J), false = split - // half (GPT-NeoX). Preserve Rs4's bool neox, not the old wrapper - // name. + // half (GPT-NeoX). Preserve the public bool neox convention, not the + // old wrapper name. auto launch = [&](auto interleaved) { LaunchRopeKernel( diff --git a/third_party/vllm b/third_party/vllm deleted file mode 160000 index ffc4f08c8..000000000 --- a/third_party/vllm +++ /dev/null @@ -1 +0,0 @@ -Subproject commit ffc4f08c8ee130d4ea6347c1bf31ffd4f8af28ab From 0f2d4fadab055fb404b436c6f05d8ac247adc039 Mon Sep 17 00:00:00 2001 From: whjthu Date: Tue, 20 Oct 2026 02:22:58 +0800 Subject: [PATCH 21/25] refactor: share native helpers and unify source archives --- .gitmodules | 3 - CMakeLists.txt | 9 + LICENSES/Apache-2.0-components.txt | 12 +- THIRD_PARTY.txt | 20 +- cmake/FlashAttentionSource.cmake | 41 +++ cmake/NativeFlashAttention.cmake | 21 +- docs/build-nvidia-inference.md | 102 +++---- pyproject.toml | 12 +- scripts/prepare_flash_attention_inference.py | 18 +- src/CMakeLists.txt | 1 + src/base/fused_gated_delta_net_gating.h | 1 - src/native/cpu/dense_view.h | 62 ++++ src/native/cpu/ops/mul_scalar/mul_scalar.h | 54 +--- .../ops/row_l2_normalize/row_l2_normalize.h | 50 +--- .../ops/flash_attn_varlen_func/kernel.cc | 280 +----------------- .../ops/flash_attn_with_kvcache/kernel.cc | 167 +---------- src/native/cuda/nvidia/ops/flash_common.cc | 124 ++++++++ src/native/cuda/nvidia/ops/flash_common.h | 26 ++ .../recurrent_delta_rule_common.cuh | 11 +- third_party/flash-attention | 1 - 20 files changed, 360 insertions(+), 655 deletions(-) create mode 100644 cmake/FlashAttentionSource.cmake create mode 100644 src/native/cpu/dense_view.h create mode 100644 src/native/cuda/nvidia/ops/flash_common.cc create mode 100644 src/native/cuda/nvidia/ops/flash_common.h delete mode 160000 third_party/flash-attention diff --git a/.gitmodules b/.gitmodules index c313414c1..be99e8a8f 100644 --- a/.gitmodules +++ b/.gitmodules @@ -1,6 +1,3 @@ [submodule ".ci"] path = .ci url = https://github.com/InfiniTensor/ci.git -[submodule "third_party/flash-attention"] - path = third_party/flash-attention - url = https://github.com/Dao-AILab/flash-attention.git diff --git a/CMakeLists.txt b/CMakeLists.txt index 198127e79..baa5ca2b2 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -7,6 +7,15 @@ if(SKBUILD_STATE STREQUAL "sdist") configure_file("${vllm_marlin_SOURCE_DIR}/${_marlin_file}" "${CMAKE_BINARY_DIR}/marlin-source/${_marlin_file}" COPYONLY) endforeach() + include(cmake/FlashAttentionSource.cmake) + foreach(_flash_file IN LISTS _flash_source_files) + configure_file("${flash_attention_SOURCE_DIR}/${_flash_file}" + "${CMAKE_BINARY_DIR}/flash-attention-source/${_flash_file}" COPYONLY) + endforeach() + configure_file("${flash_attention_cutlass_SOURCE_DIR}/LICENSE.txt" + "${CMAKE_BINARY_DIR}/flash-attention-cutlass-source/LICENSE.txt" COPYONLY) + file(COPY "${flash_attention_cutlass_SOURCE_DIR}/include" + DESTINATION "${CMAKE_BINARY_DIR}/flash-attention-cutlass-source") return() endif() project(InfiniOps LANGUAGES CXX) diff --git a/LICENSES/Apache-2.0-components.txt b/LICENSES/Apache-2.0-components.txt index 72477a894..f4f57a92a 100644 --- a/LICENSES/Apache-2.0-components.txt +++ b/LICENSES/Apache-2.0-components.txt @@ -13,21 +13,12 @@ Copyright 2025-2026 NVIDIA Copyright 2023-2026 FlashInfer community (https://flashinfer.ai/) Native source headers also retain Copyright (c) 2023, 2024, 2025 and 2026 by FlashInfer team. Shared terms: LICENSE. Original NOTICE: FlashInfer-0.6.7/NOTICE.txt. -Sampling's three TUs are csrc/sampling.cu, csrc/renorm.cu and -csrc/flashinfer_sampling_binding.cu (flashinfer/jit/sampling.py). -The fixed source include closure does not use the release's separate -FA3/CUTLASS/fmt/spdlog components. Complete FlashInfer distributions retain their -own additional applicable materials, which this native-only list does not replace. Apache TVM FFI 0.1.10 https://github.com/apache/tvm-ffi/tree/v0.1.10 Commit: 1fed0ae0421e614d45662e8ee6bcae353d3ab2ea Copyright 2024-present The Apache Software Foundation Shared terms: LICENSE. Original NOTICE: TVM-FFI-0.1.10/NOTICE.txt. -Native object source lists: CMakeLists.txt (_tvm_ffi_objs_sources and extras). -The separate python/tvm_ffi/_optional_torch_c_dlpack.py contains PyTorch-derived -code, but is not in the native runtime source list. Its notices still apply to -distributions including that addon; this list does not describe those distributions. DLPack: Copyright 2017 by Contributors Shared Apache-2.0 terms: LICENSE. Its interfaces are used by the native runtime. @@ -36,5 +27,4 @@ TVM-FFI-0.1.10/LICENSE.libbacktrace.txt; retain them for the optional compiled s CUDA toolkit supplied headers/libraries remain subject to their own applicable redistribution terms; this list does not replace those materials. -Actual resolved-package license inventories are collected separately and can -include conservative full-package materials outside the two native DSOs. +Source scope and resolved-package redistribution: repository THIRD_PARTY.txt. diff --git a/THIRD_PARTY.txt b/THIRD_PARTY.txt index 5afd54862..e588b1443 100644 --- a/THIRD_PARTY.txt +++ b/THIRD_PARTY.txt @@ -27,10 +27,8 @@ vLLM revision is not asserted to be the unique historical origin of that port. Preparation applies patches/marlin to build-directory copies, retaining the existing Torch-free includes, repack bindings and activation-order index behavior. CUDA template instances and their selector are generated together in the build -directory. The upstream marlin.cu native pointer entry and scalar_type.hpp are -prepared alongside the kernels; no copied host/type/helper subtree is retained -in Ops. A private declaration and small local check helper adapt the native -entry without Torch types, while existing CUDA error propagation is retained. +directory. The native marlin.cu host entry and scalar_type.hpp use a private +declaration/check adapter, retaining CUDA error propagation without Torch types. Marlin notices retain Copyright 2024 Elias Frantar, modifications by Neural Magic, Apache-2.0. The original generator identifies Apache-2.0 / Copyright contributors to the vLLM project. Marlin original Apache headers remain in the prepared source, including marlin.cu; @@ -63,13 +61,17 @@ Namespace, layout/backend binding and resource adaptations are not described as FlashAttention / CUTLASS Projects: https://github.com/Dao-AILab/flash-attention and https://github.com/NVIDIA/cutlass -Source: third_party/flash-attention, pinned gitlink -060c9188beec3a8b62b33a3bfa6d5d2d44975fab (v2.8.3). -Its csrc/cutlass gitlink is dc4817921edda44a549197ff3a9dcf5df0636e7b. +Official fixed source archives, acquired source-only by cmake/FlashAttentionSource.cmake: +FlashAttention 060c9188beec3a8b62b33a3bfa6d5d2d44975fab (v2.8.3), +SHA256 64391cfec73fc5a25932b2c0c2fb3b4d9f97af001164e5220af21049fc17e04e. +Its private CUTLASS dc4817921edda44a549197ff3a9dcf5df0636e7b, +SHA256 f2a3a9df5e6f010c8b02716aa2644a6f071827fafa606fac5f5241cab6a1ab56. scripts/prepare_flash_attention_inference.py prepares csrc/flash_attn/src forward headers with patches/flash-attention, preserving source notices, and copies upstream LICENSE and AUTHORS. CUTLASS headers are included directly -from the nested source dependency; its LICENSE.txt is copied for distribution. +from that separate private source; its LICENSE.txt is copied for distribution. +Complete sdists bundle the 17 original FA inputs and private CUTLASS include/ +plus LICENSE.txt under third_party/flash-attention{,-cutlass}-source. The adapted scheduling heuristic comes from csrc/flash_attn/flash_api.cpp and retains Copyright 2023 Tri Dao. The original FlashAttention/CUTLASS notices govern those portions; the Core MIT notice does not relicense them. Native launch templates, descriptors and workspace/lifetime integration are glue around the imported forward mathematics. @@ -87,7 +89,7 @@ retain their Huawei/InfiniTensor notices and vLLM-Ascend source links. Distribution SDK installation includes LICENSE, this attribution list and LICENSES under share/licenses/InfiniOps. Native FlashAttention installation also includes the exact -upstream LICENSE and AUTHORS and the nested CUTLASS LICENSE.txt. +upstream LICENSE and AUTHORS and the private CUTLASS LICENSE.txt. CUTLASS-BSD-3-Clause.txt records the separate CMake dependency at 087c84df83d254b5fb295a7a408f1a1d554085cf. CUTLASS notices apply separately to its compiled headers. Linked libraries are not relicensed by Ops: distributions bundling sampling.so or diff --git a/cmake/FlashAttentionSource.cmake b/cmake/FlashAttentionSource.cmake new file mode 100644 index 000000000..7609fd2cb --- /dev/null +++ b/cmake/FlashAttentionSource.cmake @@ -0,0 +1,41 @@ +# Source inputs only; neither upstream root project is configured or installed. +include(FetchContent) +set(_flash_bundled "${CMAKE_CURRENT_LIST_DIR}/../third_party/flash-attention-source") +if(NOT FETCHCONTENT_SOURCE_DIR_FLASH_ATTENTION AND IS_DIRECTORY "${_flash_bundled}") + set(FETCHCONTENT_SOURCE_DIR_FLASH_ATTENTION "${_flash_bundled}") +endif() +set(_flash_cutlass_bundled "${CMAKE_CURRENT_LIST_DIR}/../third_party/flash-attention-cutlass-source") +if(NOT FETCHCONTENT_SOURCE_DIR_FLASH_ATTENTION_CUTLASS AND IS_DIRECTORY "${_flash_cutlass_bundled}") + set(FETCHCONTENT_SOURCE_DIR_FLASH_ATTENTION_CUTLASS "${_flash_cutlass_bundled}") +endif() +cmake_policy(PUSH) +if(POLICY CMP0135) + cmake_policy(SET CMP0135 NEW) +endif() +FetchContent_Declare(flash_attention + URL "https://github.com/Dao-AILab/flash-attention/archive/060c9188beec3a8b62b33a3bfa6d5d2d44975fab.tar.gz" + URL_HASH SHA256=64391cfec73fc5a25932b2c0c2fb3b4d9f97af001164e5220af21049fc17e04e + SOURCE_SUBDIR csrc/flash_attn/src) +FetchContent_Declare(flash_attention_cutlass + URL "https://github.com/NVIDIA/cutlass/archive/dc4817921edda44a549197ff3a9dcf5df0636e7b.tar.gz" + URL_HASH SHA256=f2a3a9df5e6f010c8b02716aa2644a6f071827fafa606fac5f5241cab6a1ab56 + SOURCE_SUBDIR include) +FetchContent_MakeAvailable(flash_attention flash_attention_cutlass) +cmake_policy(POP) + +set(_flash_source_files LICENSE AUTHORS csrc/flash_attn/flash_api.cpp) +foreach(_flash_header block_info.h dropout.h flash.h flash_fwd_kernel.h + flash_fwd_launch_template.h hardware_info.h kernel_traits.h mask.h + namespace_config.h philox.cuh rotary.h softmax.h static_switch.h utils.h) + list(APPEND _flash_source_files "csrc/flash_attn/src/${_flash_header}") +endforeach() +foreach(_flash_file IN LISTS _flash_source_files) + if(NOT EXISTS "${flash_attention_SOURCE_DIR}/${_flash_file}") + message(FATAL_ERROR "Pinned FlashAttention source is missing ${_flash_file}") + endif() +endforeach() +foreach(_flash_cutlass_file LICENSE.txt include/cutlass/cutlass.h) + if(NOT EXISTS "${flash_attention_cutlass_SOURCE_DIR}/${_flash_cutlass_file}") + message(FATAL_ERROR "Pinned FlashAttention CUTLASS source is missing ${_flash_cutlass_file}") + endif() +endforeach() diff --git a/cmake/NativeFlashAttention.cmake b/cmake/NativeFlashAttention.cmake index 99148b920..1c2e2e64d 100644 --- a/cmake/NativeFlashAttention.cmake +++ b/cmake/NativeFlashAttention.cmake @@ -7,32 +7,28 @@ function(infiniops_add_native_flash target) set(_provider "${CMAKE_CURRENT_FUNCTION_LIST_DIR}/../src/native/cuda/nvidia/ops/flash_attn_with_kvcache") set(_varlen_provider "${CMAKE_CURRENT_FUNCTION_LIST_DIR}/../src/native/cuda/nvidia/ops/flash_attn_varlen_func") set(_source "${CMAKE_CURRENT_FUNCTION_LIST_DIR}/..") - set(_upstream "${_source}/third_party/flash-attention") - set(_kernel_source "${_upstream}/csrc/flash_attn/src") - set(_cutlass_include "${_upstream}/csrc/cutlass/include") + include("${CMAKE_CURRENT_FUNCTION_LIST_DIR}/FlashAttentionSource.cmake") + set(_upstream "${flash_attention_SOURCE_DIR}") + set(_cutlass_include "${flash_attention_cutlass_SOURCE_DIR}/include") set(_prepared "${CMAKE_CURRENT_BINARY_DIR}/native-flash-inference") set(_inputs "${_source}/scripts/prepare_flash_attention_inference.py" "${_source}/patches/flash-attention/inference.patch" - "${_upstream}/csrc/flash_attn/flash_api.cpp" - "${_upstream}/LICENSE" "${_upstream}/AUTHORS" - "${_upstream}/csrc/cutlass/LICENSE.txt" + "${flash_attention_cutlass_SOURCE_DIR}/LICENSE.txt" "${_cutlass_include}/cutlass/cutlass.h") - # Fixed forward closure: no recursive source scanning or copied CUTLASS tree. - foreach(_header block_info.h dropout.h flash.h flash_fwd_kernel.h - flash_fwd_launch_template.h hardware_info.h kernel_traits.h mask.h - namespace_config.h philox.cuh rotary.h softmax.h static_switch.h utils.h) - list(APPEND _inputs "${_kernel_source}/${_header}") + foreach(_flash_file IN LISTS _flash_source_files) + list(APPEND _inputs "${_upstream}/${_flash_file}") endforeach() foreach(_input IN LISTS _inputs) if(NOT EXISTS "${_input}") message(FATAL_ERROR - "Native FlashAttention is missing ${_input}. Initialize third_party/flash-attention and its CUTLASS submodule.") + "Native FlashAttention is missing ${_input}. Check the fixed source archives or FetchContent SOURCE_DIR overrides.") endif() endforeach() set_property(DIRECTORY APPEND PROPERTY CMAKE_CONFIGURE_DEPENDS ${_inputs}) execute_process(COMMAND "${Python_EXECUTABLE}" "${_source}/scripts/prepare_flash_attention_inference.py" --source-dir "${_upstream}" --output "${_prepared}" + --cutlass-source-dir "${flash_attention_cutlass_SOURCE_DIR}" --patch-executable "${_flash_patch}" RESULT_VARIABLE _status ERROR_VARIABLE _prepare_error) if(NOT _status EQUAL 0) @@ -57,6 +53,7 @@ function(infiniops_add_native_flash target) endforeach() endforeach() add_library(${target} OBJECT "${_provider}/kernel.cc" + "${_source}/src/native/cuda/nvidia/ops/flash_common.cc" "${_varlen_provider}/kernel.cc" ${_launches}) set_target_properties(${target} PROPERTIES POSITION_INDEPENDENT_CODE ON CUDA_STANDARD 17) target_compile_features(${target} PRIVATE cxx_std_17) diff --git a/docs/build-nvidia-inference.md b/docs/build-nvidia-inference.md index c537c0682..155129751 100644 --- a/docs/build-nvidia-inference.md +++ b/docs/build-nvidia-inference.md @@ -1,7 +1,7 @@ # Reproducible NVIDIA inference SDK inputs -This recipe describes the Linux x86-64, CPython 3.12, CUDA 12.8, SM80, -`_GLIBCXX_USE_CXX11_ABI=0` combination used by the Torch-free inference SDK. +This recipe uses Linux x86-64, CPython 3.12, CUDA 12.8 wheel inputs, an SM80 +target and `_GLIBCXX_USE_CXX11_ABI=0`. It is not a toolkit qualification report. Use the source revisions pinned by the consuming InfiniCore checkout, not floating branches. Other architectures/toolkits require separate validation. The Python packages below supply compiled libraries and metadata; LM does not @@ -15,34 +15,35 @@ execute their Python inference APIs. No `flashinfer-python` or Torch is required | TVM FFI | `apache_tvm_ffi-0.1.10-cp312-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl` | `73d478395a8625dd92fde7b7fd92b4719f18f480b78336e422cb66cc7985213d` | | CUTLASS | commit `087c84df83d254b5fb295a7a408f1a1d554085cf` archive | `99340f7a8c87f5295b91aecf8295df6693ad954807672c3869a7fb3d0084f709` | -Native Attention uses the fixed submodule revisions below. W4 kernels use the -fixed vLLM source archive through CMake FetchContent. They -are source inputs only; neither the Torch extension nor the vLLM Python package -is built or installed. - -| Source | Fixed commit | -|---|---| -| `third_party/flash-attention` | `060c9188beec3a8b62b33a3bfa6d5d2d44975fab` | -| Its `csrc/cutlass` submodule | `dc4817921edda44a549197ff3a9dcf5df0636e7b` | -| vLLM source archive (Marlin only) | `ffc4f08c8ee130d4ea6347c1bf31ffd4f8af28ab` | - -The vLLM archive SHA256 is -`d8f3d663347dff35ea8123d3eb570f67d2dd6c4258ee8b8a8ec2996c42ada678`. -It is a build source input, not a Git submodule or an installed package. -The upstream project is never configured. For an offline checkout build, pass -`-DFETCHCONTENT_SOURCE_DIR_VLLM_MARLIN=/absolute/path/to/pinned/extracted/vllm`; -the standard `FETCHCONTENT_FULLY_DISCONNECTED=ON` can then disable downloading. -Complete sdists bundle LICENSE and the eleven original Marlin inputs under -`third_party/marlin-source`, which CMake selects automatically without network. -This includes the native `marlin.cu` host entry and `csrc/core/scalar_type.hpp`; -Ops keeps only its private declaration and small error adaptation, not copies. -Building an sdist uses the standard scikit-build `sdist.cmake` flow and -`build/sdist` staging; an offline packager can pass the same source override. -Keep that sdist build directory unless also overriding its force-include path. -No generated Marlin CUDA files are stored in the source package. - -The other CUTLASS input above remains the existing dependency of ordinary Ops -providers. Its version is not substituted for FlashAttention's private headers. +Native Attention and W4 use official fixed source archives through CMake +FetchContent. Only source files are consumed; upstream projects and Torch/vLLM +runtimes are not built or installed. Attention's private CUTLASS is independent +of the ordinary Ops CUTLASS above. + +| Source | Official fixed archive | SHA256 | +|---|---|---| +| FlashAttention | [060c9188](https://github.com/Dao-AILab/flash-attention/archive/060c9188beec3a8b62b33a3bfa6d5d2d44975fab.tar.gz) | `64391cfec73fc5a25932b2c0c2fb3b4d9f97af001164e5220af21049fc17e04e` | +| FA-private CUTLASS | [dc481792](https://github.com/NVIDIA/cutlass/archive/dc4817921edda44a549197ff3a9dcf5df0636e7b.tar.gz) | `f2a3a9df5e6f010c8b02716aa2644a6f071827fafa606fac5f5241cab6a1ab56` | +| Marlin/vLLM | [ffc4f08c](https://github.com/vllm-project/vllm/archive/ffc4f08c8ee130d4ea6347c1bf31ffd4f8af28ab.tar.gz) | `d8f3d663347dff35ea8123d3eb570f67d2dd6c4258ee8b8a8ec2996c42ada678` | + +For an offline checkout build or source packager, supply the pinned extracted +trees with standard CMake overrides, alongside the ordinary CUTLASS override +shown in the build command below: + +```text +-DFETCHCONTENT_SOURCE_DIR_FLASH_ATTENTION=/absolute/path/to/flash-attention +-DFETCHCONTENT_SOURCE_DIR_FLASH_ATTENTION_CUTLASS=/absolute/path/to/private-cutlass +-DFETCHCONTENT_SOURCE_DIR_VLLM_MARLIN=/absolute/path/to/vllm +-DFETCHCONTENT_FULLY_DISCONNECTED=ON +``` + +Complete sdists bundle the original FA forward closure, private CUTLASS headers +and Marlin sources with their licenses under `third_party/*-source`; CMake uses +them automatically without Git or source downloads. `sdist.cmake` stages these +files in `build/sdist` without probing compilers, CUDA or installed SDKs. Keep +that staging path unless also overriding the `force-include` mappings. Creating +a complete sdist needs each source once, even if a later build disables it. +Generated CUDA files and prepared source copies are not stored in the sdist. Installed library hashes for the two wheels above: @@ -54,14 +55,6 @@ Installed library hashes for the two wheels above: The YAML version ranges describe the provider's accepted packages, not a lock for a release. A release using a different wheel must record its resolved version/hash and run the corresponding direct operator tests again. -`LICENSES/Apache-2.0-components.txt` records the release component copyrights -and maps their shared Apache terms to `LICENSE`. -`LICENSES/FlashInfer-0.6.7` and `LICENSES/TVM-FFI-0.1.10` retain the original -NOTICE files and native-runtime libbacktrace terms for these exact inputs. -If replacing either wheel, update those redistribution materials from the -actual release as well. SDK installation puts these materials, the native -FlashAttention notices and the Ops attribution list in `share/licenses/InfiniOps`; -downstream binary packages must retain that directory's contents. ## Acquire and install the fixed inputs @@ -74,10 +67,6 @@ export CORE=/absolute/path/to/InfiniCore export INFINI_ROOT="$PWD/sdk" export CUDA_HOME=/usr/local/cuda export INFINI_RT_ROOT="$INFINI_ROOT" -git -C "$CORE/submodules/InfiniOps" submodule update --init \ - third_party/flash-attention -git -C "$CORE/submodules/InfiniOps/third_party/flash-attention" \ - submodule update --init csrc/cutlass mkdir artifacts cd artifacts curl -fL -o 'flashinfer_jit_cache-0.6.7+cu128-cp39-abi3-manylinux_2_28_x86_64.whl' \ @@ -145,12 +134,12 @@ its explicit `public_api_only` entries preserve the three W4 API symbols without providers. Other consumers can omit those entries entirely. Disabled Marlin has no CUDA >=12.3/C++20 check or provider registration and cannot execute W4. Both native components require the standard `patch` command when enabled. -Their preparation scripts patch build-directory copies, leaving submodules +Their preparation scripts patch build-directory copies, leaving upstream sources unchanged. Jinja2 is used only to generate Marlin instances at build time. Source distributions include the consumed upstream files and can be built -without Git metadata or network access. Component-disabled builds do not need -these CUDA source submodules. An ordinary Git source archive does not include -submodule contents; use the complete source distribution for offline builds. +without Git metadata or source network access. A disabled native component does +not fetch its source during normal configure. An ordinary Git source archive +does not bundle external inputs; use the complete sdist or the overrides above. Build CCL separately using the Core instructions and the same ABI/prefix before building LM. Pin and record the installed NCCL version as part of that deployment. @@ -162,21 +151,12 @@ FlashInfer 0.6.7 and TVM-FFI 0.1.10 inputs into SDK license directory when bundling the native libraries. No additional download or manual copy into the LM source tree is needed for these inputs. -The fixed fallback is installed only for a linked TVM-FFI transport, not for -every linked backend. `sampling.so` is one JIT-cache artifact, not the entire -FlashInfer wheel; `libtvm_ffi.so` is the native runtime, not its Python package. -The fixed source closures described in `THIRD_PARTY.txt` exclude FlashInfer's -FA3/CUTLASS/fmt/spdlog components and TVM-FFI's separate PyTorch-derived Python -addon. Their six secondary license/notice copies are therefore not part of the -checked-in native fallback. Common Apache terms, original project NOTICEs, -DLPack attribution and libbacktrace terms remain. - -The resolver also preserves real license inventories from the distributions -actually selected. They can conservatively include full-package materials -beyond these DSOs; the final installed SDK is not claimed to be reduced to the -fixed fallback alone. For different binaries, validate source correspondence -and preserve the actual applicable materials. Complete wheel/SDK distributions -retain their own applicable notices independently of this native-only recipe. +The TVM-FFI transport installs a native-only fallback: common Apache terms, +component copyrights, original project NOTICEs, DLPack attribution and +libbacktrace terms. See `THIRD_PARTY.txt` for its source scope. The resolver also +preserves actual package license inventories, which may include broader wheel +materials. Replacement binaries and complete wheel/SDK distributions require +their own applicable notices; the fixed fallback is not a substitute. Preserve the separate FlashAttention/CUTLASS and Marlin notices described in `THIRD_PARTY.txt` when distributing those compiled providers. Do not treat this diff --git a/pyproject.toml b/pyproject.toml index ae7a52600..afd5defb9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -14,17 +14,7 @@ install-dir = "infini" [tool.scikit-build.sdist] cmake = true -exclude = ["third_party/flash-attention/**"] -force-include = { "build/sdist/marlin-source" = "third_party/marlin-source" } -include = [ - "third_party/flash-attention/LICENSE", - "third_party/flash-attention/AUTHORS", - "third_party/flash-attention/csrc/flash_attn/src/*.h", - "third_party/flash-attention/csrc/flash_attn/src/*.cuh", - "third_party/flash-attention/csrc/flash_attn/flash_api.cpp", - "third_party/flash-attention/csrc/cutlass/include/**", - "third_party/flash-attention/csrc/cutlass/LICENSE.txt", -] +force-include = { "build/sdist/marlin-source" = "third_party/marlin-source", "build/sdist/flash-attention-source" = "third_party/flash-attention-source", "build/sdist/flash-attention-cutlass-source" = "third_party/flash-attention-cutlass-source" } [[tool.scikit-build.overrides]] if.state = "sdist" diff --git a/scripts/prepare_flash_attention_inference.py b/scripts/prepare_flash_attention_inference.py index c42867692..4df934315 100644 --- a/scripts/prepare_flash_attention_inference.py +++ b/scripts/prepare_flash_attention_inference.py @@ -1,6 +1,6 @@ """Prepare the pinned FlashAttention forward headers for native inference. -Apply the reviewed adaptation only in build staging, never in the submodule. +Apply the reviewed adaptation only in build staging, never in upstream sources. The prepared parameter layout must not be linked to the upstream Torch extension. """ @@ -46,22 +46,23 @@ def read_source(path): if not path.is_file(): raise ValueError( f"missing FlashAttention source dependency: {path}; " - "initialize third_party/flash-attention and its CUTLASS submodule" + "check the fixed source archives or FetchContent SOURCE_DIR overrides" ) return path.read_text(encoding="utf-8") -def prepare(source, output, patch_executable="patch"): +def prepare(source, cutlass_source, output, patch_executable="patch"): source, output = pathlib.Path(source).resolve(), pathlib.Path(output).resolve() - if output == source or source in output.parents: + cutlass_source = pathlib.Path(cutlass_source).resolve() + if output.is_relative_to(source) or output.is_relative_to(cutlass_source): raise ValueError( - "FlashAttention preparation output must be outside its source tree" + "FlashAttention preparation output must be outside upstream source trees" ) kernel_source = source / "csrc/flash_attn/src" headers = {name: read_source(kernel_source / name) for name in HEADERS} api = read_source(source / "csrc/flash_attn/flash_api.cpp") notices = {name: read_source(source / name) for name in ("LICENSE", "AUTHORS")} - notices["CUTLASS-LICENSE.txt"] = read_source(source / "csrc/cutlass/LICENSE.txt") + notices["CUTLASS-LICENSE.txt"] = read_source(cutlass_source / "LICENSE.txt") if not PATCH.is_file(): raise ValueError(f"missing FlashAttention inference adaptation: {PATCH}") @@ -122,10 +123,13 @@ def prepare(source, output, patch_executable="patch"): if __name__ == "__main__": parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--source-dir", required=True) + parser.add_argument("--cutlass-source-dir", required=True) parser.add_argument("--output", required=True) parser.add_argument("--patch-executable", default="patch") args = parser.parse_args() try: - prepare(args.source_dir, args.output, args.patch_executable) + prepare( + args.source_dir, args.cutlass_source_dir, args.output, args.patch_executable + ) except (ValueError, OSError) as exc: parser.exit(1, f"FlashAttention preparation failed: {exc}\n") diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index e6147ea01..9d4bb2e85 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -161,6 +161,7 @@ if(WITH_NVIDIA) # Optional pinned kernels have their own include closure and CUDA TUs. list(FILTER NVIDIA_SOURCES EXCLUDE REGEX "/ops/flash_attn_with_kvcache/") list(FILTER NVIDIA_SOURCES EXCLUDE REGEX "/ops/flash_attn_varlen_func/") + list(FILTER NVIDIA_SOURCES EXCLUDE REGEX "/ops/flash_common\\.cc$") # The CUTLASS INT8 Tensor Core kernel needs SM75 device compilation. Scope # the extra code objects to this provider source so other CUDA operators diff --git a/src/base/fused_gated_delta_net_gating.h b/src/base/fused_gated_delta_net_gating.h index 14e214e3a..44d2f45dd 100644 --- a/src/base/fused_gated_delta_net_gating.h +++ b/src/base/fused_gated_delta_net_gating.h @@ -6,7 +6,6 @@ #include "operator.h" namespace infini::ops { -// Candidate destination: src/base/fused_gated_delta_net_gating.h. class FusedGatedDeltaNetGating : public Operator { public: FusedGatedDeltaNetGating(Tensor A_log, Tensor a, Tensor b, Tensor dt_bias, diff --git a/src/native/cpu/dense_view.h b/src/native/cpu/dense_view.h new file mode 100644 index 000000000..01732d334 --- /dev/null +++ b/src/native/cpu/dense_view.h @@ -0,0 +1,62 @@ +#ifndef INFINI_OPS_CPU_DENSE_VIEW_H_ +#define INFINI_OPS_CPU_DENSE_VIEW_H_ + +#include +#include +#include + +#include "tensor.h" + +namespace infini::ops::cpu_dense_view { + +// Storage checks only; callers retain their dtype, shape and device contracts. +inline bool ContiguousBytes(const Tensor t, std::size_t* bytes) { + if (t.shape().size() != t.strides().size()) return false; + const auto limit = + static_cast(std::numeric_limits::max()) / + t.element_size(); + bool empty = false; + for (std::size_t axis = 0; axis < t.ndim(); ++axis) { + if (t.stride(axis) < 0 || t.size(axis) > limit) return false; + empty = empty || t.size(axis) == 0; + } + if (empty) { + *bytes = 0; + return true; + } + std::size_t count = 1; + for (std::size_t axis = t.ndim(); axis > 0; --axis) { + const auto size = t.size(axis - 1); + if ((size > 1 && + t.stride(axis - 1) != static_cast(count)) || + size > limit / count) + return false; + count *= size; + } + *bytes = count * t.element_size(); + return true; +} + +inline bool ValidAddress(const Tensor t, std::size_t bytes) { + if (bytes == 0) return true; + const auto start = reinterpret_cast(t.data()); + return start && start % t.element_size() == 0 && + start <= std::numeric_limits::max() - bytes; +} + +inline bool DisjointOrSameContiguous(const Tensor input, Tensor out) { + std::size_t input_bytes = 0, out_bytes = 0; + if (!ContiguousBytes(input, &input_bytes) || + !ContiguousBytes(out, &out_bytes) || !ValidAddress(input, input_bytes) || + !ValidAddress(out, out_bytes)) + return false; + const auto ip = reinterpret_cast(input.data()); + const auto op = reinterpret_cast(out.data()); + const bool overlap = + input_bytes && out_bytes && ip < op + out_bytes && op < ip + input_bytes; + return !overlap || (ip == op && input.strides() == out.strides()); +} + +} // namespace infini::ops::cpu_dense_view + +#endif // INFINI_OPS_CPU_DENSE_VIEW_H_ diff --git a/src/native/cpu/ops/mul_scalar/mul_scalar.h b/src/native/cpu/ops/mul_scalar/mul_scalar.h index 49cb092a5..7b90cde7c 100644 --- a/src/native/cpu/ops/mul_scalar/mul_scalar.h +++ b/src/native/cpu/ops/mul_scalar/mul_scalar.h @@ -10,6 +10,7 @@ #include "base/mul_scalar.h" #include "native/cpu/caster_.h" +#include "native/cpu/dense_view.h" namespace infini::ops { @@ -22,23 +23,14 @@ class Operator : public MulScalar, input.device() != out.device() || input.dtype() != out.dtype() || input.shape() != out.shape() || !SupportedDtype(input.dtype())) return false; - std::size_t input_bytes = 0, out_bytes = 0; - if (!ContiguousBytes(input, &input_bytes) || - !ContiguousBytes(out, &out_bytes) || - !ValidAddress(input, input_bytes) || !ValidAddress(out, out_bytes)) - return false; - const auto ip = reinterpret_cast(input.data()); - const auto op = reinterpret_cast(out.data()); - const bool overlap = input_bytes && out_bytes && ip < op + out_bytes && - op < ip + input_bytes; - return !overlap || (ip == op && input.strides() == out.strides()); + return cpu_dense_view::DisjointOrSameContiguous(input, out); } Operator(const Tensor input, const double alpha, Tensor out) : MulScalar{Checked(input, alpha, out), alpha, out} { std::size_t bytes = 0; - ContiguousBytes(input, &bytes); - count_ = bytes / Width(dtype_); + cpu_dense_view::ContiguousBytes(input, &bytes); + count_ = bytes / input.element_size(); } void operator()(const Tensor input, const double alpha, @@ -86,44 +78,6 @@ class Operator : public MulScalar, type == DataType::kBFloat16; } - static std::size_t Width(DataType type) { - return type == DataType::kFloat32 ? 4u : 2u; - } - - static bool ContiguousBytes(const Tensor t, std::size_t* bytes) { - if (t.shape().size() != t.strides().size()) return false; - const auto limit = - static_cast(std::numeric_limits::max()) / - Width(t.dtype()); - bool empty = false; - for (std::size_t axis = 0; axis < t.ndim(); ++axis) { - if (t.stride(axis) < 0 || t.size(axis) > limit) return false; - empty = empty || t.size(axis) == 0; - } - if (empty) { - *bytes = 0; - return true; - } - std::size_t count = 1; - for (std::size_t axis = t.ndim(); axis > 0; --axis) { - const auto size = t.size(axis - 1); - if ((size > 1 && - t.stride(axis - 1) != static_cast(count)) || - size > limit / count) - return false; - count *= size; - } - *bytes = count * Width(t.dtype()); - return true; - } - - static bool ValidAddress(const Tensor t, std::size_t bytes) { - if (bytes == 0) return true; - const auto start = reinterpret_cast(t.data()); - return start && start % Width(t.dtype()) == 0 && - start <= std::numeric_limits::max() - bytes; - } - static float FloatAlpha(double alpha) { if (std::isfinite(alpha) && std::fabs(alpha) > std::numeric_limits::max()) diff --git a/src/native/cpu/ops/row_l2_normalize/row_l2_normalize.h b/src/native/cpu/ops/row_l2_normalize/row_l2_normalize.h index e9facbc26..912c538a7 100644 --- a/src/native/cpu/ops/row_l2_normalize/row_l2_normalize.h +++ b/src/native/cpu/ops/row_l2_normalize/row_l2_normalize.h @@ -10,6 +10,7 @@ #include "base/row_l2_normalize.h" #include "native/cpu/caster_.h" +#include "native/cpu/dense_view.h" namespace infini::ops { @@ -23,16 +24,7 @@ class Operator input.shape() != out.shape() || input.ndim() != 2 || !SupportedDtype(input.dtype()) || !std::isfinite(eps) || eps < 0) return false; - std::size_t input_bytes = 0, out_bytes = 0; - if (!ContiguousBytes(input, &input_bytes) || - !ContiguousBytes(out, &out_bytes) || - !ValidAddress(input, input_bytes) || !ValidAddress(out, out_bytes)) - return false; - const auto ip = reinterpret_cast(input.data()); - const auto op = reinterpret_cast(out.data()); - const bool overlap = input_bytes && out_bytes && ip < op + out_bytes && - op < ip + input_bytes; - return !overlap || (ip == op && input.strides() == out.strides()); + return cpu_dense_view::DisjointOrSameContiguous(input, out); } Operator(const Tensor input, const double eps, Tensor out) @@ -118,44 +110,6 @@ class Operator type == DataType::kBFloat16; } - static std::size_t Width(DataType type) { - return type == DataType::kFloat32 ? 4u : 2u; - } - - static bool ContiguousBytes(const Tensor t, std::size_t* bytes) { - if (t.shape().size() != t.strides().size()) return false; - const auto limit = - static_cast(std::numeric_limits::max()) / - Width(t.dtype()); - bool empty = false; - for (std::size_t axis = 0; axis < t.ndim(); ++axis) { - if (t.stride(axis) < 0 || t.size(axis) > limit) return false; - empty = empty || t.size(axis) == 0; - } - if (empty) { - *bytes = 0; - return true; - } - std::size_t count = 1; - for (std::size_t axis = t.ndim(); axis > 0; --axis) { - const auto size = t.size(axis - 1); - if ((size > 1 && - t.stride(axis - 1) != static_cast(count)) || - size > limit / count) - return false; - count *= size; - } - *bytes = count * Width(t.dtype()); - return true; - } - - static bool ValidAddress(const Tensor t, std::size_t bytes) { - if (bytes == 0) return true; - const auto start = reinterpret_cast(t.data()); - return start && start % Width(t.dtype()) == 0 && - start <= std::numeric_limits::max() - bytes; - } - static Tensor Checked(const Tensor input, const double eps, Tensor out) { if (!supports(input, eps, out)) throw std::invalid_argument( diff --git a/src/native/cuda/nvidia/ops/flash_attn_varlen_func/kernel.cc b/src/native/cuda/nvidia/ops/flash_attn_varlen_func/kernel.cc index e08efaa99..96920aae2 100644 --- a/src/native/cuda/nvidia/ops/flash_attn_varlen_func/kernel.cc +++ b/src/native/cuda/nvidia/ops/flash_attn_varlen_func/kernel.cc @@ -1,120 +1,22 @@ #include "kernel.h" -#include #include #include #include -#include #include #include -#include "flash.h" #include "native/cuda/nvidia/ops/flash_attn_with_kvcache/infini_flash_cuda.h" +#include "native/cuda/nvidia/ops/flash_common.h" #include "native/cuda/nvidia/ops/flash_head_dim.h" #include "native/cuda/nvidia/ops/flash_workspace.h" -namespace infini::ops::native_flash { -void launch_32_fp16_false(infini_flash::Flash_fwd_params&, cudaStream_t); -void launch_32_fp16_true(infini_flash::Flash_fwd_params&, cudaStream_t); -void launch_32_bf16_false(infini_flash::Flash_fwd_params&, cudaStream_t); -void launch_32_bf16_true(infini_flash::Flash_fwd_params&, cudaStream_t); -void launch_64_fp16_false(infini_flash::Flash_fwd_params&, cudaStream_t); -void launch_64_fp16_true(infini_flash::Flash_fwd_params&, cudaStream_t); -void launch_64_bf16_false(infini_flash::Flash_fwd_params&, cudaStream_t); -void launch_64_bf16_true(infini_flash::Flash_fwd_params&, cudaStream_t); -void launch_96_fp16_false(infini_flash::Flash_fwd_params&, cudaStream_t); -void launch_96_fp16_true(infini_flash::Flash_fwd_params&, cudaStream_t); -void launch_96_bf16_false(infini_flash::Flash_fwd_params&, cudaStream_t); -void launch_96_bf16_true(infini_flash::Flash_fwd_params&, cudaStream_t); -void launch_128_fp16_false(infini_flash::Flash_fwd_params&, cudaStream_t); -void launch_128_fp16_true(infini_flash::Flash_fwd_params&, cudaStream_t); -void launch_128_bf16_false(infini_flash::Flash_fwd_params&, cudaStream_t); -void launch_128_bf16_true(infini_flash::Flash_fwd_params&, cudaStream_t); -void launch_192_fp16_false(infini_flash::Flash_fwd_params&, cudaStream_t); -void launch_192_fp16_true(infini_flash::Flash_fwd_params&, cudaStream_t); -void launch_192_bf16_false(infini_flash::Flash_fwd_params&, cudaStream_t); -void launch_192_bf16_true(infini_flash::Flash_fwd_params&, cudaStream_t); -void launch_256_fp16_false(infini_flash::Flash_fwd_params&, cudaStream_t); -void launch_256_fp16_true(infini_flash::Flash_fwd_params&, cudaStream_t); -void launch_256_bf16_false(infini_flash::Flash_fwd_params&, cudaStream_t); -void launch_256_bf16_true(infini_flash::Flash_fwd_params&, cudaStream_t); -void launch_varlen_32_fp16_false(infini_flash::Flash_fwd_params&, cudaStream_t); -void launch_varlen_32_fp16_true(infini_flash::Flash_fwd_params&, cudaStream_t); -void launch_varlen_32_bf16_false(infini_flash::Flash_fwd_params&, cudaStream_t); -void launch_varlen_32_bf16_true(infini_flash::Flash_fwd_params&, cudaStream_t); -void launch_varlen_64_fp16_false(infini_flash::Flash_fwd_params&, cudaStream_t); -void launch_varlen_64_fp16_true(infini_flash::Flash_fwd_params&, cudaStream_t); -void launch_varlen_64_bf16_false(infini_flash::Flash_fwd_params&, cudaStream_t); -void launch_varlen_64_bf16_true(infini_flash::Flash_fwd_params&, cudaStream_t); -void launch_varlen_96_fp16_false(infini_flash::Flash_fwd_params&, cudaStream_t); -void launch_varlen_96_fp16_true(infini_flash::Flash_fwd_params&, cudaStream_t); -void launch_varlen_96_bf16_false(infini_flash::Flash_fwd_params&, cudaStream_t); -void launch_varlen_96_bf16_true(infini_flash::Flash_fwd_params&, cudaStream_t); -void launch_varlen_128_fp16_false(infini_flash::Flash_fwd_params&, - cudaStream_t); -void launch_varlen_128_fp16_true(infini_flash::Flash_fwd_params&, cudaStream_t); -void launch_varlen_128_bf16_false(infini_flash::Flash_fwd_params&, - cudaStream_t); -void launch_varlen_128_bf16_true(infini_flash::Flash_fwd_params&, cudaStream_t); -void launch_varlen_192_fp16_false(infini_flash::Flash_fwd_params&, - cudaStream_t); -void launch_varlen_192_fp16_true(infini_flash::Flash_fwd_params&, cudaStream_t); -void launch_varlen_192_bf16_false(infini_flash::Flash_fwd_params&, - cudaStream_t); -void launch_varlen_192_bf16_true(infini_flash::Flash_fwd_params&, cudaStream_t); -void launch_varlen_256_fp16_false(infini_flash::Flash_fwd_params&, - cudaStream_t); -void launch_varlen_256_fp16_true(infini_flash::Flash_fwd_params&, cudaStream_t); -void launch_varlen_256_bf16_false(infini_flash::Flash_fwd_params&, - cudaStream_t); -void launch_varlen_256_bf16_true(infini_flash::Flash_fwd_params&, cudaStream_t); -} // namespace infini::ops::native_flash - namespace infini::ops { namespace { -struct DeviceScope { - int previous; - explicit DeviceScope(int device) { - INFINI_FLASH_CUDA_CHECK(cudaGetDevice(&previous)); - INFINI_FLASH_CUDA_CHECK(cudaSetDevice(device)); - } - ~DeviceScope() { cudaSetDevice(previous); } -}; - -bool plain_layout(Tensor t, int dimensions) { - if (t.ndim() != dimensions || t.stride(-1) != 1) return false; - std::array axes{}; - int count = 0; - for (int axis = 0; axis < dimensions; ++axis) { - if (t.size(axis) <= 0 || - t.size(axis) > std::numeric_limits::max() - 127 || - t.stride(axis) < 0) - return false; - if (t.size(axis) != 1) axes[count++] = axis; - } - std::sort(axes.begin(), axes.begin() + count, - [&](int a, int b) { return t.stride(a) < t.stride(b); }); - std::size_t extent = 1; - for (int index = 0; index < count; ++index) { - const auto axis = axes[index]; - const auto size = static_cast(t.size(axis)); - const auto stride = static_cast(t.stride(axis)); - if (stride < extent || - (size - 1) > - (std::numeric_limits::max() - extent) / stride) - return false; - extent += (size - 1) * stride; - } - return true; -} - -bool vector_aligned(Tensor t) { - if (t.data() && reinterpret_cast(t.data()) % 16) return false; - for (int axis = 0; axis < t.ndim() - 1; ++axis) - if (t.size(axis) > 1 && t.stride(axis) % 8) return false; - return true; -} +using native_flash::DeviceScope; +using native_flash::plain_layout; +using native_flash::vector_aligned; bool int_aligned(Tensor t) { return !t.data() || @@ -198,8 +100,7 @@ bool supports_metadata(VARLEN_ARGS) { bool Operator::supports( VARLEN_ARGS) { if (!supports_metadata(VARLEN_VALUES)) return false; - // Qualification must not pick this provider on pre-Ampere NVIDIA devices. - // This is a device-attribute query only; no provider or scratch is created. + // These native kernels require SM80+. The attribute query allocates no scratch. int major = 0; return cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, q.device().index()) == cudaSuccess && @@ -337,173 +238,10 @@ void Operator::operator()( : nullptr; // Pinned flash_api.cpp forces split-KV for paged KV; non-split kernels do // not read block_table. This one-split case needs no accumulation scratch. - if (impl_->paged) { - switch (p.d_rounded) { - case 32: - if (p.is_bf16) { - if (p.is_causal) - native_flash::launch_32_bf16_true(p, stream); - else - native_flash::launch_32_bf16_false(p, stream); - } else { - if (p.is_causal) - native_flash::launch_32_fp16_true(p, stream); - else - native_flash::launch_32_fp16_false(p, stream); - } - break; - case 64: - if (p.is_bf16) { - if (p.is_causal) - native_flash::launch_64_bf16_true(p, stream); - else - native_flash::launch_64_bf16_false(p, stream); - } else { - if (p.is_causal) - native_flash::launch_64_fp16_true(p, stream); - else - native_flash::launch_64_fp16_false(p, stream); - } - break; - case 96: - if (p.is_bf16) { - if (p.is_causal) - native_flash::launch_96_bf16_true(p, stream); - else - native_flash::launch_96_bf16_false(p, stream); - } else { - if (p.is_causal) - native_flash::launch_96_fp16_true(p, stream); - else - native_flash::launch_96_fp16_false(p, stream); - } - break; - case 128: - if (p.is_bf16) { - if (p.is_causal) - native_flash::launch_128_bf16_true(p, stream); - else - native_flash::launch_128_bf16_false(p, stream); - } else { - if (p.is_causal) - native_flash::launch_128_fp16_true(p, stream); - else - native_flash::launch_128_fp16_false(p, stream); - } - break; - case 192: - if (p.is_bf16) { - if (p.is_causal) - native_flash::launch_192_bf16_true(p, stream); - else - native_flash::launch_192_bf16_false(p, stream); - } else { - if (p.is_causal) - native_flash::launch_192_fp16_true(p, stream); - else - native_flash::launch_192_fp16_false(p, stream); - } - break; - case 256: - if (p.is_bf16) { - if (p.is_causal) - native_flash::launch_256_bf16_true(p, stream); - else - native_flash::launch_256_bf16_false(p, stream); - } else { - if (p.is_causal) - native_flash::launch_256_fp16_true(p, stream); - else - native_flash::launch_256_fp16_false(p, stream); - } - break; - default: - throw std::logic_error("unqualified native Flash head dimension"); - } - } else { - switch (p.d_rounded) { - case 32: - if (p.is_bf16) { - if (p.is_causal) - native_flash::launch_varlen_32_bf16_true(p, stream); - else - native_flash::launch_varlen_32_bf16_false(p, stream); - } else { - if (p.is_causal) - native_flash::launch_varlen_32_fp16_true(p, stream); - else - native_flash::launch_varlen_32_fp16_false(p, stream); - } - break; - case 64: - if (p.is_bf16) { - if (p.is_causal) - native_flash::launch_varlen_64_bf16_true(p, stream); - else - native_flash::launch_varlen_64_bf16_false(p, stream); - } else { - if (p.is_causal) - native_flash::launch_varlen_64_fp16_true(p, stream); - else - native_flash::launch_varlen_64_fp16_false(p, stream); - } - break; - case 96: - if (p.is_bf16) { - if (p.is_causal) - native_flash::launch_varlen_96_bf16_true(p, stream); - else - native_flash::launch_varlen_96_bf16_false(p, stream); - } else { - if (p.is_causal) - native_flash::launch_varlen_96_fp16_true(p, stream); - else - native_flash::launch_varlen_96_fp16_false(p, stream); - } - break; - case 128: - if (p.is_bf16) { - if (p.is_causal) - native_flash::launch_varlen_128_bf16_true(p, stream); - else - native_flash::launch_varlen_128_bf16_false(p, stream); - } else { - if (p.is_causal) - native_flash::launch_varlen_128_fp16_true(p, stream); - else - native_flash::launch_varlen_128_fp16_false(p, stream); - } - break; - case 192: - if (p.is_bf16) { - if (p.is_causal) - native_flash::launch_varlen_192_bf16_true(p, stream); - else - native_flash::launch_varlen_192_bf16_false(p, stream); - } else { - if (p.is_causal) - native_flash::launch_varlen_192_fp16_true(p, stream); - else - native_flash::launch_varlen_192_fp16_false(p, stream); - } - break; - case 256: - if (p.is_bf16) { - if (p.is_causal) - native_flash::launch_varlen_256_bf16_true(p, stream); - else - native_flash::launch_varlen_256_bf16_false(p, stream); - } else { - if (p.is_causal) - native_flash::launch_varlen_256_fp16_true(p, stream); - else - native_flash::launch_varlen_256_fp16_false(p, stream); - } - break; - default: - throw std::logic_error("unqualified native Flash head dimension"); - } - } + if (impl_->paged) + native_flash::launch_splitkv(p, stream); + else + native_flash::launch_varlen(p, stream); scratch.finish(); } diff --git a/src/native/cuda/nvidia/ops/flash_attn_with_kvcache/kernel.cc b/src/native/cuda/nvidia/ops/flash_attn_with_kvcache/kernel.cc index 8697c3b79..7297772eb 100644 --- a/src/native/cuda/nvidia/ops/flash_attn_with_kvcache/kernel.cc +++ b/src/native/cuda/nvidia/ops/flash_attn_with_kvcache/kernel.cc @@ -1,89 +1,21 @@ #include "kernel.h" -#include #include #include -#include #include #include -#include "flash.h" #include "infini_flash_cuda.h" +#include "native/cuda/nvidia/ops/flash_common.h" #include "native/cuda/nvidia/ops/flash_head_dim.h" #include "native/cuda/nvidia/ops/flash_workspace.h" #include "num_splits.h" -namespace infini::ops::native_flash { -void launch_32_fp16_false(infini_flash::Flash_fwd_params&, cudaStream_t); -void launch_32_fp16_true(infini_flash::Flash_fwd_params&, cudaStream_t); -void launch_32_bf16_false(infini_flash::Flash_fwd_params&, cudaStream_t); -void launch_32_bf16_true(infini_flash::Flash_fwd_params&, cudaStream_t); -void launch_64_fp16_false(infini_flash::Flash_fwd_params&, cudaStream_t); -void launch_64_fp16_true(infini_flash::Flash_fwd_params&, cudaStream_t); -void launch_64_bf16_false(infini_flash::Flash_fwd_params&, cudaStream_t); -void launch_64_bf16_true(infini_flash::Flash_fwd_params&, cudaStream_t); -void launch_96_fp16_false(infini_flash::Flash_fwd_params&, cudaStream_t); -void launch_96_fp16_true(infini_flash::Flash_fwd_params&, cudaStream_t); -void launch_96_bf16_false(infini_flash::Flash_fwd_params&, cudaStream_t); -void launch_96_bf16_true(infini_flash::Flash_fwd_params&, cudaStream_t); -void launch_128_fp16_false(infini_flash::Flash_fwd_params&, cudaStream_t); -void launch_128_fp16_true(infini_flash::Flash_fwd_params&, cudaStream_t); -void launch_128_bf16_false(infini_flash::Flash_fwd_params&, cudaStream_t); -void launch_128_bf16_true(infini_flash::Flash_fwd_params&, cudaStream_t); -void launch_192_fp16_false(infini_flash::Flash_fwd_params&, cudaStream_t); -void launch_192_fp16_true(infini_flash::Flash_fwd_params&, cudaStream_t); -void launch_192_bf16_false(infini_flash::Flash_fwd_params&, cudaStream_t); -void launch_192_bf16_true(infini_flash::Flash_fwd_params&, cudaStream_t); -void launch_256_fp16_false(infini_flash::Flash_fwd_params&, cudaStream_t); -void launch_256_fp16_true(infini_flash::Flash_fwd_params&, cudaStream_t); -void launch_256_bf16_false(infini_flash::Flash_fwd_params&, cudaStream_t); -void launch_256_bf16_true(infini_flash::Flash_fwd_params&, cudaStream_t); -} // namespace infini::ops::native_flash - namespace infini::ops { namespace { -struct DeviceScope { - int previous; - explicit DeviceScope(int device) { - INFINI_FLASH_CUDA_CHECK(cudaGetDevice(&previous)); - INFINI_FLASH_CUDA_CHECK(cudaSetDevice(device)); - } - ~DeviceScope() { cudaSetDevice(previous); } -}; - -bool plain_layout(Tensor t, int dimensions) { - if (dimensions < 1 || dimensions > 4 || t.ndim() != dimensions || - t.stride(-1) != 1) - return false; - std::array axes{}; - int count = 0; - for (int axis = 0; axis < dimensions; ++axis) { - if (t.size(axis) <= 0 || - t.size(axis) > std::numeric_limits::max() - 127 || - t.stride(axis) < 0) - return false; - // A singleton axis is never indexed away from zero. LM unsqueeze views - // legitimately use stride zero here; this is not a broadcasted data axis. - if (t.size(axis) == 1) continue; - axes[count++] = axis; - } - // Logical dimension order need not match storage order. Static KV is a - // head-major view; the upstream kernel already receives its row/head strides. - std::sort(axes.begin(), axes.begin() + count, - [&](int a, int b) { return t.stride(a) < t.stride(b); }); - std::size_t extent = 1; - for (int index = 0; index < count; ++index) { - const int axis = axes[index]; - const auto size = static_cast(t.size(axis)); - const auto stride = static_cast(t.stride(axis)); - if (stride < extent || - (size - 1) > - (std::numeric_limits::max() - extent) / stride) - return false; - extent += (size - 1) * stride; - } - return true; -} +using native_flash::DeviceScope; +using native_flash::plain_layout; +using native_flash::vector_aligned; std::size_t product(std::initializer_list sizes) { std::size_t result = 1; @@ -117,12 +49,8 @@ bool Operator::supported_layout( return false; // Upstream vectorized loads/stores require 16-byte-aligned rows and heads. // Metadata-only descriptor views may have null data until invocation. - for (auto tensor : {q, k, v, out}) { - if (tensor.data() && reinterpret_cast(tensor.data()) % 16) - return false; - for (int axis = 0; axis < 3; ++axis) - if (tensor.size(axis) > 1 && tensor.stride(axis) % 8 != 0) return false; - } + for (auto tensor : {q, k, v, out}) + if (!vector_aligned(tensor)) return false; if ((lengths && !plain_layout(*lengths, 1)) || (blocks && (!plain_layout(*blocks, 2) || @@ -282,88 +210,7 @@ void Operator::execute( p.o_ptr = out.data(); p.cu_seqlens_k = lengths ? static_cast(lengths->data()) : nullptr; p.block_table = blocks ? static_cast(blocks->data()) : nullptr; - switch (p.d_rounded) { - case 32: - if (p.is_bf16) { - if (p.is_causal) - native_flash::launch_32_bf16_true(p, stream); - else - native_flash::launch_32_bf16_false(p, stream); - } else { - if (p.is_causal) - native_flash::launch_32_fp16_true(p, stream); - else - native_flash::launch_32_fp16_false(p, stream); - } - break; - case 64: - if (p.is_bf16) { - if (p.is_causal) - native_flash::launch_64_bf16_true(p, stream); - else - native_flash::launch_64_bf16_false(p, stream); - } else { - if (p.is_causal) - native_flash::launch_64_fp16_true(p, stream); - else - native_flash::launch_64_fp16_false(p, stream); - } - break; - case 96: - if (p.is_bf16) { - if (p.is_causal) - native_flash::launch_96_bf16_true(p, stream); - else - native_flash::launch_96_bf16_false(p, stream); - } else { - if (p.is_causal) - native_flash::launch_96_fp16_true(p, stream); - else - native_flash::launch_96_fp16_false(p, stream); - } - break; - case 128: - if (p.is_bf16) { - if (p.is_causal) - native_flash::launch_128_bf16_true(p, stream); - else - native_flash::launch_128_bf16_false(p, stream); - } else { - if (p.is_causal) - native_flash::launch_128_fp16_true(p, stream); - else - native_flash::launch_128_fp16_false(p, stream); - } - break; - case 192: - if (p.is_bf16) { - if (p.is_causal) - native_flash::launch_192_bf16_true(p, stream); - else - native_flash::launch_192_bf16_false(p, stream); - } else { - if (p.is_causal) - native_flash::launch_192_fp16_true(p, stream); - else - native_flash::launch_192_fp16_false(p, stream); - } - break; - case 256: - if (p.is_bf16) { - if (p.is_causal) - native_flash::launch_256_bf16_true(p, stream); - else - native_flash::launch_256_bf16_false(p, stream); - } else { - if (p.is_causal) - native_flash::launch_256_fp16_true(p, stream); - else - native_flash::launch_256_fp16_false(p, stream); - } - break; - default: - throw std::logic_error("unqualified native Flash head dimension"); - } + native_flash::launch_splitkv(p, stream); scratch.finish(); } } // namespace infini::ops diff --git a/src/native/cuda/nvidia/ops/flash_common.cc b/src/native/cuda/nvidia/ops/flash_common.cc new file mode 100644 index 000000000..50c558e4e --- /dev/null +++ b/src/native/cuda/nvidia/ops/flash_common.cc @@ -0,0 +1,124 @@ +#include "flash_common.h" + +#include +#include +#include +#include +#include + +#include "native/cuda/nvidia/ops/flash_attn_with_kvcache/infini_flash_cuda.h" + +namespace infini::ops::native_flash { + +DeviceScope::DeviceScope(int device) { + INFINI_FLASH_CUDA_CHECK(cudaGetDevice(&previous)); + INFINI_FLASH_CUDA_CHECK(cudaSetDevice(device)); +} + +DeviceScope::~DeviceScope() { cudaSetDevice(previous); } + +bool plain_layout(Tensor t, int dimensions) { + if (dimensions < 1 || dimensions > 4 || t.ndim() != dimensions || + t.stride(-1) != 1) + return false; + std::array axes{}; + int count = 0; + for (int axis = 0; axis < dimensions; ++axis) { + if (t.size(axis) <= 0 || + t.size(axis) > std::numeric_limits::max() - 127 || + t.stride(axis) < 0) + return false; + // A singleton is never indexed away from zero, including LM unsqueeze + // views with stride zero. Other zero-stride axes must still be rejected. + if (t.size(axis) != 1) axes[count++] = axis; + } + std::sort(axes.begin(), axes.begin() + count, + [&](int a, int b) { return t.stride(a) < t.stride(b); }); + std::size_t extent = 1; + for (int index = 0; index < count; ++index) { + const auto axis = axes[index]; + const auto size = static_cast(t.size(axis)); + const auto stride = static_cast(t.stride(axis)); + if (stride < extent || + (size - 1) > + (std::numeric_limits::max() - extent) / stride) + return false; + extent += (size - 1) * stride; + } + return true; +} + +bool vector_aligned(Tensor t) { + if (t.data() && reinterpret_cast(t.data()) % 16) return false; + for (int axis = 0; axis < t.ndim() - 1; ++axis) + if (t.size(axis) > 1 && t.stride(axis) % 8) return false; + return true; +} + +#define DECLARE_FLASH_LAUNCHES(prefix, dimension) \ + void prefix##_##dimension##_fp16_false(infini_flash::Flash_fwd_params&, \ + cudaStream_t); \ + void prefix##_##dimension##_fp16_true(infini_flash::Flash_fwd_params&, \ + cudaStream_t); \ + void prefix##_##dimension##_bf16_false(infini_flash::Flash_fwd_params&, \ + cudaStream_t); \ + void prefix##_##dimension##_bf16_true(infini_flash::Flash_fwd_params&, \ + cudaStream_t) + +DECLARE_FLASH_LAUNCHES(launch, 32); +DECLARE_FLASH_LAUNCHES(launch, 64); +DECLARE_FLASH_LAUNCHES(launch, 96); +DECLARE_FLASH_LAUNCHES(launch, 128); +DECLARE_FLASH_LAUNCHES(launch, 192); +DECLARE_FLASH_LAUNCHES(launch, 256); +DECLARE_FLASH_LAUNCHES(launch_varlen, 32); +DECLARE_FLASH_LAUNCHES(launch_varlen, 64); +DECLARE_FLASH_LAUNCHES(launch_varlen, 96); +DECLARE_FLASH_LAUNCHES(launch_varlen, 128); +DECLARE_FLASH_LAUNCHES(launch_varlen, 192); +DECLARE_FLASH_LAUNCHES(launch_varlen, 256); +#undef DECLARE_FLASH_LAUNCHES + +#define FLASH_LAUNCH_CASE(prefix, dimension) \ + case dimension: \ + if (p.is_bf16) { \ + if (p.is_causal) \ + prefix##_##dimension##_bf16_true(p, stream); \ + else \ + prefix##_##dimension##_bf16_false(p, stream); \ + } else { \ + if (p.is_causal) \ + prefix##_##dimension##_fp16_true(p, stream); \ + else \ + prefix##_##dimension##_fp16_false(p, stream); \ + } \ + break + +void launch_splitkv(infini_flash::Flash_fwd_params& p, cudaStream_t stream) { + switch (p.d_rounded) { + FLASH_LAUNCH_CASE(launch, 32); + FLASH_LAUNCH_CASE(launch, 64); + FLASH_LAUNCH_CASE(launch, 96); + FLASH_LAUNCH_CASE(launch, 128); + FLASH_LAUNCH_CASE(launch, 192); + FLASH_LAUNCH_CASE(launch, 256); + default: + throw std::logic_error("unsupported native Flash head dimension"); + } +} + +void launch_varlen(infini_flash::Flash_fwd_params& p, cudaStream_t stream) { + switch (p.d_rounded) { + FLASH_LAUNCH_CASE(launch_varlen, 32); + FLASH_LAUNCH_CASE(launch_varlen, 64); + FLASH_LAUNCH_CASE(launch_varlen, 96); + FLASH_LAUNCH_CASE(launch_varlen, 128); + FLASH_LAUNCH_CASE(launch_varlen, 192); + FLASH_LAUNCH_CASE(launch_varlen, 256); + default: + throw std::logic_error("unsupported native Flash head dimension"); + } +} +#undef FLASH_LAUNCH_CASE + +} // namespace infini::ops::native_flash diff --git a/src/native/cuda/nvidia/ops/flash_common.h b/src/native/cuda/nvidia/ops/flash_common.h new file mode 100644 index 000000000..7fdbb4fd9 --- /dev/null +++ b/src/native/cuda/nvidia/ops/flash_common.h @@ -0,0 +1,26 @@ +#pragma once + +#include + +#include "flash.h" +#include "tensor.h" + +namespace infini::ops::native_flash { + +struct DeviceScope { + int previous; + explicit DeviceScope(int device); + ~DeviceScope(); +}; + +// Non-overlapping strided rows are valid regardless of logical axis order. +// Singleton axes may have zero stride; the innermost element stride is one. +bool plain_layout(Tensor t, int dimensions); +bool vector_aligned(Tensor t); + +// Fixed FA specializations only: paged varlen and KV-cache share split-KV; +// packed, nonpaged varlen must retain the separate forward kernel family. +void launch_splitkv(infini_flash::Flash_fwd_params& p, cudaStream_t stream); +void launch_varlen(infini_flash::Flash_fwd_params& p, cudaStream_t stream); + +} // namespace infini::ops::native_flash diff --git a/src/native/cuda/nvidia/ops/recurrent_gated_delta_rule/recurrent_delta_rule_common.cuh b/src/native/cuda/nvidia/ops/recurrent_gated_delta_rule/recurrent_delta_rule_common.cuh index 37db9ef83..ff9e5bdfa 100644 --- a/src/native/cuda/nvidia/ops/recurrent_gated_delta_rule/recurrent_delta_rule_common.cuh +++ b/src/native/cuda/nvidia/ops/recurrent_gated_delta_rule/recurrent_delta_rule_common.cuh @@ -8,7 +8,7 @@ #include #include -// NVIDIA-only candidate: retain the original NVIDIA warp width. +// Retain the original NVIDIA warp width. #define INFINIOP_RECURRENT_DELTA_RULE_WARP_SIZE 32 #define INFINIOP_RECURRENT_DELTA_RULE_FULL_MASK 0xffffffffU @@ -67,15 +67,6 @@ __device__ inline Tcompute blockReduceSum(Tcompute value, Tcompute* scratch) { return result; } -__device__ inline float sigmoid(float value) { - if (value >= 0.0f) { - const float exp_neg = expf(-value); - return 1.0f / (1.0f + exp_neg); - } - const float exp_pos = expf(value); - return exp_pos / (1.0f + exp_pos); -} - // GatePolicy prepares one decay value per key dimension and a scalar beta for // the current token. This keeps the state-update implementation shared while // allowing scalar GDR and vector-decay KDA to retain their own gate semantics. diff --git a/third_party/flash-attention b/third_party/flash-attention deleted file mode 160000 index 060c9188b..000000000 --- a/third_party/flash-attention +++ /dev/null @@ -1 +0,0 @@ -Subproject commit 060c9188beec3a8b62b33a3bfa6d5d2d44975fab From c86b4f229e5a61331e281964df8b42f100fdd178 Mon Sep 17 00:00:00 2001 From: whjthu Date: Tue, 20 Oct 2026 04:03:13 +0800 Subject: [PATCH 22/25] refactor: simplify native operators and public SDK dispatch --- CMakeLists.txt | 13 +- cmake/NativeFlashAttention.cmake | 12 +- cmake/NativeW4Marlin.cmake | 6 +- patches/flash-attention/inference.patch | 50 ----- patches/marlin/w4a16.patch | 176 +----------------- scripts/generate_public_operator_header.py | 21 ++- scripts/generate_wrappers.py | 39 ++++ scripts/resolve_linked_ops.py | 9 +- src/CMakeLists.txt | 8 + src/base/causal_conv1d.h | 39 ++-- src/base/chunk_gated_delta_rule.h | 51 ++--- src/base/fused_gated_delta_net_gating.h | 15 +- src/base/gemm.h | 25 --- src/base/layer_norm.h | 23 +-- src/base/mul_scalar.h | 12 +- src/base/pad_constant.h | 12 +- src/base/quick_gelu.h | 13 +- src/base/recurrent_gated_delta_rule.h | 48 ++--- src/base/row_l2_normalize.h | 12 +- src/native/cambricon/ops/argmax/kernel.h | 5 +- src/native/cambricon/ops/tril/kernel.h | 4 +- src/native/cambricon/ops/triu/kernel.h | 4 +- src/native/cpu/ops/fill/fill.h | 4 +- src/native/cpu/ops/mul_scalar/mul_scalar.h | 18 +- .../ops/row_l2_normalize/row_l2_normalize.h | 22 ++- src/native/cuda/nvidia/ops/argmax/kernel.h | 1 - .../cuda/nvidia/ops/causal_conv1d/kernel.h | 19 +- .../chunk_gated_delta_rule_legacy_kernel.cuh | 2 +- .../ops/chunk_gated_delta_rule/kernel.h | 41 ++-- .../cuda/nvidia/ops/deepseek_moe/kernel.h | 1 - .../infini_flash_cuda.h | 4 + .../ops/fused_gated_delta_net_gating/kernel.h | 12 +- src/native/cuda/nvidia/ops/gemm/cublas.h | 9 - src/native/cuda/nvidia/ops/gemm/cublaslt.h | 9 - .../cuda/nvidia/ops/layer_norm/kernel.h | 19 +- .../cuda/nvidia/ops/moe_fused_dense/kernel.h | 1 - .../cuda/nvidia/ops/moe_fused_gate/kernel.cu | 3 +- .../cuda/nvidia/ops/moe_fused_gate/kernel.h | 3 +- .../cuda/nvidia/ops/mul_scalar/kernel.h | 11 +- .../cuda/nvidia/ops/pad_constant/kernel.h | 30 ++- .../cuda/nvidia/ops/per_tensor_i8/kernel.h | 62 ++++-- .../cuda/nvidia/ops/quick_gelu/kernel.h | 52 ++++-- .../ops/recurrent_gated_delta_rule/kernel.h | 11 +- .../recurrent_delta_rule_common.cuh | 58 ++---- ...current_gated_delta_rule_legacy_kernel.cuh | 17 +- src/native/cuda/nvidia/ops/sum/kernel.h | 6 +- src/native/cuda/nvidia/ops/take/kernel.h | 1 + .../nvidia/ops/upsample_bilinear2d/kernel.h | 1 + src/native/cuda/nvidia/ops/w4_marlin/kernel.h | 1 - src/native/cuda/nvidia/ops/w4_marlin/launch.h | 2 +- .../cuda/nvidia/ops/w4_marlin/marlin_check.h | 3 + src/operator.h | 59 ++---- tests/test_generate_public_operator_header.py | 4 +- tests/test_generate_wrappers.py | 28 +++ tests/test_resolve_linked_ops.py | 3 + 55 files changed, 450 insertions(+), 664 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index baa5ca2b2..02c416958 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -47,12 +47,6 @@ option(WITH_LINKED "Enable linked third-party operator backends" OFF) option(BUILD_NVIDIA_FLASH_INFERENCE "Build pinned Torch-free NVIDIA inference Attention" OFF) option(BUILD_NVIDIA_W4_MARLIN "Build NVIDIA AWQ/GPTQ Marlin providers (CUDA >= 12.3)" OFF) -if(BUILD_NVIDIA_FLASH_INFERENCE AND NOT WITH_NVIDIA) - message(FATAL_ERROR "Native FlashAttention requires WITH_NVIDIA") -endif() -if(BUILD_NVIDIA_W4_MARLIN AND NOT WITH_NVIDIA) - message(FATAL_ERROR "Marlin requires WITH_NVIDIA") -endif() option(WITH_NINETOOTHED "Enable NineToothed-generated kernels" OFF) @@ -216,6 +210,13 @@ if(AUTO_DETECT_DEVICES) endif() endif() +if(BUILD_NVIDIA_FLASH_INFERENCE AND NOT WITH_NVIDIA) + message(FATAL_ERROR "Native FlashAttention requires WITH_NVIDIA") +endif() +if(BUILD_NVIDIA_W4_MARLIN AND NOT WITH_NVIDIA) + message(FATAL_ERROR "Marlin requires WITH_NVIDIA") +endif() + if(AUTO_DETECT_BACKENDS) message(STATUS "Auto-detecting available backends...") diff --git a/cmake/NativeFlashAttention.cmake b/cmake/NativeFlashAttention.cmake index 1c2e2e64d..5c013c6a9 100644 --- a/cmake/NativeFlashAttention.cmake +++ b/cmake/NativeFlashAttention.cmake @@ -12,18 +12,18 @@ function(infiniops_add_native_flash target) set(_cutlass_include "${flash_attention_cutlass_SOURCE_DIR}/include") set(_prepared "${CMAKE_CURRENT_BINARY_DIR}/native-flash-inference") set(_inputs "${_source}/scripts/prepare_flash_attention_inference.py" - "${_source}/patches/flash-attention/inference.patch" - "${flash_attention_cutlass_SOURCE_DIR}/LICENSE.txt" - "${_cutlass_include}/cutlass/cutlass.h") - foreach(_flash_file IN LISTS _flash_source_files) - list(APPEND _inputs "${_upstream}/${_flash_file}") - endforeach() + "${_source}/patches/flash-attention/inference.patch") foreach(_input IN LISTS _inputs) if(NOT EXISTS "${_input}") message(FATAL_ERROR "Native FlashAttention is missing ${_input}. Check the fixed source archives or FetchContent SOURCE_DIR overrides.") endif() endforeach() + list(APPEND _inputs "${flash_attention_cutlass_SOURCE_DIR}/LICENSE.txt" + "${_cutlass_include}/cutlass/cutlass.h") + foreach(_flash_file IN LISTS _flash_source_files) + list(APPEND _inputs "${_upstream}/${_flash_file}") + endforeach() set_property(DIRECTORY APPEND PROPERTY CMAKE_CONFIGURE_DEPENDS ${_inputs}) execute_process(COMMAND "${Python_EXECUTABLE}" "${_source}/scripts/prepare_flash_attention_inference.py" diff --git a/cmake/NativeW4Marlin.cmake b/cmake/NativeW4Marlin.cmake index e7cc9ed57..9fdacb87c 100644 --- a/cmake/NativeW4Marlin.cmake +++ b/cmake/NativeW4Marlin.cmake @@ -10,15 +10,15 @@ set(_w4_prepared "${CMAKE_CURRENT_BINARY_DIR}/native-w4-marlin") set(_w4_inputs "${PROJECT_SOURCE_DIR}/scripts/prepare_marlin.py" "${PROJECT_SOURCE_DIR}/patches/marlin/w4a16.patch") -foreach(_w4_input IN LISTS _marlin_source_files) - list(APPEND _w4_inputs "${_w4_upstream}/${_w4_input}") -endforeach() foreach(_w4_input IN LISTS _w4_inputs) if(NOT EXISTS "${_w4_input}") message(FATAL_ERROR "Native Marlin is missing ${_w4_input}. Check the pinned source archive or FETCHCONTENT_SOURCE_DIR_VLLM_MARLIN.") endif() endforeach() +foreach(_w4_input IN LISTS _marlin_source_files) + list(APPEND _w4_inputs "${_w4_upstream}/${_w4_input}") +endforeach() set_property(DIRECTORY APPEND PROPERTY CMAKE_CONFIGURE_DEPENDS ${_w4_inputs}) execute_process( COMMAND "${Python_EXECUTABLE}" "${PROJECT_SOURCE_DIR}/scripts/prepare_marlin.py" diff --git a/patches/flash-attention/inference.patch b/patches/flash-attention/inference.patch index 8e25d67d3..9372b837f 100644 --- a/patches/flash-attention/inference.patch +++ b/patches/flash-attention/inference.patch @@ -70,56 +70,6 @@ #include "static_switch.h" #include "hardware_info.h" -@@ -81,7 +84,7 @@ - // printf("IsEvenMNConst = %d, IsEvenKConst = %d, Is_local = %d, Is_causal = %d, ReturnSoftmaxConst = %d, Is_dropout = %d\n", int(IsEvenMNConst), int(IsEvenKConst), int(Is_local), int(Is_causal), int(ReturnSoftmaxConst), int(Is_dropout)); - // auto kernel = &flash_fwd_kernel; - if (smem_size >= 48 * 1024) { -- C10_CUDA_CHECK(cudaFuncSetAttribute( -+ INFINI_FLASH_CUDA_CHECK(cudaFuncSetAttribute( - kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); - } - // int ctas_per_sm; -@@ -89,7 +92,7 @@ - // &ctas_per_sm, kernel, Kernel_traits::kNThreads, smem_size); - // printf("smem_size = %d, CTAs per SM = %d\n", int(smem_size), ctas_per_sm); - kernel<<>>(params); -- C10_CUDA_KERNEL_LAUNCH_CHECK(); -+ INFINI_FLASH_KERNEL_CHECK(); - }); - }); - }); -@@ -121,11 +124,11 @@ - // auto kernel = &flash_fwd_splitkv_kernel; - // auto kernel = &flash_fwd_splitkv_kernel; - if (smem_size >= 48 * 1024) { -- C10_CUDA_CHECK(cudaFuncSetAttribute( -+ INFINI_FLASH_CUDA_CHECK(cudaFuncSetAttribute( - kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); - } - kernel<<>>(params); -- C10_CUDA_KERNEL_LAUNCH_CHECK(); -+ INFINI_FLASH_KERNEL_CHECK(); - }); - }); - }); -@@ -155,7 +158,7 @@ - } else if (params.num_splits <= 128) { - flash_fwd_splitkv_combine_kernel<<>>(params); - } -- C10_CUDA_KERNEL_LAUNCH_CHECK(); -+ INFINI_FLASH_KERNEL_CHECK(); - }); - } - } -@@ -284,7 +287,7 @@ - status_ = cudaDeviceGetAttribute( - &max_smem_per_block, cudaDevAttrMaxSharedMemoryPerBlockOptin, device); - if (status_ != cudaSuccess) { -- C10_CUDA_CHECK(status_); -+ INFINI_FLASH_CUDA_CHECK(status_); - } - // printf("max_smem_per_sm = %d, max_smem_per_block = %d\n", max_smem_per_sm, max_smem_per_block); - DROPOUT_SWITCH(params.p_dropout < 1.f, Is_dropout, [&] { --- a/hardware_info.h +++ b/hardware_info.h @@ -1,3 +1,6 @@ diff --git a/patches/marlin/w4a16.patch b/patches/marlin/w4a16.patch index e3bfcd0b3..4a1d73d7b 100644 --- a/patches/marlin/w4a16.patch +++ b/patches/marlin/w4a16.patch @@ -209,117 +209,19 @@ #else // For a given "a" of size [M,K] performs a permutation of the K columns based -@@ -334,18 +341,18 @@ - int thread_n_init, int sms, bool use_atomic_add, - bool use_fp32_reduce, bool is_zp_float) { - bool is_a_8bit = a_type.size_bits() == 8; -- STD_TORCH_CHECK(prob_m > 0 && prob_n > 0 && prob_k > 0, "Invalid MNK = [", -+ INFINI_MARLIN_CHECK(prob_m > 0 && prob_n > 0 && prob_k > 0, "Invalid MNK = [", - prob_m, ", ", prob_n, ", ", prob_k, "]"); - - int group_blocks = 0; - if (has_act_order) { - if (is_k_full) { -- STD_TORCH_CHECK(group_size != -1); -+ INFINI_MARLIN_CHECK(group_size != -1); - group_blocks = group_size / 16; -- STD_TORCH_CHECK(prob_k % group_blocks == 0, "prob_k = ", prob_k, -+ INFINI_MARLIN_CHECK(prob_k % group_blocks == 0, "prob_k = ", prob_k, - " is not divisible by group_blocks = ", group_blocks); - } else { -- STD_TORCH_CHECK(group_size == 0); -+ INFINI_MARLIN_CHECK(group_size == 0); - group_blocks = 0; - } - } else { -@@ -353,7 +360,7 @@ - group_blocks = -1; - } else { - group_blocks = group_size / 16; -- STD_TORCH_CHECK(prob_k % group_blocks == 0, "prob_k = ", prob_k, -+ INFINI_MARLIN_CHECK(prob_k % group_blocks == 0, "prob_k = ", prob_k, - " is not divisible by group_blocks = ", group_blocks); - } - } -@@ -395,25 +402,25 @@ - int max_shared_mem = 0; - cudaDeviceGetAttribute(&max_shared_mem, - cudaDevAttrMaxSharedMemoryPerBlockOptin, dev); -- STD_TORCH_CHECK(max_shared_mem > 0); -+ INFINI_MARLIN_CHECK(max_shared_mem > 0); - - int major_capability, minor_capability; - cudaDeviceGetAttribute(&major_capability, cudaDevAttrComputeCapabilityMajor, - dev); - cudaDeviceGetAttribute(&minor_capability, cudaDevAttrComputeCapabilityMinor, - dev); -- STD_TORCH_CHECK(major_capability * 10 + minor_capability >= 75, -+ INFINI_MARLIN_CHECK(major_capability * 10 + minor_capability >= 75, - "marlin kernel only support Turing or newer GPUs."); - int stages = 4; - if (major_capability == 7 && minor_capability == 5) { - stages = 2; -- STD_TORCH_CHECK(a_type == vllm::kFloat16 || a_type == vllm::kS8, -+ INFINI_MARLIN_CHECK(a_type == vllm::kFloat16 || a_type == vllm::kS8, - "Turing only support FP16 or INT8 activation."); - } - if (a_type == vllm::kFE4M3fn) { -- STD_TORCH_CHECK(major_capability * 10 + minor_capability >= 89, -+ INFINI_MARLIN_CHECK(major_capability * 10 + minor_capability >= 89, - "FP8 only support Ada Lovelace or newer GPUs."); -- STD_TORCH_CHECK( -+ INFINI_MARLIN_CHECK( - major_capability * 10 + minor_capability == 89 || - major_capability == 12, - "Marlin W4A8-FP8 only support SM89 or SM12x device (It is slower than " -@@ -443,9 +450,9 @@ - if (thread_k != -1 && thread_n != -1) { - thread_tfg = thread_config_t{thread_k, thread_n, default_threads}; - exec_cfg = exec_config_t{1, thread_tfg}; -- STD_TORCH_CHECK(prob_n % thread_n == 0, "prob_n = ", prob_n, -+ INFINI_MARLIN_CHECK(prob_n % thread_n == 0, "prob_n = ", prob_n, - " is not divisible by thread_n = ", thread_n); -- STD_TORCH_CHECK(prob_k % thread_k == 0, "prob_k = ", prob_k, -+ INFINI_MARLIN_CHECK(prob_k % thread_k == 0, "prob_k = ", prob_k, - " is not divisible by thread_k = ", thread_k); - } else { - // Auto config -@@ -485,7 +492,7 @@ - int thread_k_blocks = thread_k / 16; - int thread_n_blocks = thread_n / 16; - -- STD_TORCH_CHECK( -+ INFINI_MARLIN_CHECK( - is_valid_config(thread_tfg, thread_m_blocks, prob_m_split, prob_n, - prob_k, num_bits, group_size, has_act_order, is_k_full, - has_zp, is_zp_float, is_a_8bit, stages, -@@ -506,7 +513,7 @@ - num_threads, is_zp_float, stages); - - if (kernel == MarlinDefault) { -- STD_TORCH_CHECK( -+ INFINI_MARLIN_CHECK( - false, "Unsupported shapes: MNK = [", prob_m, ", ", prob_n, ", ", - prob_k, "]", ", has_act_order = ", has_act_order, - ", num_groups = ", num_groups, ", group_size = ", group_size, -@@ -517,8 +524,14 @@ +@@ -517,8 +524,9 @@ ", num_threads = ", num_threads, ", num_bits = ", num_bits); } - cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, - max_shared_mem_new); -+ cudaStreamCaptureStatus capture_status = cudaStreamCaptureStatusNone; + ::infini::ops::native_w4::check( -+ cudaStreamIsCapturing(stream, &capture_status), "cudaStreamIsCapturing"); -+ if (capture_status == cudaStreamCaptureStatusNone) { -+ ::infini::ops::native_w4::check( -+ cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, -+ max_shared_mem_new), "cudaFuncSetAttribute"); -+ } ++ cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, ++ max_shared_mem_new), "cudaFuncSetAttribute"); bool part_use_atomic_add = use_atomic_add && div_ceil(prob_m_split, 64) * prob_n <= 2048; -@@ -542,6 +555,7 @@ +@@ -542,6 +550,7 @@ } // namespace marlin @@ -327,7 +229,7 @@ torch::stable::Tensor marlin_gemm( torch::stable::Tensor& a, std::optional c_or_none, torch::stable::Tensor& b_q_weight, -@@ -893,8 +907,12 @@ +@@ -893,8 +902,12 @@ return c; } @@ -356,71 +258,3 @@ +#include "marlin_check.h" namespace vllm { - -@@ -51,7 +51,7 @@ - // IEEE 754 compliant floating point type - static constexpr ScalarType float_IEEE754(uint8_t exponent, - uint8_t mantissa) { -- STD_TORCH_CHECK(mantissa > 0 && exponent > 0); -+ INFINI_MARLIN_CHECK(mantissa > 0 && exponent > 0); - return ScalarType(exponent, mantissa, true, 0, false, NAN_IEEE_754); - } - -@@ -59,9 +59,9 @@ - static constexpr ScalarType float_(uint8_t exponent, uint8_t mantissa, - bool finite_values_only, - NanRepr nan_repr) { -- STD_TORCH_CHECK(nan_repr < NAN_REPR_ID_MAX, "Invalid NanRepr"); -- STD_TORCH_CHECK(mantissa > 0 && exponent > 0); -- STD_TORCH_CHECK( -+ INFINI_MARLIN_CHECK(nan_repr < NAN_REPR_ID_MAX, "Invalid NanRepr"); -+ INFINI_MARLIN_CHECK(mantissa > 0 && exponent > 0); -+ INFINI_MARLIN_CHECK( - nan_repr != NAN_IEEE_754, - "use `float_IEEE754` constructor for floating point types that " - "follow IEEE 754 conventions"); -@@ -183,7 +183,7 @@ - - private: - double _floating_point_max() const { -- STD_TORCH_CHECK(mantissa <= 52 && exponent <= 11, -+ INFINI_MARLIN_CHECK(mantissa <= 52 && exponent <= 11, - "Cannot represent max/min as a double for type ", str()); - - uint64_t max_mantissa = (uint64_t(1) << mantissa) - 1; -@@ -193,7 +193,7 @@ - - uint64_t max_exponent = (uint64_t(1) << exponent) - 2; - if (nan_repr == NAN_EXTD_RANGE_MAX_MIN || nan_repr == NAN_NONE) { -- STD_TORCH_CHECK(exponent < 11, -+ INFINI_MARLIN_CHECK(exponent < 11, - "Cannot represent max/min as a double for type ", str()); - max_exponent += 1; - } -@@ -223,7 +223,7 @@ - if (is_floating_point()) { - return {_floating_point_max()}; - } else { -- STD_TORCH_CHECK(size_bits() < 64 || size_bits() == 64 && is_signed(), -+ INFINI_MARLIN_CHECK(size_bits() < 64 || size_bits() == 64 && is_signed(), - "Cannot represent max as a int64_t"); - return {(int64_t(1) << mantissa) - 1}; - } -@@ -231,7 +231,7 @@ - - constexpr std::variant _raw_min() const { - if (is_floating_point()) { -- STD_TORCH_CHECK( -+ INFINI_MARLIN_CHECK( - is_signed(), - "We currently assume all floating point types are signed"); - constexpr uint64_t sign_bit_double = (uint64_t(1) << 63); -@@ -241,7 +241,7 @@ - uint64_t min_raw = max_raw | sign_bit_double; - return {*reinterpret_cast(&min_raw)}; - } else { -- STD_TORCH_CHECK(!is_signed() || size_bits() <= 64, -+ INFINI_MARLIN_CHECK(!is_signed() || size_bits() <= 64, - "Cannot represent min as a int64_t"); - if (is_signed()) { - // set the top bit to 1 (i.e. INT64_MIN) and the rest to 0 diff --git a/scripts/generate_public_operator_header.py b/scripts/generate_public_operator_header.py index 9dad53155..9c2f724b6 100644 --- a/scripts/generate_public_operator_header.py +++ b/scripts/generate_public_operator_header.py @@ -206,7 +206,9 @@ def _find_matching_brace(text: str, open_brace: int) -> int: raise RuntimeError("Could not find the end of `Operator::Call`.") -def _hide_definition(text: str, signature: re.Pattern[str], name: str) -> str: +def _hide_definition( + text: str, signature: re.Pattern[str], name: str, replacement: str = ";" +) -> str: match = signature.search(text) if match is None: @@ -223,15 +225,26 @@ def _hide_definition(text: str, signature: re.Pattern[str], name: str) -> str: body_end = _find_matching_brace(text, body_start) signature = text[match.start() : match.end()].rstrip() - return f"{text[: match.start()]}{signature};{text[body_end + 1 :]}" + return f"{text[: match.start()]}{signature}{replacement}{text[body_end + 1 :]}" def generate_public_operator_header(source: Path, output: Path) -> None: public_text = _hide_definition( - source.read_text(), _CALL_SIGNATURE_RE, "Operator::Call" + source.read_text(), + _CALL_SIGNATURE_RE, + "Operator::Call", + " {\n return detail::OperatorFactory::Call(handle, config, args...);\n }", ) for name, signature in _PROVIDER_SIGNATURES: - public_text = _hide_definition(public_text, signature, name) + replacement = ";" + if name.startswith("Operator::Make(") or name.startswith("Operator::CanMake("): + method = "CanMake" if "CanMake" in name else "Make" + first = "tensors" if "vector" in name else "tensor" + replacement = ( + " {\n return detail::OperatorFactory::" + f"{method}(config, {first}, args...);\n }}" + ) + public_text = _hide_definition(public_text, signature, name, replacement) public_text = public_text.replace('#include "host_range_profiler.h"\n', "") output.parent.mkdir(parents=True, exist_ok=True) output.write_text(public_text) diff --git a/scripts/generate_wrappers.py b/scripts/generate_wrappers.py index 1d73dcc68..c2984b09d 100644 --- a/scripts/generate_wrappers.py +++ b/scripts/generate_wrappers.py @@ -1426,6 +1426,7 @@ def _generate_generated_dispatch_source(impl_paths, definitions): impl_includes = "\n".join(f'#include "{impl_path}"' for impl_path in impl_paths) return f"""#include "generated_dispatch.h" +#include "infini/operator_call_instantiations.h" #include "host_range_profiler.h" #include @@ -1604,6 +1605,8 @@ def _append_optional_params(prefix, params): op_type = _op_cpp_type(operator.name) declarations = [] definitions = [] + factory_declarations = [] + factory_definitions = [] seen_entries = set() @@ -1687,6 +1690,42 @@ def _append_unique(declaration, definition): f"template void {operator_instantiation};", ) + factory_declarations.append(f" static void Call({function_params});") + factory_definitions.append( + f"void detail::OperatorFactory<{op_type}>::Call({function_params}) {{\n" + f" return Operator<{op_type}>::Call(handle, config, " + f"{', '.join(arg.spelling for arg in args)});\n}}" + ) + + for constructor in operator.constructors or operator.calls: + args = [arg for arg in constructor.get_arguments() if arg.spelling != "stream"] + params = ", ".join( + f"const {_normalized_type(arg)}& {arg.spelling}" for arg in args + ) + formal_params = _append_optional_params("const Config& config", params) + arguments = ", ".join(arg.spelling for arg in args) + factory_declarations.extend( + [ + f" static std::unique_ptr> Make({formal_params});", + f" static bool CanMake({formal_params});", + ] + ) + factory_definitions.extend( + [ + f"std::unique_ptr> detail::OperatorFactory<{op_type}>::Make({formal_params}) {{\n" + f" return Operator<{op_type}>::Make(config, {arguments});\n}}", + f"bool detail::OperatorFactory<{op_type}>::CanMake({formal_params}) {{\n" + f" return Operator<{op_type}>::CanMake(config, {arguments});\n}}", + ] + ) + + declarations.append( + f"template <>\nstruct detail::OperatorFactory<{op_type}> {{\n" + + "\n".join(dict.fromkeys(factory_declarations)) + + "\n};" + ) + definitions.extend(dict.fromkeys(factory_definitions)) + return declarations, definitions diff --git a/scripts/resolve_linked_ops.py b/scripts/resolve_linked_ops.py index 3de837db9..67ce8d9c3 100644 --- a/scripts/resolve_linked_ops.py +++ b/scripts/resolve_linked_ops.py @@ -388,14 +388,7 @@ def _load_distribution(config): def _locate_distribution_notices(config): - try: - distribution = importlib.metadata.distribution( - config.python_distribution_package - ) - except importlib.metadata.PackageNotFoundError: - # A custom library locator can supply a DSO without package metadata. - # Retain the static conservative redistribution notices in that case. - return [] + distribution = importlib.metadata.distribution(config.python_distribution_package) distribution_root = pathlib.Path(distribution.locate_file("")).resolve() notices = set() # Use only this resolved package's inventory. Sparse RECORDs and editable diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 9d4bb2e85..47ef0320a 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -269,6 +269,8 @@ if(WITH_ILUVATAR) ) file(GLOB_RECURSE ILUVATAR_SOURCES CONFIGURE_DEPENDS ${ILUVATAR_PATTERNS}) + list(FILTER ILUVATAR_SOURCES EXCLUDE REGEX + ".*/native/cuda/(hygon|mars|metax|moore|nvidia|thead)/.*") target_compile_definitions(infiniops PUBLIC WITH_ILUVATAR=1) target_sources(infiniops PRIVATE ${ILUVATAR_SOURCES}) @@ -348,6 +350,8 @@ if(WITH_METAX) ) file(GLOB_RECURSE METAX_SOURCES CONFIGURE_DEPENDS ${METAX_PATTERNS}) + list(FILTER METAX_SOURCES EXCLUDE REGEX + ".*/native/cuda/(hygon|iluvatar|mars|moore|nvidia|thead)/.*") set_source_files_properties(${METAX_SOURCES} PROPERTIES LANGUAGE CXX) @@ -374,6 +378,8 @@ if(WITH_MARS) ) file(GLOB_RECURSE MARS_SOURCES CONFIGURE_DEPENDS ${MARS_PATTERNS}) + list(FILTER MARS_SOURCES EXCLUDE REGEX + ".*/native/cuda/(hygon|iluvatar|metax|moore|nvidia|thead)/.*") set_source_files_properties(${MARS_SOURCES} PROPERTIES LANGUAGE CXX) @@ -407,6 +413,8 @@ if(WITH_MOORE) ) file(GLOB_RECURSE MOORE_SOURCES CONFIGURE_DEPENDS ${MOORE_PATTERNS}) + list(FILTER MOORE_SOURCES EXCLUDE REGEX + ".*/native/cuda/(hygon|iluvatar|mars|metax|nvidia|thead)/.*") set_source_files_properties(${MOORE_SOURCES} PROPERTIES LANGUAGE CXX) diff --git a/src/base/causal_conv1d.h b/src/base/causal_conv1d.h index cd339364f..023805e32 100644 --- a/src/base/causal_conv1d.h +++ b/src/base/causal_conv1d.h @@ -1,6 +1,7 @@ #pragma once #include #include +#include #include #include #include @@ -15,16 +16,17 @@ class CausalConv1d : public Operator { const std::optional cu_seqlens, const std::optional initial_state_indices, const std::optional final_state_indices, Tensor out, - std::optional final_conv_state) { - const Views views{qkv, - weight, - conv_state, - bias, - cu_seqlens, - initial_state_indices, - final_state_indices, - out, - final_conv_state}; + std::optional final_conv_state) + : metadata_{qkv, + weight, + conv_state, + bias, + cu_seqlens, + initial_state_indices, + final_state_indices, + out, + final_conv_state} { + const auto& views = metadata_; Validate(views); device_ = qkv.device(); dtype_ = qkv.dtype(); @@ -32,15 +34,6 @@ class CausalConv1d : public Operator { channels_ = qkv.shape()[2]; pool_ = conv_state.shape()[0]; requests_ = cu_seqlens ? cu_seqlens->numel() - 1 : qkv.shape()[0]; - for (size_t i = 0; i < views.size(); ++i) { - metadata_[i].present = views[i].has_value(); - if (!views[i]) continue; - metadata_[i].shape.assign(views[i]->shape().begin(), - views[i]->shape().end()); - metadata_[i].strides.assign(views[i]->strides().begin(), - views[i]->strides().end()); - metadata_[i].dtype = views[i]->dtype(); - } } virtual void operator()(const Tensor qkv, const Tensor weight, Tensor conv_state, const std::optional bias, @@ -52,16 +45,10 @@ class CausalConv1d : public Operator { protected: using Views = std::array, 9>; - struct Metadata { - bool present{false}; - Tensor::Shape shape; - Tensor::Strides strides; - DataType dtype{DataType::kFloat32}; - }; Device device_; DataType dtype_; size_t tokens_, channels_, pool_, requests_; - std::array metadata_; + const Views metadata_; static void Validate(const Views& v) { for (const auto& view : v) if (view && view->shape().size() != view->strides().size()) diff --git a/src/base/chunk_gated_delta_rule.h b/src/base/chunk_gated_delta_rule.h index 220fda061..1a2399279 100644 --- a/src/base/chunk_gated_delta_rule.h +++ b/src/base/chunk_gated_delta_rule.h @@ -1,6 +1,7 @@ #pragma once #include #include +#include #include #include #include @@ -16,18 +17,20 @@ class ChunkGatedDeltaRule : public Operator { const std::optional final_state_indices, Tensor out, std::optional final_state, bool use_qk_l2norm, size_t chunk_size) - : norm_(use_qk_l2norm), chunk_(chunk_size) { - const Views views{q, - k, - v, - g, - beta, - initial_state, - cu_seqlens, - initial_state_indices, - final_state_indices, - out, - final_state}; + : norm_(use_qk_l2norm), + chunk_(chunk_size), + metadata_{q, + k, + v, + g, + beta, + initial_state, + cu_seqlens, + initial_state_indices, + final_state_indices, + out, + final_state} { + const auto& views = metadata_; Validate(views, chunk_size); device_ = q.device(); data_type_ = q.dtype(); @@ -38,15 +41,6 @@ class ChunkGatedDeltaRule : public Operator { Hv_ = v.shape()[2]; D_ = q.shape()[3]; pool_ = initial_state.shape()[0]; - for (size_t i = 0; i < views.size(); ++i) - if (views[i]) { - metadata_[i].present = true; - metadata_[i].shape.assign(views[i]->shape().begin(), - views[i]->shape().end()); - metadata_[i].strides.assign(views[i]->strides().begin(), - views[i]->strides().end()); - metadata_[i].dtype = views[i]->dtype(); - } } virtual void operator()(const Tensor q, const Tensor k, const Tensor v, const Tensor g, const Tensor beta, @@ -59,17 +53,11 @@ class ChunkGatedDeltaRule : public Operator { protected: using Views = std::array, 11>; - struct Metadata { - bool present{false}; - Tensor::Shape shape; - Tensor::Strides strides; - DataType dtype{DataType::kFloat32}; - }; Device device_; DataType data_type_, gate_type_; bool norm_; size_t chunk_, B_, T_, Hk_, Hv_, D_, pool_; - std::array metadata_; + const Views metadata_; static bool floating(DataType dtype) { return dtype == DataType::kFloat16 || dtype == DataType::kBFloat16 || dtype == DataType::kFloat32; @@ -143,11 +131,8 @@ class ChunkGatedDeltaRule : public Operator { if (norm != norm_ || chunk != chunk_) throw std::invalid_argument("chunk delta attributes changed after Make"); for (size_t i = 0; i < views.size(); ++i) { - if (views[i].has_value() != metadata_[i].present || - (views[i] && (views[i]->device() != device_ || - views[i]->shape() != metadata_[i].shape || - views[i]->strides() != metadata_[i].strides || - views[i]->dtype() != metadata_[i].dtype))) + if (views[i].has_value() != metadata_[i].has_value() || + (views[i] && !std::equal_to{}(*views[i], *metadata_[i]))) throw std::invalid_argument("chunk delta metadata changed after Make"); if (views[i] && views[i]->numel() && !views[i]->data()) throw std::invalid_argument( diff --git a/src/base/fused_gated_delta_net_gating.h b/src/base/fused_gated_delta_net_gating.h index 44d2f45dd..9fe9a6c95 100644 --- a/src/base/fused_gated_delta_net_gating.h +++ b/src/base/fused_gated_delta_net_gating.h @@ -1,5 +1,6 @@ #pragma once #include +#include #include #include "data_type.h" @@ -11,8 +12,10 @@ class FusedGatedDeltaNetGating : public Operator { FusedGatedDeltaNetGating(Tensor A_log, Tensor a, Tensor b, Tensor dt_bias, float beta, float threshold, Tensor g, Tensor beta_output) - : beta_(beta), threshold_(threshold) { - std::array tensors{A_log, a, b, dt_bias, g, beta_output}; + : beta_(beta), + threshold_(threshold), + metadata_{A_log, a, b, dt_bias, g, beta_output} { + const auto& tensors = metadata_; if (a.ndim() != 3 || b.shape() != a.shape() || g.shape() != a.shape() || beta_output.shape() != a.shape() || A_log.ndim() != 1 || dt_bias.ndim() != 1 || A_log.numel() != a.shape()[2] || @@ -32,10 +35,6 @@ class FusedGatedDeltaNetGating : public Operator { if (tensors[i].device() != device_ || (i >= 4 && tensors[i].HasBroadcastDim())) throw std::invalid_argument("gating device/output layout mismatch"); - shapes_[i].assign(tensors[i].shape().begin(), tensors[i].shape().end()); - strides_[i].assign(tensors[i].strides().begin(), - tensors[i].strides().end()); - types_[i] = tensors[i].dtype(); } } virtual void operator()(Tensor A_log, Tensor a, Tensor b, Tensor dt_bias, @@ -46,8 +45,6 @@ class FusedGatedDeltaNetGating : public Operator { Device device_; size_t total_, seq_, hidden_; float beta_, threshold_; - std::array shapes_; - std::array strides_; - std::array types_; + const std::array metadata_; }; } // namespace infini::ops diff --git a/src/base/gemm.h b/src/base/gemm.h index 0bfc7c730..dfd3d20dc 100644 --- a/src/base/gemm.h +++ b/src/base/gemm.h @@ -51,14 +51,6 @@ class GemmConfig : public Config { y.dtype() == DataType::kFloat32; } - static void RequireSupported(const Config& config, const Tensor& a, - const Tensor& b, const Tensor& y) { - if (!SupportsMetadata(config, a, b, y)) { - throw std::invalid_argument( - "FullFp32 GEMM requires float32 A, B, and output"); - } - } - private: static bool IsValid(Fp32Precision precision) { return precision == Fp32Precision::DefaultFast || @@ -72,18 +64,6 @@ class Gemm : public Operator { public: using Operator::Make; - template - static std::unique_ptr> Make(const Config& config, - const Tensor a, const Tensor b, - Tail&&... tail) { - static_assert(sizeof...(Tail) == 1 || sizeof...(Tail) == 6, - "GEMM requires short or full argument signature"); - const auto& y = - std::get(std::forward_as_tuple(tail...)); - GemmConfig::RequireSupported(config, a, b, y); - return Operator::Make(config, a, b, std::forward(tail)...); - } - Gemm(const Tensor a, const Tensor b, const std::optional c, std::optional alpha, std::optional beta, std::optional trans_a, std::optional trans_b, Tensor y) @@ -205,11 +185,6 @@ struct CacheKeyBuilder { template detail::CacheKey operator()(const Config& config, const Tensor a, const Tensor b, const Tail&... tail) const { - static_assert(sizeof...(Tail) == 1 || sizeof...(Tail) == 6, - "GEMM requires short or full argument signature"); - const auto& y = - std::get(std::forward_as_tuple(tail...)); - GemmConfig::RequireSupported(config, a, b, y); return detail::CacheKey::Build( config.implementation_index(), static_cast(GemmConfig::PrecisionFrom(config)), a, b, tail...); diff --git a/src/base/layer_norm.h b/src/base/layer_norm.h index 59c513bd3..57e357646 100644 --- a/src/base/layer_norm.h +++ b/src/base/layer_norm.h @@ -1,6 +1,7 @@ #ifndef INFINI_OPS_BASE_LAYER_NORM_H_ #define INFINI_OPS_BASE_LAYER_NORM_H_ +#include #include #include "operator.h" @@ -15,18 +16,11 @@ class LayerNorm : public Operator { LayerNorm(const Tensor input, const Tensor weight, const std::optional bias, float eps, Tensor out) - : input_shape_{input.shape()}, - input_strides_{input.strides()}, - weight_shape_{weight.shape()}, - weight_strides_{weight.strides()}, - bias_shape_{bias ? Tensor::Shape{bias->shape()} : Tensor::Shape{}}, - bias_strides_{bias ? Tensor::Strides{bias->strides()} - : Tensor::Strides{}}, - out_shape_{out.shape()}, - out_strides_{out.strides()}, - dtype_{input.dtype()}, + : input_meta_{input}, + weight_meta_{weight}, + bias_meta_{bias}, + out_meta_{out}, device_{input.device()}, - has_bias_{bias.has_value()}, eps_{eps} {} virtual void operator()(const Tensor input, const Tensor weight, @@ -34,11 +28,10 @@ class LayerNorm : public Operator { Tensor out) const = 0; protected: - Tensor::Shape input_shape_, weight_shape_, bias_shape_, out_shape_; - Tensor::Strides input_strides_, weight_strides_, bias_strides_, out_strides_; - DataType dtype_; + const Tensor input_meta_, weight_meta_; + const std::optional bias_meta_; + const Tensor out_meta_; Device device_; - bool has_bias_; float eps_; }; diff --git a/src/base/mul_scalar.h b/src/base/mul_scalar.h index d3ec51e06..1af4fa879 100644 --- a/src/base/mul_scalar.h +++ b/src/base/mul_scalar.h @@ -3,6 +3,7 @@ #include #include +#include #include "operator.h" @@ -13,11 +14,8 @@ class MulScalar : public Operator { static constexpr bool kAutomaticDefault = true; MulScalar(const Tensor input, const double alpha, Tensor out) - : input_shape_(input.shape()), - input_strides_(input.strides()), - out_shape_(out.shape()), - out_strides_(out.strides()), - dtype_(input.dtype()), + : input_meta_(input), + out_meta_(out), device_(input.device()), alpha_(alpha) {} @@ -25,9 +23,7 @@ class MulScalar : public Operator { Tensor out) const = 0; protected: - Tensor::Shape input_shape_, out_shape_; - Tensor::Strides input_strides_, out_strides_; - DataType dtype_; + const Tensor input_meta_, out_meta_; Device device_; double alpha_; }; diff --git a/src/base/pad_constant.h b/src/base/pad_constant.h index 5bf074b08..ee01d9837 100644 --- a/src/base/pad_constant.h +++ b/src/base/pad_constant.h @@ -2,6 +2,7 @@ #define INFINI_OPS_BASE_PAD_CONSTANT_H_ #include +#include #include #include "operator.h" @@ -16,11 +17,8 @@ class PadConstant : public Operator { PadConstant(const Tensor input, const std::vector padding, const double value, Tensor out) - : input_shape_(input.shape()), - input_strides_(input.strides()), - out_shape_(out.shape()), - out_strides_(out.strides()), - dtype_(input.dtype()), + : input_meta_(input), + out_meta_(out), device_(input.device()), padding_(padding), value_(value) {} @@ -30,9 +28,7 @@ class PadConstant : public Operator { const double value, Tensor out) const = 0; protected: - Tensor::Shape input_shape_, out_shape_; - Tensor::Strides input_strides_, out_strides_; - DataType dtype_; + const Tensor input_meta_, out_meta_; Device device_; std::vector padding_; double value_; diff --git a/src/base/quick_gelu.h b/src/base/quick_gelu.h index 15daa14f2..d976b199e 100644 --- a/src/base/quick_gelu.h +++ b/src/base/quick_gelu.h @@ -1,6 +1,8 @@ #ifndef INFINI_OPS_BASE_QUICK_GELU_H_ #define INFINI_OPS_BASE_QUICK_GELU_H_ +#include + #include "operator.h" namespace infini::ops { @@ -12,19 +14,12 @@ class QuickGelu : public Operator { static constexpr bool kAutomaticDefault = true; QuickGelu(const Tensor input, Tensor out) - : input_shape_{input.shape()}, - input_strides_{input.strides()}, - out_shape_{out.shape()}, - out_strides_{out.strides()}, - dtype_{input.dtype()}, - device_{input.device()} {} + : input_meta_{input}, out_meta_{out}, device_{input.device()} {} virtual void operator()(const Tensor input, Tensor out) const = 0; protected: - Tensor::Shape input_shape_, out_shape_; - Tensor::Strides input_strides_, out_strides_; - DataType dtype_; + const Tensor input_meta_, out_meta_; Device device_; }; diff --git a/src/base/recurrent_gated_delta_rule.h b/src/base/recurrent_gated_delta_rule.h index 61c991420..74865fd8c 100644 --- a/src/base/recurrent_gated_delta_rule.h +++ b/src/base/recurrent_gated_delta_rule.h @@ -1,6 +1,7 @@ #pragma once #include #include +#include #include #include #include @@ -16,17 +17,18 @@ class RecurrentGatedDeltaRule : public Operator { const std::optional final_state_indices, Tensor out, std::optional final_state, bool use_qk_l2norm) - : norm_(use_qk_l2norm) { - const Views views{q, - k, - v, - g, - beta, - initial_state, - initial_state_indices, - final_state_indices, - out, - final_state}; + : norm_(use_qk_l2norm), + metadata_{q, + k, + v, + g, + beta, + initial_state, + initial_state_indices, + final_state_indices, + out, + final_state} { + const auto& views = metadata_; Validate(views); device_ = q.device(); data_type_ = q.dtype(); @@ -36,15 +38,6 @@ class RecurrentGatedDeltaRule : public Operator { Hv_ = v.shape()[2]; D_ = q.shape()[3]; pool_ = initial_state.shape()[0]; - for (size_t i = 0; i < views.size(); ++i) - if (views[i]) { - metadata_[i].present = true; - metadata_[i].shape.assign(views[i]->shape().begin(), - views[i]->shape().end()); - metadata_[i].strides.assign(views[i]->strides().begin(), - views[i]->strides().end()); - metadata_[i].dtype = views[i]->dtype(); - } } virtual void operator()(const Tensor q, const Tensor k, const Tensor v, const Tensor g, const Tensor beta, @@ -56,17 +49,11 @@ class RecurrentGatedDeltaRule : public Operator { protected: using Views = std::array, 10>; - struct Metadata { - bool present{false}; - Tensor::Shape shape; - Tensor::Strides strides; - DataType dtype{DataType::kFloat32}; - }; Device device_; DataType data_type_, gate_type_; bool norm_; size_t B_, Hk_, Hv_, D_, pool_; - std::array metadata_; + const Views metadata_; static bool floating(DataType dtype) { return dtype == DataType::kFloat16 || dtype == DataType::kBFloat16 || dtype == DataType::kFloat32; @@ -130,11 +117,8 @@ class RecurrentGatedDeltaRule : public Operator { throw std::invalid_argument( "recurrent delta normalization changed after Make"); for (size_t i = 0; i < views.size(); ++i) { - if (views[i].has_value() != metadata_[i].present || - (views[i] && (views[i]->device() != device_ || - views[i]->shape() != metadata_[i].shape || - views[i]->strides() != metadata_[i].strides || - views[i]->dtype() != metadata_[i].dtype))) + if (views[i].has_value() != metadata_[i].has_value() || + (views[i] && !std::equal_to{}(*views[i], *metadata_[i]))) throw std::invalid_argument( "recurrent delta metadata changed after Make"); if (views[i] && views[i]->numel() && !views[i]->data()) diff --git a/src/base/row_l2_normalize.h b/src/base/row_l2_normalize.h index 0a91f2dc1..3dde4d929 100644 --- a/src/base/row_l2_normalize.h +++ b/src/base/row_l2_normalize.h @@ -3,6 +3,7 @@ #include #include +#include #include "operator.h" @@ -13,11 +14,8 @@ class RowL2Normalize : public Operator { static constexpr bool kAutomaticDefault = true; RowL2Normalize(const Tensor input, const double eps, Tensor out) - : input_shape_(input.shape()), - input_strides_(input.strides()), - out_shape_(out.shape()), - out_strides_(out.strides()), - dtype_(input.dtype()), + : input_meta_(input), + out_meta_(out), device_(input.device()), eps_(eps) {} @@ -25,9 +23,7 @@ class RowL2Normalize : public Operator { Tensor out) const = 0; protected: - Tensor::Shape input_shape_, out_shape_; - Tensor::Strides input_strides_, out_strides_; - DataType dtype_; + const Tensor input_meta_, out_meta_; Device device_; double eps_; }; diff --git a/src/native/cambricon/ops/argmax/kernel.h b/src/native/cambricon/ops/argmax/kernel.h index 2a173bc37..ae31554eb 100644 --- a/src/native/cambricon/ops/argmax/kernel.h +++ b/src/native/cambricon/ops/argmax/kernel.h @@ -16,9 +16,8 @@ class Operator : public Argmax { bool keepdim, const Tensor y) { if (x.device().type() != Device::Type::kCambricon || x.device().index() < 0 || x.dtype() != DataType::kFloat32 || - x.ndim() != 1 || !x.numel() || x.numel() > INT_MAX || - !x.IsContiguous() || !x.data() || dim || keepdim || - x.device() != y.device() || y.ndim() != 0 || + !x.numel() || x.numel() > INT_MAX || !x.IsContiguous() || !x.data() || + dim || keepdim || x.device() != y.device() || y.ndim() != 0 || y.dtype() != DataType::kInt64 || !y.IsContiguous() || !y.data()) return false; const auto a = reinterpret_cast(x.data()), diff --git a/src/native/cambricon/ops/tril/kernel.h b/src/native/cambricon/ops/tril/kernel.h index 0f6398bdc..02db31b52 100644 --- a/src/native/cambricon/ops/tril/kernel.h +++ b/src/native/cambricon/ops/tril/kernel.h @@ -15,8 +15,8 @@ class Operator : public Tril { if (x.device().type() != Device::Type::kCambricon || x.device().index() < 0 || x.dtype() != DataType::kFloat32 || y.dtype() != x.dtype() || x.shape() != y.shape() || - x.device() != y.device() || x.ndim() < 2 || x.ndim() > 3 || - !x.IsContiguous() || !y.IsContiguous() || !x.data() || !y.data()) + x.device() != y.device() || x.ndim() < 2 || !x.IsContiguous() || + !y.IsContiguous() || !x.data() || !y.data()) return false; size_t count = 1; for (auto n : x.shape()) { diff --git a/src/native/cambricon/ops/triu/kernel.h b/src/native/cambricon/ops/triu/kernel.h index 81eb83de4..4dc04a53c 100644 --- a/src/native/cambricon/ops/triu/kernel.h +++ b/src/native/cambricon/ops/triu/kernel.h @@ -15,8 +15,8 @@ class Operator : public Triu { if (x.device().type() != Device::Type::kCambricon || x.device().index() < 0 || x.dtype() != DataType::kFloat32 || y.dtype() != x.dtype() || x.shape() != y.shape() || - x.device() != y.device() || x.ndim() < 2 || x.ndim() > 3 || - !x.IsContiguous() || !y.IsContiguous() || !x.data() || !y.data()) + x.device() != y.device() || x.ndim() < 2 || !x.IsContiguous() || + !y.IsContiguous() || !x.data() || !y.data()) return false; size_t count = 1; for (auto n : x.shape()) { diff --git a/src/native/cpu/ops/fill/fill.h b/src/native/cpu/ops/fill/fill.h index 4ae669cf2..c6bf17e9f 100644 --- a/src/native/cpu/ops/fill/fill.h +++ b/src/native/cpu/ops/fill/fill.h @@ -1,6 +1,7 @@ #ifndef INFINI_OPS_CPU_FILL_FILL_H_ #define INFINI_OPS_CPU_FILL_FILL_H_ +#include #include #include @@ -35,7 +36,8 @@ class Operator : public Fill, if (input.shape() != input_shape_ || input.strides() != input_strides_ || input.dtype() != input_type_ || out.shape() != out_shape_ || out.strides() != out_strides_ || out.dtype() != out_type_ || - out.device().index() != device_index_ || value != value_ || + out.device().index() != device_index_ || + !(value == value_ || (std::isnan(value) && std::isnan(value_))) || !supports(input, value, out)) throw std::invalid_argument("CPU Fill metadata/value changed after Make"); if (out.numel() == 0) return; diff --git a/src/native/cpu/ops/mul_scalar/mul_scalar.h b/src/native/cpu/ops/mul_scalar/mul_scalar.h index 7b90cde7c..8c7d592be 100644 --- a/src/native/cpu/ops/mul_scalar/mul_scalar.h +++ b/src/native/cpu/ops/mul_scalar/mul_scalar.h @@ -35,23 +35,27 @@ class Operator : public MulScalar, void operator()(const Tensor input, const double alpha, Tensor out) const override { - if (input.shape() != input_shape_ || out.shape() != out_shape_ || - input.strides() != input_strides_ || out.strides() != out_strides_ || - input.dtype() != dtype_ || out.dtype() != dtype_ || - input.device() != device_ || out.device() != device_ || + if (!std::equal_to{}(input, input_meta_) || + !std::equal_to{}(out, out_meta_) || std::memcmp(&alpha, &alpha_, sizeof(alpha))) throw std::invalid_argument( "CPU MulScalar: call metadata or alpha changed after Make"); - if (!supports(input, alpha, out)) + const auto bytes = count_ * input_meta_.element_size(); + const auto ip = reinterpret_cast(input.data()); + const auto op = reinterpret_cast(out.data()); + if (!cpu_dense_view::ValidAddress(input, bytes) || + !cpu_dense_view::ValidAddress(out, bytes) || + (bytes && ip < op + bytes && op < ip + bytes && + !(ip == op && input_meta_.strides() == out_meta_.strides()))) throw std::invalid_argument( - "CPU MulScalar: call pointer, layout or alias changed after Make"); + "CPU MulScalar: call pointer or alias changed after Make"); if (count_ == 0) return; // CPU scalar arithmetic retains float opmath precision even for half/BF16. const float multiplier = FloatAlpha(alpha_); DispatchFunc< Device::Type::kCpu, List>( - dtype_, + input_meta_.dtype(), [&](auto tag) { using T = typename decltype(tag)::type; const auto* src = static_cast(input.data()); diff --git a/src/native/cpu/ops/row_l2_normalize/row_l2_normalize.h b/src/native/cpu/ops/row_l2_normalize/row_l2_normalize.h index 912c538a7..ab5ab1f2e 100644 --- a/src/native/cpu/ops/row_l2_normalize/row_l2_normalize.h +++ b/src/native/cpu/ops/row_l2_normalize/row_l2_normalize.h @@ -28,20 +28,25 @@ class Operator } Operator(const Tensor input, const double eps, Tensor out) - : RowL2Normalize{Checked(input, eps, out), eps, out} {} + : RowL2Normalize{Checked(input, eps, out), eps, out} { + cpu_dense_view::ContiguousBytes(input, &bytes_); + } void operator()(const Tensor input, const double eps, Tensor out) const override { - if (input.shape() != input_shape_ || out.shape() != out_shape_ || - input.strides() != input_strides_ || out.strides() != out_strides_ || - input.dtype() != dtype_ || out.dtype() != dtype_ || - input.device() != device_ || out.device() != device_ || + if (!std::equal_to{}(input, input_meta_) || + !std::equal_to{}(out, out_meta_) || std::memcmp(&eps, &eps_, sizeof(eps))) throw std::invalid_argument( "CPU RowL2Normalize: call metadata or eps changed after Make"); - if (!supports(input, eps, out)) + const auto ip = reinterpret_cast(input.data()); + const auto op = reinterpret_cast(out.data()); + if (!cpu_dense_view::ValidAddress(input, bytes_) || + !cpu_dense_view::ValidAddress(out, bytes_) || + (bytes_ && ip < op + bytes_ && op < ip + bytes_ && + !(ip == op && input_meta_.strides() == out_meta_.strides()))) throw std::invalid_argument( - "CPU RowL2Normalize: call pointer, layout or alias changed after " + "CPU RowL2Normalize: call pointer or alias changed after " "Make"); if (input.size(0) == 0 || input.size(1) == 0) return; const float epsilon = eps_ > std::numeric_limits::max() @@ -50,7 +55,7 @@ class Operator DispatchFunc< Device::Type::kCpu, List>( - dtype_, + input_meta_.dtype(), [&](auto tag) { using T = typename decltype(tag)::type; const auto* src = static_cast(input.data()); @@ -75,6 +80,7 @@ class Operator } private: + std::size_t bytes_ = 0; template static T Round(float value) { if constexpr (IsBFloat16) { diff --git a/src/native/cuda/nvidia/ops/argmax/kernel.h b/src/native/cuda/nvidia/ops/argmax/kernel.h index 32f461d67..24c7fae2d 100644 --- a/src/native/cuda/nvidia/ops/argmax/kernel.h +++ b/src/native/cuda/nvidia/ops/argmax/kernel.h @@ -43,7 +43,6 @@ class Operator : public Argmax { return workspace_bytes_; } - bool requires_workspace() const override { return workspace_bytes_ != 0; } void operator()(const Tensor input, const std::optional dim, const bool keepdim, Tensor out) const override { diff --git a/src/native/cuda/nvidia/ops/causal_conv1d/kernel.h b/src/native/cuda/nvidia/ops/causal_conv1d/kernel.h index c1023f653..276fd41d7 100644 --- a/src/native/cuda/nvidia/ops/causal_conv1d/kernel.h +++ b/src/native/cuda/nvidia/ops/causal_conv1d/kernel.h @@ -124,11 +124,12 @@ class Operator : public CausalConv1d { p_.channels = channels_; p_.pool = pool_; for (size_t i = 0; i < metadata_.size(); ++i) - for (size_t j = 0; j < metadata_[i].strides.size(); ++j) - p_.strides[i][j] = metadata_[i].strides[j]; - p_.cu_i64 = metadata_[4].dtype == DataType::kInt64; - p_.initial_i64 = metadata_[5].dtype == DataType::kInt64; - p_.final_i64 = metadata_[6].dtype == DataType::kInt64; + if (metadata_[i]) + for (size_t j = 0; j < metadata_[i]->strides().size(); ++j) + p_.strides[i][j] = metadata_[i]->strides()[j]; + p_.cu_i64 = metadata_[4] && metadata_[4]->dtype() == DataType::kInt64; + p_.initial_i64 = metadata_[5] && metadata_[5]->dtype() == DataType::kInt64; + p_.final_i64 = metadata_[6] && metadata_[6]->dtype() == DataType::kInt64; } // Device capture/state-pool lifetime is not supported for this stateful path. // Zero device allocations does not alone establish capture eligibility. @@ -149,11 +150,9 @@ class Operator : public CausalConv1d { final_conv_state}; for (size_t i = 0; i < views.size(); ++i) { const auto& t = views[i]; - if (t.has_value() != metadata_[i].present || - (t && - (t->device() != device_ || t->shape() != metadata_[i].shape || - t->strides() != metadata_[i].strides || - t->dtype() != metadata_[i].dtype || (t->numel() && !t->data())))) + if (t.has_value() != metadata_[i].has_value() || + (t && (!std::equal_to{}(*t, *metadata_[i]) || + (t->numel() && !t->data())))) throw std::invalid_argument( "causal conv metadata changed or nonempty storage is null"); } diff --git a/src/native/cuda/nvidia/ops/chunk_gated_delta_rule/chunk_gated_delta_rule_legacy_kernel.cuh b/src/native/cuda/nvidia/ops/chunk_gated_delta_rule/chunk_gated_delta_rule_legacy_kernel.cuh index 5f5e513dd..813188ed1 100644 --- a/src/native/cuda/nvidia/ops/chunk_gated_delta_rule/chunk_gated_delta_rule_legacy_kernel.cuh +++ b/src/native/cuda/nvidia/ops/chunk_gated_delta_rule/chunk_gated_delta_rule_legacy_kernel.cuh @@ -60,7 +60,7 @@ __device__ void chunkGatedDeltaRuleRecurrentKernel( const void* initial_state_indices, const void* final_state_indices, bool cu_seqlens_i64, bool initial_state_indices_i64, bool final_state_indices_i64, bool use_qk_l2norm, bool has_cu_seqlens, - bool indexed_state_pool, size_t T, size_t, size_t pool_size, size_t Hk, + bool indexed_state_pool, size_t T, size_t pool_size, size_t Hk, size_t value_heads_per_key_head, ptrdiff_t out_s0, ptrdiff_t out_s1, ptrdiff_t out_s2, ptrdiff_t initial_s0, ptrdiff_t initial_s1, ptrdiff_t initial_s2, ptrdiff_t initial_s3, ptrdiff_t final_s0, diff --git a/src/native/cuda/nvidia/ops/chunk_gated_delta_rule/kernel.h b/src/native/cuda/nvidia/ops/chunk_gated_delta_rule/kernel.h index 267ddf400..79fdc30c2 100644 --- a/src/native/cuda/nvidia/ops/chunk_gated_delta_rule/kernel.h +++ b/src/native/cuda/nvidia/ops/chunk_gated_delta_rule/kernel.h @@ -14,7 +14,7 @@ namespace chunk_native_detail { struct Layout { ptrdiff_t s[11][4]{}; bool cu_i64, initial_i64, final_i64, norm, indexed, packed; - size_t T, chunk, pool, Hk, ratio; + size_t T, pool, Hk, ratio; }; template __global__ void kernel(float* workspace, const Data* q, const Data* k, @@ -25,11 +25,11 @@ __global__ void kernel(float* workspace, const Data* q, const Data* k, chunkGatedDeltaRuleRecurrentKernel( workspace, out, state, final_state, q, k, v, g, beta, cu, initial_indices, final_indices, p.cu_i64, p.initial_i64, p.final_i64, p.norm, p.packed, - p.indexed, p.T, p.chunk, p.pool, p.Hk, p.ratio, p.s[9][0], p.s[9][1], - p.s[9][2], p.s[5][0], p.s[5][1], p.s[5][2], p.s[5][3], p.s[10][0], - p.s[10][1], p.s[10][2], p.s[10][3], p.s[0][0], p.s[0][1], p.s[0][2], - p.s[1][0], p.s[1][1], p.s[1][2], p.s[2][0], p.s[2][1], p.s[2][2], - p.s[3][0], p.s[3][1], p.s[3][2], p.s[4][0], p.s[4][1], p.s[4][2]); + p.indexed, p.T, p.pool, p.Hk, p.ratio, p.s[9][0], p.s[9][1], p.s[9][2], + p.s[5][0], p.s[5][1], p.s[5][2], p.s[5][3], p.s[10][0], p.s[10][1], + p.s[10][2], p.s[10][3], p.s[0][0], p.s[0][1], p.s[0][2], p.s[1][0], + p.s[1][1], p.s[1][2], p.s[2][0], p.s[2][1], p.s[2][2], p.s[3][0], + p.s[3][1], p.s[3][2], p.s[4][0], p.s[4][1], p.s[4][2]); } } // namespace chunk_native_detail template <> @@ -73,15 +73,15 @@ class Operator p_.ratio = Hv_ / Hk_; p_.norm = norm_; p_.T = T_; - p_.chunk = chunk_; - p_.packed = metadata_[6].present; - p_.indexed = metadata_[7].present || metadata_[8].present; - p_.cu_i64 = metadata_[6].dtype == DataType::kInt64; - p_.initial_i64 = metadata_[7].dtype == DataType::kInt64; - p_.final_i64 = metadata_[8].dtype == DataType::kInt64; + p_.packed = metadata_[6].has_value(); + p_.indexed = metadata_[7].has_value() || metadata_[8].has_value(); + p_.cu_i64 = metadata_[6] && metadata_[6]->dtype() == DataType::kInt64; + p_.initial_i64 = metadata_[7] && metadata_[7]->dtype() == DataType::kInt64; + p_.final_i64 = metadata_[8] && metadata_[8]->dtype() == DataType::kInt64; for (size_t i = 0; i < metadata_.size(); ++i) - for (size_t j = 0; j < metadata_[i].strides.size(); ++j) - p_.s[i][j] = metadata_[i].strides[j]; + if (metadata_[i]) + for (size_t j = 0; j < metadata_[i]->strides().size(); ++j) + p_.s[i][j] = metadata_[i]->strides()[j]; } size_t workspace_size_in_bytes() const override { return bytes_; } void operator()(const Tensor q, const Tensor k, const Tensor v, @@ -164,13 +164,12 @@ class Operator }; dim3 grid(static_cast(B_), static_cast(Hv_), 1); const auto stream = static_cast(stream_ ? stream_ : nullptr); - chunk_native_detail::kernel - <<>>( - static_cast(workspace_), static_cast(ptr(0)), - static_cast(ptr(1)), static_cast(ptr(2)), - static_cast(ptr(3)), static_cast(ptr(4)), - static_cast(ptr(5)), ptr(6), ptr(7), ptr(8), - static_cast(ptr(9)), static_cast(ptr(10)), p_); + chunk_native_detail::kernel<<>>( + static_cast(workspace_), static_cast(ptr(0)), + static_cast(ptr(1)), static_cast(ptr(2)), + static_cast(ptr(3)), static_cast(ptr(4)), + static_cast(ptr(5)), ptr(6), ptr(7), ptr(8), + static_cast(ptr(9)), static_cast(ptr(10)), p_); } }; } // namespace infini::ops diff --git a/src/native/cuda/nvidia/ops/deepseek_moe/kernel.h b/src/native/cuda/nvidia/ops/deepseek_moe/kernel.h index cf924e3a5..571b7c688 100644 --- a/src/native/cuda/nvidia/ops/deepseek_moe/kernel.h +++ b/src/native/cuda/nvidia/ops/deepseek_moe/kernel.h @@ -10,7 +10,6 @@ class Operator : public DeepseekMoe { const std::vector up, const std::vector down, size_t intermediate, size_t num_experts); size_t workspace_size_in_bytes() const override; - bool requires_workspace() const override { return true; } private: std::shared_ptr blas_; diff --git a/src/native/cuda/nvidia/ops/flash_attn_with_kvcache/infini_flash_cuda.h b/src/native/cuda/nvidia/ops/flash_attn_with_kvcache/infini_flash_cuda.h index f1d9df3e6..68e1a0452 100644 --- a/src/native/cuda/nvidia/ops/flash_attn_with_kvcache/infini_flash_cuda.h +++ b/src/native/cuda/nvidia/ops/flash_attn_with_kvcache/infini_flash_cuda.h @@ -14,3 +14,7 @@ inline void infini_flash_check(cudaError_t error, const char* operation) { #define INFINI_FLASH_CUDA_CHECK(call) infini_flash_check((call), #call) #define INFINI_FLASH_KERNEL_CHECK() INFINI_FLASH_CUDA_CHECK(cudaGetLastError()) + +// Preserve the upstream call sites in the private inference source target. +#define C10_CUDA_CHECK INFINI_FLASH_CUDA_CHECK +#define C10_CUDA_KERNEL_LAUNCH_CHECK INFINI_FLASH_KERNEL_CHECK diff --git a/src/native/cuda/nvidia/ops/fused_gated_delta_net_gating/kernel.h b/src/native/cuda/nvidia/ops/fused_gated_delta_net_gating/kernel.h index e33833dc8..41699f41f 100644 --- a/src/native/cuda/nvidia/ops/fused_gated_delta_net_gating/kernel.h +++ b/src/native/cuda/nvidia/ops/fused_gated_delta_net_gating/kernel.h @@ -1,6 +1,4 @@ #pragma once -// Candidate destination: -// src/native/cuda/nvidia/ops/fused_gated_delta_net_gating/kernel.h. #include #include #include @@ -63,8 +61,8 @@ class Operator p_.beta = beta; p_.threshold = threshold; for (size_t i = 0; i < 6; ++i) - for (size_t j = 0; j < strides_[i].size(); ++j) - p_.strides[i][j] = strides_[i][j]; + for (size_t j = 0; j < metadata_[i].strides().size(); ++j) + p_.strides[i][j] = metadata_[i].strides()[j]; } // All metadata is a kernel value argument. There is no allocation, host copy, // synchronization or data-dependent host work in operator(); stream launch @@ -75,9 +73,7 @@ class Operator check_device(); std::array tensors{A, a, b, dt, g, betao}; for (size_t i = 0; i < 6; ++i) - if (tensors[i].device() != device_ || tensors[i].shape() != shapes_[i] || - tensors[i].strides() != strides_[i] || - tensors[i].dtype() != types_[i]) + if (!std::equal_to{}(tensors[i], metadata_[i])) throw std::invalid_argument("gating metadata changed after Make"); if (std::memcmp(&beta, &beta_, sizeof(float)) || std::memcmp(&threshold, &threshold_, sizeof(float))) @@ -86,7 +82,7 @@ class Operator const dim3 grid( static_cast(total_ / 256 + (total_ % 256 != 0))); const auto stream = static_cast(stream_ ? stream_ : nullptr); - switch (types_[1]) { + switch (metadata_[1].dtype()) { case DataType::kFloat32: launch(A, a, b, dt, g, betao, grid, stream); break; diff --git a/src/native/cuda/nvidia/ops/gemm/cublas.h b/src/native/cuda/nvidia/ops/gemm/cublas.h index feeedb9e5..795b71e9c 100644 --- a/src/native/cuda/nvidia/ops/gemm/cublas.h +++ b/src/native/cuda/nvidia/ops/gemm/cublas.h @@ -11,15 +11,6 @@ template <> class Operator : public BlasGemm> { public: - Operator(const Config&, const Tensor a, const Tensor b, - const std::optional c, std::optional alpha, - std::optional beta, std::optional trans_a, - std::optional trans_b, Tensor y) - : Operator(a, b, c, alpha, beta, trans_a, trans_b, y) {} - - Operator(const Config&, const Tensor a, const Tensor b, Tensor y) - : Operator(a, b, y) {} - template static bool supports(const Config& config, const Tensor a, const Tensor b, const Tail&... tail) { diff --git a/src/native/cuda/nvidia/ops/gemm/cublaslt.h b/src/native/cuda/nvidia/ops/gemm/cublaslt.h index 4ae15fd18..540a1beb4 100644 --- a/src/native/cuda/nvidia/ops/gemm/cublaslt.h +++ b/src/native/cuda/nvidia/ops/gemm/cublaslt.h @@ -18,15 +18,6 @@ namespace infini::ops { template <> class Operator : public Gemm { public: - Operator(const Config&, const Tensor a, const Tensor b, - const std::optional c, std::optional alpha, - std::optional beta, std::optional trans_a, - std::optional trans_b, Tensor y) - : Operator(a, b, c, alpha, beta, trans_a, trans_b, y) {} - - Operator(const Config&, const Tensor a, const Tensor b, Tensor y) - : Operator(a, b, y) {} - template static bool supports(const Config& config, const Tensor a, const Tensor b, const Tail&... tail) { diff --git a/src/native/cuda/nvidia/ops/layer_norm/kernel.h b/src/native/cuda/nvidia/ops/layer_norm/kernel.h index 51ac98b64..e8e7bd21e 100644 --- a/src/native/cuda/nvidia/ops/layer_norm/kernel.h +++ b/src/native/cuda/nvidia/ops/layer_norm/kernel.h @@ -109,7 +109,7 @@ class Operator : public LayerNorm { SpanBytes(out, &out_bytes_); if (bias) SpanBytes(*bias, &bias_bytes_); // An identical input layout inherits the validated output non-overlap. - exact_inplace_layout_ = input_strides_ == out_strides_; + exact_inplace_layout_ = input_meta_.strides() == out_meta_.strides(); layout_.rank = static_cast(input.ndim()); layout_.dim = input.size(-1); layout_.weight_stride = static_cast(weight.stride(0)); @@ -128,16 +128,11 @@ class Operator : public LayerNorm { void operator()(const Tensor input, const Tensor weight, const std::optional bias, float eps, Tensor out) const override { - if (input.shape() != input_shape_ || input.strides() != input_strides_ || - weight.shape() != weight_shape_ || - weight.strides() != weight_strides_ || out.shape() != out_shape_ || - out.strides() != out_strides_ || input.dtype() != dtype_ || - weight.dtype() != dtype_ || out.dtype() != dtype_ || - input.device() != device_ || weight.device() != device_ || - out.device() != device_ || bias.has_value() != has_bias_ || - (bias && - (bias->shape() != bias_shape_ || bias->strides() != bias_strides_ || - bias->dtype() != dtype_ || bias->device() != device_)) || + if (!std::equal_to{}(input, input_meta_) || + !std::equal_to{}(weight, weight_meta_) || + !std::equal_to{}(out, out_meta_) || + bias.has_value() != bias_meta_.has_value() || + (bias && !std::equal_to{}(*bias, *bias_meta_)) || bits(eps) != bits(eps_)) throw std::invalid_argument( "NVIDIA LayerNorm call metadata or alias changed after Make"); @@ -163,7 +158,7 @@ class Operator : public LayerNorm { if (!rows_) return; nvidia_layer_norm_detail::Launch( input.data(), weight.data(), bias ? bias->data() : nullptr, out.data(), - dtype_, layout_, rows_, eps_, + input_meta_.dtype(), layout_, rows_, eps_, static_cast(stream_ ? stream_ : 0)); } diff --git a/src/native/cuda/nvidia/ops/moe_fused_dense/kernel.h b/src/native/cuda/nvidia/ops/moe_fused_dense/kernel.h index 5bd6dfb71..76bad5c2f 100644 --- a/src/native/cuda/nvidia/ops/moe_fused_dense/kernel.h +++ b/src/native/cuda/nvidia/ops/moe_fused_dense/kernel.h @@ -10,7 +10,6 @@ class Operator : public MoeFusedDense { size_t block_size, const std::optional bias1 = std::nullopt, const std::optional bias2 = std::nullopt); size_t workspace_size_in_bytes() const override; - bool requires_workspace() const override { return true; } static bool supports(Tensor output, const Tensor hidden, const Tensor w13, const Tensor w2, const Tensor weights, const Tensor ids, const Tensor sorted, const Tensor experts, diff --git a/src/native/cuda/nvidia/ops/moe_fused_gate/kernel.cu b/src/native/cuda/nvidia/ops/moe_fused_gate/kernel.cu index f6f49f021..bfd364dcf 100644 --- a/src/native/cuda/nvidia/ops/moe_fused_gate/kernel.cu +++ b/src/native/cuda/nvidia/ops/moe_fused_gate/kernel.cu @@ -220,7 +220,8 @@ __global__ void moe_fused_gate_kernel_dynamic( void Operator::operator()(Tensor weights, Tensor indices, const Tensor input, const Tensor bias, size_t groups, size_t selected_groups, size_t shared, float scale, bool apply_scale) const { MoeDeviceGuard device_guard{metadata_[0].device().index()}; Invocation(weights, indices, input, bias, groups, selected_groups, shared, scale, apply_scale); if(input.size(0)==0) return; - moe_detail::Require(!(input.size(1)&(input.size(1)-1)) && groups<=32 && !(groups&(groups-1)) && input.size(1)/groups<=32); + moe_detail::Require(groups <= 32 && !(groups & (groups - 1)) && + input.size(1) / groups <= 32); const auto stream=static_cast(stream_); auto run=[&](auto value) { using T=decltype(value); moe_fused_gate_kernel_dynamic<<>>( diff --git a/src/native/cuda/nvidia/ops/moe_fused_gate/kernel.h b/src/native/cuda/nvidia/ops/moe_fused_gate/kernel.h index 8ed752d2c..e4734e793 100644 --- a/src/native/cuda/nvidia/ops/moe_fused_gate/kernel.h +++ b/src/native/cuda/nvidia/ops/moe_fused_gate/kernel.h @@ -12,8 +12,7 @@ class Operator : public MoeFusedGate { return MoeFusedGate::supports(weights, indices, input, bias, groups, selected_groups, shared, scale, apply_scale) && - !(e & (e - 1)) && groups <= 32 && !(groups & (groups - 1)) && - e / groups <= 32; + groups <= 32 && !(groups & (groups - 1)) && e / groups <= 32; } void operator()(Tensor weights, Tensor indices, const Tensor input, const Tensor bias, size_t groups, size_t selected_groups, diff --git a/src/native/cuda/nvidia/ops/mul_scalar/kernel.h b/src/native/cuda/nvidia/ops/mul_scalar/kernel.h index e77ddd41d..1b90f6041 100644 --- a/src/native/cuda/nvidia/ops/mul_scalar/kernel.h +++ b/src/native/cuda/nvidia/ops/mul_scalar/kernel.h @@ -106,10 +106,8 @@ class Operator : public MulScalar { void operator()(const Tensor input, const double alpha, Tensor out) const override { - if (input.shape() != input_shape_ || out.shape() != out_shape_ || - input.strides() != input_strides_ || out.strides() != out_strides_ || - input.dtype() != dtype_ || out.dtype() != dtype_ || - input.device() != device_ || out.device() != device_ || + if (!std::equal_to{}(input, input_meta_) || + !std::equal_to{}(out, out_meta_) || std::memcmp(&alpha, &alpha_, sizeof(alpha))) throw std::invalid_argument( "NVIDIA MulScalar metadata or alpha changed after Make"); @@ -121,12 +119,13 @@ class Operator : public MulScalar { reinterpret_cast(out.data()) % out.element_size() || (nvidia_binary_view::Overlaps(input, input_span_bytes_, out, out_span_bytes_) && - !(input.data() == out.data() && input_strides_ == out_strides_))) + !(input.data() == out.data() && + input_meta_.strides() == out_meta_.strides()))) throw std::invalid_argument( "NVIDIA MulScalar current device, pointer or alias changed after " "Make"); nvidia_mul_scalar_detail::Launch( - input.data(), out.data(), dtype_, alpha_, + input.data(), out.data(), input_meta_.dtype(), alpha_, layout_.contiguous ? nullptr : &layout_, count_, static_cast(stream_ ? stream_ : 0)); } diff --git a/src/native/cuda/nvidia/ops/pad_constant/kernel.h b/src/native/cuda/nvidia/ops/pad_constant/kernel.h index b90142374..55eac7a43 100644 --- a/src/native/cuda/nvidia/ops/pad_constant/kernel.h +++ b/src/native/cuda/nvidia/ops/pad_constant/kernel.h @@ -94,27 +94,35 @@ class Operator : public PadConstant { const double value, Tensor out) : PadConstant(Checked(input, padding, value, out), padding, value, out) { BuildLayout(input, padding, out); + nvidia_binary_view::SpanBytes(input, &input_bytes_); + nvidia_binary_view::SpanBytes(out, &out_bytes_); + exact_inplace_layout_ = + input.shape() == out.shape() && input.strides() == out.strides(); } bool capture_safe() const override { return true; } void operator()(const Tensor input, const std::vector padding, const double value, Tensor out) const override { - if (input.shape() != input_shape_ || input.strides() != input_strides_ || - out.shape() != out_shape_ || out.strides() != out_strides_ || - input.dtype() != dtype_ || out.dtype() != dtype_ || - input.device() != device_ || out.device() != device_ || - padding != padding_ || std::memcmp(&value, &value_, sizeof(value)) || - !supports(input, padding, value, out)) + if (!std::equal_to{}(input, input_meta_) || + !std::equal_to{}(out, out_meta_) || padding != padding_ || + std::memcmp(&value, &value_, sizeof(value))) throw std::invalid_argument( "NVIDIA PadConstant metadata or alias changed after Make"); + int current = -1; + if (cudaGetDevice(¤t) != cudaSuccess || current != device_.index() || + (input_bytes_ && !input.data()) || (out_bytes_ && !out.data()) || + (nvidia_binary_view::Overlaps(input, input_bytes_, out, out_bytes_) && + !(input.data() == out.data() && exact_inplace_layout_))) + throw std::invalid_argument( + "NVIDIA PadConstant current device, pointer or alias changed after " + "Make"); if (out.numel() == 0 || - (input.data() == out.data() && input.shape() == out.shape() && - input.strides() == out.strides())) + (input.data() == out.data() && exact_inplace_layout_)) return; nvidia_pad_constant_detail::Launch( - input.data(), out.data(), dtype_, value_, layout_, out.numel(), - static_cast(stream_ ? stream_ : 0)); + input.data(), out.data(), input_meta_.dtype(), value_, layout_, + out.numel(), static_cast(stream_ ? stream_ : 0)); } private: @@ -140,6 +148,8 @@ class Operator : public PadConstant { } } nvidia_pad_constant_detail::Layout layout_{}; + std::size_t input_bytes_ = 0, out_bytes_ = 0; + bool exact_inplace_layout_ = false; }; } // namespace infini::ops diff --git a/src/native/cuda/nvidia/ops/per_tensor_i8/kernel.h b/src/native/cuda/nvidia/ops/per_tensor_i8/kernel.h index fe57cd753..40a3437ac 100644 --- a/src/native/cuda/nvidia/ops/per_tensor_i8/kernel.h +++ b/src/native/cuda/nvidia/ops/per_tensor_i8/kernel.h @@ -18,7 +18,8 @@ void Quantize(const void*, std::int8_t*, float*, DataType, Layout, std::size_t, void Dequantize(void*, const std::int8_t*, const float*, DataType, Layout, std::size_t, cudaStream_t); inline bool Supports(Tensor x, Tensor packed, Tensor scale, - const std::optional zero) { + const std::optional zero, bool quant, + bool is_static = true) { using namespace nvidia_binary_view; if (zero || (x.dtype() != DataType::kFloat16 && x.dtype() != DataType::kFloat32 && @@ -27,8 +28,11 @@ inline bool Supports(Tensor x, Tensor packed, Tensor scale, scale.dtype() != DataType::kFloat32 || scale.numel() != 1 || x.device().type() != Device::Type::kNvidia || x.device() != packed.device() || x.device() != scale.device() || - x.shape() != packed.shape() || x.ndim() > 4 || !NonOverlappingOutput(x) || - !NonOverlappingOutput(packed)) + x.shape() != packed.shape() || x.ndim() > 4 || + x.strides().size() != x.ndim() || + packed.strides().size() != packed.ndim() || + !(quant ? NonOverlappingOutput(packed) : NonOverlappingOutput(x)) || + (quant && !is_static && !NonOverlappingOutput(scale))) return false; int device = -1; if (cudaGetDevice(&device) != cudaSuccess || device != x.device().index()) @@ -51,14 +55,25 @@ inline bool Supports(Tensor x, Tensor packed, Tensor scale, if ((xb && reinterpret_cast(x.data()) % x.element_size()) || reinterpret_cast(scale.data()) % sizeof(float)) return false; - return !Overlaps(x, xb, packed, pb) && !Overlaps(x, xb, scale, sb) && - !Overlaps(packed, pb, scale, sb); + return !Overlaps(x, xb, packed, pb) && + !(quant && !is_static && Overlaps(x, xb, scale, sb)) && + !(quant ? Overlaps(packed, pb, scale, sb) + : Overlaps(x, xb, scale, sb)); +} +inline Tensor Checked(Tensor x, Tensor packed, Tensor scale, + const std::optional zero, bool quant, + bool is_static = true) { + if (!Supports(x, packed, scale, zero, quant, is_static)) + throw std::invalid_argument("unsupported per-tensor INT8 operands"); + return x; } struct Metadata { DataType dtype; Layout layout{}; std::size_t count; - Metadata(Tensor x, Tensor packed) : dtype(x.dtype()), count(x.numel()) { + Device device; + Metadata(Tensor x, Tensor packed) + : dtype(x.dtype()), count(x.numel()), device(x.device()) { layout.ndim = x.ndim(); for (std::size_t d = 0; d < x.ndim(); ++d) { layout.shape[d] = x.shape()[d]; @@ -66,6 +81,17 @@ struct Metadata { layout.packed_stride[d] = packed.strides()[d]; } } + bool Matches(Tensor x, Tensor packed) const { + if (x.dtype() != dtype || x.device() != device || + x.ndim() != static_cast(layout.ndim)) + return false; + for (std::size_t d = 0; d < x.ndim(); ++d) + if (x.shape()[d] != static_cast(layout.shape[d]) || + x.strides()[d] != layout.x_stride[d] || + packed.strides()[d] != layout.packed_stride[d]) + return false; + return true; + } }; } // namespace infini::ops::nvidia_per_tensor_i8_detail namespace infini::ops { @@ -74,17 +100,23 @@ class Operator : public PerTensorQuantI8 { public: static bool supports(const Tensor x, Tensor packed, Tensor scale, - const std::optional zero, const bool) { - return nvidia_per_tensor_i8_detail::Supports(x, packed, scale, zero); + const std::optional zero, const bool is_static) { + return nvidia_per_tensor_i8_detail::Supports(x, packed, scale, zero, true, + is_static); } Operator(const Tensor x, Tensor packed, Tensor scale, const std::optional zero, const bool is_static) - : PerTensorQuantI8(x, packed, scale, zero, is_static), meta_(x, packed) {} + : PerTensorQuantI8(x, packed, scale, zero, is_static), + meta_(nvidia_per_tensor_i8_detail::Checked(x, packed, scale, zero, true, + is_static), + packed), + is_static_(is_static) {} bool capture_safe() const override { return true; } void operator()(const Tensor x, Tensor packed, Tensor scale, const std::optional zero, const bool is_static) const override { - if (!supports(x, packed, scale, zero, is_static)) + if (!supports(x, packed, scale, zero, is_static) || + !meta_.Matches(x, packed) || is_static != is_static_) throw std::invalid_argument("unsupported per-tensor INT8 operands"); nvidia_per_tensor_i8_detail::Quantize( x.data(), static_cast(packed.data()), @@ -95,6 +127,7 @@ class Operator private: nvidia_per_tensor_i8_detail::Metadata meta_; + bool is_static_; }; template <> class Operator @@ -102,15 +135,18 @@ class Operator public: static bool supports(Tensor x, const Tensor packed, const Tensor scale, const std::optional zero) { - return nvidia_per_tensor_i8_detail::Supports(x, packed, scale, zero); + return nvidia_per_tensor_i8_detail::Supports(x, packed, scale, zero, false); } Operator(Tensor x, const Tensor packed, const Tensor scale, const std::optional zero) - : PerTensorDequantI8(x, packed, scale, zero), meta_(x, packed) {} + : PerTensorDequantI8(x, packed, scale, zero), + meta_( + nvidia_per_tensor_i8_detail::Checked(x, packed, scale, zero, false), + packed) {} bool capture_safe() const override { return true; } void operator()(Tensor x, const Tensor packed, const Tensor scale, const std::optional zero) const override { - if (!supports(x, packed, scale, zero)) + if (!supports(x, packed, scale, zero) || !meta_.Matches(x, packed)) throw std::invalid_argument("unsupported per-tensor INT8 operands"); nvidia_per_tensor_i8_detail::Dequantize( x.data(), static_cast(packed.data()), diff --git a/src/native/cuda/nvidia/ops/quick_gelu/kernel.h b/src/native/cuda/nvidia/ops/quick_gelu/kernel.h index 04e06d3b2..e5668f654 100644 --- a/src/native/cuda/nvidia/ops/quick_gelu/kernel.h +++ b/src/native/cuda/nvidia/ops/quick_gelu/kernel.h @@ -31,7 +31,7 @@ class Operator : public QuickGelu { using namespace nvidia_binary_view; const auto rank = input.shape().size(); if (input.device().type() != Device::Type::kNvidia || - input.device() != out.device() || rank < 1 || rank > 8 || + input.device() != out.device() || rank > 8 || input.strides().size() != rank || out.strides().size() != rank || input.shape() != out.shape() || input.dtype() != out.dtype() || (input.dtype() != DataType::kFloat16 && @@ -79,36 +79,52 @@ class Operator : public QuickGelu { } Operator(const Tensor input, Tensor out) - : QuickGelu(checked(input, out), out) {} + : QuickGelu(checked(input, out), out) { + nvidia_binary_view::SpanBytes(input, &input_bytes_); + nvidia_binary_view::SpanBytes(out, &out_bytes_); + exact_inplace_layout_ = input.strides() == out.strides(); + layout_.rank = static_cast(input.ndim()); + for (std::size_t axis = 0; axis < input.ndim(); ++axis) { + layout_.shape[axis] = input.shape()[axis]; + layout_.input_strides[axis] = + static_cast(input.strides()[axis]); + layout_.out_strides[axis] = static_cast(out.strides()[axis]); + } + } // Replay submits a kernel with by-value layout and no device allocation. bool capture_safe() const override { return true; } void operator()(const Tensor input, Tensor out) const override { - if (input.shape() != input_shape_ || input.strides() != input_strides_ || - out.shape() != out_shape_ || out.strides() != out_strides_ || - input.dtype() != dtype_ || out.dtype() != dtype_ || - input.device() != device_ || out.device() != device_ || - !supports(input, out)) + if (!std::equal_to{}(input, input_meta_) || + !std::equal_to{}(out, out_meta_)) throw std::invalid_argument( "NVIDIA QuickGelu call metadata or alias changed after Make"); - std::size_t count = 1; - nvidia_quick_gelu_detail::Layout layout{}; - layout.rank = static_cast(input.ndim()); - for (std::size_t axis = 0; axis < input.ndim(); ++axis) { - layout.shape[axis] = input.shape()[axis]; - layout.input_strides[axis] = - static_cast(input.strides()[axis]); - layout.out_strides[axis] = static_cast(out.strides()[axis]); - count *= input.shape()[axis]; - } + int current = -1; + const auto bound = [](Tensor tensor, std::size_t bytes) { + return !bytes || + (tensor.data() && reinterpret_cast(tensor.data()) % + tensor.element_size() == + 0); + }; + if (cudaGetDevice(¤t) != cudaSuccess || current != device_.index() || + !bound(input, input_bytes_) || !bound(out, out_bytes_) || + (nvidia_binary_view::Overlaps(input, input_bytes_, out, out_bytes_) && + !(input.data() == out.data() && exact_inplace_layout_))) + throw std::invalid_argument( + "NVIDIA QuickGelu current device, pointer or alias changed after " + "Make"); + const auto count = input_meta_.numel(); if (!count) return; nvidia_quick_gelu_detail::Launch( - input.data(), out.data(), dtype_, layout, count, + input.data(), out.data(), input_meta_.dtype(), layout_, count, static_cast(stream_ ? stream_ : 0)); } private: + nvidia_quick_gelu_detail::Layout layout_{}; + std::size_t input_bytes_ = 0, out_bytes_ = 0; + bool exact_inplace_layout_ = false; static Tensor checked(const Tensor input, Tensor out) { if (!supports(input, out)) throw std::invalid_argument( diff --git a/src/native/cuda/nvidia/ops/recurrent_gated_delta_rule/kernel.h b/src/native/cuda/nvidia/ops/recurrent_gated_delta_rule/kernel.h index 78b11cf74..ddd5c8c71 100644 --- a/src/native/cuda/nvidia/ops/recurrent_gated_delta_rule/kernel.h +++ b/src/native/cuda/nvidia/ops/recurrent_gated_delta_rule/kernel.h @@ -62,12 +62,13 @@ class Operator p_.Hk = Hk_; p_.ratio = Hv_ / Hk_; p_.norm = norm_; - p_.indexed = metadata_[6].present || metadata_[7].present; - p_.initial_i64 = metadata_[6].dtype == DataType::kInt64; - p_.final_i64 = metadata_[7].dtype == DataType::kInt64; + p_.indexed = metadata_[6].has_value() || metadata_[7].has_value(); + p_.initial_i64 = metadata_[6] && metadata_[6]->dtype() == DataType::kInt64; + p_.final_i64 = metadata_[7] && metadata_[7]->dtype() == DataType::kInt64; for (size_t i = 0; i < metadata_.size(); ++i) - for (size_t j = 0; j < metadata_[i].strides.size(); ++j) - p_.s[i][j] = metadata_[i].strides[j]; + if (metadata_[i]) + for (size_t j = 0; j < metadata_[i]->strides().size(); ++j) + p_.s[i][j] = metadata_[i]->strides()[j]; } void operator()(const Tensor q, const Tensor k, const Tensor v, const Tensor g, const Tensor beta, Tensor initial_state, diff --git a/src/native/cuda/nvidia/ops/recurrent_gated_delta_rule/recurrent_delta_rule_common.cuh b/src/native/cuda/nvidia/ops/recurrent_gated_delta_rule/recurrent_delta_rule_common.cuh index ff9e5bdfa..da7e64162 100644 --- a/src/native/cuda/nvidia/ops/recurrent_gated_delta_rule/recurrent_delta_rule_common.cuh +++ b/src/native/cuda/nvidia/ops/recurrent_gated_delta_rule/recurrent_delta_rule_common.cuh @@ -1,6 +1,6 @@ // Ported from InfiniCore 1ab85ef1d3c2bce49edf1fc73a5b48becaf7552e: // src/infiniop/ops/recurrent_gated_delta_rule/cuda/recurrent_delta_rule_common.cuh. -// Modified for InfiniOps: NVIDIA-only warp definitions and backend integration. +// Modified for InfiniOps: NVIDIA single-token decode and backend integration. // Copyright 2025 InfiniTensor; MIT notice: LICENSES/InfiniCore-MIT.txt. #ifndef __RECURRENT_DELTA_RULE_COMMON_CUH__ #define __RECURRENT_DELTA_RULE_COMMON_CUH__ @@ -67,25 +67,22 @@ __device__ inline Tcompute blockReduceSum(Tcompute value, Tcompute* scratch) { return result; } -// GatePolicy prepares one decay value per key dimension and a scalar beta for -// the current token. This keeps the state-update implementation shared while -// allowing scalar GDR and vector-decay KDA to retain their own gate semantics. +// The NVIDIA provider submits one decode token per batch row. template __device__ void recurrentDeltaRuleWarpSequence( Tdata* out, Tdata* initial_state, Tdata* final_state, const Tdata* q, - const Tdata* k, const Tdata* v, const void* cu_seqlens, - const void* initial_state_indices, const void* final_state_indices, - bool cu_seqlens_i64, bool initial_state_indices_i64, - bool final_state_indices_i64, bool use_qk_l2norm, bool has_cu_seqlens, - bool indexed_state_pool, size_t total_tokens, size_t pool_size, - size_t num_key_heads, size_t value_heads_per_key_head, Tcompute query_scale, - ptrdiff_t out_s0, ptrdiff_t out_s1, ptrdiff_t out_s2, ptrdiff_t initial_s0, - ptrdiff_t initial_s1, ptrdiff_t initial_s2, ptrdiff_t initial_s3, - ptrdiff_t final_s0, ptrdiff_t final_s1, ptrdiff_t final_s2, - ptrdiff_t final_s3, ptrdiff_t q_s0, ptrdiff_t q_s1, ptrdiff_t q_s2, - ptrdiff_t k_s0, ptrdiff_t k_s1, ptrdiff_t k_s2, ptrdiff_t v_s0, - ptrdiff_t v_s1, ptrdiff_t v_s2, GatePolicy gate_policy, Tcompute* shared) { + const Tdata* k, const Tdata* v, const void* initial_state_indices, + const void* final_state_indices, bool initial_state_indices_i64, + bool final_state_indices_i64, bool use_qk_l2norm, bool indexed_state_pool, + size_t pool_size, size_t num_key_heads, size_t value_heads_per_key_head, + Tcompute query_scale, ptrdiff_t out_s0, ptrdiff_t out_s1, ptrdiff_t out_s2, + ptrdiff_t initial_s0, ptrdiff_t initial_s1, ptrdiff_t initial_s2, + ptrdiff_t initial_s3, ptrdiff_t final_s0, ptrdiff_t final_s1, + ptrdiff_t final_s2, ptrdiff_t final_s3, ptrdiff_t q_s0, ptrdiff_t q_s1, + ptrdiff_t q_s2, ptrdiff_t k_s0, ptrdiff_t k_s1, ptrdiff_t k_s2, + ptrdiff_t v_s0, ptrdiff_t v_s1, ptrdiff_t v_s2, GatePolicy gate_policy, + Tcompute* shared) { constexpr int WARP_SIZE = INFINIOP_RECURRENT_DELTA_RULE_WARP_SIZE; constexpr int NUM_THREADS = WARPS_PER_BLOCK * WARP_SIZE; constexpr int STATE_VALUES_PER_LANE = (Dk + WARP_SIZE - 1) / WARP_SIZE; @@ -103,17 +100,6 @@ __device__ void recurrentDeltaRuleWarpSequence( return; } - int64_t token_begin = 0; - int64_t token_end = static_cast(total_tokens); - if (has_cu_seqlens) { - token_begin = loadOptionalIndex(cu_seqlens, cu_seqlens_i64, batch_idx, 0); - token_end = loadOptionalIndex(cu_seqlens, cu_seqlens_i64, batch_idx + 1, 0); - if (token_begin < 0 || token_end < token_begin || - token_end > static_cast(total_tokens)) { - return; - } - } - int64_t read_slot = batch_idx; int64_t write_slot = batch_idx; if (indexed_state_pool) { @@ -129,15 +115,10 @@ __device__ void recurrentDeltaRuleWarpSequence( (final_state_indices != nullptr && (write_slot < 0 || write_slot >= static_cast(pool_size)))) { if (valid_value_dim && lane_idx == 0) { - const int token_batch = has_cu_seqlens ? 0 : batch_idx; - for (int64_t token_idx = token_begin; token_idx < token_end; - ++token_idx) { - const ptrdiff_t out_base = - static_cast(token_batch) * out_s0 + - static_cast(token_idx) * out_s1 + - static_cast(value_head_idx) * out_s2; - out[out_base + value_dim_idx] = static_cast(0.0f); - } + const ptrdiff_t out_base = + static_cast(batch_idx) * out_s0 + + static_cast(value_head_idx) * out_s2; + out[out_base + value_dim_idx] = static_cast(0.0f); } return; } @@ -179,8 +160,9 @@ __device__ void recurrentDeltaRuleWarpSequence( Tcompute* reduction_scratch = decay_local + Dk; Tcompute* beta_shared = reduction_scratch + NUM_THREADS; - const int token_batch = has_cu_seqlens ? 0 : batch_idx; - for (int64_t token_idx = token_begin; token_idx < token_end; ++token_idx) { + const int token_batch = batch_idx; + { + constexpr int64_t token_idx = 0; const ptrdiff_t q_base = static_cast(token_batch) * q_s0 + static_cast(token_idx) * q_s1 + static_cast(key_head_idx) * q_s2; diff --git a/src/native/cuda/nvidia/ops/recurrent_gated_delta_rule/recurrent_gated_delta_rule_legacy_kernel.cuh b/src/native/cuda/nvidia/ops/recurrent_gated_delta_rule/recurrent_gated_delta_rule_legacy_kernel.cuh index a50ceaa1f..d5dd6e14b 100644 --- a/src/native/cuda/nvidia/ops/recurrent_gated_delta_rule/recurrent_gated_delta_rule_legacy_kernel.cuh +++ b/src/native/cuda/nvidia/ops/recurrent_gated_delta_rule/recurrent_gated_delta_rule_legacy_kernel.cuh @@ -1,6 +1,6 @@ // Copied from InfiniCore 1ab85ef1d3c2bce49edf1fc73a5b48becaf7552e: // src/infiniop/ops/recurrent_gated_delta_rule/cuda/kernel.cuh. -// Modified for InfiniOps: formatting and this source notice only. +// Modified for InfiniOps: retain the scalar, single-token decode path. // Copyright 2025 InfiniTensor; MIT notice: LICENSES/InfiniCore-MIT.txt. #ifndef __RECURRENT_GATED_DELTA_RULE_KERNEL_CUH__ #define __RECURRENT_GATED_DELTA_RULE_KERNEL_CUH__ @@ -66,14 +66,13 @@ __device__ void recurrentGatedDeltaRuleIndexedPoolWarpKernel( }; op::recurrent_gated_delta_rule::cuda::recurrentDeltaRuleWarpSequence< Tdata, Tcompute, Dk, Dv, WARPS_PER_BLOCK>( - out, initial_state, final_state, q, k, v, nullptr, initial_state_indices, - final_state_indices, false, initial_state_indices_i64, - final_state_indices_i64, use_qk_l2norm, false, indexed_state_pool, 1, - pool_size, num_key_heads, value_heads_per_key_head, - rsqrtf(static_cast(Dk)), out_s0, out_s1, out_s2, initial_s0, - initial_s1, initial_s2, initial_s3, final_s0, final_s1, final_s2, - final_s3, q_s0, q_s1, q_s2, k_s0, k_s1, k_s2, v_s0, v_s1, v_s2, - gate_policy, shared); + out, initial_state, final_state, q, k, v, initial_state_indices, + final_state_indices, initial_state_indices_i64, final_state_indices_i64, + use_qk_l2norm, indexed_state_pool, pool_size, num_key_heads, + value_heads_per_key_head, rsqrtf(static_cast(Dk)), out_s0, + out_s1, out_s2, initial_s0, initial_s1, initial_s2, initial_s3, final_s0, + final_s1, final_s2, final_s3, q_s0, q_s1, q_s2, k_s0, k_s1, k_s2, v_s0, + v_s1, v_s2, gate_policy, shared); } #endif // __RECURRENT_GATED_DELTA_RULE_KERNEL_CUH__ diff --git a/src/native/cuda/nvidia/ops/sum/kernel.h b/src/native/cuda/nvidia/ops/sum/kernel.h index 34f1685de..e0d1132be 100644 --- a/src/native/cuda/nvidia/ops/sum/kernel.h +++ b/src/native/cuda/nvidia/ops/sum/kernel.h @@ -43,8 +43,9 @@ class Operator : public Sum { Tensor out, bool invocation) { if (input.device().type() != Device::Type::kNvidia || input.device() != out.device() || input.dtype() != out.dtype() || - !dim || dim->size() != 1 || keepdim || dtype || - input.shape().size() < 2 || input.shape().size() > 4 || + !dim || dim->size() != 1 || keepdim || + (dtype && *dtype != input.dtype()) || input.shape().size() < 1 || + input.shape().size() > 4 || out.shape().size() + 1 != input.shape().size()) return false; if (input.dtype() != DataType::kFloat16 && @@ -151,6 +152,7 @@ class Operator : public Sum { static bool CheckedContiguous(const Tensor tensor) { std::size_t expected = 1; for (std::size_t i = tensor.shape().size(); i > 0; --i) { + if (tensor.shape()[i - 1] == 1) continue; if (expected > static_cast( std::numeric_limits::max()) || tensor.strides()[i - 1] != static_cast(expected)) diff --git a/src/native/cuda/nvidia/ops/take/kernel.h b/src/native/cuda/nvidia/ops/take/kernel.h index c44bb9bba..26d7bfb28 100644 --- a/src/native/cuda/nvidia/ops/take/kernel.h +++ b/src/native/cuda/nvidia/ops/take/kernel.h @@ -123,6 +123,7 @@ class Operator : public Take { if (shape.size() != strides.size()) return false; std::size_t expected = 1; for (std::size_t i = shape.size(); i > 0; --i) { + if (shape[i - 1] == 1) continue; if (expected > static_cast( std::numeric_limits::max()) || strides[i - 1] != static_cast(expected)) diff --git a/src/native/cuda/nvidia/ops/upsample_bilinear2d/kernel.h b/src/native/cuda/nvidia/ops/upsample_bilinear2d/kernel.h index 3590ba063..5c3893f2f 100644 --- a/src/native/cuda/nvidia/ops/upsample_bilinear2d/kernel.h +++ b/src/native/cuda/nvidia/ops/upsample_bilinear2d/kernel.h @@ -175,6 +175,7 @@ class Operator static bool CheckedContiguous(const Tensor tensor) { std::size_t expected = 1; for (std::size_t i = 4; i > 0; --i) { + if (tensor.shape()[i - 1] == 1) continue; if (expected > static_cast( std::numeric_limits::max()) || tensor.strides()[i - 1] != static_cast(expected) || diff --git a/src/native/cuda/nvidia/ops/w4_marlin/kernel.h b/src/native/cuda/nvidia/ops/w4_marlin/kernel.h index 7fd2b47a3..9df1d31a7 100644 --- a/src/native/cuda/nvidia/ops/w4_marlin/kernel.h +++ b/src/native/cuda/nvidia/ops/w4_marlin/kernel.h @@ -48,7 +48,6 @@ class Operator : public W4MarlinGemm { const int64_t group_size, const bool is_k_full, const float alpha, Tensor out) const override; std::size_t workspace_size_in_bytes() const override { return bytes_; } - bool requires_workspace() const override { return true; } // Original host selector still sets CUDA function attributes and queries the // device at invocation. No capture promise without a qualified frozen launch. diff --git a/src/native/cuda/nvidia/ops/w4_marlin/launch.h b/src/native/cuda/nvidia/ops/w4_marlin/launch.h index d0a9f1dfc..c0e0e2a8f 100644 --- a/src/native/cuda/nvidia/ops/w4_marlin/launch.h +++ b/src/native/cuda/nvidia/ops/w4_marlin/launch.h @@ -40,7 +40,7 @@ inline void info(int device, int& sms, int& shared) { cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, device), "W4 SM query"); if (major < 8) - throw std::runtime_error("W4 candidate builds SM80+ kernels only"); + throw std::runtime_error("W4 Marlin builds SM80+ kernels only"); check(cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, device), "W4 SM count"); check(cudaDeviceGetAttribute(&shared, cudaDevAttrMaxSharedMemoryPerBlockOptin, diff --git a/src/native/cuda/nvidia/ops/w4_marlin/marlin_check.h b/src/native/cuda/nvidia/ops/w4_marlin/marlin_check.h index 085902de1..a4d0536d9 100644 --- a/src/native/cuda/nvidia/ops/w4_marlin/marlin_check.h +++ b/src/native/cuda/nvidia/ops/w4_marlin/marlin_check.h @@ -24,3 +24,6 @@ template ::infini::ops::native_w4::marlin_fail(__VA_ARGS__); \ } \ } while (false) + +// Keep upstream check call sites intact inside the private prepared sources. +#define STD_TORCH_CHECK INFINI_MARLIN_CHECK diff --git a/src/operator.h b/src/operator.h index 7edfa78f9..3ec0c6733 100644 --- a/src/operator.h +++ b/src/operator.h @@ -29,6 +29,10 @@ namespace infini::ops::detail { +// Generated formal overloads normalize public calls before crossing the ABI. +template +struct OperatorFactory; + // A key with providers covering different domains may opt into metadata-based // defaults. Explicit implementation indices and other keys remain unchanged. template @@ -291,10 +295,6 @@ class OperatorBase { virtual std::size_t workspace_size_in_bytes() const { return 0; } - // Existing providers may treat workspace as optional. Only providers that - // opt in here make an absent/undersized Call workspace an error. - virtual bool requires_workspace() const { return false; } - // Capture is opt-in for each prepared implementation. In particular, an // operator's device type alone does not establish resource lifetime safety. virtual bool capture_safe() const { return false; } @@ -482,23 +482,14 @@ class Operator : public OperatorBase { if (required_workspace == 0) { return (*op)(handle, args...); } - // An explicit workspace takes precedence. Preserve legacy optional - // workspace behavior unless this provider declares it mandatory. + // An explicit workspace takes precedence. Providers validate their own + // workspace requirements before accessing it. if (handle.workspace() != nullptr || handle.workspace_size_in_bytes() != 0) { - if (op->requires_workspace() && - (handle.workspace() == nullptr || - handle.workspace_size_in_bytes() < required_workspace)) { - throw std::invalid_argument( - "operator workspace is missing or too small"); - } return (*op)(handle, args...); } auto lease = handle.acquire_workspace(required_workspace); if (!lease.data) { - if (op->requires_workspace()) { - throw std::invalid_argument("operator workspace lease is missing"); - } return (*op)(handle, args...); } if (lease.size_in_bytes < required_workspace) { @@ -582,11 +573,14 @@ class Operator : public OperatorBase { } static std::size_t FirstActiveImplementationIndex(List<>) { - assert(false && "operator has no active implementation for this device"); - std::abort(); + throw std::invalid_argument( + "operator has no active implementation for this device"); } static std::size_t DefaultImplementationIndex(Device::Type dev_type) { + if (!detail::ListContains(dev_type, ActiveDevices{})) { + throw std::invalid_argument("operator device is not enabled"); + } std::size_t default_index{0}; DispatchFunc>( @@ -649,7 +643,6 @@ class Operator : public OperatorBase { Device::Type dispatch_device_type, bool strict, Selected&& selected, Args&&... args) { bool matched = false; - auto cache_args = std::forward_as_tuple(args...); if (!detail::ListContains(dispatch_device_type, ActiveDevices{})) { if (strict) { @@ -670,25 +663,18 @@ class Operator : public OperatorBase { constexpr std::size_t kImplementationIndex = decltype(implementation_tag)::value; using Provider = Operator; + constexpr bool kConfigConstructor = + std::is_constructible_v; constexpr bool kConfigAware = - std::is_constructible_v && detail::HasConfigAwareSupports::value; - if constexpr (kConfigAware || + if constexpr (kConfigConstructor || std::is_constructible_v) { bool supported = false; if constexpr (kConfigAware) { - supported = std::apply( - [&](const auto&... cached_args) { - return Provider::supports(config, cached_args...); - }, - cache_args); + supported = Provider::supports(config, std::as_const(args)...); } else if (!config.requires_config_aware_provider()) { - supported = std::apply( - [](const auto&... cached_args) { - return detail::ProviderSupports(0, - cached_args...); - }, - cache_args); + supported = detail::ProviderSupports( + 0, std::as_const(args)...); } if (!supported) { if (automatic) return; @@ -699,14 +685,9 @@ class Operator : public OperatorBase { } return; } - std::apply( - [&](auto&... cached_args) { - matched = true; - selected(device_tag, implementation_tag, - std::bool_constant{}, - cached_args...); - }, - cache_args); + matched = true; + selected(device_tag, implementation_tag, + std::bool_constant{}, args...); } else { if (!automatic && strict) { throw std::invalid_argument( diff --git a/tests/test_generate_public_operator_header.py b/tests/test_generate_public_operator_header.py index a9a56954e..091e626ae 100644 --- a/tests/test_generate_public_operator_header.py +++ b/tests/test_generate_public_operator_header.py @@ -31,8 +31,6 @@ def test_public_operator_header_omits_host_range_profiler(tmp_path): "static void Call(const Handle& handle, const Config& config," in public_header ) for name in ( - "Make", - "CanMake", "active_implementation_indices", "DefaultConfig", "DefaultImplementationIndex", @@ -42,6 +40,8 @@ def test_public_operator_header_omits_host_range_profiler(tmp_path): assert re.search(rf"\b{name}\s*\([^{{;]*\)\s*;", public_header) assert "DispatchFunc>(\n" not in public_header assert "Config::Automatic()" not in public_header + for name in ("Make", "CanMake", "Call"): + assert f"detail::OperatorFactory::{name}(" in public_header public_section = public_header.split("class Operator : public OperatorBase {", 1)[1] public_section = public_section.split(" protected:", 1)[0] assert "static Config DefaultConfig(Device::Type dev_type);" in public_section diff --git a/tests/test_generate_wrappers.py b/tests/test_generate_wrappers.py index b8672e852..b0254aa35 100644 --- a/tests/test_generate_wrappers.py +++ b/tests/test_generate_wrappers.py @@ -135,6 +135,33 @@ def test_operator_call_instantiations_externalize_active_implementation_query(): assert f"template {signature}" in definitions +def test_public_factory_normalizes_formal_scalar_and_tensor_list_arguments(): + module = _load_generator_module() + signature = module._ParsedFunction( + [ + module._ParsedArgument("const std::vector", "inputs"), + module._ParsedArgument("const double", "alpha"), + module._ParsedArgument("Tensor", "out"), + ] + ) + operator = module._Operator("cat", constructors=[signature], calls=[signature]) + declarations, definitions = module._generate_operator_call_instantiation_entries( + operator + ) + header = "\n".join(declarations) + source = "\n".join(definitions) + assert "struct detail::OperatorFactory<::infini::ops::Cat>" in header + assert ( + "Make(const Config& config, const std::vector& inputs, const double& alpha, const Tensor& out)" + in header + ) + assert "Operator<::infini::ops::Cat>::Make(config, inputs, alpha, out)" in source + assert ( + "Operator<::infini::ops::Cat>::Call(handle, config, inputs, alpha, out)" + in source + ) + + def test_operator_call_instantiations_keep_scalar_and_optional_tensor_overloads_distinct( monkeypatch, tmp_path ): @@ -681,6 +708,7 @@ def test_generated_dispatch_calls_start_with_dispatch_profile_scope( " HostRangeLayer::kDispatchCall};" ) assert '#include "host_range_profiler.h"' in text + assert '#include "infini/operator_call_instantiations.h"' in text dispatch_call_bodies = re.findall( r"void CallProfileOp\([^)]*\) \{\n(.*?)\n\}", text, diff --git a/tests/test_resolve_linked_ops.py b/tests/test_resolve_linked_ops.py index 90dcee9c3..6d89b9481 100644 --- a/tests/test_resolve_linked_ops.py +++ b/tests/test_resolve_linked_ops.py @@ -18,6 +18,9 @@ def _load_resolver_module(): assert spec.loader is not None sys.modules[spec.name] = module spec.loader.exec_module(module) + # Resolver tests replace package discovery with fixture DSO paths; notices + # need the same explicit fixture boundary instead of installed metadata. + module._locate_distribution_notices = lambda config: [] return module From d696f4c162e7ca5eced8797c44f617c7f0c71fee Mon Sep 17 00:00:00 2001 From: whjthu Date: Tue, 20 Oct 2026 11:38:30 +0800 Subject: [PATCH 23/25] style: format native operator sources --- src/native/cuda/nvidia/ops/argmax/kernel.h | 1 - .../chunk_gated_delta_rule_legacy_kernel.cuh | 3 +-- src/native/cuda/nvidia/ops/flash_attn_varlen_func/kernel.cc | 3 ++- .../recurrent_gated_delta_rule/recurrent_delta_rule_common.cuh | 3 +-- 4 files changed, 4 insertions(+), 6 deletions(-) diff --git a/src/native/cuda/nvidia/ops/argmax/kernel.h b/src/native/cuda/nvidia/ops/argmax/kernel.h index 24c7fae2d..1a507ad26 100644 --- a/src/native/cuda/nvidia/ops/argmax/kernel.h +++ b/src/native/cuda/nvidia/ops/argmax/kernel.h @@ -43,7 +43,6 @@ class Operator : public Argmax { return workspace_bytes_; } - void operator()(const Tensor input, const std::optional dim, const bool keepdim, Tensor out) const override { if (input.shape() != input_shape_ || input.strides() != input_strides_ || diff --git a/src/native/cuda/nvidia/ops/chunk_gated_delta_rule/chunk_gated_delta_rule_legacy_kernel.cuh b/src/native/cuda/nvidia/ops/chunk_gated_delta_rule/chunk_gated_delta_rule_legacy_kernel.cuh index 813188ed1..0c8ac1a52 100644 --- a/src/native/cuda/nvidia/ops/chunk_gated_delta_rule/chunk_gated_delta_rule_legacy_kernel.cuh +++ b/src/native/cuda/nvidia/ops/chunk_gated_delta_rule/chunk_gated_delta_rule_legacy_kernel.cuh @@ -103,8 +103,7 @@ __device__ void chunkGatedDeltaRuleRecurrentKernel( : loadOptionalIndex(final_state_indices, final_state_indices_i64, batch_idx, batch_idx); - if (read_slot < 0 || - read_slot >= static_cast(pool_size) || + if (read_slot < 0 || read_slot >= static_cast(pool_size) || (final_state_indices != nullptr && (write_slot < 0 || write_slot >= static_cast(pool_size)))) { return; diff --git a/src/native/cuda/nvidia/ops/flash_attn_varlen_func/kernel.cc b/src/native/cuda/nvidia/ops/flash_attn_varlen_func/kernel.cc index 96920aae2..a45ddfd0f 100644 --- a/src/native/cuda/nvidia/ops/flash_attn_varlen_func/kernel.cc +++ b/src/native/cuda/nvidia/ops/flash_attn_varlen_func/kernel.cc @@ -100,7 +100,8 @@ bool supports_metadata(VARLEN_ARGS) { bool Operator::supports( VARLEN_ARGS) { if (!supports_metadata(VARLEN_VALUES)) return false; - // These native kernels require SM80+. The attribute query allocates no scratch. + // These native kernels require SM80+. The attribute query allocates no + // scratch. int major = 0; return cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, q.device().index()) == cudaSuccess && diff --git a/src/native/cuda/nvidia/ops/recurrent_gated_delta_rule/recurrent_delta_rule_common.cuh b/src/native/cuda/nvidia/ops/recurrent_gated_delta_rule/recurrent_delta_rule_common.cuh index da7e64162..9a715fdc7 100644 --- a/src/native/cuda/nvidia/ops/recurrent_gated_delta_rule/recurrent_delta_rule_common.cuh +++ b/src/native/cuda/nvidia/ops/recurrent_gated_delta_rule/recurrent_delta_rule_common.cuh @@ -110,8 +110,7 @@ __device__ void recurrentDeltaRuleWarpSequence( ? static_cast(batch_idx) : loadOptionalIndex(final_state_indices, final_state_indices_i64, batch_idx, batch_idx); - if (read_slot < 0 || - read_slot >= static_cast(pool_size) || + if (read_slot < 0 || read_slot >= static_cast(pool_size) || (final_state_indices != nullptr && (write_slot < 0 || write_slot >= static_cast(pool_size)))) { if (valid_value_dim && lane_idx == 0) { From bcfe96daf338b9f9992049f7965b27b6afe88b68 Mon Sep 17 00:00:00 2001 From: whjthu Date: Tue, 20 Oct 2026 19:07:02 +0800 Subject: [PATCH 24/25] refactor: unify scalar multiplication and padding interfaces --- cmake/NativeFlashAttention.cmake | 2 +- ..._ops_nvidia_dense_sampling_torch_free.json | 3 +- ...a_dense_sampling_torch_free_no_marlin.json | 3 +- src/base/causal_conv1d.h | 70 ++++++++--------- src/base/chunk_gated_delta_rule.h | 77 ++++++++----------- src/base/fused_gated_delta_net_gating.h | 29 ++++--- src/base/mrope.h | 38 ++++----- src/base/mul.h | 34 ++++++++ src/base/mul_scalar.h | 47 ----------- src/base/pad.h | 44 +++++++++++ src/base/pad_constant.h | 39 ---------- src/base/recurrent_gated_delta_rule.h | 67 +++++++--------- src/base/separate_table_rope.h | 7 +- src/native/cpu/ops/mul/mul.h | 16 ++++ .../{mul_scalar/mul_scalar.h => mul/scalar.h} | 30 ++++---- .../ops/chunk_gated_delta_rule/kernel.h | 5 ++ .../ops/flash_attn_varlen_func/kernel.cc | 2 +- .../ops/flash_attn_with_kvcache/kernel.cc | 2 +- src/native/cuda/nvidia/ops/flash_common.cc | 2 +- .../infini_flash_cuda.h | 3 + src/native/cuda/nvidia/ops/mul/kernel.h | 18 +++++ .../{mul_scalar/kernel.cu => mul/scalar.cu} | 25 +++--- .../ops/{mul_scalar/kernel.h => mul/scalar.h} | 35 +++++---- .../ops/{pad_constant => pad}/kernel.cu | 36 +++++---- .../nvidia/ops/{pad_constant => pad}/kernel.h | 59 +++++++------- src/native/cuda/ops/mul/kernel.h | 6 ++ 26 files changed, 362 insertions(+), 337 deletions(-) delete mode 100644 src/base/mul_scalar.h create mode 100644 src/base/pad.h delete mode 100644 src/base/pad_constant.h rename src/native/cpu/ops/{mul_scalar/mul_scalar.h => mul/scalar.h} (79%) rename src/native/cuda/nvidia/ops/{flash_attn_with_kvcache => }/infini_flash_cuda.h (80%) rename src/native/cuda/nvidia/ops/{mul_scalar/kernel.cu => mul/scalar.cu} (82%) rename src/native/cuda/nvidia/ops/{mul_scalar/kernel.h => mul/scalar.h} (86%) rename src/native/cuda/nvidia/ops/{pad_constant => pad}/kernel.cu (71%) rename src/native/cuda/nvidia/ops/{pad_constant => pad}/kernel.h (74%) diff --git a/cmake/NativeFlashAttention.cmake b/cmake/NativeFlashAttention.cmake index 5c013c6a9..36d25c185 100644 --- a/cmake/NativeFlashAttention.cmake +++ b/cmake/NativeFlashAttention.cmake @@ -61,7 +61,7 @@ function(infiniops_add_native_flash target) FLASH_NAMESPACE=infini_flash FLASHATTENTION_DISABLE_DROPOUT) target_include_directories(${target} PRIVATE "${_source}/src" "${_source}/include" "${_prepared}" - "${_cutlass_include}" "${_provider}" "${_varlen_provider}" + "${_cutlass_include}" "${_provider}/.." ${INFINI_RT_INCLUDE_DIRS}) target_compile_options(${target} PRIVATE # Match the pinned upstream setup.py numerical compilation mode. diff --git a/scripts/configs/infiniops_ops_nvidia_dense_sampling_torch_free.json b/scripts/configs/infiniops_ops_nvidia_dense_sampling_torch_free.json index f0e4c3b07..a08aa92fd 100644 --- a/scripts/configs/infiniops_ops_nvidia_dense_sampling_torch_free.json +++ b/scripts/configs/infiniops_ops_nvidia_dense_sampling_torch_free.json @@ -24,8 +24,7 @@ "moe_fused_gate": {"implementations": [0]}, "mrope": {"implementations": [0]}, "mul": {"implementations": [0]}, - "mul_scalar": {"implementations": [0]}, - "pad_constant": {"implementations": [0]}, + "pad": {"implementations": [0]}, "per_channel_quant_i8": {"implementations": [0]}, "per_tensor_dequant_i8": {"implementations": [0]}, "per_tensor_quant_i8": {"implementations": [0]}, diff --git a/scripts/configs/infiniops_ops_nvidia_dense_sampling_torch_free_no_marlin.json b/scripts/configs/infiniops_ops_nvidia_dense_sampling_torch_free_no_marlin.json index 0b76aa23d..32fb6c5a3 100644 --- a/scripts/configs/infiniops_ops_nvidia_dense_sampling_torch_free_no_marlin.json +++ b/scripts/configs/infiniops_ops_nvidia_dense_sampling_torch_free_no_marlin.json @@ -24,8 +24,7 @@ "moe_fused_gate": {"implementations": [0]}, "mrope": {"implementations": [0]}, "mul": {"implementations": [0]}, - "mul_scalar": {"implementations": [0]}, - "pad_constant": {"implementations": [0]}, + "pad": {"implementations": [0]}, "per_channel_quant_i8": {"implementations": [0]}, "per_tensor_dequant_i8": {"implementations": [0]}, "per_tensor_quant_i8": {"implementations": [0]}, diff --git a/src/base/causal_conv1d.h b/src/base/causal_conv1d.h index 023805e32..ff723e2a8 100644 --- a/src/base/causal_conv1d.h +++ b/src/base/causal_conv1d.h @@ -1,5 +1,6 @@ #pragma once #include +#include #include #include #include @@ -51,58 +52,53 @@ class CausalConv1d : public Operator { const Views metadata_; static void Validate(const Views& v) { for (const auto& view : v) - if (view && view->shape().size() != view->strides().size()) - throw std::invalid_argument("causal conv shape/stride rank mismatch"); + assert((!view || view->shape().size() == view->strides().size()) && + "causal conv shape/stride rank mismatch"); const auto &q = *v[0], &w = *v[1], &state = *v[2], &out = *v[7]; const auto dtype = q.dtype(); - if (dtype != DataType::kFloat32 && dtype != DataType::kFloat16 && - dtype != DataType::kBFloat16) - throw std::invalid_argument("causal conv supports F16/BF16/F32"); - if (q.ndim() != 3 || w.ndim() != 3 || state.ndim() != 3 || - out.shape() != q.shape()) - throw std::invalid_argument("causal conv rank/shape mismatch"); + assert((dtype == DataType::kFloat32 || dtype == DataType::kFloat16 || + dtype == DataType::kBFloat16) && + "causal conv supports F16/BF16/F32"); + assert(q.ndim() == 3 && w.ndim() == 3 && state.ndim() == 3 && + out.shape() == q.shape() && "causal conv rank/shape mismatch"); const size_t channels = q.shape()[2]; - if (w.shape()[0] != channels || w.shape()[1] != 1 || w.shape()[2] == 0 || - state.shape()[1] != channels || state.shape()[2] != w.shape()[2] - 1) - throw std::invalid_argument("causal conv weight/state shape mismatch"); - if (v[6].has_value() == v[8].has_value()) - throw std::invalid_argument( - "causal conv requires exactly one final state destination mode"); + assert(w.shape()[0] == channels && w.shape()[1] == 1 && w.shape()[2] > 0 && + state.shape()[1] == channels && + state.shape()[2] == w.shape()[2] - 1 && + "causal conv weight/state shape mismatch"); + assert(v[6].has_value() != v[8].has_value() && + "causal conv requires exactly one final state destination mode"); size_t requests = q.shape()[0]; if (v[4]) { - if (q.shape()[0] != 1 || v[4]->ndim() != 1 || v[4]->numel() < 2) - throw std::invalid_argument( - "causal conv packed sequence shape mismatch"); + assert(q.shape()[0] == 1 && v[4]->ndim() == 1 && v[4]->numel() >= 2 && + "causal conv packed sequence shape mismatch"); requests = v[4]->numel() - 1; } - if (!v[5] && !v[6] && state.shape()[0] != requests) - throw std::invalid_argument("causal conv state batch mismatch"); - if (v[8] && - v[8]->shape() != Tensor::Shape{requests, channels, state.shape()[2]}) - throw std::invalid_argument("causal conv final state shape mismatch"); - if (v[3] && v[3]->shape() != Tensor::Shape{channels}) - throw std::invalid_argument("causal conv bias shape mismatch"); + assert((v[5] || v[6] || state.shape()[0] == requests) && + "causal conv state batch mismatch"); + assert((!v[8] || v[8]->shape() == + Tensor::Shape{requests, channels, state.shape()[2]}) && + "causal conv final state shape mismatch"); + assert((!v[3] || v[3]->shape() == Tensor::Shape{channels}) && + "causal conv bias shape mismatch"); for (size_t i = 0; i < v.size(); ++i) { if (!v[i]) continue; - if (v[i]->device() != q.device()) - throw std::invalid_argument("causal conv device mismatch"); + assert(v[i]->device() == q.device() && "causal conv device mismatch"); if (i >= 4 && i <= 6) { - if (v[i]->ndim() != 1 || !v[i]->IsContiguous() || - (v[i]->dtype() != DataType::kInt32 && - v[i]->dtype() != DataType::kInt64) || - (i != 4 && v[i]->numel() != requests)) - throw std::invalid_argument( - "causal conv requires contiguous I32/I64 sequence/slot indices"); - } else if (v[i]->dtype() != dtype) { - throw std::invalid_argument("causal conv data dtype mismatch"); + assert(v[i]->ndim() == 1 && v[i]->IsContiguous() && + (v[i]->dtype() == DataType::kInt32 || + v[i]->dtype() == DataType::kInt64) && + (i == 4 || v[i]->numel() == requests) && + "causal conv requires contiguous I32/I64 sequence/slot indices"); + } else { + assert(v[i]->dtype() == dtype && "causal conv data dtype mismatch"); } // Preserve ordinary positive strided data and reject broadcast // destinations. if (i == 2 || i == 7 || i == 8) for (size_t axis = 0; axis < v[i]->ndim(); ++axis) - if (v[i]->shape()[axis] != 1 && v[i]->strides()[axis] == 0) - throw std::invalid_argument( - "causal conv broadcast destination unsupported"); + assert((v[i]->shape()[axis] == 1 || v[i]->strides()[axis] != 0) && + "causal conv broadcast destination unsupported"); } if (q.shape()[1] > static_cast(std::numeric_limits::max()) || diff --git a/src/base/chunk_gated_delta_rule.h b/src/base/chunk_gated_delta_rule.h index 1a2399279..c91b7dadb 100644 --- a/src/base/chunk_gated_delta_rule.h +++ b/src/base/chunk_gated_delta_rule.h @@ -1,5 +1,6 @@ #pragma once #include +#include #include #include #include @@ -63,63 +64,51 @@ class ChunkGatedDeltaRule : public Operator { dtype == DataType::kFloat32; } static void Validate(const Views& a, size_t chunk) { - if (!chunk) - throw std::invalid_argument("chunk delta chunk_size must be positive"); + assert(chunk > 0 && "chunk delta chunk_size must be positive"); const auto &q = *a[0], &k = *a[1], &v = *a[2], &g = *a[3], &beta = *a[4], &state = *a[5], &out = *a[9]; - if (!floating(q.dtype()) || !floating(g.dtype()) || - beta.dtype() != g.dtype()) - throw std::invalid_argument( - "chunk delta data/gates support F16/BF16/F32"); - if (q.ndim() != 4 || k.shape() != q.shape() || v.ndim() != 4 || - out.shape() != v.shape() || state.ndim() != 4 || g.ndim() != 3 || - beta.shape() != g.shape()) - throw std::invalid_argument("chunk delta rank/shape mismatch"); + assert(floating(q.dtype()) && floating(g.dtype()) && + beta.dtype() == g.dtype() && + "chunk delta data/gates support F16/BF16/F32"); + assert(q.ndim() == 4 && k.shape() == q.shape() && v.ndim() == 4 && + out.shape() == v.shape() && state.ndim() == 4 && g.ndim() == 3 && + beta.shape() == g.shape() && "chunk delta rank/shape mismatch"); const auto physical_B = q.shape()[0], T = q.shape()[1], Hk = q.shape()[2], Hv = v.shape()[2], D = q.shape()[3], Dv = v.shape()[3]; size_t B = physical_B; if (a[6]) { - if (physical_B != 1 || a[6]->ndim() != 1 || a[6]->numel() < 2) - throw std::invalid_argument( - "chunk delta packed cu_seqlens shape mismatch"); + assert(physical_B == 1 && a[6]->ndim() == 1 && a[6]->numel() >= 2 && + "chunk delta packed cu_seqlens shape mismatch"); B = a[6]->numel() - 1; } - if (v.shape()[0] != physical_B || v.shape()[1] != T || !Hk || !Hv || - Hv % Hk || !D || !Dv || g.shape() != Tensor::Shape{physical_B, T, Hv}) - throw std::invalid_argument( - "chunk delta data/gate shape or head grouping mismatch"); - if (a[8].has_value() == a[10].has_value()) - throw std::invalid_argument( - "chunk delta requires exactly one final destination mode"); - if (state.shape()[1] != Hv || state.shape()[2] != Dv || - state.shape()[3] != D || (!a[7] && !a[8] && state.shape()[0] != B) || - (a[10] && a[10]->shape() != Tensor::Shape{B, Hv, Dv, D})) - throw std::invalid_argument( - "chunk delta state layout must be pool/B,Hv,Dv,Dk"); + assert((v.shape()[0] == physical_B && v.shape()[1] == T && Hk && Hv && + Hv % Hk == 0 && D && Dv && + g.shape() == Tensor::Shape{physical_B, T, Hv}) && + "chunk delta data/gate shape or head grouping mismatch"); + assert(a[8].has_value() != a[10].has_value() && + "chunk delta requires exactly one final destination mode"); + assert((state.shape()[1] == Hv && state.shape()[2] == Dv && + state.shape()[3] == D && (a[7] || a[8] || state.shape()[0] == B) && + (!a[10] || a[10]->shape() == Tensor::Shape{B, Hv, Dv, D})) && + "chunk delta state layout must be pool/B,Hv,Dv,Dk"); for (size_t i = 0; i < a.size(); ++i) { if (!a[i]) continue; - if (a[i]->device() != q.device()) - throw std::invalid_argument("chunk delta device mismatch"); + assert(a[i]->device() == q.device() && "chunk delta device mismatch"); for (auto stride : a[i]->strides()) - if (stride < 0) - throw std::invalid_argument( - "chunk delta negative stride unsupported"); + assert(stride >= 0 && "chunk delta negative stride unsupported"); if (i >= 6 && i <= 8) { - if (a[i]->shape() != Tensor::Shape{i == 6 ? B + 1 : B} || - !a[i]->IsContiguous() || - (a[i]->dtype() != DataType::kInt32 && - a[i]->dtype() != DataType::kInt64)) - throw std::invalid_argument( - "chunk delta requires contiguous I32/I64 cu/slot vectors"); - } else if (i != 3 && i != 4 && a[i]->dtype() != q.dtype()) { - throw std::invalid_argument("chunk delta data dtype mismatch"); + assert((a[i]->shape() == Tensor::Shape{i == 6 ? B + 1 : B} && + a[i]->IsContiguous() && + (a[i]->dtype() == DataType::kInt32 || + a[i]->dtype() == DataType::kInt64)) && + "chunk delta requires contiguous I32/I64 cu/slot vectors"); + } else if (i != 3 && i != 4) { + assert(a[i]->dtype() == q.dtype() && "chunk delta data dtype mismatch"); } - if ((i <= 2 || i == 9) && a[i]->strides()[3] != 1) - throw std::invalid_argument( - "chunk delta innermost data dimension must be contiguous"); - if ((i == 5 || i == 9 || i == 10) && a[i]->HasBroadcastDim()) - throw std::invalid_argument( - "chunk delta broadcast destination unsupported"); + assert((!(i <= 2 || i == 9) || a[i]->strides()[3] == 1) && + "chunk delta innermost data dimension must be contiguous"); + assert((!(i == 5 || i == 9 || i == 10) || !a[i]->HasBroadcastDim()) && + "chunk delta broadcast destination unsupported"); } if (T > static_cast(std::numeric_limits::max()) || state.shape()[0] > diff --git a/src/base/fused_gated_delta_net_gating.h b/src/base/fused_gated_delta_net_gating.h index 9fe9a6c95..afbf6377c 100644 --- a/src/base/fused_gated_delta_net_gating.h +++ b/src/base/fused_gated_delta_net_gating.h @@ -1,7 +1,7 @@ #pragma once #include +#include #include -#include #include "data_type.h" #include "operator.h" @@ -16,25 +16,24 @@ class FusedGatedDeltaNetGating : public Operator { threshold_(threshold), metadata_{A_log, a, b, dt_bias, g, beta_output} { const auto& tensors = metadata_; - if (a.ndim() != 3 || b.shape() != a.shape() || g.shape() != a.shape() || - beta_output.shape() != a.shape() || A_log.ndim() != 1 || - dt_bias.ndim() != 1 || A_log.numel() != a.shape()[2] || - dt_bias.numel() != a.shape()[2]) - throw std::invalid_argument("gating shape mismatch"); - if ((a.dtype() != DataType::kFloat32 && a.dtype() != DataType::kFloat16 && - a.dtype() != DataType::kBFloat16) || - b.dtype() != a.dtype() || A_log.dtype() != a.dtype() || - dt_bias.dtype() != a.dtype() || g.dtype() != DataType::kFloat32 || - beta_output.dtype() != DataType::kFloat32) - throw std::invalid_argument("gating dtype mismatch"); + assert(a.ndim() == 3 && b.shape() == a.shape() && g.shape() == a.shape() && + beta_output.shape() == a.shape() && A_log.ndim() == 1 && + dt_bias.ndim() == 1 && A_log.numel() == a.shape()[2] && + dt_bias.numel() == a.shape()[2] && "gating shape mismatch"); + assert( + (a.dtype() == DataType::kFloat32 || a.dtype() == DataType::kFloat16 || + a.dtype() == DataType::kBFloat16) && + b.dtype() == a.dtype() && A_log.dtype() == a.dtype() && + dt_bias.dtype() == a.dtype() && g.dtype() == DataType::kFloat32 && + beta_output.dtype() == DataType::kFloat32 && "gating dtype mismatch"); device_ = a.device(); total_ = a.numel(); seq_ = a.shape()[1]; hidden_ = a.shape()[2]; for (size_t i = 0; i < tensors.size(); ++i) { - if (tensors[i].device() != device_ || - (i >= 4 && tensors[i].HasBroadcastDim())) - throw std::invalid_argument("gating device/output layout mismatch"); + assert(tensors[i].device() == device_ && + (i < 4 || !tensors[i].HasBroadcastDim()) && + "gating device/output layout mismatch"); } } virtual void operator()(Tensor A_log, Tensor a, Tensor b, Tensor dt_bias, diff --git a/src/base/mrope.h b/src/base/mrope.h index 19028afd6..359b28e46 100644 --- a/src/base/mrope.h +++ b/src/base/mrope.h @@ -1,6 +1,6 @@ #pragma once #include -#include +#include #include "operator.h" namespace infini::ops { @@ -9,26 +9,26 @@ class Mrope : public Operator { Mrope(const Tensor q, const Tensor k, const Tensor cos, const Tensor sin, const Tensor p, int head, int rot, int st, int sh, int sw, bool interleaved, Tensor qo, Tensor ko) { - if (head <= 0 || rot <= 0 || rot > head || rot % 2 || st <= 0 || sh <= 0 || - sw <= 0 || int64_t(st) + sh + sw != rot / 2) - throw std::invalid_argument("Mrope invalid sections"); + assert(head > 0 && rot > 0 && rot <= head && rot % 2 == 0 && st > 0 && + sh > 0 && sw > 0 && int64_t(st) + sh + sw == rot / 2 && + "Mrope invalid sections"); for (const auto& t : {k, cos, sin, p, qo, ko}) - if (t.device() != q.device()) - throw std::invalid_argument("Mrope devices differ"); + assert(t.device() == q.device() && "Mrope devices differ"); for (const auto& t : {q, k, qo, ko}) - if (t.ndim() != 2 || t.stride(1) != 1 || t.size(0) != q.size(0) || - t.dtype() != q.dtype() || t.size(1) % head) - throw std::invalid_argument("Mrope invalid q/k metadata"); - if (q.shape() != qo.shape() || k.shape() != ko.shape() || - (q.dtype() != DataType::kFloat16 && q.dtype() != DataType::kBFloat16 && - q.dtype() != DataType::kFloat32 && q.dtype() != DataType::kFloat64) || - cos.ndim() != 2 || sin.ndim() != 2 || cos.shape() != sin.shape() || - cos.size(1) != rot / 2 || cos.stride(1) != 1 || sin.stride(1) != 1 || - cos.dtype() != q.dtype() || sin.dtype() != q.dtype() || - (p.dtype() != DataType::kInt32 && p.dtype() != DataType::kInt64) || - !((p.ndim() == 1 && p.size(0) == q.size(0)) || - (p.ndim() == 2 && p.size(0) == 3 && p.size(1) == q.size(0)))) - throw std::invalid_argument("Mrope invalid tables/positions/outputs"); + assert(t.ndim() == 2 && t.stride(1) == 1 && t.size(0) == q.size(0) && + t.dtype() == q.dtype() && t.size(1) % head == 0 && + "Mrope invalid q/k metadata"); + assert( + q.shape() == qo.shape() && k.shape() == ko.shape() && + (q.dtype() == DataType::kFloat16 || q.dtype() == DataType::kBFloat16 || + q.dtype() == DataType::kFloat32 || q.dtype() == DataType::kFloat64) && + cos.ndim() == 2 && sin.ndim() == 2 && cos.shape() == sin.shape() && + cos.size(1) == rot / 2 && cos.stride(1) == 1 && sin.stride(1) == 1 && + cos.dtype() == q.dtype() && sin.dtype() == q.dtype() && + (p.dtype() == DataType::kInt32 || p.dtype() == DataType::kInt64) && + ((p.ndim() == 1 && p.size(0) == q.size(0)) || + (p.ndim() == 2 && p.size(0) == 3 && p.size(1) == q.size(0))) && + "Mrope invalid tables/positions/outputs"); } virtual void operator()(const Tensor q, const Tensor k, const Tensor cos, const Tensor sin, const Tensor p, int head, int rot, diff --git a/src/base/mul.h b/src/base/mul.h index 7ac019f73..52f3b5e94 100644 --- a/src/base/mul.h +++ b/src/base/mul.h @@ -1,6 +1,10 @@ #ifndef INFINI_OPS_BASE_MUL_H_ #define INFINI_OPS_BASE_MUL_H_ +#include +#include +#include + #include "operator.h" namespace infini::ops { @@ -32,8 +36,20 @@ class Mul : public Operator { ValidateBroadcast(input, other, out); } + Mul(const Tensor input, const double other, Tensor out) + : output_size_(out.numel()), + input_type_(input.dtype()), + other_type_(input.dtype()), + out_type_(out.dtype()), + scalar_mode_(true) {} + virtual void operator()(const Tensor input, const Tensor other, Tensor out) const = 0; + // Backends implementing only tensor multiplication remain concrete. + virtual void operator()(const Tensor input, const double other, + Tensor out) const { + throw std::invalid_argument("scalar Mul is unsupported by this provider"); + } protected: static Tensor::Strides BroadcastStrides(const Tensor input, @@ -98,6 +114,24 @@ class Mul : public Operator { bool is_other_contiguous_{false}; bool is_out_contiguous_{false}; + bool scalar_mode_{false}; +}; + +template <> +struct CacheKeyBuilder { + detail::CacheKey operator()(const Config& config, const Tensor input, + const Tensor other, Tensor out) const { + return detail::CacheKey::Build(config.implementation_index(), input, other, + out); + } + detail::CacheKey operator()(const Config& config, const Tensor input, + const double alpha, Tensor out) const { + std::uint64_t bits; + static_assert(sizeof(bits) == sizeof(alpha)); + std::memcpy(&bits, &alpha, sizeof(bits)); + return detail::CacheKey::Build(config.implementation_index(), input, bits, + out); + } }; } // namespace infini::ops diff --git a/src/base/mul_scalar.h b/src/base/mul_scalar.h deleted file mode 100644 index 1af4fa879..000000000 --- a/src/base/mul_scalar.h +++ /dev/null @@ -1,47 +0,0 @@ -#ifndef INFINI_OPS_BASE_MUL_SCALAR_H_ -#define INFINI_OPS_BASE_MUL_SCALAR_H_ - -#include -#include -#include - -#include "operator.h" - -namespace infini::ops { - -class MulScalar : public Operator { - public: - static constexpr bool kAutomaticDefault = true; - - MulScalar(const Tensor input, const double alpha, Tensor out) - : input_meta_(input), - out_meta_(out), - device_(input.device()), - alpha_(alpha) {} - - virtual void operator()(const Tensor input, const double alpha, - Tensor out) const = 0; - - protected: - const Tensor input_meta_, out_meta_; - Device device_; - double alpha_; -}; - -// The scalar is an execution argument, not a device tensor. Hash its exact -// bits so a cached instance cannot substitute +0 for -0 or another NaN payload. -template <> -struct CacheKeyBuilder { - detail::CacheKey operator()(const Config& config, const Tensor input, - const double alpha, Tensor out) const { - std::uint64_t bits; - static_assert(sizeof(bits) == sizeof(alpha)); - std::memcpy(&bits, &alpha, sizeof(bits)); - return detail::CacheKey::Build(config.implementation_index(), input, bits, - out); - } -}; - -} // namespace infini::ops - -#endif // INFINI_OPS_BASE_MUL_SCALAR_H_ diff --git a/src/base/pad.h b/src/base/pad.h new file mode 100644 index 000000000..1837b248e --- /dev/null +++ b/src/base/pad.h @@ -0,0 +1,44 @@ +#ifndef INFINI_OPS_BASE_PAD_H_ +#define INFINI_OPS_BASE_PAD_H_ + +#include +#include +#include +#include +#include + +#include "operator.h" + +namespace infini::ops { + +// Preserve the ATen pad.out public signature, including its optional fill. +// Individual providers decide which padding modes and layouts they support. +class Pad : public Operator { + public: + static constexpr bool kAutomaticDefault = true; + + Pad(const Tensor input, const std::vector pad, + const std::string mode, const std::optional value, Tensor out) + : input_meta_(input), + out_meta_(out), + device_(input.device()), + padding_(pad), + mode_(mode), + value_(value.value_or(0.0)) {} + + virtual void operator()(const Tensor input, const std::vector pad, + const std::string mode, + const std::optional value, + Tensor out) const = 0; + + protected: + const Tensor input_meta_, out_meta_; + Device device_; + std::vector padding_; + std::string mode_; + double value_; +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_BASE_PAD_H_ diff --git a/src/base/pad_constant.h b/src/base/pad_constant.h deleted file mode 100644 index ee01d9837..000000000 --- a/src/base/pad_constant.h +++ /dev/null @@ -1,39 +0,0 @@ -#ifndef INFINI_OPS_BASE_PAD_CONSTANT_H_ -#define INFINI_OPS_BASE_PAD_CONSTANT_H_ - -#include -#include -#include - -#include "operator.h" - -namespace infini::ops { - -// Constant padding in Torch's trailing-dimension pair order. This is an -// inference primitive, not the old generic reflect/replicate/circular Pad. -class PadConstant : public Operator { - public: - static constexpr bool kAutomaticDefault = true; - - PadConstant(const Tensor input, const std::vector padding, - const double value, Tensor out) - : input_meta_(input), - out_meta_(out), - device_(input.device()), - padding_(padding), - value_(value) {} - - virtual void operator()(const Tensor input, - const std::vector padding, - const double value, Tensor out) const = 0; - - protected: - const Tensor input_meta_, out_meta_; - Device device_; - std::vector padding_; - double value_; -}; - -} // namespace infini::ops - -#endif // INFINI_OPS_BASE_PAD_CONSTANT_H_ diff --git a/src/base/recurrent_gated_delta_rule.h b/src/base/recurrent_gated_delta_rule.h index 74865fd8c..b9d7cf417 100644 --- a/src/base/recurrent_gated_delta_rule.h +++ b/src/base/recurrent_gated_delta_rule.h @@ -1,5 +1,6 @@ #pragma once #include +#include #include #include #include @@ -61,51 +62,41 @@ class RecurrentGatedDeltaRule : public Operator { static void Validate(const Views& a) { const auto &q = *a[0], &k = *a[1], &v = *a[2], &g = *a[3], &beta = *a[4], &state = *a[5], &out = *a[8]; - if (!floating(q.dtype()) || !floating(g.dtype()) || - beta.dtype() != g.dtype()) - throw std::invalid_argument( - "recurrent delta data/gates support F16/BF16/F32"); - if (q.ndim() != 4 || k.shape() != q.shape() || v.ndim() != 4 || - out.shape() != v.shape() || state.ndim() != 4 || g.ndim() != 3 || - beta.shape() != g.shape()) - throw std::invalid_argument("recurrent delta rank/shape mismatch"); + assert(floating(q.dtype()) && floating(g.dtype()) && + beta.dtype() == g.dtype() && + "recurrent delta data/gates support F16/BF16/F32"); + assert(q.ndim() == 4 && k.shape() == q.shape() && v.ndim() == 4 && + out.shape() == v.shape() && state.ndim() == 4 && g.ndim() == 3 && + beta.shape() == g.shape() && "recurrent delta rank/shape mismatch"); const auto B = q.shape()[0], Hk = q.shape()[2], Hv = v.shape()[2], D = q.shape()[3], Dv = v.shape()[3], T = q.shape()[1]; - if (v.shape()[0] != B || v.shape()[1] != T || !Hk || !Hv || Hv % Hk || !D || - !Dv || g.shape() != Tensor::Shape{B, T, Hv}) - throw std::invalid_argument( - "recurrent delta data/gate shape or head grouping mismatch"); - if (a[7].has_value() == a[9].has_value()) - throw std::invalid_argument( - "recurrent delta requires exactly one final destination mode"); - if (state.shape()[1] != Hv || state.shape()[2] != Dv || - state.shape()[3] != D || (!a[6] && !a[7] && state.shape()[0] != B) || - (a[9] && a[9]->shape() != Tensor::Shape{B, Hv, Dv, D})) - throw std::invalid_argument( - "recurrent delta state layout must be pool/B,Hv,Dv,Dk"); + assert((v.shape()[0] == B && v.shape()[1] == T && Hk && Hv && + Hv % Hk == 0 && D && Dv && g.shape() == Tensor::Shape{B, T, Hv}) && + "recurrent delta data/gate shape or head grouping mismatch"); + assert(a[7].has_value() != a[9].has_value() && + "recurrent delta requires exactly one final destination mode"); + assert((state.shape()[1] == Hv && state.shape()[2] == Dv && + state.shape()[3] == D && (a[6] || a[7] || state.shape()[0] == B) && + (!a[9] || a[9]->shape() == Tensor::Shape{B, Hv, Dv, D})) && + "recurrent delta state layout must be pool/B,Hv,Dv,Dk"); for (size_t i = 0; i < a.size(); ++i) { if (!a[i]) continue; - if (a[i]->device() != q.device()) - throw std::invalid_argument("recurrent delta device mismatch"); + assert(a[i]->device() == q.device() && "recurrent delta device mismatch"); for (auto stride : a[i]->strides()) - if (stride < 0) - throw std::invalid_argument( - "recurrent delta negative stride unsupported"); + assert(stride >= 0 && "recurrent delta negative stride unsupported"); if (i == 6 || i == 7) { - if (a[i]->shape() != Tensor::Shape{B} || !a[i]->IsContiguous() || - (a[i]->dtype() != DataType::kInt32 && - a[i]->dtype() != DataType::kInt64)) - throw std::invalid_argument( - "recurrent delta indices require contiguous I32/I64 B-vector"); - } else if (i != 3 && i != 4 && a[i]->dtype() != q.dtype()) { - throw std::invalid_argument("recurrent delta data dtype mismatch"); + assert((a[i]->shape() == Tensor::Shape{B} && a[i]->IsContiguous() && + (a[i]->dtype() == DataType::kInt32 || + a[i]->dtype() == DataType::kInt64)) && + "recurrent delta indices require contiguous I32/I64 B-vector"); + } else if (i != 3 && i != 4) { + assert(a[i]->dtype() == q.dtype() && + "recurrent delta data dtype mismatch"); } - if ((i <= 2 || i == 8) && a[i]->strides()[3] != 1) - throw std::invalid_argument( - "recurrent delta innermost data dimension must be contiguous"); - if ((i == 5 || i == 8 || i == 9) && a[i]->HasBroadcastDim()) - throw std::invalid_argument( - "recurrent delta broadcast destination unsupported"); + assert((!(i <= 2 || i == 8) || a[i]->strides()[3] == 1) && + "recurrent delta innermost data dimension must be contiguous"); + assert((!(i == 5 || i == 8 || i == 9) || !a[i]->HasBroadcastDim()) && + "recurrent delta broadcast destination unsupported"); } if (state.shape()[0] > static_cast(std::numeric_limits::max())) diff --git a/src/base/separate_table_rope.h b/src/base/separate_table_rope.h index 02e4f861a..a04bfd082 100644 --- a/src/base/separate_table_rope.h +++ b/src/base/separate_table_rope.h @@ -2,9 +2,9 @@ #define INFINI_OPS_BASE_SEPARATE_TABLE_ROPE_H_ #include #include +#include #include #include -#include #include "data_type.h" #include "operator.h" @@ -79,9 +79,8 @@ class SeparateTableRope : public Operator { SeparateTableRope(const Tensor input, const Tensor positions, const Tensor sin_table, const Tensor cos_table, bool is_neox, Tensor out) { - if (!supports(input, positions, sin_table, cos_table, is_neox, out)) - throw std::invalid_argument( - "SeparateTableRope: unsupported metadata/alias"); + assert(supports(input, positions, sin_table, cos_table, is_neox, out) && + "SeparateTableRope: unsupported metadata/alias"); } virtual void operator()(const Tensor input, const Tensor positions, const Tensor sin_table, const Tensor cos_table, diff --git a/src/native/cpu/ops/mul/mul.h b/src/native/cpu/ops/mul/mul.h index c3039b005..385f4e283 100644 --- a/src/native/cpu/ops/mul/mul.h +++ b/src/native/cpu/ops/mul/mul.h @@ -6,6 +6,7 @@ #include "base/mul.h" #include "common/generic_utils.h" #include "native/cpu/caster_.h" +#include "native/cpu/ops/mul/scalar.h" namespace infini::ops { @@ -15,9 +16,23 @@ class Operator : public Mul, public: Operator(const Tensor input, const Tensor other, Tensor out) : Mul{input, other, out} {} + Operator(const Tensor input, const double alpha, Tensor out) + : Mul{input, alpha, out}, scalar_(std::in_place, input, alpha, out) {} + + static bool supports(const Tensor input, const double alpha, Tensor out) { + return CpuScalarMul::supports(input, alpha, out); + } + void operator()(const Tensor input, const double alpha, + Tensor out) const override { + if (!scalar_) + throw std::invalid_argument("Mul prepared for tensor operands"); + scalar_->Run(input, alpha, out); + } void operator()(const Tensor input, const Tensor other, Tensor out) const override { + if (scalar_mode_) + throw std::invalid_argument("Mul prepared for a scalar operand"); DispatchFunc( out_type_, [&](auto tag) { @@ -28,6 +43,7 @@ class Operator : public Mul, } private: + std::optional scalar_; template void Compute(const Tensor input, const Tensor other, Tensor out) const { using ComputeType = std::conditional_t || diff --git a/src/native/cpu/ops/mul_scalar/mul_scalar.h b/src/native/cpu/ops/mul/scalar.h similarity index 79% rename from src/native/cpu/ops/mul_scalar/mul_scalar.h rename to src/native/cpu/ops/mul/scalar.h index 8c7d592be..03e99aea3 100644 --- a/src/native/cpu/ops/mul_scalar/mul_scalar.h +++ b/src/native/cpu/ops/mul/scalar.h @@ -1,22 +1,21 @@ -#ifndef INFINI_OPS_CPU_MUL_SCALAR_MUL_SCALAR_H_ -#define INFINI_OPS_CPU_MUL_SCALAR_MUL_SCALAR_H_ +#ifndef INFINI_OPS_CPU_MUL_SCALAR_H_ +#define INFINI_OPS_CPU_MUL_SCALAR_H_ #include #include #include #include +#include #include #include -#include "base/mul_scalar.h" +#include "base/mul.h" #include "native/cpu/caster_.h" #include "native/cpu/dense_view.h" namespace infini::ops { -template <> -class Operator : public MulScalar, - Caster { +class CpuScalarMul : Caster { public: static bool supports(const Tensor input, const double, Tensor out) { if (input.device().type() != Device::Type::kCpu || @@ -26,20 +25,19 @@ class Operator : public MulScalar, return cpu_dense_view::DisjointOrSameContiguous(input, out); } - Operator(const Tensor input, const double alpha, Tensor out) - : MulScalar{Checked(input, alpha, out), alpha, out} { + CpuScalarMul(const Tensor input, const double alpha, Tensor out) + : input_meta_(Checked(input, alpha, out)), out_meta_(out), alpha_(alpha) { std::size_t bytes = 0; cpu_dense_view::ContiguousBytes(input, &bytes); count_ = bytes / input.element_size(); } - void operator()(const Tensor input, const double alpha, - Tensor out) const override { + void Run(const Tensor input, const double alpha, Tensor out) const { if (!std::equal_to{}(input, input_meta_) || !std::equal_to{}(out, out_meta_) || std::memcmp(&alpha, &alpha_, sizeof(alpha))) throw std::invalid_argument( - "CPU MulScalar: call metadata or alpha changed after Make"); + "CPU scalar Mul: call metadata or alpha changed after Make"); const auto bytes = count_ * input_meta_.element_size(); const auto ip = reinterpret_cast(input.data()); const auto op = reinterpret_cast(out.data()); @@ -48,7 +46,7 @@ class Operator : public MulScalar, (bytes && ip < op + bytes && op < ip + bytes && !(ip == op && input_meta_.strides() == out_meta_.strides()))) throw std::invalid_argument( - "CPU MulScalar: call pointer or alias changed after Make"); + "CPU scalar Mul: call pointer or alias changed after Make"); if (count_ == 0) return; // CPU scalar arithmetic retains float opmath precision even for half/BF16. const float multiplier = FloatAlpha(alpha_); @@ -73,7 +71,7 @@ class Operator : public MulScalar, } } }, - "CPU MulScalar"); + "CPU scalar Mul"); } private: @@ -92,12 +90,14 @@ class Operator : public MulScalar, static Tensor Checked(const Tensor input, const double alpha, Tensor out) { if (!supports(input, alpha, out)) throw std::invalid_argument( - "CPU MulScalar: unsupported dtype/device/layout or alias"); + "CPU scalar Mul: unsupported dtype/device/layout or alias"); return input; } std::size_t count_ = 0; + const Tensor input_meta_, out_meta_; + double alpha_; }; } // namespace infini::ops -#endif // INFINI_OPS_CPU_MUL_SCALAR_MUL_SCALAR_H_ +#endif // INFINI_OPS_CPU_MUL_SCALAR_H_ diff --git a/src/native/cuda/nvidia/ops/chunk_gated_delta_rule/kernel.h b/src/native/cuda/nvidia/ops/chunk_gated_delta_rule/kernel.h index 79fdc30c2..000fa36ca 100644 --- a/src/native/cuda/nvidia/ops/chunk_gated_delta_rule/kernel.h +++ b/src/native/cuda/nvidia/ops/chunk_gated_delta_rule/kernel.h @@ -12,6 +12,9 @@ namespace infini::ops { namespace chunk_native_detail { struct Layout { + // Rows follow Views: q, k, v, g, beta, initial_state, cu_seqlens, + // initial_state_indices, final_state_indices, out, final_state (0..10). + // Each row stores that tensor's axis strides, up to the four state axes. ptrdiff_t s[11][4]{}; bool cu_i64, initial_i64, final_i64, norm, indexed, packed; size_t T, pool, Hk, ratio; @@ -22,6 +25,8 @@ __global__ void kernel(float* workspace, const Data* q, const Data* k, Data* state, const void* cu, const void* initial_indices, const void* final_indices, Data* out, Data* final_state, Layout p) { + // q/k/v/out have a contiguous last axis, so only their first three strides + // are passed. Initial/final state uses all axes of [pool,Hv,Dv,Dk]. chunkGatedDeltaRuleRecurrentKernel( workspace, out, state, final_state, q, k, v, g, beta, cu, initial_indices, final_indices, p.cu_i64, p.initial_i64, p.final_i64, p.norm, p.packed, diff --git a/src/native/cuda/nvidia/ops/flash_attn_varlen_func/kernel.cc b/src/native/cuda/nvidia/ops/flash_attn_varlen_func/kernel.cc index a45ddfd0f..7486c1e58 100644 --- a/src/native/cuda/nvidia/ops/flash_attn_varlen_func/kernel.cc +++ b/src/native/cuda/nvidia/ops/flash_attn_varlen_func/kernel.cc @@ -6,10 +6,10 @@ #include #include -#include "native/cuda/nvidia/ops/flash_attn_with_kvcache/infini_flash_cuda.h" #include "native/cuda/nvidia/ops/flash_common.h" #include "native/cuda/nvidia/ops/flash_head_dim.h" #include "native/cuda/nvidia/ops/flash_workspace.h" +#include "native/cuda/nvidia/ops/infini_flash_cuda.h" namespace infini::ops { namespace { diff --git a/src/native/cuda/nvidia/ops/flash_attn_with_kvcache/kernel.cc b/src/native/cuda/nvidia/ops/flash_attn_with_kvcache/kernel.cc index 7297772eb..c278c79d0 100644 --- a/src/native/cuda/nvidia/ops/flash_attn_with_kvcache/kernel.cc +++ b/src/native/cuda/nvidia/ops/flash_attn_with_kvcache/kernel.cc @@ -5,10 +5,10 @@ #include #include -#include "infini_flash_cuda.h" #include "native/cuda/nvidia/ops/flash_common.h" #include "native/cuda/nvidia/ops/flash_head_dim.h" #include "native/cuda/nvidia/ops/flash_workspace.h" +#include "native/cuda/nvidia/ops/infini_flash_cuda.h" #include "num_splits.h" namespace infini::ops { diff --git a/src/native/cuda/nvidia/ops/flash_common.cc b/src/native/cuda/nvidia/ops/flash_common.cc index 50c558e4e..d1921a423 100644 --- a/src/native/cuda/nvidia/ops/flash_common.cc +++ b/src/native/cuda/nvidia/ops/flash_common.cc @@ -6,7 +6,7 @@ #include #include -#include "native/cuda/nvidia/ops/flash_attn_with_kvcache/infini_flash_cuda.h" +#include "native/cuda/nvidia/ops/infini_flash_cuda.h" namespace infini::ops::native_flash { diff --git a/src/native/cuda/nvidia/ops/flash_attn_with_kvcache/infini_flash_cuda.h b/src/native/cuda/nvidia/ops/infini_flash_cuda.h similarity index 80% rename from src/native/cuda/nvidia/ops/flash_attn_with_kvcache/infini_flash_cuda.h rename to src/native/cuda/nvidia/ops/infini_flash_cuda.h index 68e1a0452..c72ab625b 100644 --- a/src/native/cuda/nvidia/ops/flash_attn_with_kvcache/infini_flash_cuda.h +++ b/src/native/cuda/nvidia/ops/infini_flash_cuda.h @@ -1,5 +1,8 @@ #pragma once +// Private CUDA checks shared by both native FlashAttention providers and the +// pinned forward kernels; this header does not depend on either provider. + #include #include diff --git a/src/native/cuda/nvidia/ops/mul/kernel.h b/src/native/cuda/nvidia/ops/mul/kernel.h index 579e5e093..c62f55b44 100644 --- a/src/native/cuda/nvidia/ops/mul/kernel.h +++ b/src/native/cuda/nvidia/ops/mul/kernel.h @@ -4,6 +4,7 @@ #include #include "native/cuda/nvidia/caster.cuh" +#include "native/cuda/nvidia/ops/mul/scalar.h" #include "native/cuda/nvidia/runtime_.h" #include "native/cuda/ops/mul/kernel.h" @@ -14,8 +15,25 @@ class Operator : public CudaMul> { public: using CudaMul>::CudaMul; + using CudaMul>::operator(); + Operator(const Tensor input, const double alpha, Tensor out) + : CudaMul>{input, alpha, out}, + scalar_(std::in_place, input, alpha, out) {} + static bool supports(const Tensor input, const double alpha, Tensor out) { + return NvidiaScalarMul::supports(input, alpha, out); + } + void operator()(const Tensor input, const double alpha, + Tensor out) const override { + if (!scalar_) + throw std::invalid_argument("Mul prepared for tensor operands"); + scalar_->Run(input, alpha, out, + static_cast(stream_ ? stream_ : 0)); + } // All dispatched types use kernels and constructor-owned metadata on replay. bool capture_safe() const override { return true; } + + private: + std::optional scalar_; }; } // namespace infini::ops diff --git a/src/native/cuda/nvidia/ops/mul_scalar/kernel.cu b/src/native/cuda/nvidia/ops/mul/scalar.cu similarity index 82% rename from src/native/cuda/nvidia/ops/mul_scalar/kernel.cu rename to src/native/cuda/nvidia/ops/mul/scalar.cu index 2d2d28604..6deaa0b74 100644 --- a/src/native/cuda/nvidia/ops/mul_scalar/kernel.cu +++ b/src/native/cuda/nvidia/ops/mul/scalar.cu @@ -1,11 +1,11 @@ -#include "native/cuda/nvidia/ops/mul_scalar/kernel.h" - #include #include #include #include +#include "native/cuda/nvidia/ops/mul/scalar.h" + namespace infini::ops::nvidia_mul_scalar_detail { namespace { template @@ -19,15 +19,16 @@ __device__ __forceinline__ __half Multiply(__half value, __half alpha) { } template <> -__device__ __forceinline__ __nv_bfloat16 Multiply( - __nv_bfloat16 value, __nv_bfloat16 alpha) { +__device__ __forceinline__ __nv_bfloat16 Multiply(__nv_bfloat16 value, + __nv_bfloat16 alpha) { return __hmul(value, alpha); } template __global__ void MulScalarContiguousKernel(const T* input, T* out, T alpha, std::size_t count) { - const auto start = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const auto start = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; const auto step = static_cast(gridDim.x) * blockDim.x; for (std::size_t linear = start; linear < count; linear += step) { out[linear] = Multiply(input[linear], alpha); @@ -37,7 +38,8 @@ __global__ void MulScalarContiguousKernel(const T* input, T* out, T alpha, template __global__ void MulScalarStridedKernel(const T* input, T* out, Layout layout, T alpha, std::size_t count) { - const auto start = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const auto start = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; const auto step = static_cast(gridDim.x) * blockDim.x; for (std::size_t linear = start; linear < count; linear += step) { std::size_t remaining = linear; @@ -61,8 +63,7 @@ void Submit(const void* input, void* out, const Layout* layout, T alpha, std::min((count - 1) / kThreads + 1, 65535)); if (!layout) MulScalarContiguousKernel<<>>( - static_cast(input), static_cast(out), alpha, - count); + static_cast(input), static_cast(out), alpha, count); else MulScalarStridedKernel<<>>( static_cast(input), static_cast(out), *layout, alpha, @@ -79,9 +80,9 @@ void Launch(const void* input, void* out, DataType dtype, double alpha, __float2half_rn(static_cast(alpha)), count, stream); break; case DataType::kBFloat16: - Submit<__nv_bfloat16>( - input, out, layout, - __float2bfloat16_rn(static_cast(alpha)), count, stream); + Submit<__nv_bfloat16>(input, out, layout, + __float2bfloat16_rn(static_cast(alpha)), + count, stream); break; case DataType::kFloat32: Submit(input, out, layout, static_cast(alpha), count, @@ -91,7 +92,7 @@ void Launch(const void* input, void* out, DataType dtype, double alpha, Submit(input, out, layout, alpha, count, stream); break; default: - throw std::invalid_argument("NVIDIA MulScalar unsupported dtype"); + throw std::invalid_argument("NVIDIA scalar Mul unsupported dtype"); } } } // namespace infini::ops::nvidia_mul_scalar_detail diff --git a/src/native/cuda/nvidia/ops/mul_scalar/kernel.h b/src/native/cuda/nvidia/ops/mul/scalar.h similarity index 86% rename from src/native/cuda/nvidia/ops/mul_scalar/kernel.h rename to src/native/cuda/nvidia/ops/mul/scalar.h index 1b90f6041..c96553aed 100644 --- a/src/native/cuda/nvidia/ops/mul_scalar/kernel.h +++ b/src/native/cuda/nvidia/ops/mul/scalar.h @@ -1,14 +1,15 @@ -#ifndef INFINI_OPS_NVIDIA_MUL_SCALAR_KERNEL_H_ -#define INFINI_OPS_NVIDIA_MUL_SCALAR_KERNEL_H_ +#ifndef INFINI_OPS_NVIDIA_MUL_SCALAR_H_ +#define INFINI_OPS_NVIDIA_MUL_SCALAR_H_ #include #include #include +#include #include #include -#include "base/mul_scalar.h" +#include "base/mul.h" #include "native/cuda/nvidia/ops/binary_view_capture.h" #include "native/cuda/nvidia/runtime_.h" @@ -26,8 +27,7 @@ void Launch(const void* input, void* out, DataType dtype, double alpha, } // namespace infini::ops::nvidia_mul_scalar_detail namespace infini::ops { -template <> -class Operator : public MulScalar { +class NvidiaScalarMul { public: static bool supports(const Tensor input, const double, Tensor out) { using namespace nvidia_binary_view; @@ -88,8 +88,11 @@ class Operator : public MulScalar { return !Overlaps(input, input_bytes, out, out_bytes) || exact_inplace; } - Operator(const Tensor input, const double alpha, Tensor out) - : MulScalar(Checked(input, alpha, out), alpha, out) { + NvidiaScalarMul(const Tensor input, const double alpha, Tensor out) + : input_meta_(Checked(input, alpha, out)), + out_meta_(out), + device_(input.device()), + alpha_(alpha) { nvidia_binary_view::SpanBytes(input, &input_span_bytes_); nvidia_binary_view::SpanBytes(out, &out_span_bytes_); count_ = out.numel(); @@ -102,15 +105,13 @@ class Operator : public MulScalar { } } - bool capture_safe() const override { return true; } - - void operator()(const Tensor input, const double alpha, - Tensor out) const override { + void Run(const Tensor input, const double alpha, Tensor out, + cudaStream_t stream) const { if (!std::equal_to{}(input, input_meta_) || !std::equal_to{}(out, out_meta_) || std::memcmp(&alpha, &alpha_, sizeof(alpha))) throw std::invalid_argument( - "NVIDIA MulScalar metadata or alpha changed after Make"); + "NVIDIA scalar Mul metadata or alpha changed after Make"); if (count_ == 0) return; int current_device = -1; if (cudaGetDevice(¤t_device) != cudaSuccess || @@ -122,22 +123,24 @@ class Operator : public MulScalar { !(input.data() == out.data() && input_meta_.strides() == out_meta_.strides()))) throw std::invalid_argument( - "NVIDIA MulScalar current device, pointer or alias changed after " + "NVIDIA scalar Mul current device, pointer or alias changed after " "Make"); nvidia_mul_scalar_detail::Launch( input.data(), out.data(), input_meta_.dtype(), alpha_, - layout_.contiguous ? nullptr : &layout_, count_, - static_cast(stream_ ? stream_ : 0)); + layout_.contiguous ? nullptr : &layout_, count_, stream); } private: static Tensor Checked(const Tensor input, const double alpha, Tensor out) { if (!supports(input, alpha, out)) throw std::invalid_argument( - "unsupported NVIDIA MulScalar metadata, device or alias"); + "unsupported NVIDIA scalar Mul metadata, device or alias"); return input; } nvidia_mul_scalar_detail::Layout layout_{}; + const Tensor input_meta_, out_meta_; + Device device_; + double alpha_; std::size_t input_span_bytes_ = 0, out_span_bytes_ = 0, count_ = 0; }; } // namespace infini::ops diff --git a/src/native/cuda/nvidia/ops/pad_constant/kernel.cu b/src/native/cuda/nvidia/ops/pad/kernel.cu similarity index 71% rename from src/native/cuda/nvidia/ops/pad_constant/kernel.cu rename to src/native/cuda/nvidia/ops/pad/kernel.cu index 3d7f52ffa..0a4f99d63 100644 --- a/src/native/cuda/nvidia/ops/pad_constant/kernel.cu +++ b/src/native/cuda/nvidia/ops/pad/kernel.cu @@ -1,17 +1,18 @@ -#include "native/cuda/nvidia/ops/pad_constant/kernel.h" - #include #include #include #include -namespace infini::ops::nvidia_pad_constant_detail { +#include "native/cuda/nvidia/ops/pad/kernel.h" + +namespace infini::ops::nvidia_pad_detail { namespace { template -__global__ void PadConstantKernel(const T* input, T* out, Layout layout, - T fill, std::size_t count) { - const auto start = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; +__global__ void PadKernel(const T* input, T* out, Layout layout, T fill, + std::size_t count) { + const auto start = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; const auto step = static_cast(gridDim.x) * blockDim.x; for (std::size_t linear = start; linear < count; linear += step) { std::size_t remaining = linear; @@ -23,8 +24,10 @@ __global__ void PadConstantKernel(const T* input, T* out, Layout layout, remaining /= static_cast(layout.output_shape[d]); output_offset += coord * layout.output_strides[d]; const auto source = coord - layout.left_padding[d]; - if (source < 0 || source >= layout.input_shape[d]) inside = false; - else input_offset += source * layout.input_strides[d]; + if (source < 0 || source >= layout.input_shape[d]) + inside = false; + else + input_offset += source * layout.input_strides[d]; } out[output_offset] = inside ? input[input_offset] : fill; } @@ -36,7 +39,7 @@ void Submit(const void* input, void* out, Layout layout, T value, constexpr unsigned kThreads = 256; const auto blocks = static_cast( std::min((count - 1) / kThreads + 1, 65535)); - PadConstantKernel<<>>( + PadKernel<<>>( static_cast(input), static_cast(out), layout, value, count); } } // namespace @@ -47,8 +50,8 @@ void Launch(const void* input, void* out, DataType dtype, double value, switch (dtype) { case DataType::kFloat16: Submit<__half>(input, out, layout, - __float2half_rn(static_cast(value)), - output_count, stream); + __float2half_rn(static_cast(value)), output_count, + stream); break; case DataType::kBFloat16: Submit<__nv_bfloat16>(input, out, layout, @@ -56,16 +59,17 @@ void Launch(const void* input, void* out, DataType dtype, double value, output_count, stream); break; case DataType::kFloat32: - Submit(input, out, layout, static_cast(value), - output_count, stream); + Submit(input, out, layout, static_cast(value), output_count, + stream); break; case DataType::kFloat64: Submit(input, out, layout, value, output_count, stream); break; default: - throw std::invalid_argument("NVIDIA PadConstant unsupported dtype"); + throw std::invalid_argument("NVIDIA Pad unsupported dtype"); } const auto status = cudaPeekAtLastError(); - if (status != cudaSuccess) throw std::runtime_error(cudaGetErrorString(status)); + if (status != cudaSuccess) + throw std::runtime_error(cudaGetErrorString(status)); } -} // namespace infini::ops::nvidia_pad_constant_detail +} // namespace infini::ops::nvidia_pad_detail diff --git a/src/native/cuda/nvidia/ops/pad_constant/kernel.h b/src/native/cuda/nvidia/ops/pad/kernel.h similarity index 74% rename from src/native/cuda/nvidia/ops/pad_constant/kernel.h rename to src/native/cuda/nvidia/ops/pad/kernel.h index 55eac7a43..ad6c54733 100644 --- a/src/native/cuda/nvidia/ops/pad_constant/kernel.h +++ b/src/native/cuda/nvidia/ops/pad/kernel.h @@ -1,5 +1,5 @@ -#ifndef INFINI_OPS_NVIDIA_PAD_CONSTANT_KERNEL_H_ -#define INFINI_OPS_NVIDIA_PAD_CONSTANT_KERNEL_H_ +#ifndef INFINI_OPS_NVIDIA_PAD_KERNEL_H_ +#define INFINI_OPS_NVIDIA_PAD_KERNEL_H_ #include @@ -9,11 +9,11 @@ #include #include -#include "base/pad_constant.h" +#include "base/pad.h" #include "native/cuda/nvidia/ops/binary_view_capture.h" #include "native/cuda/nvidia/runtime_.h" -namespace infini::ops::nvidia_pad_constant_detail { +namespace infini::ops::nvidia_pad_detail { struct Layout { std::int64_t ndim = 0; std::int64_t input_shape[16]{}; @@ -25,16 +25,18 @@ struct Layout { void Launch(const void* input, void* out, DataType dtype, double value, Layout layout, std::size_t output_count, cudaStream_t stream); -} // namespace infini::ops::nvidia_pad_constant_detail +} // namespace infini::ops::nvidia_pad_detail namespace infini::ops { template <> -class Operator : public PadConstant { +class Operator : public Pad { public: - static bool supports(const Tensor input, - const std::vector padding, const double, + static bool supports(const Tensor input, const std::vector pad, + const std::string mode, const std::optional, Tensor out) { using namespace nvidia_binary_view; + if (mode != "constant") return false; + const auto& padding = pad; const auto ndim = input.shape().size(); if (input.device().type() != Device::Type::kNvidia || input.device() != out.device() || input.dtype() != out.dtype() || @@ -90,10 +92,11 @@ class Operator : public PadConstant { return !Overlaps(input, input_bytes, out, output_bytes) || exact_inplace; } - Operator(const Tensor input, const std::vector padding, - const double value, Tensor out) - : PadConstant(Checked(input, padding, value, out), padding, value, out) { - BuildLayout(input, padding, out); + Operator(const Tensor input, const std::vector pad, + const std::string mode, const std::optional value, + Tensor out) + : Pad(Checked(input, pad, mode, value, out), pad, mode, value, out) { + BuildLayout(input, pad, out); nvidia_binary_view::SpanBytes(input, &input_bytes_); nvidia_binary_view::SpanBytes(out, &out_bytes_); exact_inplace_layout_ = @@ -102,36 +105,38 @@ class Operator : public PadConstant { bool capture_safe() const override { return true; } - void operator()(const Tensor input, const std::vector padding, - const double value, Tensor out) const override { + void operator()(const Tensor input, const std::vector pad, + const std::string mode, const std::optional value, + Tensor out) const override { + const double fill = value.value_or(0.0); if (!std::equal_to{}(input, input_meta_) || - !std::equal_to{}(out, out_meta_) || padding != padding_ || - std::memcmp(&value, &value_, sizeof(value))) + !std::equal_to{}(out, out_meta_) || pad != padding_ || + mode != mode_ || std::memcmp(&fill, &value_, sizeof(fill))) throw std::invalid_argument( - "NVIDIA PadConstant metadata or alias changed after Make"); + "NVIDIA Pad metadata or alias changed after Make"); int current = -1; if (cudaGetDevice(¤t) != cudaSuccess || current != device_.index() || (input_bytes_ && !input.data()) || (out_bytes_ && !out.data()) || (nvidia_binary_view::Overlaps(input, input_bytes_, out, out_bytes_) && !(input.data() == out.data() && exact_inplace_layout_))) throw std::invalid_argument( - "NVIDIA PadConstant current device, pointer or alias changed after " + "NVIDIA Pad current device, pointer or alias changed after " "Make"); if (out.numel() == 0 || (input.data() == out.data() && exact_inplace_layout_)) return; - nvidia_pad_constant_detail::Launch( - input.data(), out.data(), input_meta_.dtype(), value_, layout_, - out.numel(), static_cast(stream_ ? stream_ : 0)); + nvidia_pad_detail::Launch(input.data(), out.data(), input_meta_.dtype(), + value_, layout_, out.numel(), + static_cast(stream_ ? stream_ : 0)); } private: - static Tensor Checked(const Tensor input, - const std::vector& padding, double value, + static Tensor Checked(const Tensor input, const std::vector& pad, + const std::string& mode, std::optional value, Tensor out) { - if (!supports(input, padding, value, out)) + if (!supports(input, pad, mode, value, out)) throw std::invalid_argument( - "unsupported NVIDIA PadConstant metadata, device or alias"); + "unsupported NVIDIA Pad metadata, mode, device or alias"); return input; } void BuildLayout(const Tensor input, const std::vector& padding, @@ -147,10 +152,10 @@ class Operator : public PadConstant { pair < padding.size() / 2 ? padding[2 * pair] : 0; } } - nvidia_pad_constant_detail::Layout layout_{}; + nvidia_pad_detail::Layout layout_{}; std::size_t input_bytes_ = 0, out_bytes_ = 0; bool exact_inplace_layout_ = false; }; } // namespace infini::ops -#endif // INFINI_OPS_NVIDIA_PAD_CONSTANT_KERNEL_H_ +#endif // INFINI_OPS_NVIDIA_PAD_KERNEL_H_ diff --git a/src/native/cuda/ops/mul/kernel.h b/src/native/cuda/ops/mul/kernel.h index 032d44022..7cb08d0be 100644 --- a/src/native/cuda/ops/mul/kernel.h +++ b/src/native/cuda/ops/mul/kernel.h @@ -63,6 +63,8 @@ class CudaMul : public Mul { void operator()(const Tensor input, const Tensor other, Tensor out) const override { + if (scalar_mode_) + throw std::invalid_argument("Mul prepared for a scalar operand"); if (output_size_ == 0) { return; } @@ -93,6 +95,10 @@ class CudaMul : public Mul { "CudaMul::operator()"); } + protected: + CudaMul(const Tensor input, const double alpha, Tensor out) + : Mul{input, alpha, out}, block_size_(0) {} + private: const int block_size_ = RuntimeUtils::GetOptimalBlockSize(); From f998e4d6b445e29a950d048d0bc31f494ca7cce9 Mon Sep 17 00:00:00 2001 From: whjthu Date: Tue, 20 Oct 2026 21:47:04 +0800 Subject: [PATCH 25/25] fix: relax state pool and rotary metadata constraints --- src/base/causal_conv1d.h | 6 ++---- src/base/chunk_gated_delta_rule.h | 5 +++-- src/base/mrope.h | 4 ++-- src/base/recurrent_gated_delta_rule.h | 5 +++-- src/native/cpu/ops/separate_table_rope/kernel.h | 5 +---- 5 files changed, 11 insertions(+), 14 deletions(-) diff --git a/src/base/causal_conv1d.h b/src/base/causal_conv1d.h index ff723e2a8..309e87559 100644 --- a/src/base/causal_conv1d.h +++ b/src/base/causal_conv1d.h @@ -74,7 +74,7 @@ class CausalConv1d : public Operator { "causal conv packed sequence shape mismatch"); requests = v[4]->numel() - 1; } - assert((v[5] || v[6] || state.shape()[0] == requests) && + assert((v[5] || state.shape()[0] >= requests) && "causal conv state batch mismatch"); assert((!v[8] || v[8]->shape() == Tensor::Shape{requests, channels, state.shape()[2]}) && @@ -93,9 +93,7 @@ class CausalConv1d : public Operator { } else { assert(v[i]->dtype() == dtype && "causal conv data dtype mismatch"); } - // Preserve ordinary positive strided data and reject broadcast - // destinations. - if (i == 2 || i == 7 || i == 8) + if ((i == 2 && v[6]) || i == 7 || i == 8) for (size_t axis = 0; axis < v[i]->ndim(); ++axis) assert((v[i]->shape()[axis] == 1 || v[i]->strides()[axis] != 0) && "causal conv broadcast destination unsupported"); diff --git a/src/base/chunk_gated_delta_rule.h b/src/base/chunk_gated_delta_rule.h index c91b7dadb..ab2a189c2 100644 --- a/src/base/chunk_gated_delta_rule.h +++ b/src/base/chunk_gated_delta_rule.h @@ -88,7 +88,7 @@ class ChunkGatedDeltaRule : public Operator { assert(a[8].has_value() != a[10].has_value() && "chunk delta requires exactly one final destination mode"); assert((state.shape()[1] == Hv && state.shape()[2] == Dv && - state.shape()[3] == D && (a[7] || a[8] || state.shape()[0] == B) && + state.shape()[3] == D && (a[7] || state.shape()[0] >= B) && (!a[10] || a[10]->shape() == Tensor::Shape{B, Hv, Dv, D})) && "chunk delta state layout must be pool/B,Hv,Dv,Dk"); for (size_t i = 0; i < a.size(); ++i) { @@ -107,7 +107,8 @@ class ChunkGatedDeltaRule : public Operator { } assert((!(i <= 2 || i == 9) || a[i]->strides()[3] == 1) && "chunk delta innermost data dimension must be contiguous"); - assert((!(i == 5 || i == 9 || i == 10) || !a[i]->HasBroadcastDim()) && + assert((!((i == 5 && a[8]) || i == 9 || i == 10) || + !a[i]->HasBroadcastDim()) && "chunk delta broadcast destination unsupported"); } if (T > static_cast(std::numeric_limits::max()) || diff --git a/src/base/mrope.h b/src/base/mrope.h index 359b28e46..36fc3abe3 100644 --- a/src/base/mrope.h +++ b/src/base/mrope.h @@ -9,8 +9,8 @@ class Mrope : public Operator { Mrope(const Tensor q, const Tensor k, const Tensor cos, const Tensor sin, const Tensor p, int head, int rot, int st, int sh, int sw, bool interleaved, Tensor qo, Tensor ko) { - assert(head > 0 && rot > 0 && rot <= head && rot % 2 == 0 && st > 0 && - sh > 0 && sw > 0 && int64_t(st) + sh + sw == rot / 2 && + assert(head > 0 && rot >= 0 && rot <= head && rot % 2 == 0 && st >= 0 && + sh >= 0 && sw >= 0 && int64_t(st) + sh + sw == rot / 2 && "Mrope invalid sections"); for (const auto& t : {k, cos, sin, p, qo, ko}) assert(t.device() == q.device() && "Mrope devices differ"); diff --git a/src/base/recurrent_gated_delta_rule.h b/src/base/recurrent_gated_delta_rule.h index b9d7cf417..0df174b80 100644 --- a/src/base/recurrent_gated_delta_rule.h +++ b/src/base/recurrent_gated_delta_rule.h @@ -76,7 +76,7 @@ class RecurrentGatedDeltaRule : public Operator { assert(a[7].has_value() != a[9].has_value() && "recurrent delta requires exactly one final destination mode"); assert((state.shape()[1] == Hv && state.shape()[2] == Dv && - state.shape()[3] == D && (a[6] || a[7] || state.shape()[0] == B) && + state.shape()[3] == D && (a[6] || state.shape()[0] >= B) && (!a[9] || a[9]->shape() == Tensor::Shape{B, Hv, Dv, D})) && "recurrent delta state layout must be pool/B,Hv,Dv,Dk"); for (size_t i = 0; i < a.size(); ++i) { @@ -95,7 +95,8 @@ class RecurrentGatedDeltaRule : public Operator { } assert((!(i <= 2 || i == 8) || a[i]->strides()[3] == 1) && "recurrent delta innermost data dimension must be contiguous"); - assert((!(i == 5 || i == 8 || i == 9) || !a[i]->HasBroadcastDim()) && + assert((!((i == 5 && a[7]) || i == 8 || i == 9) || + !a[i]->HasBroadcastDim()) && "recurrent delta broadcast destination unsupported"); } if (state.shape()[0] > diff --git a/src/native/cpu/ops/separate_table_rope/kernel.h b/src/native/cpu/ops/separate_table_rope/kernel.h index fb7be02b6..be9944c75 100644 --- a/src/native/cpu/ops/separate_table_rope/kernel.h +++ b/src/native/cpu/ops/separate_table_rope/kernel.h @@ -15,10 +15,7 @@ class Operator } Operator(const Tensor x, const Tensor p, const Tensor s, const Tensor c, bool neox, Tensor y) - : SeparateTableRope(x, p, s, c, neox, y) { - if (!supports(x, p, s, c, neox, y)) - throw std::invalid_argument("SeparateTableRope: CPU only"); - } + : SeparateTableRope(x, p, s, c, neox, y) {} void operator()(const Tensor x, const Tensor p, const Tensor s, const Tensor c, bool neox, Tensor y) const override { if (x.numel() == 0) return;