diff --git a/CMakeLists.txt b/CMakeLists.txt index 0d8c3e5b3..02c416958 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1,4 +1,23 @@ 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() + 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) if(POLICY CMP0116) @@ -26,6 +45,9 @@ 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) +option(BUILD_NVIDIA_W4_MARLIN "Build NVIDIA AWQ/GPTQ Marlin providers (CUDA >= 12.3)" OFF) + option(WITH_NINETOOTHED "Enable NineToothed-generated kernels" OFF) option(WITH_TRITON "Enable the Triton JIT backend" OFF) @@ -188,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...") @@ -239,7 +268,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() @@ -267,7 +299,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( @@ -360,6 +392,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/LICENSES/Apache-2.0-components.txt b/LICENSES/Apache-2.0-components.txt new file mode 100644 index 000000000..f4f57a92a --- /dev/null +++ b/LICENSES/Apache-2.0-components.txt @@ -0,0 +1,30 @@ +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 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/) +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. + +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. + +DLPack: Copyright 2017 by Contributors +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. + +Source scope and resolved-package redistribution: repository THIRD_PARTY.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/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/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/InfiniCore-MIT.txt b/LICENSES/InfiniCore-MIT.txt new file mode 100644 index 000000000..5b2b72465 --- /dev/null +++ b/LICENSES/InfiniCore-MIT.txt @@ -0,0 +1,8 @@ +The MIT License (MIT) +Copyright © 2025 InfiniTensor + +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. 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/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/README.md b/README.md index cb1c2d844..0d784b317 100644 --- a/README.md +++ b/README.md @@ -49,6 +49,7 @@ InfiniOps extension so `import infini.ops` can load its runtime dependency. | `-DWITH_CAMBRICON=[ON\|OFF]` | Compile the Cambricon implementation | OFF | | `-DWITH_ASCEND=[ON\|OFF]` | Compile the Ascend implementation | OFF | | `-DWITH_TORCH=[ON\|OFF]` | Compile generated PyTorch ATen-backed operators | OFF | +| `-DBUILD_NVIDIA_W4_MARLIN=[ON\|OFF]` | Build AWQ/GPTQ Marlin providers; requires CUDA >= 12.3 and C++20 | OFF | | `-DAUTO_DETECT_DEVICES=[ON\|OFF]` | Auto-detect available platforms | ON | | `-DINFINI_RT_ROOT=` | InfiniRT install prefix containing `include/` and `lib/` | `$INFINI_RT_ROOT` | diff --git a/THIRD_PARTY.txt b/THIRD_PARTY.txt new file mode 100644 index 000000000..e588b1443 --- /dev/null +++ b/THIRD_PARTY.txt @@ -0,0 +1,115 @@ +THIRD-PARTY SOURCE ATTRIBUTION + +Imported/ported material is identified below; new Ops integration glue is not claimed as third-party code. +Direct donor: https://github.com/InfiniTensor/InfiniCore +Pinned direct donor: 1ab85ef1d3c2bce49edf1fc73a5b48becaf7552e +Core material: MIT, Copyright 2025 InfiniTensor; original text: LICENSES/InfiniCore-MIT.txt. +The Core MIT notice does not replace the original licenses or notices of the upstream projects below. +Earlier upstream revisions are not asserted where only the fixed direct Core donor is recorded. + +Path prefixes: +P = src/native/cuda/nvidia/ops/ +C = src/infiniop/ops/ in the pinned direct Core donor + +W4 / Marlin +Upstream projects: https://github.com/IST-DASLab/marlin and https://github.com/vllm-project/vllm +Current kernel source: the fixed vLLM source archive at commit +ffc4f08c8ee130d4ea6347c1bf31ffd4f8af28ab, +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 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; +Apache text: LICENSE. Original source notices remain in prepared files. +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 +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. +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 +ADAPTED P/moe_fused_dense/kernel.cu <- C/moe_fused_dense/nvidia/moe_fused_dense_nvidia.cu +ADAPTED P/moe_fused_gate/kernel.cu <- C/moe_fused_gate/nvidia/moe_fused_gate_nvidia.cu +SGLang project: https://github.com/sgl-project/sglang; earlier upstream path: sgl-kernel/csrc/moe/moe_fused_gate.cu. +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. +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. + +FlashAttention / CUTLASS +Projects: https://github.com/Dao-AILab/flash-attention and https://github.com/NVIDIA/cutlass +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 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. + +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 +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 +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 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/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/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/NativeFlashAttention.cmake b/cmake/NativeFlashAttention.cmake new file mode 100644 index 000000000..36d25c185 --- /dev/null +++ b/cmake/NativeFlashAttention.cmake @@ -0,0 +1,79 @@ +# 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(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") + set(_varlen_provider "${CMAKE_CURRENT_FUNCTION_LIST_DIR}/../src/native/cuda/nvidia/ops/flash_attn_varlen_func") + set(_source "${CMAKE_CURRENT_FUNCTION_LIST_DIR}/..") + 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") + 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" + --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) + message(FATAL_ERROR "Pinned FlashAttention preparation failed: ${_prepare_error}") + 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" + "${_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) + target_compile_definitions(${target} PRIVATE + FLASH_NAMESPACE=infini_flash FLASHATTENTION_DISABLE_DROPOUT) + target_include_directories(${target} PRIVATE + "${_source}/src" "${_source}/include" "${_prepared}" + "${_cutlass_include}" "${_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) + 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/cmake/NativeW4Marlin.cmake b/cmake/NativeW4Marlin.cmake new file mode 100644 index 000000000..9fdacb87c --- /dev/null +++ b/cmake/NativeW4Marlin.cmake @@ -0,0 +1,56 @@ +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(Python REQUIRED COMPONENTS Interpreter) +find_program(_w4_patch NAMES patch REQUIRED) +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 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" + --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_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_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 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 new file mode 100644 index 000000000..155129751 --- /dev/null +++ b/docs/build-nvidia-inference.md @@ -0,0 +1,166 @@ +# Reproducible NVIDIA inference SDK inputs + +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 +execute their Python inference APIs. No `flashinfer-python` or Torch is required. + +## Resolved third-party inputs + +| Input | Exact artifact | SHA256 | +|---|---|---| +| 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` | +| CUTLASS | commit `087c84df83d254b5fb295a7a408f1a1d554085cf` archive | `99340f7a8c87f5295b91aecf8295df6693ad954807672c3869a7fb3d0084f709` | + +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: + +```text +4aa059811da929fac4a4b25b26153649965d9d84ce35f29f06119fa874846c77 sampling.so +7657c74066e959a038324d618d92647a23c79433b3f7a9fbf4fd6354c4d3f9be libtvm_ffi.so +``` + +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. + +## Acquire and install the fixed inputs + +Run in a new working directory. Set `CORE` to an initialized recursive +InfiniCore checkout and use a new, absolute installation prefix: + +```bash +set -euo pipefail +export CORE=/absolute/path/to/InfiniCore +export INFINI_ROOT="$PWD/sdk" +export CUDA_HOME=/usr/local/cuda +export INFINI_RT_ROOT="$INFINI_ROOT" +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 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 +99340f7a8c87f5295b91aecf8295df6693ad954807672c3869a7fb3d0084f709 cutlass.tar.gz +SHA256 +tar -xzf cutlass.tar.gz +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 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" +install -d "$INFINI_ROOT/lib" +install -m644 "$INFINI_ROOT/python/flashinfer_jit_cache/jit_cache/sampling/sampling.so" "$INFINI_ROOT/lib/" +install -m644 "$INFINI_ROOT/python/tvm_ffi/lib/libtvm_ffi.so" "$INFINI_ROOT/lib/" +export PYTHONPATH="$INFINI_ROOT/python${PYTHONPATH:+:$PYTHONPATH}" +export LD_LIBRARY_PATH="$INFINI_ROOT/lib:$CUDA_HOME/lib64${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}" +``` + +Retain both complete package directories and their `.dist-info` directories; +the Ops resolver uses `importlib.metadata`, not just the two copied libraries. +The selected FlashInfer wheel is large because it includes other precompiled +kernels; those are not added as Ops providers by this recipe. + +## Build the SDK + +```bash +cmake -S "$CORE/submodules/InfiniRT" -B build/rt \ + -DCMAKE_BUILD_TYPE=Release -DCMAKE_INSTALL_PREFIX="$INFINI_ROOT" \ + -DCMAKE_INSTALL_LIBDIR=lib -DWITH_CPU=ON -DWITH_NVIDIA=ON \ + -DCMAKE_CUDA_ARCHITECTURES=80 \ + -DCMAKE_CXX_FLAGS=-D_GLIBCXX_USE_CXX11_ABI=0 \ + -DCMAKE_CUDA_FLAGS=-D_GLIBCXX_USE_CXX11_ABI=0 +cmake --build build/rt -j4 +cmake --install build/rt +cmake -S "$CORE/submodules/InfiniOps" -B build/ops \ + -DCMAKE_BUILD_TYPE=Release -DCMAKE_INSTALL_PREFIX="$INFINI_ROOT" \ + -DCMAKE_INSTALL_LIBDIR=lib -DINFINI_RT_ROOT="$INFINI_ROOT" \ + -DWITH_CPU=ON -DWITH_NVIDIA=ON -DWITH_TORCH=OFF -DWITH_LINKED=ON \ + -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 \ + -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 \ + -DCMAKE_CXX_FLAGS=-D_GLIBCXX_USE_CXX11_ABI=0 \ + -DCMAKE_CUDA_FLAGS=-D_GLIBCXX_USE_CXX11_ABI=0 \ + -DPython_EXECUTABLE="$(command -v python)" +cmake --build build/ops -j4 +cmake --install build/ops +``` + +Marlin is enabled here because the LM profile explicitly contains W4 operators. +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 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 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. + +## Third-party notices and validation + +`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 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 +dependency recipe as model qualification: validate the installed SDK with a tiny +sampling call and an attention call, then import the consuming `_infinilm`. +Record source revisions, wheel/library hashes, compiler/driver versions, ABI, +architecture, configure commands and actual library origins with the result. diff --git a/patches/flash-attention/inference.patch b/patches/flash-attention/inference.patch new file mode 100644 index 000000000..9372b837f --- /dev/null +++ b/patches/flash-attention/inference.patch @@ -0,0 +1,98 @@ +--- 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" +--- 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..4a1d73d7b --- /dev/null +++ b/patches/marlin/w4a16.patch @@ -0,0 +1,260 @@ +--- a/dequant.h ++++ b/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/marlin_template.h ++++ b/marlin_template.h +@@ -1035,7 +1035,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) +--- 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 { + +@@ -213,6 +215,7 @@ + + } // namespace marlin + ++#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 { + +@@ -269,6 +271,7 @@ + + } // namespace marlin + ++#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 = """ +-// 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() + +@@ -162,12 +164,16 @@ + ] + + ++# 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]) ++ 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) + +- 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) + +--- 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 +@@ -517,8 +524,9 @@ + ", num_threads = ", num_threads, ", num_bits = ", num_bits); + } + +- cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, +- max_shared_mem_new); ++ ::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 +550,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 +902,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 { diff --git a/pyproject.toml b/pyproject.toml index 2dd4ba2db..afd5defb9 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>=1.1.1", "pybind11", "libclang", "packaging", "pyyaml", "jinja2"] build-backend = "scikit_build_core.build" [project] @@ -12,6 +12,14 @@ dev = ["pytest", "pytest-cov", "pytest-xdist", "ruff==0.15.22", "torch", "pyyaml [tool.scikit-build.wheel] install-dir = "infini" +[tool.scikit-build.sdist] +cmake = true +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" +build-dir = "build/sdist" + [tool.scikit-build.cmake.define] AUTO_DETECT_DEVICES = "ON" AUTO_DETECT_BACKENDS = "ON" 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..a08aa92fd --- /dev/null +++ b/scripts/configs/infiniops_ops_nvidia_dense_sampling_torch_free.json @@ -0,0 +1,55 @@ +{ + "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]}, + "pad": {"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": {"implementations": [0]}, + "w4_marlin_parameters": {"implementations": [0]}, + "w4_marlin_repack": {"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 new file mode 100644 index 000000000..32fb6c5a3 --- /dev/null +++ b/scripts/configs/infiniops_ops_nvidia_dense_sampling_torch_free_no_marlin.json @@ -0,0 +1,55 @@ +{ + "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]}, + "pad": {"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_public_operator_header.py b/scripts/generate_public_operator_header.py index eb6a3add4..9c2f724b6 100644 --- a/scripts/generate_public_operator_header.py +++ b/scripts/generate_public_operator_header.py @@ -16,6 +16,101 @@ 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(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( + 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,28 +206,45 @@ 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, replacement: 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() - 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_operator_call_definition(source.read_text()) + public_text = _hide_definition( + source.read_text(), + _CALL_SIGNATURE_RE, + "Operator::Call", + " {\n return detail::OperatorFactory::Call(handle, config, args...);\n }", + ) + for name, signature in _PROVIDER_SIGNATURES: + 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 4b98dfb9b..c2984b09d 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: @@ -1421,8 +1426,11 @@ 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 + // clang-format off {impl_includes} // clang-format on @@ -1597,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() @@ -1615,6 +1625,11 @@ def _append_unique(declaration, definition): f"Operator<{op_type}>::active_implementation_indices(Device::Type);", ) + _append_unique( + f"extern template Config Operator<{op_type}>::DefaultConfig(Device::Type);", + f"template Config Operator<{op_type}>::DefaultConfig(Device::Type);", + ) + _append_unique( f"extern template std::size_t " f"Operator<{op_type}>::DefaultImplementationIndex(Device::Type);", @@ -1646,6 +1661,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,11 +1681,51 @@ def _append_unique(declaration, definition): f"extern template std::unique_ptr> {make_instantiation};", f"template std::unique_ptr> {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};", ) + 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 @@ -1861,12 +1923,35 @@ def _select_ops_from_config(ops, config, config_path): selected = {} for op_name, selection in config.items(): + if selection.get("public_api_only") is True: + base_header = _BASE_DIR / f"{op_name}.h" + if ( + base_header.resolve().parent != _BASE_DIR.resolve() + or not base_header.is_file() + ): + raise ValueError( + f"{config_path}: public API-only operator {op_name!r} " + "has no hand-written base header" + ) + selected[op_name] = [] + continue 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: @@ -1909,6 +1994,8 @@ def _get_all_ops( with_ninetoothed=False, with_linked=False, with_triton=False, + with_native_flash=False, + with_native_w4=False, ): scan_dirs = set(devices) @@ -1939,6 +2026,30 @@ def _get_all_ops( impl_roots.append(_GENERATION_DIR) impl_headers_by_operator = _index_impl_headers(impl_roots, scan_dirs) + excluded_headers = set() + if not with_native_flash: + excluded_headers.update( + { + ( + _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_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.resolve() not in excluded_headers + ] for base_dir in base_dirs: for file_path in base_dir.iterdir(): @@ -2134,6 +2245,16 @@ 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( + "--with-native-w4", + action="store_true", + help="Include the compiled NVIDIA Marlin providers.", + ) parser.add_argument( "--ops", nargs="+", @@ -2170,6 +2291,8 @@ 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, + with_native_w4=args.with_native_w4, ) if config_path is not None: diff --git a/scripts/ops_config.py b/scripts/ops_config.py index b1e764fb8..9d967cd70 100644 --- a/scripts/ops_config.py +++ b/scripts/ops_config.py @@ -86,6 +86,14 @@ def _normalize_selection(path, op_name, value): f"{path}: {op_name!r} must map to implementation path(s) or an object" ) + if "public_api_only" in value: + if set(value) != {"public_api_only"} or value["public_api_only"] is not True: + raise OpsConfigError( + f"{path}: {op_name!r} public_api_only must be exactly true " + "without implementation selections" + ) + return {"headers": None, "implementations": (), "public_api_only": True} + unknown_keys = sorted(set(value) - {"implementations"}) if unknown_keys: raise OpsConfigError( diff --git a/scripts/prepare_flash_attention_inference.py b/scripts/prepare_flash_attention_inference.py new file mode 100644 index 000000000..4df934315 --- /dev/null +++ b/scripts/prepare_flash_attention_inference.py @@ -0,0 +1,135 @@ +"""Prepare the pinned FlashAttention forward headers for native inference. + +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. +""" + +import argparse +import pathlib +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): + # 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 read_source(path): + if not path.is_file(): + raise ValueError( + f"missing FlashAttention source dependency: {path}; " + "check the fixed source archives or FetchContent SOURCE_DIR overrides" + ) + return path.read_text(encoding="utf-8") + + +def prepare(source, cutlass_source, output, patch_executable="patch"): + source, output = pathlib.Path(source).resolve(), pathlib.Path(output).resolve() + 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 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(cutlass_source / "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(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("--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.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/scripts/prepare_marlin.py b/scripts/prepare_marlin.py new file mode 100644 index 000000000..c831d771c --- /dev/null +++ b/scripts/prepare_marlin.py @@ -0,0 +1,110 @@ +#!/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(name): UPSTREAM / name + for name in ( + "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", + ) +} +SOURCES[pathlib.Path("core/scalar_type.hpp")] = pathlib.Path( + "csrc/core/scalar_type.hpp" +) +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}; " + "check the pinned source archive or offline FetchContent override" + ) + 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 / "generate_kernels.py"), "8.0"], + cwd=staging, + check=True, + ) + prepared = [path for path in SOURCES if path.suffix != ".py"] + prepared.extend(pathlib.Path(name) for name in GENERATED) + for destination in prepared: + write_if_changed(output / 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..67ce8d9c3 100644 --- a/scripts/resolve_linked_ops.py +++ b/scripts/resolve_linked_ops.py @@ -387,6 +387,32 @@ def _load_distribution(config): return distribution +def _locate_distribution_notices(config): + 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 + # 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 +663,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 +850,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 ce9f75770..47ef0320a 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,10 @@ 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/") + 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 @@ -158,7 +172,17 @@ if(WITH_NVIDIA) COMPILE_OPTIONS "--generate-code=arch=compute_75,code=sm_75;--generate-code=arch=compute_75,code=compute_75") + list(FILTER NVIDIA_SOURCES EXCLUDE REGEX "/ops/w4_marlin/") enable_language(CUDA) + if(BUILD_NVIDIA_W4_MARLIN) + include("${PROJECT_SOURCE_DIR}/cmake/NativeW4Marlin.cmake") + endif() + + 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}) @@ -245,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}) @@ -324,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) @@ -350,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) @@ -383,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) @@ -401,10 +433,12 @@ if(WITH_CAMBRICON) find_program(CNCC_COMPILER cncc HINTS "${NEUWARE_HOME}/bin" "$ENV{NEUWARE_HOME}/bin" /usr/local/neuware/bin) if(CNCC_COMPILER) message(STATUS "Found cncc: ${CNCC_COMPILER}") + separate_arguments(INFINI_MLU_HOST_ABI_FLAGS NATIVE_COMMAND "${CMAKE_CXX_FLAGS}") set(MLU_COMPILE_OPTS -c --bang-mlu-arch=mtp_592 -O3 -fPIC -Wall -Werror -std=c++17 -pthread -I${CMAKE_CURRENT_SOURCE_DIR} -I${NEUWARE_HOME}/include ${INFINI_RT_INCLUDE_FLAGS} + ${INFINI_MLU_HOST_ABI_FLAGS} -idirafter /usr/local/neuware/lib/clang/11.1.0/include ) function(compile_mlu_file src_file) @@ -670,6 +704,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.") @@ -716,7 +756,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 ) @@ -739,6 +779,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) @@ -1220,6 +1264,12 @@ 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(BUILD_NVIDIA_W4_MARLIN) + list(APPEND GENERATOR_ARGS --with-native-w4) + endif() if(WITH_NINETOOTHED) list(APPEND GENERATOR_ARGS --with-ninetoothed) endif() @@ -1649,6 +1699,39 @@ 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(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) + 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}") + install(FILES "${_notice}" + DESTINATION "${CMAKE_INSTALL_DATADIR}/licenses/InfiniOps/resolved/${_notice_group}") + endforeach() +endif() + set(INFINI_OPS_INSTALL_CMAKEDIR "${CMAKE_INSTALL_LIBDIR}/cmake/InfiniOps") install(EXPORT InfiniOpsTargets 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/cat.h b/src/base/cat.h index b7f2c9f4c..0e9f190d8 100644 --- a/src/base/cat.h +++ b/src/base/cat.h @@ -9,6 +9,25 @@ 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/causal_conv1d.h b/src/base/causal_conv1d.h new file mode 100644 index 000000000..309e87559 --- /dev/null +++ b/src/base/causal_conv1d.h @@ -0,0 +1,108 @@ +#pragma once +#include +#include +#include +#include +#include +#include +#include + +#include "operator.h" + +namespace infini::ops { +class CausalConv1d : public Operator { + public: + CausalConv1d(const Tensor qkv, const Tensor weight, Tensor conv_state, + const std::optional bias, + const std::optional cu_seqlens, + const std::optional initial_state_indices, + const std::optional final_state_indices, Tensor out, + 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(); + tokens_ = qkv.shape()[1]; + channels_ = qkv.shape()[2]; + pool_ = conv_state.shape()[0]; + requests_ = cu_seqlens ? cu_seqlens->numel() - 1 : qkv.shape()[0]; + } + virtual void operator()(const Tensor qkv, const Tensor weight, + Tensor conv_state, const std::optional bias, + 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 = 0; + + protected: + using Views = std::array, 9>; + Device device_; + DataType dtype_; + size_t tokens_, channels_, pool_, requests_; + const Views metadata_; + static void Validate(const Views& v) { + for (const auto& view : v) + 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(); + 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]; + 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]) { + assert(q.shape()[0] == 1 && v[4]->ndim() == 1 && v[4]->numel() >= 2 && + "causal conv packed sequence shape mismatch"); + requests = v[4]->numel() - 1; + } + 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]}) && + "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; + assert(v[i]->device() == q.device() && "causal conv device mismatch"); + if (i >= 4 && i <= 6) { + 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"); + } + 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"); + } + if (q.shape()[1] > + 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"); + } +}; +} // namespace infini::ops diff --git a/src/base/chunk_gated_delta_rule.h b/src/base/chunk_gated_delta_rule.h new file mode 100644 index 000000000..ab2a189c2 --- /dev/null +++ b/src/base/chunk_gated_delta_rule.h @@ -0,0 +1,133 @@ +#pragma once +#include +#include +#include +#include +#include +#include +#include + +#include "operator.h" +namespace infini::ops { +class ChunkGatedDeltaRule : public Operator { + public: + ChunkGatedDeltaRule(const Tensor q, const Tensor k, const Tensor v, + const Tensor g, const Tensor beta, Tensor initial_state, + const std::optional cu_seqlens, + const std::optional initial_state_indices, + 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), + 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(); + gate_type_ = g.dtype(); + B_ = cu_seqlens ? cu_seqlens->numel() - 1 : q.shape()[0]; + T_ = q.shape()[1]; + Hk_ = q.shape()[2]; + Hv_ = v.shape()[2]; + D_ = q.shape()[3]; + pool_ = initial_state.shape()[0]; + } + virtual 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, + const std::optional initial_state_indices, + const std::optional final_state_indices, + Tensor out, std::optional final_state, + bool use_qk_l2norm, size_t chunk_size) const = 0; + + protected: + using Views = std::array, 11>; + Device device_; + DataType data_type_, gate_type_; + bool norm_; + size_t chunk_, B_, T_, Hk_, Hv_, D_, pool_; + const Views metadata_; + static bool floating(DataType dtype) { + return dtype == DataType::kFloat16 || dtype == DataType::kBFloat16 || + dtype == DataType::kFloat32; + } + static void Validate(const Views& a, size_t chunk) { + 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]; + 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]) { + assert(physical_B == 1 && a[6]->ndim() == 1 && a[6]->numel() >= 2 && + "chunk delta packed cu_seqlens shape mismatch"); + B = a[6]->numel() - 1; + } + 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] || 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; + assert(a[i]->device() == q.device() && "chunk delta device mismatch"); + for (auto stride : a[i]->strides()) + assert(stride >= 0 && "chunk delta negative stride unsupported"); + if (i >= 6 && i <= 8) { + 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"); + } + assert((!(i <= 2 || i == 9) || a[i]->strides()[3] == 1) && + "chunk delta innermost data dimension must be contiguous"); + 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()) || + state.shape()[0] > + static_cast(std::numeric_limits::max())) + throw std::overflow_error( + "chunk delta token/pool exceeds int64 index extent"); + } + void CheckInvocation(const Views& views, bool norm, size_t chunk) const { + 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].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( + "chunk delta nonempty tensor has null data"); + } + } +}; +} // namespace infini::ops 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/deepseek_moe.h b/src/base/deepseek_moe.h new file mode 100644 index 000000000..c589fcf5c --- /dev/null +++ b/src/base/deepseek_moe.h @@ -0,0 +1,100 @@ +#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 < num_experts; ++e) { + Require(gate[e].ndim() == 2 && up[e].ndim() == 2 && down[e].ndim() == 2); + Require(gate[e].shape() == Tensor::Shape{intermediate, hidden.size(1)} && + up[e].shape() == gate[e].shape() && + down[e].shape() == Tensor::Shape{hidden.size(1), intermediate}); + Require(gate[e].dtype() == hidden.dtype() && + up[e].dtype() == hidden.dtype() && + down[e].dtype() == hidden.dtype()); + } + Views(Collect(output, hidden, ids, weights, gate, up, down), 1, true); + // Chunked prefill consumes successive token ranges. A shifted overlap could + // overwrite future hidden rows. + Require(!Overlap(output, hidden) || output.data() == hidden.data()); + } + std::vector metadata_; + size_t intermediate_, num_experts_; +}; +} // namespace infini::ops diff --git a/src/base/embedding.h b/src/base/embedding.h index 8d10243bb..101c9c90a 100644 --- a/src/base/embedding.h +++ b/src/base/embedding.h @@ -46,7 +46,10 @@ 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..3d85680ad 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()}, diff --git a/src/base/flash_attn_varlen_func.h b/src/base/flash_attn_varlen_func.h index 35ebb5aad..99400e242 100644 --- a/src/base/flash_attn_varlen_func.h +++ b/src/base/flash_attn_varlen_func.h @@ -14,6 +14,86 @@ 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..6b26fceba 100644 --- a/src/base/flash_attn_with_kvcache.h +++ b/src/base/flash_attn_with_kvcache.h @@ -4,16 +4,152 @@ #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); + } + + 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 +603,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/fused_gated_delta_net_gating.h b/src/base/fused_gated_delta_net_gating.h new file mode 100644 index 000000000..afbf6377c --- /dev/null +++ b/src/base/fused_gated_delta_net_gating.h @@ -0,0 +1,49 @@ +#pragma once +#include +#include +#include + +#include "data_type.h" +#include "operator.h" + +namespace infini::ops { +class FusedGatedDeltaNetGating : public Operator { + public: + FusedGatedDeltaNetGating(Tensor A_log, Tensor a, Tensor b, Tensor dt_bias, + float beta, float threshold, Tensor g, + Tensor beta_output) + : beta_(beta), + threshold_(threshold), + metadata_{A_log, a, b, dt_bias, g, beta_output} { + const auto& tensors = metadata_; + 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) { + 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, + float beta, float threshold, Tensor g, + Tensor beta_output) const = 0; + + protected: + Device device_; + size_t total_, seq_, hidden_; + float beta_, threshold_; + const std::array metadata_; +}; +} // namespace infini::ops diff --git a/src/base/gemm.h b/src/base/gemm.h index fe6df2655..dfd3d20dc 100644 --- a/src/base/gemm.h +++ b/src/base/gemm.h @@ -4,13 +4,66 @@ #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; + } + + 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 a.dtype() == DataType::kFloat32 && b.dtype() == DataType::kFloat32 && + y.dtype() == DataType::kFloat32; + } + + private: + static bool IsValid(Fp32Precision precision) { + return precision == Fp32Precision::DefaultFast || + precision == Fp32Precision::FullFp32; + } + + Fp32Precision fp32_precision_; +}; + class Gemm : public Operator { public: + using Operator::Make; + 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 +119,19 @@ 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 +180,17 @@ 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 { + return detail::CacheKey::Build( + config.implementation_index(), + static_cast(GemmConfig::PrecisionFrom(config)), a, b, tail...); + } +}; + } // namespace infini::ops #endif diff --git a/src/base/i8_gemm.h b/src/base/i8_gemm.h new file mode 100644 index 000000000..e6ff7ba4b --- /dev/null +++ b/src/base/i8_gemm.h @@ -0,0 +1,15 @@ +#pragma once +#include + +#include "operator.h" +namespace infini::ops { +class I8Gemm : public Operator { + public: + static constexpr bool kAutomaticDefault = true; + I8Gemm(Tensor c, const Tensor a, const Tensor as, const Tensor b, + const Tensor bs, const std::optional bias) {} + virtual void operator()(Tensor c, const Tensor a, const Tensor as, + const Tensor b, const Tensor bs, + const std::optional bias) const = 0; +}; +} // namespace infini::ops diff --git a/src/base/layer_norm.h b/src/base/layer_norm.h new file mode 100644 index 000000000..57e357646 --- /dev/null +++ b/src/base/layer_norm.h @@ -0,0 +1,40 @@ +#ifndef INFINI_OPS_BASE_LAYER_NORM_H_ +#define INFINI_OPS_BASE_LAYER_NORM_H_ + +#include +#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_meta_{input}, + weight_meta_{weight}, + bias_meta_{bias}, + out_meta_{out}, + device_{input.device()}, + eps_{eps} {} + + virtual void operator()(const Tensor input, const Tensor weight, + const std::optional bias, float eps, + Tensor out) const = 0; + + protected: + const Tensor input_meta_, weight_meta_; + const std::optional bias_meta_; + const Tensor out_meta_; + Device device_; + float eps_; +}; + +} // namespace infini::ops + +#endif diff --git a/src/base/mamba_selective_scan.h b/src/base/mamba_selective_scan.h new file mode 100644 index 000000000..3f5ae4ec6 --- /dev/null +++ b/src/base/mamba_selective_scan.h @@ -0,0 +1,141 @@ +#ifndef INFINI_OPS_BASE_MAMBA_SELECTIVE_SCAN_H_ +#define INFINI_OPS_BASE_MAMBA_SELECTIVE_SCAN_H_ + +#include +#include +#include +#include + +#include "operator.h" + +namespace infini::ops { +// Ordering deliberately matches the original infiniop descriptor/call ABI. +// data() denotes the logical origin; no storage-base offset is added again. +struct MambaScanLayout { + std::size_t batch, length, channels, states; + std::ptrdiff_t strides[10][3]{}; + bool contiguous; +}; + +class MambaSelectiveScan : public Operator { + public: + MambaSelectiveScan(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) + : layout_{ValidateMetadata( + {out, x, dt, b, c, a_log, d, gate, dt_bias, state})}, + dtype_{x.dtype()}, + device_{x.device()} {} + virtual 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 = 0; + + static MambaScanLayout ValidateMetadata(const std::array& ts) { + const auto dtype = ts[1].dtype(); + if (dtype != DataType::kFloat32 && dtype != DataType::kFloat16 && + dtype != DataType::kBFloat16) + throw std::invalid_argument( + "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)) + throw std::invalid_argument( + "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]) || + ts[i].device() != ts[1].device() || + ts[i].dtype() != (i == 9 ? DataType::kFloat32 + : (i == 5 || i == 6) ? parameter_dtype + : dtype)) + throw std::invalid_argument( + "MambaSelectiveScan: incompatible rank/device/dtype"); + SpanBytes(ts[i]); // Reject negative strides and overflow, including + // unused rows. + } + const auto batch = ts[1].size(0), length = ts[1].size(1), + channels = ts[1].size(2), states = ts[3].size(2); + for (int i : {0, 2, 7}) + if (ts[i].shape() != ts[1].shape()) + throw std::invalid_argument( + "MambaSelectiveScan: x/dt/gate/out shape mismatch"); + if (ts[3].size(0) != batch || ts[3].size(1) != length || + ts[4].shape() != ts[3].shape() || ts[5].size(0) != channels || + ts[5].size(1) != states || ts[6].size(0) != channels || + ts[8].size(0) != channels || ts[9].size(0) < batch || + ts[9].size(1) != channels || ts[9].size(2) != states) + throw std::invalid_argument( + "MambaSelectiveScan: B/C/A_log/D/bias/state shape mismatch"); + // The recurrence may use a larger state pool; only the first batch rows + // change. + ValidateWritable(ts[0]); + ValidateWritable(ts[9]); + if (channels && batch > std::numeric_limits::max() / channels) + throw std::overflow_error("MambaSelectiveScan: launch size overflow"); + MambaScanLayout layout{batch, length, channels, states, {}, true}; + for (int i = 0; i < 10; ++i) { + std::size_t expected = 1; + for (int axis = ranks[i] - 1; axis >= 0; --axis) { + layout.strides[i][axis] = ts[i].stride(axis); + if (ts[i].size(axis) > 1 && + ts[i].stride(axis) != static_cast(expected)) + layout.contiguous = false; + if (ts[i].size(axis) && + expected > + std::numeric_limits::max() / ts[i].size(axis)) + throw std::overflow_error( + "MambaSelectiveScan: shape product overflow"); + expected *= ts[i].size(axis); + } + } + return layout; + } + + protected: + MambaScanLayout layout_; + DataType dtype_; + Device device_; + + private: + static std::size_t SpanBytes(const Tensor& t) { + std::size_t last = 0; + bool empty = false; + const auto limit = + static_cast(std::numeric_limits::max()); + for (std::size_t axis = 0; axis < t.ndim(); ++axis) { + if (t.stride(axis) < 0) + throw std::invalid_argument("MambaSelectiveScan: negative stride"); + if (!t.size(axis)) { + empty = true; + continue; + } + const auto stride = static_cast(t.stride(axis)); + if (stride && t.size(axis) - 1 > (limit - last) / stride) + throw std::overflow_error("MambaSelectiveScan: span overflow"); + last += (t.size(axis) - 1) * stride; + } + const auto width = t.dtype() == DataType::kFloat32 ? 4u : 2u; + if (empty) return 0; + if (last >= limit / width) + throw std::overflow_error("MambaSelectiveScan: byte span overflow"); + return (last + 1) * width; + } + static void ValidateWritable(const Tensor& t) { + // Monotone pitched rows with arbitrary positive inner stride. Broadcast or + // overlapping writable dimensions are rejected before submitting a kernel. + for (std::size_t axis = 0; axis < t.ndim(); ++axis) + if (!t.size(axis)) return; + std::size_t extent = 1; + for (int axis = static_cast(t.ndim()) - 1; axis >= 0; --axis) { + if (t.size(axis) > 1 && static_cast(t.stride(axis)) < extent) + throw std::invalid_argument( + "MambaSelectiveScan: overlapping writable layout"); + if (t.size(axis)) + extent += (t.size(axis) - 1) * static_cast(t.stride(axis)); + } + } +}; +} // namespace infini::ops +#endif diff --git a/src/base/moe_align_block_size.h b/src/base/moe_align_block_size.h index 81069eb33..6193a3afe 100644 --- a/src/base/moe_align_block_size.h +++ b/src/base/moe_align_block_size.h @@ -135,11 +135,11 @@ class MoeAlignBlockSize : public Operator { "`MoeAlignBlockSize` output tensors are too small"); if (maybe_expert_map) { - assert(maybe_expert_map->ndim() == 1 && - maybe_expert_map->numel() == - static_cast(num_experts_) && - "`MoeAlignBlockSize` requires `expert_map` shape " - "[`num_experts`]"); + assert( + maybe_expert_map->ndim() == 1 && maybe_expert_map->numel() > 0 && + maybe_expert_map->numel() <= + static_cast(std::numeric_limits::max()) && + "`MoeAlignBlockSize` requires an int32-addressable `expert_map`"); assert(maybe_expert_map->dtype() == DataType::kInt32 && "`MoeAlignBlockSize` currently requires int32 `expert_map`"); assert(maybe_expert_map->IsContiguous() && diff --git a/src/base/moe_compute_common.h b/src/base/moe_compute_common.h new file mode 100644 index 000000000..0ce55c7ee --- /dev/null +++ b/src/base/moe_compute_common.h @@ -0,0 +1,64 @@ +#pragma once +#include +#include +#include +#include + +#include "operator.h" +namespace infini::ops::moe_detail { +inline void Require(bool ok) { + if (!ok) + throw std::invalid_argument( + "MoE metadata, parameter, pointer, or alias contract violated"); +} +inline bool Reduced(DataType d) { + return d == DataType::kFloat16 || d == DataType::kBFloat16; +} +inline size_t Product(size_t a, size_t b) { + Require(!b || a <= std::numeric_limits::max() / b); + return a * b; +} +inline void Bound(size_t n) { + Require(n <= size_t(std::numeric_limits::max())); +} +inline bool Overlap(const Tensor a, const Tensor b) { + const size_t an = Product(a.numel(), a.element_size()), + bn = Product(b.numel(), b.element_size()); + if (!an || !bn) return false; + const uintptr_t ap = reinterpret_cast(a.data()), + bp = reinterpret_cast(b.data()); + return ap <= bp ? bp - ap < an : ap - bp < bn; +} +inline void Views(const std::vector& v, size_t output_count, + bool output_hidden_alias = false) { + Require(!v.empty()); + for (size_t i = 0; i < v.size(); ++i) { + Require(v[i].device() == v[0].device() && v[i].IsContiguous()); + Product(v[i].numel(), v[i].element_size()); + Require(!v[i].numel() || v[i].data()); + Require(!v[i].numel() || + reinterpret_cast(v[i].data()) % v[i].element_size() == + 0); + for (size_t j = 0; j < i && j < output_count; ++j) { + if (!(output_hidden_alias && j == 0 && i == 1)) + Require(!Overlap(v[i], v[j])); + } + } +} +inline void Same(const std::vector& a, const std::vector& b) { + Require(a.size() == b.size()); + const std::equal_to eq; + for (size_t i = 0; i < a.size(); ++i) Require(eq(a[i], b[i])); +} +inline void Workspace(void* p, size_t available, size_t required, + const std::vector& v) { + Require(p && available >= required && + reinterpret_cast(p) % 256 == 0); + const auto ap = reinterpret_cast(p); + for (const auto& t : v) { + const auto bp = reinterpret_cast(t.data()); + const auto bytes = Product(t.numel(), t.element_size()); + if (bytes) Require(ap <= bp ? bp - ap >= required : ap - bp >= bytes); + } +} +} // namespace infini::ops::moe_detail diff --git a/src/base/moe_fused_dense.h b/src/base/moe_fused_dense.h new file mode 100644 index 000000000..ca72275e8 --- /dev/null +++ b/src/base/moe_fused_dense.h @@ -0,0 +1,119 @@ +#pragma once +#include + +#include "base/moe_compute_common.h" +namespace infini::ops { +class MoeFusedDense : public Operator { + public: + MoeFusedDense(Tensor output, const Tensor hidden, const Tensor w13, + const Tensor w2, const Tensor weights, const Tensor ids, + const Tensor sorted, const Tensor experts, const Tensor total, + size_t block_size, + const std::optional bias1 = std::nullopt, + const std::optional bias2 = std::nullopt) + : metadata_{output, hidden, w13, w2, weights, + ids, sorted, experts, total}, + block_size_(block_size), + bias1_(bias1), + bias2_(bias2) { + Validate(output, hidden, w13, w2, weights, ids, sorted, experts, total, + block_size, bias1, bias2); + } + 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, + const Tensor total, size_t block_size, + const std::optional bias1 = std::nullopt, + const std::optional bias2 = std::nullopt) { + try { + Validate(output, hidden, w13, w2, weights, ids, sorted, experts, total, + block_size, bias1, bias2); + return true; + } catch (const std::invalid_argument&) { + return false; + } + } + virtual void operator()( + Tensor output, const Tensor hidden, const Tensor w13, const Tensor w2, + const Tensor weights, const Tensor ids, const Tensor sorted, + const Tensor experts, const Tensor total, size_t block_size, + const std::optional bias1 = std::nullopt, + const std::optional bias2 = std::nullopt) const = 0; + + protected: + void Invocation(Tensor output, const Tensor hidden, const Tensor w13, + const Tensor w2, const Tensor weights, const Tensor ids, + const Tensor sorted, const Tensor experts, const Tensor total, + size_t block_size, const std::optional bias1, + const std::optional bias2) const { + moe_detail::Require(block_size == block_size_); + moe_detail::Same(metadata_, {output, hidden, w13, w2, weights, ids, sorted, + experts, total}); + moe_detail::Require(bias1.has_value() == bias1_.has_value() && + bias2.has_value() == bias2_.has_value()); + if (bias1) moe_detail::Require(std::equal_to{}(*bias1, *bias1_)); + if (bias2) moe_detail::Require(std::equal_to{}(*bias2, *bias2_)); + Validate(output, hidden, w13, w2, weights, ids, sorted, experts, total, + block_size, bias1, bias2); + } + static void Validate(Tensor output, const Tensor hidden, const Tensor w13, + const Tensor w2, const Tensor weights, const Tensor ids, + const Tensor sorted, const Tensor experts, + const Tensor total, size_t block_size, + const std::optional bias1, + const std::optional bias2) { + using namespace moe_detail; + Require(hidden.ndim() == 2 && output.ndim() == 2 && w13.ndim() == 3 && + w2.ndim() == 3 && weights.ndim() == 2 && ids.ndim() == 2 && + sorted.ndim() == 1 && experts.ndim() == 1 && total.ndim() == 1); + Require(Reduced(hidden.dtype()) && output.dtype() == hidden.dtype() && + w13.dtype() == hidden.dtype() && w2.dtype() == hidden.dtype()); + Require(weights.dtype() == DataType::kFloat32 && + ids.dtype() == DataType::kInt32 && + sorted.dtype() == DataType::kInt32 && + experts.dtype() == DataType::kInt32 && + total.dtype() == DataType::kInt32); + const auto n = hidden.size(0), h = hidden.size(1), e = w13.size(0), + m = w2.size(2), k = ids.size(1); + Bound(n); + Bound(h); + Bound(e); + Bound(m); + Bound(k); + Bound(block_size); + Bound(sorted.numel()); + Bound(experts.numel()); + Require(n && h && e && m && k && block_size && + output.shape() == hidden.shape()); + Require(w13.size(1) == 2 * m && w13.size(2) == h && w2.size(0) == e && + w2.size(1) == h && weights.shape() == ids.shape() && + ids.size(0) == n && total.numel() == 1); + const auto pairs = Product(n, k); + Bound(pairs); + const auto capacity = pairs < e ? Product(pairs, block_size) + : pairs + Product(e, block_size - 1); + Bound(capacity); + Require(sorted.numel() >= capacity && + experts.numel() >= (capacity + block_size - 1) / block_size); + Bound(Product(sorted.numel(), m)); + Bound(Product(sorted.numel(), h)); + Bound(Product(sorted.numel(), 2 * m)); + std::vector views{output, hidden, w13, w2, weights, + ids, sorted, experts, total}; + if (bias1) { + Require(bias1->ndim() == 2 && bias1->size(0) == e && + bias1->size(1) == 2 * m && bias1->dtype() == hidden.dtype()); + views.push_back(*bias1); + } + if (bias2) { + Require(bias2->ndim() == 2 && bias2->size(0) == e && + bias2->size(1) == h && bias2->dtype() == hidden.dtype()); + views.push_back(*bias2); + } + Views(views, 1, true); + } + std::vector metadata_; + size_t block_size_; + std::optional bias1_, bias2_; +}; +} // namespace infini::ops diff --git a/src/base/moe_fused_gate.h b/src/base/moe_fused_gate.h new file mode 100644 index 000000000..f1191add0 --- /dev/null +++ b/src/base/moe_fused_gate.h @@ -0,0 +1,72 @@ +#pragma once +#include "base/moe_compute_common.h" +namespace infini::ops { +class MoeFusedGate : public Operator { + public: + MoeFusedGate(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) + : metadata_{weights, indices, input, bias}, + groups_(groups), + selected_groups_(selected_groups), + shared_(shared), + scale_(scale), + apply_scale_(apply_scale) { + Validate(weights, indices, input, bias, groups, selected_groups, shared, + scale, apply_scale); + } + static bool supports(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) { + try { + Validate(weights, indices, input, bias, groups, selected_groups, shared, + scale, apply_scale); + return true; + } catch (const std::invalid_argument&) { + return false; + } + } + virtual void 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 = 0; + + protected: + void Invocation(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 { + moe_detail::Require( + groups == groups_ && selected_groups == selected_groups_ && + shared == shared_ && scale == scale_ && apply_scale == apply_scale_); + moe_detail::Same(metadata_, {weights, indices, input, bias}); + Validate(weights, indices, input, bias, groups, selected_groups, shared, + scale, apply_scale); + } + static void Validate(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) { + using namespace moe_detail; + Require(input.ndim() == 2 && weights.ndim() == 2 && indices.ndim() == 2 && + bias.ndim() == 1); + Require(input.dtype() == DataType::kFloat32 || Reduced(input.dtype())); + Require(weights.dtype() == DataType::kFloat32 && + bias.dtype() == DataType::kFloat32 && + indices.dtype() == DataType::kInt32); + const auto n = input.size(0), e = input.size(1), k = weights.size(1); + Bound(n); + Bound(e); + Bound(k); + Bound(Product(n, k)); + Require(e && groups && e % groups == 0 && e / groups >= 2); + Require(selected_groups && selected_groups <= groups && k > shared && + k - shared <= selected_groups * (e / groups)); + Require(std::isfinite(scale) && scale > 0 && weights.size(0) == n && + indices.shape() == weights.shape() && bias.size(0) == e); + Views({weights, indices, input, bias}, 2); + } + std::vector metadata_; + size_t groups_, selected_groups_, shared_; + float scale_; + bool apply_scale_; +}; +} // namespace infini::ops diff --git a/src/base/mrope.h b/src/base/mrope.h new file mode 100644 index 000000000..36fc3abe3 --- /dev/null +++ b/src/base/mrope.h @@ -0,0 +1,38 @@ +#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) { + 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"); + for (const auto& t : {q, k, qo, ko}) + 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, + int st, int sh, int sw, bool interleaved, Tensor qo, + Tensor ko) const = 0; +}; +} // namespace infini::ops 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/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/per_channel_quant_i8.h b/src/base/per_channel_quant_i8.h new file mode 100644 index 000000000..bc14868e5 --- /dev/null +++ b/src/base/per_channel_quant_i8.h @@ -0,0 +1,11 @@ +#pragma once +#include "operator.h" +namespace infini::ops { +class PerChannelQuantI8 : public Operator { + public: + static constexpr bool kAutomaticDefault = true; + PerChannelQuantI8(const Tensor x, Tensor packed, Tensor scale) {} + virtual void operator()(const Tensor x, Tensor packed, + Tensor scale) const = 0; +}; +} // namespace infini::ops diff --git a/src/base/per_tensor_dequant_i8.h b/src/base/per_tensor_dequant_i8.h new file mode 100644 index 000000000..6fa20cc83 --- /dev/null +++ b/src/base/per_tensor_dequant_i8.h @@ -0,0 +1,14 @@ +#pragma once +#include + +#include "operator.h" +namespace infini::ops { +class PerTensorDequantI8 : public Operator { + public: + static constexpr bool kAutomaticDefault = true; + PerTensorDequantI8(Tensor x, const Tensor packed, const Tensor scale, + const std::optional zero) {} + virtual void operator()(Tensor x, const Tensor packed, const Tensor scale, + const std::optional zero) const = 0; +}; +} // namespace infini::ops diff --git a/src/base/per_tensor_quant_i8.h b/src/base/per_tensor_quant_i8.h new file mode 100644 index 000000000..3f48b962a --- /dev/null +++ b/src/base/per_tensor_quant_i8.h @@ -0,0 +1,15 @@ +#pragma once +#include + +#include "operator.h" +namespace infini::ops { +class PerTensorQuantI8 : public Operator { + public: + static constexpr bool kAutomaticDefault = true; + PerTensorQuantI8(const Tensor x, Tensor packed, Tensor scale, + const std::optional zero, const bool is_static) {} + virtual void operator()(const Tensor x, Tensor packed, Tensor scale, + const std::optional zero, + const bool is_static) const = 0; +}; +} // namespace infini::ops diff --git a/src/base/quick_gelu.h b/src/base/quick_gelu.h new file mode 100644 index 000000000..d976b199e --- /dev/null +++ b/src/base/quick_gelu.h @@ -0,0 +1,27 @@ +#ifndef INFINI_OPS_BASE_QUICK_GELU_H_ +#define INFINI_OPS_BASE_QUICK_GELU_H_ + +#include + +#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_meta_{input}, out_meta_{out}, device_{input.device()} {} + + virtual void operator()(const Tensor input, Tensor out) const = 0; + + protected: + const Tensor input_meta_, out_meta_; + Device device_; +}; + +} // namespace infini::ops +#endif diff --git a/src/base/recurrent_gated_delta_rule.h b/src/base/recurrent_gated_delta_rule.h new file mode 100644 index 000000000..0df174b80 --- /dev/null +++ b/src/base/recurrent_gated_delta_rule.h @@ -0,0 +1,122 @@ +#pragma once +#include +#include +#include +#include +#include +#include +#include + +#include "operator.h" +namespace infini::ops { +class RecurrentGatedDeltaRule : public Operator { + public: + RecurrentGatedDeltaRule(const Tensor q, const Tensor k, const Tensor v, + const Tensor g, const Tensor beta, + Tensor initial_state, + const std::optional initial_state_indices, + const std::optional final_state_indices, + Tensor out, std::optional final_state, + bool use_qk_l2norm) + : 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(); + gate_type_ = g.dtype(); + B_ = q.shape()[0]; + Hk_ = q.shape()[2]; + Hv_ = v.shape()[2]; + D_ = q.shape()[3]; + pool_ = initial_state.shape()[0]; + } + virtual 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, + const std::optional final_state_indices, + Tensor out, std::optional final_state, + bool use_qk_l2norm) const = 0; + + protected: + using Views = std::array, 10>; + Device device_; + DataType data_type_, gate_type_; + bool norm_; + size_t B_, Hk_, Hv_, D_, pool_; + const Views metadata_; + static bool floating(DataType dtype) { + return dtype == DataType::kFloat16 || dtype == DataType::kBFloat16 || + dtype == DataType::kFloat32; + } + 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]; + 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]; + 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] || 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; + assert(a[i]->device() == q.device() && "recurrent delta device mismatch"); + for (auto stride : a[i]->strides()) + assert(stride >= 0 && "recurrent delta negative stride unsupported"); + if (i == 6 || i == 7) { + 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"); + } + assert((!(i <= 2 || i == 8) || a[i]->strides()[3] == 1) && + "recurrent delta innermost data dimension must be contiguous"); + assert((!((i == 5 && a[7]) || i == 8 || i == 9) || + !a[i]->HasBroadcastDim()) && + "recurrent delta broadcast destination unsupported"); + } + if (state.shape()[0] > + static_cast(std::numeric_limits::max())) + throw std::overflow_error( + "recurrent delta pool exceeds int64 index extent"); + } + void CheckInvocation(const Views& views, bool norm) const { + if (norm != norm_) + 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].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()) + throw std::invalid_argument( + "recurrent delta nonempty tensor has null data"); + } + } +}; +} // namespace infini::ops diff --git a/src/base/reshape_and_cache_flash.h b/src/base/reshape_and_cache_flash.h index a11c997b6..4ff7fdce8 100644 --- a/src/base/reshape_and_cache_flash.h +++ b/src/base/reshape_and_cache_flash.h @@ -14,6 +14,33 @@ 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/base/row_l2_normalize.h b/src/base/row_l2_normalize.h new file mode 100644 index 000000000..3dde4d929 --- /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 + +#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_meta_(input), + out_meta_(out), + device_(input.device()), + eps_(eps) {} + + virtual void operator()(const Tensor input, const double eps, + Tensor out) const = 0; + + protected: + const Tensor input_meta_, out_meta_; + 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..a04bfd082 --- /dev/null +++ b/src/base/separate_table_rope.h @@ -0,0 +1,90 @@ +#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 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() || + 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; + return true; + } + SeparateTableRope(const Tensor input, const Tensor positions, + const Tensor sin_table, const Tensor cos_table, + bool is_neox, Tensor out) { + 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, + bool is_neox, Tensor out) const = 0; +}; +} // namespace infini::ops +#endif 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/topk_softmax.h b/src/base/topk_softmax.h index 43aaac734..0b5efc327 100644 --- a/src/base/topk_softmax.h +++ b/src/base/topk_softmax.h @@ -2,15 +2,56 @@ #define INFINI_OPS_BASE_TOPK_SOFTMAX_H_ #include +#include #include #include #include #include +#include #include "operator.h" namespace infini::ops { +class TopkSoftmaxConfig : public Config { + public: + explicit TopkSoftmaxConfig(float softcap = 0.0f, float min_sum = 0.0f) + : softcap_{softcap}, min_sum_{min_sum} { + if (!std::isfinite(softcap) || softcap < 0.0f) { + throw std::invalid_argument( + "TopkSoftmax softcap must be finite and nonnegative"); + } + if (!std::isfinite(min_sum) || min_sum < 0.0f) { + throw std::invalid_argument( + "TopkSoftmax min_sum must be finite and nonnegative"); + } + } + + std::unique_ptr Clone() const override { + return std::make_unique(*this); + } + + bool requires_config_aware_provider() const override { + return softcap_ != 0.0f || min_sum_ != 0.0f; + } + float softcap() const { return softcap_; } + float min_sum() const { return min_sum_; } + + static float SoftcapFrom(const Config& config) { + const auto* typed = dynamic_cast(&config); + return typed ? typed->softcap() : 0.0f; + } + + static float MinSumFrom(const Config& config) { + const auto* typed = dynamic_cast(&config); + return typed ? typed->min_sum() : 0.0f; + } + + private: + float softcap_; + float min_sum_; +}; + // Aligned with vLLM `_moe_C::topk_softmax`. class TopkSoftmax : public Operator { public: @@ -43,6 +84,10 @@ class TopkSoftmax : public Operator { Tensor token_expert_indices) const = 0; protected: + float softcap_{0.0f}; + + float min_sum_{0.0f}; + void ValidateCallMetadata(const Tensor gating_output, std::optional bias, std::optional is_padding, @@ -166,6 +211,16 @@ class TopkSoftmax : public Operator { Tensor token_expert_indices_metadata_; }; +template <> +struct CacheKeyBuilder { + template + detail::CacheKey operator()(const Config& config, const Args&... args) const { + return detail::CacheKey::Build( + config.implementation_index(), TopkSoftmaxConfig::SoftcapFrom(config), + TopkSoftmaxConfig::MinSumFrom(config), args...); + } +}; + } // namespace infini::ops #endif // INFINI_OPS_BASE_TOPK_SOFTMAX_H_ 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/base/w4_marlin_common.h b/src/base/w4_marlin_common.h new file mode 100644 index 000000000..beb085e47 --- /dev/null +++ b/src/base/w4_marlin_common.h @@ -0,0 +1,75 @@ +#pragma once +#include +#include +#include +#include + +#include "operator.h" + +namespace infini::ops::w4 { +inline void require(bool value, const char* message) { + if (!value) throw std::invalid_argument(message); +} +inline std::size_t mul(std::size_t a, std::size_t b) { + if (b && a > std::numeric_limits::max() / b) + throw std::overflow_error("W4 size overflow"); + return a * b; +} +inline bool fp16(const Tensor& t) { + return t.dtype() == DataType::kFloat16 || t.dtype() == DataType::kBFloat16; +} +inline void shape(const Tensor& t, const Tensor::Shape& s, DataType type) { + require(t.shape() == s && t.dtype() == type && t.IsContiguous(), + "W4 requires the exact contiguous shape and dtype"); +} +inline void device(const Tensor& ref, const Tensor& t) { + require(ref.device() == t.device(), "W4 requires one device"); +} +inline void dims(int64_t k, int64_t n) { + require(k > 0 && n > 0 && k % 128 == 0 && n % 64 == 0 && + k <= std::numeric_limits::max() && + n <= std::numeric_limits::max() / 2, + "W4 Marlin requires K%128=0, N%64=0 and bounded dimensions"); + require(mul(static_cast(k), static_cast(n)) / 8 <= + static_cast(std::numeric_limits::max()), + "W4 packed indexing exceeds original kernel int32 limits"); +} +inline int64_t groups(int64_t k, int64_t group) { + require(group == -1 || group == 32 || group == 64 || group == 128, + "W4 supports groups -1,32,64,128"); + require(group == -1 || k % group == 0, "W4 group must divide K"); + return group == -1 ? 1 : k / group; +} +inline std::size_t span(const Tensor& t) { + if (!t.numel()) return 0; + std::size_t extent = 1; + for (int64_t i = 0; i < static_cast(t.ndim()); ++i) { + require(t.stride(i) >= 0, "W4 negative stride is unsupported"); + const auto part = mul(t.size(i) - 1, static_cast(t.stride(i))); + if (part > std::numeric_limits::max() - extent) + throw std::overflow_error("W4 tensor span overflow"); + extent += part; + } + return mul(extent, t.element_size()); +} +inline void disjoint(const void* a, std::size_t ab, const void* b, + std::size_t bb) { + if (!ab || !bb || !a || !b) return; // Make may use address-free metadata. + const auto x = reinterpret_cast(a), + y = reinterpret_cast(b); + require(ab <= std::numeric_limits::max() - x && + bb <= std::numeric_limits::max() - y, + "W4 address range overflow"); + require(x + ab <= y || y + bb <= x, + "W4 writable storage overlaps an input/output"); +} +inline void disjoint(const Tensor& a, const Tensor& b) { + disjoint(a.data(), span(a), b.data(), span(b)); +} +inline void address(const Tensor& t, std::size_t alignment) { + if (!t.numel()) return; + require( + t.data() && reinterpret_cast(t.data()) % alignment == 0, + "W4 device pointer is null or misaligned"); +} +} // namespace infini::ops::w4 diff --git a/src/base/w4_marlin_gemm.h b/src/base/w4_marlin_gemm.h new file mode 100644 index 000000000..37d49b1b2 --- /dev/null +++ b/src/base/w4_marlin_gemm.h @@ -0,0 +1,83 @@ +#pragma once +#include "base/w4_marlin_common.h" + +namespace infini::ops { +// Already repacked B. Y = alpha * Marlin(A,B,S,Z) + ordinary-order bias. +// No dense weight allocation; one stream, one external workspace per +// invocation. +class W4MarlinGemm : public Operator { + public: + W4MarlinGemm(const Tensor a, const Tensor b, const Tensor scales, + const Tensor zeros, const Tensor g_idx, const Tensor perm, + const Tensor bias, const int64_t scheme, + const int64_t group_size, const bool is_k_full, + const float alpha, Tensor out) + : scheme_{scheme}, group_{group_size}, full_{is_k_full}, alpha_{alpha} { + w4::require(a.ndim() == 2 && out.ndim() == 2, "W4 GEMM requires matrices"); + k_ = a.size(1); + n_ = out.size(1); + m_ = a.size(0); + w4::dims(k_, n_); + const auto groups = w4::groups(k_, group_); + w4::require(m_ > 0 && m_ <= std::numeric_limits::max() && w4::fp16(a), + "W4 activation must be nonempty F16/BF16"); + w4::require(a.stride(1) == 1 && a.stride(0) >= k_ && a.stride(0) % 8 == 0 && + a.stride(0) <= std::numeric_limits::max(), + "W4 A requires aligned contiguous rows"); + w4::require(scheme == 0 || scheme == 1, "unknown W4 GEMM scheme"); + w4::require(std::isfinite(alpha), "W4 alpha must be finite"); + w4::shape(out, {m_, n_}, a.dtype()); + w4::shape(b, {k_ / 16, n_ * 2}, DataType::kInt32); + w4::require(scales.ndim() == 2 && scales.size(0) > 0 && + scales.size(0) <= std::numeric_limits::max(), + "W4 scales rows are invalid"); + const auto g = scales.size(0); + w4::shape(scales, {g, n_}, a.dtype()); + w4::shape(zeros, scheme == 0 ? Tensor::Shape{g, n_ / 8} : Tensor::Shape{0}, + DataType::kInt32); + w4::shape(g_idx, {g_idx.numel() ? k_ : 0}, DataType::kInt32); + w4::shape(perm, {perm.numel() ? k_ : 0}, DataType::kInt32); + w4::shape(bias, {bias.numel() ? n_ : 0}, a.dtype()); + const bool act = perm.numel() != 0; + w4::require(act == (g_idx.numel() != 0) && (scheme == 1 || !act), + "W4 g_idx/perm pairing is invalid"); + w4::require(full_ || act, "partial-K W4 requires GPTQ activation order"); + w4::require(!full_ || g == groups, + "W4 full-K group count differs from checkpoint"); + w4::require(!act || group_ != -1, + "activation order requires grouped quantization"); + w4::require(!act || !full_ || g > 1, + "full-K activation order requires multiple groups"); + w4::require(full_ || static_cast(group_) < k_, + "partial-K scale ordering requires group_size < K"); + w4::require( + w4::mul(m_, k_) <= + static_cast(std::numeric_limits::max()) && + w4::mul(m_, n_) <= + static_cast(std::numeric_limits::max()) && + w4::mul(g, n_) <= + static_cast(std::numeric_limits::max()), + "W4 indexing exceeds original int32 kernel limits"); + for (const auto& t : {a, b, scales, zeros, out}) + if (t.data()) w4::address(t, 16); + for (const auto& t : {g_idx, perm, bias}) + if (t.data()) w4::address(t, 4); + for (const auto& t : {a, b, scales, zeros, g_idx, perm, bias, out}) + w4::device(a, t); + for (const auto& t : {a, b, scales, zeros, g_idx, perm, bias}) + w4::disjoint(out, t); + } + virtual void operator()(const Tensor a, const Tensor b, const Tensor scales, + const Tensor zeros, const Tensor g_idx, + const Tensor perm, const Tensor bias, + const int64_t scheme, const int64_t group_size, + const bool is_k_full, const float alpha, + Tensor out) const = 0; + + protected: + int64_t scheme_, group_; + Tensor::Size k_, n_, m_; + bool full_; + float alpha_; +}; +} // namespace infini::ops diff --git a/src/base/w4_marlin_parameters.h b/src/base/w4_marlin_parameters.h new file mode 100644 index 000000000..62b17de20 --- /dev/null +++ b/src/base/w4_marlin_parameters.h @@ -0,0 +1,54 @@ +#pragma once +#include "base/w4_marlin_common.h" + +namespace infini::ops { +// Scale and AWQ zero layout transformation is an Ops loading primitive. +// For GPTQ symmetric input, qzeros/zeros_out are empty after LM validates +// z+1=8. +class W4MarlinParameters : public Operator { + public: + W4MarlinParameters(const Tensor scales, const Tensor qzeros, + 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) + : scheme_{scheme}, + group_{group_size}, + k_{static_cast(size_k)}, + n_{static_cast(size_n)} { + w4::dims(size_k, size_n); + w4::groups(k_, group_); + w4::require(scheme == 0 || scheme == 1, "unknown W4 parameter scheme"); + w4::require(w4::fp16(scales) && scales.ndim() == 2 && scales.size(0) > 0, + "W4 scales must be a nonempty F16/BF16 matrix"); + const auto g = scales.size(0); + w4::require(w4::mul(g, n_) <= + static_cast(std::numeric_limits::max()), + "W4 scale indexing exceeds int32 limits"); + w4::require( + scheme != 0 || g == static_cast(w4::groups(k_, group_)), + "AWQ scales rows mismatch group_size"); + w4::shape(scales, {g, n_}, scales.dtype()); + w4::shape(scales_out, {g, n_}, scales.dtype()); + w4::shape(qzeros, scheme == 0 ? Tensor::Shape{g, n_ / 8} : Tensor::Shape{0}, + DataType::kInt32); + w4::shape(zeros_out, + scheme == 0 ? Tensor::Shape{g, n_ / 8} : Tensor::Shape{0}, + DataType::kInt32); + for (const auto& t : {scales, qzeros, scales_out, zeros_out}) + w4::device(scales, t); + for (const auto& t : {scales, qzeros}) { + w4::disjoint(scales_out, t); + w4::disjoint(zeros_out, t); + } + w4::disjoint(scales_out, zeros_out); + } + virtual void operator()(const Tensor scales, const Tensor qzeros, + 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 = 0; + + protected: + int64_t scheme_, group_; + Tensor::Size k_, n_; +}; +} // namespace infini::ops diff --git a/src/base/w4_marlin_repack.h b/src/base/w4_marlin_repack.h new file mode 100644 index 000000000..caa533c1c --- /dev/null +++ b/src/base/w4_marlin_repack.h @@ -0,0 +1,37 @@ +#pragma once +#include "base/w4_marlin_common.h" + +namespace infini::ops { +// scheme=0: AWQ GEMM nibble order; scheme=1: GPTQ K-packed order. +// perm is a verified loading-time bijection; empty means identity. +class W4MarlinRepack : public Operator { + public: + W4MarlinRepack(const Tensor qweight, const Tensor perm, const int64_t scheme, + const int64_t size_k, const int64_t size_n, Tensor out) + : scheme_{scheme}, + k_{static_cast(size_k)}, + n_{static_cast(size_n)} { + w4::dims(size_k, size_n); + w4::require(scheme == 0 || scheme == 1, "unknown W4 checkpoint packing"); + w4::shape( + qweight, + scheme == 0 ? Tensor::Shape{k_, n_ / 8} : Tensor::Shape{k_ / 8, n_}, + DataType::kInt32); + w4::shape(out, {k_ / 16, n_ * 2}, DataType::kInt32); + w4::shape(perm, {perm.numel() ? k_ : 0}, DataType::kInt32); + w4::require(scheme == 1 || perm.numel() == 0, + "AWQ cannot take a K permutation"); + w4::device(qweight, perm); + w4::device(qweight, out); + w4::disjoint(out, qweight); + w4::disjoint(out, perm); + } + virtual 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 = 0; + + protected: + int64_t scheme_; + Tensor::Size k_, n_; +}; +} // namespace infini::ops 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/ascend/ops/tril/kernel.h b/src/native/ascend/ops/tril/kernel.h new file mode 100644 index 000000000..24ec19bbf --- /dev/null +++ b/src/native/ascend/ops/tril/kernel.h @@ -0,0 +1,71 @@ +#ifndef INFINI_OPS_ASCEND_TRIL_KERNEL_H_ +#define INFINI_OPS_ASCEND_TRIL_KERNEL_H_ + +#include +#include + +#include "aclnnop/aclnn_tril.h" +#include "base/tril.h" +#include "native/ascend/common.h" + +namespace infini::ops { + +// Eager triangular masking. Each call owns fresh ACL descriptors and workspace; +// exact input/output aliasing uses the dedicated inplace API. +template <> +class Operator : public Tril { + public: + Operator(const Tensor input, const int64_t diagonal, Tensor out) + : Tril(input, diagonal, out) {} + + void operator()(const Tensor input, const int64_t diagonal, + Tensor out) const override { + auto stream = static_cast(stream_); + auto t_input = ascend::BuildAclTensor(input); + auto t_out = ascend::BuildAclTensor(out); + if (!t_input || !t_out) FailStop("descriptor creation", -1); + uint64_t workspace_size = 0; + aclOpExecutor* executor = nullptr; + const bool inplace = input.data() == out.data(); + auto ret = inplace ? aclnnInplaceTrilGetWorkspaceSize( + t_out, diagonal, &workspace_size, &executor) + : aclnnTrilGetWorkspaceSize(t_input, diagonal, t_out, + &workspace_size, &executor); + if (ret != ACL_SUCCESS || !executor) FailStop("workspace query", ret); + void* workspace = nullptr; + if (workspace_size) { + ret = aclrtMalloc(&workspace, workspace_size, ACL_MEM_MALLOC_NORMAL_ONLY); + if (ret != ACL_SUCCESS || !workspace) + FailStop("workspace allocation", ret); + } + ret = inplace + ? aclnnInplaceTril(workspace, workspace_size, executor, stream) + : aclnnTril(workspace, workspace_size, executor, stream); + if (ret != ACL_SUCCESS) FailStop("execution", ret); + // Phase two consumes a non-repeatable executor. Retain caller descriptors + // and backing until successful completion before releasing local resources. + ret = aclrtSynchronizeStream(stream); + if (ret != ACL_SUCCESS) FailStop("stream completion", ret); + if (workspace) { + ret = aclrtFree(workspace); + if (ret != ACL_SUCCESS) FailStop("workspace release", ret); + } + ret = aclDestroyTensor(t_input); + if (ret != ACL_SUCCESS) FailStop("input descriptor release", ret); + ret = aclDestroyTensor(t_out); + if (ret != ACL_SUCCESS) FailStop("output descriptor release", ret); + } + + private: + // ACL failures may leave executor ownership or queued work uncertain. + // Retain resources and stop this process without freeing in-flight backing. + [[noreturn]] static void FailStop(const char* phase, int64_t status) { + std::fprintf(stderr, "Ascend Tril %s failed (%lld)\n", phase, + static_cast(status)); + std::abort(); + } +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/ascend/ops/triu/kernel.h b/src/native/ascend/ops/triu/kernel.h new file mode 100644 index 000000000..7f4d4c8dd --- /dev/null +++ b/src/native/ascend/ops/triu/kernel.h @@ -0,0 +1,71 @@ +#ifndef INFINI_OPS_ASCEND_TRIU_KERNEL_H_ +#define INFINI_OPS_ASCEND_TRIU_KERNEL_H_ + +#include +#include + +#include "aclnnop/aclnn_triu.h" +#include "base/triu.h" +#include "native/ascend/common.h" + +namespace infini::ops { + +// Eager triangular masking. Each call owns fresh ACL descriptors and workspace; +// exact input/output aliasing uses the dedicated inplace API. +template <> +class Operator : public Triu { + public: + Operator(const Tensor input, const int64_t diagonal, Tensor out) + : Triu(input, diagonal, out) {} + + void operator()(const Tensor input, const int64_t diagonal, + Tensor out) const override { + auto stream = static_cast(stream_); + auto t_input = ascend::BuildAclTensor(input); + auto t_out = ascend::BuildAclTensor(out); + if (!t_input || !t_out) FailStop("descriptor creation", -1); + uint64_t workspace_size = 0; + aclOpExecutor* executor = nullptr; + const bool inplace = input.data() == out.data(); + auto ret = inplace ? aclnnInplaceTriuGetWorkspaceSize( + t_out, diagonal, &workspace_size, &executor) + : aclnnTriuGetWorkspaceSize(t_input, diagonal, t_out, + &workspace_size, &executor); + if (ret != ACL_SUCCESS || !executor) FailStop("workspace query", ret); + void* workspace = nullptr; + if (workspace_size) { + ret = aclrtMalloc(&workspace, workspace_size, ACL_MEM_MALLOC_NORMAL_ONLY); + if (ret != ACL_SUCCESS || !workspace) + FailStop("workspace allocation", ret); + } + ret = inplace + ? aclnnInplaceTriu(workspace, workspace_size, executor, stream) + : aclnnTriu(workspace, workspace_size, executor, stream); + if (ret != ACL_SUCCESS) FailStop("execution", ret); + // Phase two consumes a non-repeatable executor. Retain caller descriptors + // and backing until successful completion before releasing local resources. + ret = aclrtSynchronizeStream(stream); + if (ret != ACL_SUCCESS) FailStop("stream completion", ret); + if (workspace) { + ret = aclrtFree(workspace); + if (ret != ACL_SUCCESS) FailStop("workspace release", ret); + } + ret = aclDestroyTensor(t_input); + if (ret != ACL_SUCCESS) FailStop("input descriptor release", ret); + ret = aclDestroyTensor(t_out); + if (ret != ACL_SUCCESS) FailStop("output descriptor release", ret); + } + + private: + // ACL failures may leave executor ownership or queued work uncertain. + // Retain resources and stop this process without freeing in-flight backing. + [[noreturn]] static void FailStop(const char* phase, int64_t status) { + std::fprintf(stderr, "Ascend Triu %s failed (%lld)\n", phase, + static_cast(status)); + std::abort(); + } +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cambricon/data_type_.h b/src/native/cambricon/data_type_.h index 805b311ed..af029413a 100644 --- a/src/native/cambricon/data_type_.h +++ b/src/native/cambricon/data_type_.h @@ -2,5 +2,28 @@ #define INFINI_OPS_CAMBRICON_DATA_TYPE__H_ #include +// Native arithmetic is confined to genuine BANG compilation. The public RT +// type map uses the same complete storage types in host and native objects. +#if defined(__BANGCC__) +#include +#include + +namespace infini::ops::cambricon { +template +struct NativeScalar { + using type = StorageT; +}; +template <> +struct NativeScalar { + using type = __half; +}; +template <> +struct NativeScalar { + using type = __bang_bfloat16; +}; +template +using NativeScalarType = typename NativeScalar::type; +} // namespace infini::ops::cambricon +#endif #endif diff --git a/src/native/cambricon/dense_triangle.h b/src/native/cambricon/dense_triangle.h new file mode 100644 index 000000000..a68662f36 --- /dev/null +++ b/src/native/cambricon/dense_triangle.h @@ -0,0 +1,11 @@ +#ifndef INFINI_OPS_CAMBRICON_DENSE_TRIANGLE_LAUNCH_H_ +#define INFINI_OPS_CAMBRICON_DENSE_TRIANGLE_LAUNCH_H_ +#include + +#include +#include +namespace infini::ops { +void DenseTriangleLaunch(cnrtQueue_t, const void*, void*, std::size_t, + std::size_t, std::size_t, int64_t, bool, std::size_t); +} +#endif diff --git a/src/native/cambricon/ops/add/kernel.mlu b/src/native/cambricon/ops/add/kernel.mlu index 355de0bb1..2cbda223d 100644 --- a/src/native/cambricon/ops/add/kernel.mlu +++ b/src/native/cambricon/ops/add/kernel.mlu @@ -136,7 +136,7 @@ __mlu_global__ void AddKernel( } } -template +template void AddUnion(void* workspace, int core_per_cluster, int cluster_count, cnrtQueue_t queue, const void* input, const void* other, void* out, const size_t* out_shape, const ptrdiff_t* out_strides, @@ -144,6 +144,8 @@ void AddUnion(void* workspace, int core_per_cluster, int cluster_count, const size_t* other_shape, const ptrdiff_t* other_strides, size_t output_size, int ndim, bool fast_path, bool out_contiguous, double alpha) { + using T = cambricon::NativeScalarType; + cnrtDim3_t kernel_dim; cnrtFunctionType_t kernel_type; @@ -190,19 +192,19 @@ void AddUnion(void* workspace, int core_per_cluster, int cluster_count, fast_path, out_contiguous, static_cast>(alpha)); } -template void AddUnion<__half>(void*, int, int, cnrtQueue_t, const void*, - const void*, void*, const size_t*, - const ptrdiff_t*, const size_t*, - const ptrdiff_t*, const size_t*, - const ptrdiff_t*, size_t, int, bool, bool, - double); - -template void AddUnion<__bang_bfloat16>(void*, int, int, cnrtQueue_t, - const void*, const void*, void*, - const size_t*, const ptrdiff_t*, - const size_t*, const ptrdiff_t*, - const size_t*, const ptrdiff_t*, size_t, - int, bool, bool, double); +template void AddUnion(void*, int, int, cnrtQueue_t, const void*, + const void*, void*, const size_t*, + const ptrdiff_t*, const size_t*, + const ptrdiff_t*, const size_t*, + const ptrdiff_t*, size_t, int, bool, bool, + double); + +template void AddUnion(void*, int, int, cnrtQueue_t, const void*, + const void*, void*, const size_t*, + const ptrdiff_t*, const size_t*, + const ptrdiff_t*, const size_t*, + const ptrdiff_t*, size_t, int, bool, bool, + double); template void AddUnion(void*, int, int, cnrtQueue_t, const void*, const void*, void*, const size_t*, diff --git a/src/native/cambricon/ops/argmax/kernel.h b/src/native/cambricon/ops/argmax/kernel.h new file mode 100644 index 000000000..ae31554eb --- /dev/null +++ b/src/native/cambricon/ops/argmax/kernel.h @@ -0,0 +1,48 @@ +#ifndef INFINI_OPS_CAMBRICON_ARGMAX_KERNEL_H_ +#define INFINI_OPS_CAMBRICON_ARGMAX_KERNEL_H_ +#include +#include +#include + +#include "base/argmax.h" +#include "native/cambricon/common.h" +namespace infini::ops { +void CambriconArgmaxLaunch(cnrtQueue_t, const void*, void*, int32_t, + std::size_t, uint32_t); +template <> +class Operator : public Argmax { + public: + static bool supports(const Tensor x, const std::optional dim, + bool keepdim, const Tensor y) { + if (x.device().type() != Device::Type::kCambricon || + x.device().index() < 0 || x.dtype() != DataType::kFloat32 || + !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()), + b = reinterpret_cast(y.data()); + if (a % sizeof(float) || b % alignof(int64_t)) return false; + return b > a ? b - a >= x.numel() * sizeof(float) + : a - b >= sizeof(int64_t); + } + Operator(const Tensor x, const std::optional dim, bool keepdim, + Tensor y) + : Argmax(x, dim, keepdim, y) {} + void operator()(const Tensor x, const std::optional dim, + bool keepdim, Tensor y) const override { + if (!supports(x, dim, keepdim, y) || x.shape() != input_shape_ || + x.strides() != input_strides_ || x.dtype() != input_type_ || + y.shape() != out_shape_ || y.strides() != out_strides_ || + y.dtype() != out_type_ || x.device().index() != device_index_ || + dim != dim_ || keepdim != keepdim_) + throw std::invalid_argument( + "Cambricon Argmax prepared FP32 metadata changed"); + CambriconArgmaxLaunch(static_cast(stream_), x.data(), y.data(), + static_cast(x.numel()), sizeof(float), + 0x7f800000U); + CNRT_CHECK(cnrtGetLastError()); + } +}; +} // namespace infini::ops +#endif diff --git a/src/native/cambricon/ops/argmax/kernel.mlu b/src/native/cambricon/ops/argmax/kernel.mlu new file mode 100644 index 000000000..bba862691 --- /dev/null +++ b/src/native/cambricon/ops/argmax/kernel.mlu @@ -0,0 +1,55 @@ +#include + +#include "native/cambricon/ops/argmax/kernel.h" + +namespace infini::ops { +template +__mlu_global__ void CambriconArgmaxKernel(const Bits* input, int64_t* output, + int32_t count, + uint32_t positive_infinity) { + __nram__ Bits buffer[1024]; + __nram__ int64_t result; + const uint32_t sign = sizeof(Bits) == 4 ? 0x80000000U : 0x8000U; + const uint32_t mask = sizeof(Bits) == 4 ? 0xffffffffU : 0xffffU; + uint32_t best_key = 0; + int32_t best_index = 0; + bool best_nan = false; + // Storage-only ordering avoids FTZ or conversion rounding of subnormals. + // Ascending traversal and strict replacement preserve the first equal max. + for (int32_t start = 0; start < count;) { + const int32_t n = count - start < 1024 ? count - start : 1024; + __memcpy(buffer, input + start, n * sizeof(Bits), GDRAM2NRAM); + for (int32_t j = 0; j < n; ++j) { + uint32_t bits = static_cast(buffer[j]); + const uint32_t magnitude = bits & (sign - 1U); + const bool value_nan = magnitude > positive_infinity; + if (magnitude == 0) bits = 0; // -0 and +0 have equal numeric value. + const uint32_t key = bits & sign ? (~bits & mask) : (bits ^ sign); + if (start + j == 0 || + (value_nan ? !best_nan : !best_nan && key > best_key)) { + best_key = key; + best_index = start + j; + best_nan = value_nan; + } + } + start += n; // Never advances beyond INT_MAX for the validated domain. + } + result = static_cast(best_index); + __memcpy(output, &result, sizeof(result), NRAM2GDRAM); +} + +void CambriconArgmaxLaunch(cnrtQueue_t queue, const void* input, void* output, + int32_t count, std::size_t item_bytes, + uint32_t positive_infinity) { + const cnrtDim3_t dim = {1, 1, 1}; + if (item_bytes == 4) { + CambriconArgmaxKernel<<>>( + static_cast(input), static_cast(output), + count, positive_infinity); + } else { + CambriconArgmaxKernel<<>>( + static_cast(input), static_cast(output), + count, positive_infinity); + } +} +} // namespace infini::ops diff --git a/src/native/cambricon/ops/embedding/kernel.mlu b/src/native/cambricon/ops/embedding/kernel.mlu index 3d1c45fdb..a4f64adf1 100644 --- a/src/native/cambricon/ops/embedding/kernel.mlu +++ b/src/native/cambricon/ops/embedding/kernel.mlu @@ -257,7 +257,8 @@ void EmbeddingKernelLaunch( {input_dtype, weight_dtype}, [&](auto input_tag, auto weight_tag) { using IndexT = typename decltype(input_tag)::type; - using T = typename decltype(weight_tag)::type; + using T = + cambricon::NativeScalarType; if (apply_max_norm && vocab_size > 0) { (void)cnrtGetLastError(); diff --git a/src/native/cambricon/ops/fill/kernel.mlu b/src/native/cambricon/ops/fill/kernel.mlu index 0f508b9ca..bba39f1fc 100644 --- a/src/native/cambricon/ops/fill/kernel.mlu +++ b/src/native/cambricon/ops/fill/kernel.mlu @@ -13,12 +13,11 @@ __nram__ char fill_nram_buffer[NRAM_MAX_SIZE] __attribute__((aligned(128))); template __mlu_device__ void FillBuffer(T* buffer, std::size_t count, - const T* device_value) { + const T* device_value, T host_scalar) { + const T value = device_value != nullptr ? device_value[0] : host_scalar; if constexpr (std::is_same_v) { - const T value = device_value[0]; __bang_write_value(buffer, count, value); } else { - const T value = device_value[0]; for (std::size_t i = 0; i < count; ++i) { buffer[i] = value; } @@ -26,7 +25,7 @@ __mlu_device__ void FillBuffer(T* buffer, std::size_t count, } template -__mlu_global__ void FillKernel(const T* device_value, T* out, +__mlu_global__ void FillKernel(const T* device_value, T host_scalar, T* out, const std::size_t* out_shape, const ptrdiff_t* out_strides, std::size_t output_size, int ndim, @@ -44,7 +43,7 @@ __mlu_global__ void FillKernel(const T* device_value, T* out, for (std::size_t processed = range.begin; processed < range.end;) { const std::size_t current = std::min(block_size, range.end - processed); - FillBuffer(buffer, current, device_value); + FillBuffer(buffer, current, device_value, host_scalar); if (out_contiguous) { __memcpy(out + processed, buffer, current * sizeof(T), NRAM2GDRAM); @@ -108,14 +107,17 @@ void FillRaw64Union(void* workspace, cnrtQueue_t queue, int core_per_cluster, CNRT_CHECK(cnrtGetLastError()); } -template +template void FillUnion(void* workspace, cnrtQueue_t queue, int core_per_cluster, - int cluster_count, const T* device_value, double host_value, - T* out, const std::size_t* out_shape, - const ptrdiff_t* out_strides, std::size_t output_size, int ndim, - bool out_contiguous) { + int cluster_count, const StorageT* storage_device_value, + double host_value, StorageT* storage_out, + const std::size_t* out_shape, const ptrdiff_t* out_strides, + std::size_t output_size, int ndim, bool out_contiguous) { + using T = cambricon::NativeScalarType; + auto device_value = reinterpret_cast(storage_device_value); + auto out = reinterpret_cast(storage_out); + auto* bytes = static_cast(workspace); - auto* device_scalar = reinterpret_cast(bytes); bytes += sizeof(std::uint64_t); std::size_t* device_shape = nullptr; ptrdiff_t* device_strides = nullptr; @@ -132,21 +134,18 @@ void FillUnion(void* workspace, cnrtQueue_t queue, int core_per_cluster, const cnrtDim3_t kernel_dim = {static_cast(core_per_cluster), static_cast(cluster_count), 1}; + T converted_host_value{}; if (device_value == nullptr) { - T converted_host_value; if constexpr (std::is_same_v || std::is_same_v) { converted_host_value = static_cast(static_cast(host_value)); } else { converted_host_value = static_cast(host_value); } - CNRT_CHECK(cnrtMemcpy(device_scalar, &converted_host_value, sizeof(T), - cnrtMemcpyHostToDev)); - device_value = device_scalar; } FillKernel<<>>( - device_value, out, device_shape, device_strides, output_size, ndim, - out_contiguous); + device_value, converted_host_value, out, device_shape, device_strides, + output_size, ndim, out_contiguous); CNRT_CHECK(cnrtGetLastError()); } @@ -155,8 +154,8 @@ void FillUnion(void* workspace, cnrtQueue_t queue, int core_per_cluster, T*, const std::size_t*, const ptrdiff_t*, \ std::size_t, int, bool) -INSTANTIATE_FILL(__half); -INSTANTIATE_FILL(__bang_bfloat16); +INSTANTIATE_FILL(Float16); +INSTANTIATE_FILL(BFloat16); INSTANTIATE_FILL(float); INSTANTIATE_FILL(int8_t); INSTANTIATE_FILL(int16_t); diff --git a/src/native/cambricon/ops/fused_add_rms_norm/kernel.mlu b/src/native/cambricon/ops/fused_add_rms_norm/kernel.mlu index bef1592a3..fc51b1928 100644 --- a/src/native/cambricon/ops/fused_add_rms_norm/kernel.mlu +++ b/src/native/cambricon/ops/fused_add_rms_norm/kernel.mlu @@ -118,12 +118,14 @@ __mlu_global__ void FusedAddRmsNormKernel( } // namespace fused_add_rms_norm_detail -template +template void FusedAddRmsNormUnion(int core_per_cluster, int cluster_count, cnrtQueue_t queue, void* input, void* residual, const void* weight, size_t num_tokens, size_t dim, ptrdiff_t input_row_stride, ptrdiff_t residual_row_stride, float epsilon) { + using T = cambricon::NativeScalarType; + cnrtDim3_t kernel_dim; kernel_dim.x = core_per_cluster; kernel_dim.y = cluster_count; @@ -138,13 +140,12 @@ void FusedAddRmsNormUnion(int core_per_cluster, int cluster_count, CNRT_CHECK(cnrtGetLastError()); } -template void FusedAddRmsNormUnion<__half>(int, int, cnrtQueue_t, void*, void*, - const void*, size_t, size_t, - ptrdiff_t, ptrdiff_t, float); -template void FusedAddRmsNormUnion<__bang_bfloat16>(int, int, cnrtQueue_t, - void*, void*, const void*, - size_t, size_t, ptrdiff_t, - ptrdiff_t, float); +template void FusedAddRmsNormUnion(int, int, cnrtQueue_t, void*, void*, + const void*, size_t, size_t, + ptrdiff_t, ptrdiff_t, float); +template void FusedAddRmsNormUnion(int, int, cnrtQueue_t, void*, + void*, const void*, size_t, size_t, + ptrdiff_t, ptrdiff_t, float); template void FusedAddRmsNormUnion(int, int, cnrtQueue_t, void*, void*, const void*, size_t, size_t, ptrdiff_t, ptrdiff_t, float); diff --git a/src/native/cambricon/ops/gelu/kernel.mlu b/src/native/cambricon/ops/gelu/kernel.mlu index b79775873..46e563cc4 100644 --- a/src/native/cambricon/ops/gelu/kernel.mlu +++ b/src/native/cambricon/ops/gelu/kernel.mlu @@ -113,13 +113,15 @@ __mlu_global__ void GeluKernel(const T* input, T* output, } // namespace -template +template void GeluUnion(void* workspace, int core_per_cluster, int cluster_count, cnrtQueue_t queue, const void* input, void* out, const size_t* input_shape, const ptrdiff_t* input_strides, const size_t* out_shape, const ptrdiff_t* out_strides, size_t output_size, int ndim, bool approximate, bool input_contiguous, bool out_contiguous) { + using T = cambricon::NativeScalarType; + size_t* device_input_shape = nullptr; size_t* device_out_shape = nullptr; ptrdiff_t* device_input_strides = nullptr; @@ -163,8 +165,8 @@ void GeluUnion(void* workspace, int core_per_cluster, int cluster_count, const size_t*, const ptrdiff_t*, const size_t*, \ const ptrdiff_t*, size_t, int, bool, bool, bool) -INSTANTIATE_GELU(__half); -INSTANTIATE_GELU(__bang_bfloat16); +INSTANTIATE_GELU(Float16); +INSTANTIATE_GELU(BFloat16); INSTANTIATE_GELU(float); #undef INSTANTIATE_GELU diff --git a/src/native/cambricon/ops/mul/kernel.mlu b/src/native/cambricon/ops/mul/kernel.mlu index 189f13f53..10719c3dc 100644 --- a/src/native/cambricon/ops/mul/kernel.mlu +++ b/src/native/cambricon/ops/mul/kernel.mlu @@ -85,13 +85,19 @@ __mlu_global__ void MulKernel(const T* input, const T* other, T* out, } // namespace -template +template void MulUnion(void* workspace, cnrtQueue_t queue, int core_per_cluster, - int cluster_count, const T* input, const T* other, T* out, + int cluster_count, const StorageT* storage_input, + const StorageT* storage_other, StorageT* storage_out, const std::size_t* out_shape, const ptrdiff_t* input_strides, const ptrdiff_t* other_strides, const ptrdiff_t* out_strides, std::size_t output_size, int ndim, bool fast_path, bool out_contiguous) { + using T = cambricon::NativeScalarType; + auto input = reinterpret_cast(storage_input); + auto other = reinterpret_cast(storage_other); + auto out = reinterpret_cast(storage_out); + auto* bytes = static_cast(workspace); auto* device_shape = reinterpret_cast(bytes); auto* device_input_strides = @@ -129,8 +135,8 @@ void MulUnion(void* workspace, cnrtQueue_t queue, int core_per_cluster, const ptrdiff_t*, const ptrdiff_t*, std::size_t, \ int, bool, bool) -INSTANTIATE_MUL(__half); -INSTANTIATE_MUL(__bang_bfloat16); +INSTANTIATE_MUL(Float16); +INSTANTIATE_MUL(BFloat16); INSTANTIATE_MUL(float); INSTANTIATE_MUL(int16_t); INSTANTIATE_MUL(int32_t); diff --git a/src/native/cambricon/ops/relu/kernel.mlu b/src/native/cambricon/ops/relu/kernel.mlu index 4e5dd5863..beaeb3516 100644 --- a/src/native/cambricon/ops/relu/kernel.mlu +++ b/src/native/cambricon/ops/relu/kernel.mlu @@ -101,13 +101,15 @@ __mlu_global__ void ReluKernel(const T* input, T* output, const size_t* shape, } // namespace -template +template void ReluUnion(void* workspace, int core_per_cluster, int cluster_count, cnrtQueue_t queue, const void* input, void* out, const size_t* shape, const ptrdiff_t* input_strides, const ptrdiff_t* out_strides, size_t output_size, int ndim, bool input_contiguous, bool out_contiguous, bool needs_input_copy) { + using T = cambricon::NativeScalarType; + auto* workspace_bytes = static_cast(workspace); auto* device_shape = reinterpret_cast(workspace_bytes); auto* device_input_strides = @@ -162,8 +164,8 @@ void ReluUnion(void* workspace, int core_per_cluster, int cluster_count, const ptrdiff_t*, size_t, int, bool, bool, bool) INSTANTIATE_RELU(float); -INSTANTIATE_RELU(__half); -INSTANTIATE_RELU(__bang_bfloat16); +INSTANTIATE_RELU(Float16); +INSTANTIATE_RELU(BFloat16); INSTANTIATE_RELU(int64_t); INSTANTIATE_RELU(int32_t); INSTANTIATE_RELU(int16_t); diff --git a/src/native/cambricon/ops/reshape_and_cache_flash/kernel.mlu b/src/native/cambricon/ops/reshape_and_cache_flash/kernel.mlu index 5def4e0b2..bd4ca4de0 100644 --- a/src/native/cambricon/ops/reshape_and_cache_flash/kernel.mlu +++ b/src/native/cambricon/ops/reshape_and_cache_flash/kernel.mlu @@ -47,7 +47,7 @@ __mlu_global__ void ReshapeAndCacheFlashKernel( } // namespace reshape_and_cache_flash_detail -template +template void ReshapeAndCacheFlashUnion( int core_per_cluster, int cluster_count, cnrtQueue_t queue, const void* key, const void* value, const void* slot_mapping, void* key_cache, @@ -61,6 +61,8 @@ void ReshapeAndCacheFlashUnion( std::ptrdiff_t value_cache_page_stride, std::ptrdiff_t key_cache_head_stride, std::ptrdiff_t value_cache_head_stride) { + using T = cambricon::NativeScalarType; + cnrtDim3_t kernel_dim; kernel_dim.x = core_per_cluster; kernel_dim.y = cluster_count; @@ -88,8 +90,8 @@ void ReshapeAndCacheFlashUnion( std::ptrdiff_t, std::ptrdiff_t, std::ptrdiff_t, std::ptrdiff_t, \ std::ptrdiff_t, std::ptrdiff_t) -INFINI_INSTANTIATE_RESHAPE_AND_CACHE_FLASH(__half); -INFINI_INSTANTIATE_RESHAPE_AND_CACHE_FLASH(__bang_bfloat16); +INFINI_INSTANTIATE_RESHAPE_AND_CACHE_FLASH(Float16); +INFINI_INSTANTIATE_RESHAPE_AND_CACHE_FLASH(BFloat16); INFINI_INSTANTIATE_RESHAPE_AND_CACHE_FLASH(float); #undef INFINI_INSTANTIATE_RESHAPE_AND_CACHE_FLASH diff --git a/src/native/cambricon/ops/rms_norm/kernel.mlu b/src/native/cambricon/ops/rms_norm/kernel.mlu index 0c8dcf9dd..4d225039f 100644 --- a/src/native/cambricon/ops/rms_norm/kernel.mlu +++ b/src/native/cambricon/ops/rms_norm/kernel.mlu @@ -233,11 +233,14 @@ __mlu_global__ void RmsNorm(const T* input, const TW* weight, T* output, } } -template +template void RmsNormUnion(void* workspace, int core_per_cluster, int cluster_count, cnrtQueue_t queue, void* y, const void* x, const void* w, const size_t* shape, const ptrdiff_t* y_strides, const ptrdiff_t* x_strides, float eps, int ndim) { + using T = cambricon::NativeScalarType; + using TW = cambricon::NativeScalarType; + cnrtDim3_t kernel_dim; cnrtFunctionType_t kernel_type; @@ -287,40 +290,45 @@ void RmsNormUnion(void* workspace, int core_per_cluster, int cluster_count, x_, w_, y_, mlu_shape, mlu_y_strides, mlu_x_strides, eps, ndim, dim_s); } -template void RmsNormUnion<__half, __half>(void*, int, int, cnrtQueue_t, void*, +template void RmsNormUnion(void*, int, int, cnrtQueue_t, + void*, const void*, const void*, + const size_t*, const ptrdiff_t*, + const ptrdiff_t*, float, int); + +template void RmsNormUnion(void*, int, int, cnrtQueue_t, + void*, const void*, const void*, + const size_t*, const ptrdiff_t*, + const ptrdiff_t*, float, int); + +template void RmsNormUnion(void*, int, int, cnrtQueue_t, void*, const void*, const void*, const size_t*, const ptrdiff_t*, const ptrdiff_t*, float, int); -template void RmsNormUnion<__half, __bang_bfloat16>( - void*, int, int, cnrtQueue_t, void*, const void*, const void*, - const size_t*, const ptrdiff_t*, const ptrdiff_t*, float, int); - -template void RmsNormUnion<__half, float>(void*, int, int, cnrtQueue_t, void*, - const void*, const void*, - const size_t*, const ptrdiff_t*, - const ptrdiff_t*, float, int); +template void RmsNormUnion(void*, int, int, cnrtQueue_t, + void*, const void*, const void*, + const size_t*, const ptrdiff_t*, + const ptrdiff_t*, float, int); -template void RmsNormUnion<__bang_bfloat16, __half>( - void*, int, int, cnrtQueue_t, void*, const void*, const void*, - const size_t*, const ptrdiff_t*, const ptrdiff_t*, float, int); +template void RmsNormUnion(void*, int, int, cnrtQueue_t, + void*, const void*, const void*, + const size_t*, const ptrdiff_t*, + const ptrdiff_t*, float, int); -template void RmsNormUnion<__bang_bfloat16, __bang_bfloat16>( - void*, int, int, cnrtQueue_t, void*, const void*, const void*, - const size_t*, const ptrdiff_t*, const ptrdiff_t*, float, int); +template void RmsNormUnion(void*, int, int, cnrtQueue_t, void*, + const void*, const void*, + const size_t*, const ptrdiff_t*, + const ptrdiff_t*, float, int); -template void RmsNormUnion<__bang_bfloat16, float>( - void*, int, int, cnrtQueue_t, void*, const void*, const void*, - const size_t*, const ptrdiff_t*, const ptrdiff_t*, float, int); - -template void RmsNormUnion(void*, int, int, cnrtQueue_t, void*, - const void*, const void*, - const size_t*, const ptrdiff_t*, - const ptrdiff_t*, float, int); +template void RmsNormUnion(void*, int, int, cnrtQueue_t, void*, + const void*, const void*, + const size_t*, const ptrdiff_t*, + const ptrdiff_t*, float, int); -template void RmsNormUnion( - void*, int, int, cnrtQueue_t, void*, const void*, const void*, - const size_t*, const ptrdiff_t*, const ptrdiff_t*, float, int); +template void RmsNormUnion(void*, int, int, cnrtQueue_t, void*, + const void*, const void*, + const size_t*, const ptrdiff_t*, + const ptrdiff_t*, float, int); template void RmsNormUnion(void*, int, int, cnrtQueue_t, void*, const void*, const void*, diff --git a/src/native/cambricon/ops/rotary_embedding/kernel.mlu b/src/native/cambricon/ops/rotary_embedding/kernel.mlu index 0b240e83b..416ecd9aa 100644 --- a/src/native/cambricon/ops/rotary_embedding/kernel.mlu +++ b/src/native/cambricon/ops/rotary_embedding/kernel.mlu @@ -150,15 +150,19 @@ void LaunchRotaryEmbeddingCambricon( size_t rope_dim_offset, ptrdiff_t batch_stride, ptrdiff_t token_stride, ptrdiff_t head_stride, ptrdiff_t cache_token_stride, bool is_neox, bool inverse) { + using NativeT = cambricon::NativeScalarType; + using NativeCache = cambricon::NativeScalarType; + cnrtDim3_t kernel_dim; kernel_dim.x = core_per_cluster; kernel_dim.y = cluster_count; kernel_dim.z = 1; - RotaryEmbeddingCambriconKernel + RotaryEmbeddingCambriconKernel <<>>( - positions, data, cos_sin_cache, num_tokens, sequence_length, - num_heads, rot_dim, rope_dim_offset, batch_stride, token_stride, - head_stride, cache_token_stride, is_neox, inverse); + positions, reinterpret_cast(data), + reinterpret_cast(cos_sin_cache), num_tokens, + sequence_length, num_heads, rot_dim, rope_dim_offset, batch_stride, + token_stride, head_stride, cache_token_stride, is_neox, inverse); } #define INSTANTIATE_ROTARY_EMBEDDING(T, TCACHE) \ @@ -167,14 +171,14 @@ void LaunchRotaryEmbeddingCambricon( size_t, size_t, size_t, size_t, ptrdiff_t, ptrdiff_t, ptrdiff_t, \ ptrdiff_t, bool, bool) -INSTANTIATE_ROTARY_EMBEDDING(__half, __half); -INSTANTIATE_ROTARY_EMBEDDING(__half, __bang_bfloat16); -INSTANTIATE_ROTARY_EMBEDDING(__half, float); -INSTANTIATE_ROTARY_EMBEDDING(__bang_bfloat16, __half); -INSTANTIATE_ROTARY_EMBEDDING(__bang_bfloat16, __bang_bfloat16); -INSTANTIATE_ROTARY_EMBEDDING(__bang_bfloat16, float); -INSTANTIATE_ROTARY_EMBEDDING(float, __half); -INSTANTIATE_ROTARY_EMBEDDING(float, __bang_bfloat16); +INSTANTIATE_ROTARY_EMBEDDING(Float16, Float16); +INSTANTIATE_ROTARY_EMBEDDING(Float16, BFloat16); +INSTANTIATE_ROTARY_EMBEDDING(Float16, float); +INSTANTIATE_ROTARY_EMBEDDING(BFloat16, Float16); +INSTANTIATE_ROTARY_EMBEDDING(BFloat16, BFloat16); +INSTANTIATE_ROTARY_EMBEDDING(BFloat16, float); +INSTANTIATE_ROTARY_EMBEDDING(float, Float16); +INSTANTIATE_ROTARY_EMBEDDING(float, BFloat16); INSTANTIATE_ROTARY_EMBEDDING(float, float); #undef INSTANTIATE_ROTARY_EMBEDDING diff --git a/src/native/cambricon/ops/silu_and_mul/kernel.mlu b/src/native/cambricon/ops/silu_and_mul/kernel.mlu index a0a689b89..b6c4d0228 100644 --- a/src/native/cambricon/ops/silu_and_mul/kernel.mlu +++ b/src/native/cambricon/ops/silu_and_mul/kernel.mlu @@ -110,13 +110,15 @@ __mlu_global__ void SiluAndMulKernel( } // namespace -template +template void SiluAndMulUnion(void* workspace, int core_per_cluster, int cluster_count, cnrtQueue_t queue, const void* input, void* out, const size_t* input_shape, const ptrdiff_t* input_strides, const size_t* out_shape, const ptrdiff_t* out_strides, size_t output_size, size_t hidden_size, int ndim, bool input_contiguous, bool out_contiguous) { + using T = cambricon::NativeScalarType; + auto* workspace_bytes = static_cast(workspace); auto* device_input_shape = reinterpret_cast(workspace_bytes); auto* device_out_shape = device_input_shape + ndim; @@ -153,8 +155,8 @@ void SiluAndMulUnion(void* workspace, int core_per_cluster, int cluster_count, const size_t*, const ptrdiff_t*, size_t, \ size_t, int, bool, bool) -INSTANTIATE_SILU_AND_MUL(__half); -INSTANTIATE_SILU_AND_MUL(__bang_bfloat16); +INSTANTIATE_SILU_AND_MUL(Float16); +INSTANTIATE_SILU_AND_MUL(BFloat16); INSTANTIATE_SILU_AND_MUL(float); #undef INSTANTIATE_SILU_AND_MUL diff --git a/src/native/cambricon/ops/top_k_top_p_sampling_from_logits/kernel.mlu b/src/native/cambricon/ops/top_k_top_p_sampling_from_logits/kernel.mlu index 766b90bb3..3300208f0 100644 --- a/src/native/cambricon/ops/top_k_top_p_sampling_from_logits/kernel.mlu +++ b/src/native/cambricon/ops/top_k_top_p_sampling_from_logits/kernel.mlu @@ -204,12 +204,15 @@ __mlu_global__ void TopKTopPSamplingFromLogitsKernel( } // namespace -template +template void TopKTopPSamplingFromLogitsUnion( - cnrtQueue_t queue, int core_per_cluster, int cluster_count, const T* logits, - const int32_t* top_k, const float* top_p, const int32_t* row_indices, - bool has_indices, bool joint, OutputIndex* out, int32_t batch_size, - int32_t vocab_size, uint32_t seed, uint32_t offset) { + cnrtQueue_t queue, int core_per_cluster, int cluster_count, + const StorageT* storage_logits, const int32_t* top_k, const float* top_p, + const int32_t* row_indices, bool has_indices, bool joint, OutputIndex* out, + int32_t batch_size, int32_t vocab_size, uint32_t seed, uint32_t offset) { + using T = cambricon::NativeScalarType; + auto logits = reinterpret_cast(storage_logits); + const cnrtDim3_t kernel_dim = {static_cast(core_per_cluster), static_cast(cluster_count), 1}; TopKTopPSamplingFromLogitsKernel @@ -225,10 +228,10 @@ void TopKTopPSamplingFromLogitsUnion( const int32_t*, bool, bool, OUTPUT_INDEX*, int32_t, int32_t, uint32_t, \ uint32_t) -INSTANTIATE_SAMPLING(__half, int32_t); -INSTANTIATE_SAMPLING(__half, int64_t); -INSTANTIATE_SAMPLING(__bang_bfloat16, int32_t); -INSTANTIATE_SAMPLING(__bang_bfloat16, int64_t); +INSTANTIATE_SAMPLING(Float16, int32_t); +INSTANTIATE_SAMPLING(Float16, int64_t); +INSTANTIATE_SAMPLING(BFloat16, int32_t); +INSTANTIATE_SAMPLING(BFloat16, int64_t); INSTANTIATE_SAMPLING(float, int32_t); INSTANTIATE_SAMPLING(float, int64_t); diff --git a/src/native/cambricon/ops/tril/kernel.h b/src/native/cambricon/ops/tril/kernel.h new file mode 100644 index 000000000..02db31b52 --- /dev/null +++ b/src/native/cambricon/ops/tril/kernel.h @@ -0,0 +1,47 @@ +#ifndef INFINI_OPS_CAMBRICON_TRIL_KERNEL_H_ +#define INFINI_OPS_CAMBRICON_TRIL_KERNEL_H_ +#include +#include +#include + +#include "base/tril.h" +#include "native/cambricon/common.h" +#include "native/cambricon/dense_triangle.h" +namespace infini::ops { +template <> +class Operator : public Tril { + public: + static bool supports(const Tensor x, int64_t, const Tensor y) { + 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.IsContiguous() || + !y.IsContiguous() || !x.data() || !y.data()) + return false; + size_t count = 1; + for (auto n : x.shape()) { + if (!n || n > INT_MAX || count > size_t(INT_MAX) / n) return false; + count *= n; + } + const auto a = reinterpret_cast(x.data()), + b = reinterpret_cast(y.data()); + if (a % sizeof(float) || b % sizeof(float)) return false; + return a == b || (a > b ? a - b >= count * sizeof(float) + : b - a >= count * sizeof(float)); + } + Operator(const Tensor x, int64_t diagonal, Tensor y) : Tril(x, diagonal, y) {} + void operator()(const Tensor x, int64_t diagonal, Tensor y) const override { + if (!supports(x, diagonal, y) || x.shape() != input_shape_ || + x.strides() != input_strides_ || x.dtype() != input_type_ || + y.shape() != out_shape_ || y.strides() != out_strides_ || + y.dtype() != out_type_ || x.device().index() != device_index_) + throw std::invalid_argument( + "Cambricon Tril prepared FP32 metadata changed"); + DenseTriangleLaunch(static_cast(stream_), x.data(), y.data(), + y.numel(), y.size(-2), y.size(-1), diagonal, true, + sizeof(float)); + CNRT_CHECK(cnrtGetLastError()); + } +}; +} // namespace infini::ops +#endif diff --git a/src/native/cambricon/ops/tril/kernel.mlu b/src/native/cambricon/ops/tril/kernel.mlu new file mode 100644 index 000000000..061d21ec3 --- /dev/null +++ b/src/native/cambricon/ops/tril/kernel.mlu @@ -0,0 +1,62 @@ +#include + +#include "native/cambricon/common.h" +#include "native/cambricon/dense_triangle.h" +namespace infini::ops { +template +__mlu_global__ void DenseTriangleKernel(const T* src, T* dst, size_t count, + size_t rows, size_t cols, + int32_t diagonal, bool lower) { + __nram__ T buffer[1024]; + const int32_t rows32 = static_cast(rows), + cols32 = static_cast(cols); + for (size_t start = static_cast(taskId) * 1024; start < count; + start += static_cast(taskDim) * 1024) { + size_t n = count - start < 1024 ? count - start : 1024; + __memcpy(buffer, src + start, n * sizeof(T), GDRAM2NRAM); + int32_t cursor = 0; + while (cursor < static_cast(n)) { + const int32_t i = static_cast(start + cursor), col = i % cols32, + row = (i / cols32) % rows32; + const int32_t remaining = static_cast(n) - cursor; + const int32_t segment = + cols32 - col < remaining ? cols32 - col : remaining; + const int64_t boundary = + static_cast(row) + diagonal - col + (lower ? 1 : 0); + const int32_t split = boundary <= 0 ? 0 + : boundary >= segment + ? segment + : static_cast(boundary); + const int32_t begin = cursor + (lower ? split : 0), + end = cursor + (lower ? segment : split); + for (int32_t j = begin; j < end; ++j) buffer[j] = T(0); + cursor += segment; + } + __memcpy(dst + start, buffer, n * sizeof(T), NRAM2GDRAM); + } +} +void DenseTriangleLaunch(cnrtQueue_t q, const void* src, void* dst, + size_t count, size_t rows, size_t cols, + int64_t diagonal, bool lower, size_t itemsize) { + // Matrix differences are in [-(rows-1), cols-1]. Clamp on the host + // before converting to the native kernel scalar width; all int64 diagonals + // retain semantics. + const int32_t bounded = + diagonal < -static_cast(rows) ? -static_cast(rows) + : diagonal > static_cast(cols) ? static_cast(cols) + : static_cast(diagonal); + int device = -1, cores = 0, clusters = 0; + CNRT_CHECK(cnrtGetDevice(&device)); + cnrt_utils::GetLaunchConfig(Device{Device::Type::kCambricon, device}, &cores, + &clusters); + cnrtDim3_t dim = {static_cast(cores * clusters), 1, 1}; + if (itemsize == 4) + DenseTriangleKernel<<>>( + static_cast(src), static_cast(dst), count, + rows, cols, bounded, lower); + else + DenseTriangleKernel<<>>( + static_cast(src), static_cast(dst), count, + rows, cols, bounded, lower); +} +} // namespace infini::ops diff --git a/src/native/cambricon/ops/triu/kernel.h b/src/native/cambricon/ops/triu/kernel.h new file mode 100644 index 000000000..4dc04a53c --- /dev/null +++ b/src/native/cambricon/ops/triu/kernel.h @@ -0,0 +1,47 @@ +#ifndef INFINI_OPS_CAMBRICON_TRIU_KERNEL_H_ +#define INFINI_OPS_CAMBRICON_TRIU_KERNEL_H_ +#include +#include +#include + +#include "base/triu.h" +#include "native/cambricon/common.h" +#include "native/cambricon/dense_triangle.h" +namespace infini::ops { +template <> +class Operator : public Triu { + public: + static bool supports(const Tensor x, int64_t, const Tensor y) { + 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.IsContiguous() || + !y.IsContiguous() || !x.data() || !y.data()) + return false; + size_t count = 1; + for (auto n : x.shape()) { + if (!n || n > INT_MAX || count > size_t(INT_MAX) / n) return false; + count *= n; + } + const auto a = reinterpret_cast(x.data()), + b = reinterpret_cast(y.data()); + if (a % sizeof(float) || b % sizeof(float)) return false; + return a == b || (a > b ? a - b >= count * sizeof(float) + : b - a >= count * sizeof(float)); + } + Operator(const Tensor x, int64_t diagonal, Tensor y) : Triu(x, diagonal, y) {} + void operator()(const Tensor x, int64_t diagonal, Tensor y) const override { + if (!supports(x, diagonal, y) || x.shape() != input_shape_ || + x.strides() != input_strides_ || x.dtype() != input_type_ || + y.shape() != out_shape_ || y.strides() != out_strides_ || + y.dtype() != out_type_ || x.device().index() != device_index_) + throw std::invalid_argument( + "Cambricon Triu prepared FP32 metadata changed"); + DenseTriangleLaunch(static_cast(stream_), x.data(), y.data(), + y.numel(), y.size(-2), y.size(-1), diagonal, false, + sizeof(float)); + CNRT_CHECK(cnrtGetLastError()); + } +}; +} // namespace infini::ops +#endif 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/embedding/kernel.h b/src/native/cpu/ops/embedding/kernel.h new file mode 100644 index 000000000..aaa2912c7 --- /dev/null +++ b/src/native/cpu/ops/embedding/kernel.h @@ -0,0 +1,113 @@ +#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..c6bf17e9f --- /dev/null +++ b/src/native/cpu/ops/fill/fill.h @@ -0,0 +1,83 @@ +#ifndef INFINI_OPS_CPU_FILL_FILL_H_ +#define INFINI_OPS_CPU_FILL_FILL_H_ + +#include +#include +#include + +#include "base/fill.h" +#include "native/cpu/caster_.h" + +namespace infini::ops { + +// Native scalar Fill provider. Tensor-valued Fill remains 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 && + 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"); + } + + 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_ || (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; + 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 scalar Fill"); + } + + 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 requires valid scalar fill metadata"); + return input; + } +}; +} // namespace infini::ops +#endif 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.h b/src/native/cpu/ops/mul/scalar.h new file mode 100644 index 000000000..03e99aea3 --- /dev/null +++ b/src/native/cpu/ops/mul/scalar.h @@ -0,0 +1,103 @@ +#ifndef INFINI_OPS_CPU_MUL_SCALAR_H_ +#define INFINI_OPS_CPU_MUL_SCALAR_H_ + +#include +#include +#include +#include +#include +#include +#include + +#include "base/mul.h" +#include "native/cpu/caster_.h" +#include "native/cpu/dense_view.h" + +namespace infini::ops { + +class CpuScalarMul : 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; + return cpu_dense_view::DisjointOrSameContiguous(input, 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 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 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()); + 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 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_); + DispatchFunc< + Device::Type::kCpu, + List>( + input_meta_.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 scalar Mul"); + } + + private: + static bool SupportedDtype(DataType type) { + return type == DataType::kFloat32 || type == DataType::kFloat16 || + type == DataType::kBFloat16; + } + + 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 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_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..ab5ab1f2e --- /dev/null +++ b/src/native/cpu/ops/row_l2_normalize/row_l2_normalize.h @@ -0,0 +1,128 @@ +#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" +#include "native/cpu/dense_view.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; + return cpu_dense_view::DisjointOrSameContiguous(input, out); + } + + Operator(const Tensor input, const double eps, Tensor 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 (!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"); + 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 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< + Device::Type::kCpu, + List>( + input_meta_.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: + std::size_t bytes_ = 0; + 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 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..be9944c75 --- /dev/null +++ b/src/native/cpu/ops/separate_table_rope/kernel.h @@ -0,0 +1,78 @@ +#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 "math.h" +#include "native/cpu/caster_.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) {} + 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..01e7e683a --- /dev/null +++ b/src/native/cpu/ops/separate_table_rope/math.h @@ -0,0 +1,78 @@ +#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/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/blas_utils.h b/src/native/cuda/nvidia/blas_utils.h index 45644ae38..fe2134f55 100644 --- a/src/native/cuda/nvidia/blas_utils.h +++ b/src/native/cuda/nvidia/blas_utils.h @@ -5,6 +5,7 @@ #include "cublas_v2.h" // clang-format on +#include "base/gemm.h" #include "data_type.h" #include "native/cuda/blas_utils.h" @@ -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..ccd70b31f 100644 --- a/src/native/cuda/nvidia/ops/add/kernel.h +++ b/src/native/cuda/nvidia/ops/add/kernel.h @@ -14,6 +14,8 @@ class Operator : public CudaAdd> { public: using CudaAdd>::CudaAdd; + // All dispatched types use kernels and constructor-owned metadata on replay. + bool capture_safe() const override { return true; } }; } // 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..944262aa7 --- /dev/null +++ b/src/native/cuda/nvidia/ops/argmax/kernel.cuh @@ -0,0 +1,129 @@ +#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..1a507ad26 --- /dev/null +++ b/src/native/cuda/nvidia/ops/argmax/kernel.h @@ -0,0 +1,200 @@ +#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: + 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) { + 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; + 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_; + } + + 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 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 || + (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) { + 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)) { + 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)) { + 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()); + 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)) { + return false; + } + int current_device = -1; + const auto status = cudaGetDevice(¤t_device); + 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( + "unsupported NVIDIA Argmax metadata, storage or 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<<< + final_grid, nvidia_argmax_detail::kThreads, 0, stream>>>( + 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..f0f8df020 --- /dev/null +++ b/src/native/cuda/nvidia/ops/binary_view_capture.h @@ -0,0 +1,76 @@ +#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; +} + +} // 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..7bd6a4aa5 --- /dev/null +++ b/src/native/cuda/nvidia/ops/cast/kernel.h @@ -0,0 +1,147 @@ +#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()} {} + + 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, + 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/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..229bd2ec1 --- /dev/null +++ b/src/native/cuda/nvidia/ops/cat/kernel.h @@ -0,0 +1,107 @@ +#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; + 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/causal_conv1d/kernel.h b/src/native/cuda/nvidia/ops/causal_conv1d/kernel.h new file mode 100644 index 000000000..276fd41d7 --- /dev/null +++ b/src/native/cuda/nvidia/ops/causal_conv1d/kernel.h @@ -0,0 +1,202 @@ +#pragma once +#include +#include +#include + +#include + +#include "base/causal_conv1d.h" +#include "native/cuda/nvidia/caster.cuh" +#include "native/cuda/nvidia/runtime_.h" + +namespace infini::ops { +namespace causal_conv1d_detail { +struct Layout { + ptrdiff_t strides[9][3]{}; + size_t tokens, channels, pool; + bool cu_i64, initial_i64, final_i64; +}; +__device__ inline int64_t index(const void* p, bool i64, int i, int fallback) { + return p ? (i64 ? static_cast(p)[i] + : static_cast(static_cast(p)[i])) + : fallback; +} +template +__device__ inline float history(const T* state, const T* qkv, int64_t pos, + int64_t begin, int batch, int channel, + ptrdiff_t base, const Layout& p) { + if (pos < 3) return static_cast(state[base + pos * p.strides[2][2]]); + const auto off = static_cast(batch) * p.strides[0][0] + + static_cast(begin + pos - 3) * p.strides[0][1] + + static_cast(channel) * p.strides[0][2]; + return static_cast(qkv[off]); +} +// Port of baseline InfiniCore-0.2.9 causal_conv1d_k4_kernel. One thread per +// request/channel reads the entire sequence before committing its final state. +template +__global__ void kernel(const T* qkv, const T* weight, T* state, const T* bias, + const void* cu, const void* initial, + const void* final_indices, T* out, T* final_state, + Layout p) { + const int request = blockIdx.x, channel = blockIdx.y; + if (static_cast(channel) >= p.channels || threadIdx.x != 0) return; + int64_t begin = 0, end = static_cast(p.tokens); + int batch = request; + if (cu) { + begin = index(cu, p.cu_i64, request, 0); + end = index(cu, p.cu_i64, request + 1, 0); + batch = 0; + // Preserve baseline device-value behavior: malformed range leaves the + // corresponding output/state untouched. This is not host validation. + if (begin < 0 || end < begin || end > static_cast(p.tokens)) + return; + } + const int64_t length = end - begin; + const int64_t read = index(initial, p.initial_i64, request, request); + const int64_t write = index(final_indices, p.final_i64, request, request); + if (read < 0 || write < 0 || read >= static_cast(p.pool) || + (final_indices && write >= static_cast(p.pool))) + return; + const auto base = static_cast(read) * p.strides[2][0] + + static_cast(channel) * p.strides[2][1]; + const auto wb = static_cast(channel) * p.strides[1][0]; + for (int64_t t = 0; t < length; ++t) { + float acc = static_cast(weight[wb]) * + history(state, qkv, t, begin, batch, channel, base, p) + + static_cast(weight[wb + p.strides[1][2]]) * + history(state, qkv, t + 1, begin, batch, channel, base, p) + + static_cast(weight[wb + 2 * p.strides[1][2]]) * + history(state, qkv, t + 2, begin, batch, channel, base, p) + + static_cast(weight[wb + 3 * p.strides[1][2]]) * + history(state, qkv, t + 3, begin, batch, channel, base, p); + if (bias) + acc += static_cast( + bias[static_cast(channel) * p.strides[3][0]]); + const auto off = static_cast(batch) * p.strides[7][0] + + static_cast(begin + t) * p.strides[7][1] + + static_cast(channel) * p.strides[7][2]; + out[off] = static_cast(acc); + } + T* target = final_indices ? state : final_state; + const int stride_index = final_indices ? 2 : 8; + const auto target_base = + static_cast(final_indices ? write : request) * + p.strides[stride_index][0] + + static_cast(channel) * p.strides[stride_index][1]; + // Same ordered commit as baseline. Exact initial/final alias is valid for + // short sequences because copied history source positions lie ahead of + // writes. + target[target_base] = static_cast( + history(state, qkv, length, begin, batch, channel, base, p)); + target[target_base + p.strides[stride_index][2]] = static_cast( + history(state, qkv, length + 1, begin, batch, channel, base, p)); + target[target_base + 2 * p.strides[stride_index][2]] = static_cast( + history(state, qkv, length + 2, begin, batch, channel, base, p)); +} +} // namespace causal_conv1d_detail + +template <> +class Operator : public CausalConv1d { + public: + Operator(const Tensor qkv, const Tensor weight, Tensor conv_state, + const std::optional bias, + const std::optional cu_seqlens, + const std::optional initial_state_indices, + const std::optional final_state_indices, Tensor out, + std::optional final_conv_state) + : CausalConv1d(qkv, weight, conv_state, bias, cu_seqlens, + initial_state_indices, final_state_indices, out, + 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) + throw std::invalid_argument( + "causal conv request/channel launch extent unsupported"); + p_.tokens = tokens_; + p_.channels = channels_; + p_.pool = pool_; + for (size_t i = 0; i < metadata_.size(); ++i) + 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. + void operator()(const Tensor qkv, const Tensor weight, Tensor conv_state, + const std::optional bias, + 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 override { + const Views views{qkv, + weight, + conv_state, + bias, + cu_seqlens, + initial_state_indices, + 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].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"); + } + check_device(); + if (!requests_ || !channels_) return; + switch (dtype_) { + case DataType::kFloat32: + launch(views); + break; + case DataType::kFloat16: + launch(views); + break; + case DataType::kBFloat16: + launch<__nv_bfloat16>(views); + break; + default: + throw std::invalid_argument("causal conv dtype unsupported"); + } + if (cudaPeekAtLastError() != cudaSuccess) + throw std::runtime_error("causal conv CUDA launch failed"); + } + + private: + causal_conv1d_detail::Layout p_{}; + void check_device() const { + int current = -1; + if (cudaGetDevice(¤t) != cudaSuccess) + throw std::runtime_error("causal conv device query failed"); + if (current != device_.index()) + throw std::invalid_argument("causal conv owner device mismatch"); + } + template + void launch(const Views& v) const { + auto pointer = [&](size_t i) -> void* { + return v[i] ? const_cast(v[i]->data()) : nullptr; + }; + const dim3 grid(static_cast(requests_), + static_cast(channels_), 1); + const auto stream = static_cast(stream_ ? stream_ : nullptr); + causal_conv1d_detail::kernel<<>>( + static_cast(pointer(0)), static_cast(pointer(1)), + static_cast(pointer(2)), static_cast(pointer(3)), + pointer(4), pointer(5), pointer(6), static_cast(pointer(7)), + static_cast(pointer(8)), p_); + } +}; +} // namespace infini::ops 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 new file mode 100644 index 000000000..0c8ac1a52 --- /dev/null +++ b/src/native/cuda/nvidia/ops/chunk_gated_delta_rule/chunk_gated_delta_rule_legacy_kernel.cuh @@ -0,0 +1,259 @@ +// 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 +// chunked-matrix branch is hard-disabled and is not copied as a second path. +namespace infini::ops::chunk_native_detail { +__device__ inline int64_t loadOptionalIndex(const void* indices, bool is_i64, + int idx, int fallback) { + if (indices == nullptr) { + return static_cast(fallback); + } + return is_i64 + ? static_cast(indices)[idx] + : static_cast(static_cast(indices)[idx]); +} + +template +__device__ inline float loadAsFloat(const T* ptr, ptrdiff_t offset) { + return static_cast(ptr[offset]); +} + +template <> +__device__ inline float loadAsFloat(const half* ptr, ptrdiff_t offset) { + return __half2float(ptr[offset]); +} + +template <> +__device__ inline float loadAsFloat(const cuda_bfloat16* ptr, + ptrdiff_t offset) { + return __bfloat162float(ptr[offset]); +} + +#define CGDR_FOR(idx, n) \ + for (int idx = threadIdx.x; idx < static_cast(n); idx += blockDim.x) + +template +__device__ Tcompute blockReduceSum(Tcompute v) { + __shared__ Tcompute smem[NUM_THREADS]; + smem[threadIdx.x] = v; + __syncthreads(); + + for (int s = NUM_THREADS / 2; s > 0; s >>= 1) { + if (threadIdx.x < s) { + smem[threadIdx.x] += smem[threadIdx.x + s]; + } + __syncthreads(); + } + return smem[0]; +} + +template +__device__ void chunkGatedDeltaRuleRecurrentKernel( + Tcompute* state_workspace, Tdata* out, Tdata* initial_state, + Tdata* final_state, const Tdata* q, const Tdata* k, const Tdata* v, + const Tgate* g, const Tgate* beta, 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 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, + 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, + ptrdiff_t g_s0, ptrdiff_t g_s1, ptrdiff_t g_s2, ptrdiff_t beta_s0, + ptrdiff_t beta_s1, ptrdiff_t beta_s2) { + const int batch_idx = blockIdx.x; + const int value_head_idx = blockIdx.y; + const int key_head_idx = + value_head_idx / static_cast(value_heads_per_key_head); + + if (key_head_idx >= static_cast(Hk)) { + return; + } + + int64_t token_begin = 0; + int64_t token_end = static_cast(T); + + 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(T)) { + return; + } + } + + int64_t read_slot = batch_idx; + int64_t write_slot = batch_idx; + + if (indexed_state_pool) { + read_slot = loadOptionalIndex( + initial_state_indices, initial_state_indices_i64, batch_idx, batch_idx); + + write_slot = + final_state_indices == nullptr + ? 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) || + (final_state_indices != nullptr && + (write_slot < 0 || write_slot >= static_cast(pool_size)))) { + return; + } + } + + const ptrdiff_t initial_base = + indexed_state_pool + ? static_cast(read_slot) * initial_s0 + + static_cast(value_head_idx) * initial_s1 + : static_cast(batch_idx) * initial_s0 + + static_cast(value_head_idx) * initial_s1; + + Tdata* final_state_target = nullptr; + ptrdiff_t final_base = 0; + + if (indexed_state_pool && final_state_indices != nullptr) { + final_state_target = initial_state; + final_base = static_cast(write_slot) * initial_s0 + + static_cast(value_head_idx) * initial_s1; + } else { + final_state_target = final_state; + final_base = static_cast(batch_idx) * final_s0 + + static_cast(value_head_idx) * final_s1; + } + + const ptrdiff_t workspace_block = + (static_cast(batch_idx) * gridDim.y + + static_cast(value_head_idx)) * + static_cast(Dk * Dv); + + Tcompute* state_local = state_workspace + workspace_block; + __shared__ Tcompute q_vec[Dk]; + __shared__ Tcompute k_vec[Dk]; + __shared__ Tcompute v_new[Dv]; + __shared__ Tcompute scalar_buf[3]; + + const int token_batch = has_cu_seqlens ? 0 : batch_idx; + const Tcompute scale = rsqrtf(static_cast(Dk)); + + CGDR_FOR(i, Dk * Dv) { + int dk = i / Dv; + int dv = i % Dv; + + ptrdiff_t read_idx = initial_base + + static_cast(dv) * initial_s2 + + static_cast(dk) * initial_s3; + + state_local[i] = + static_cast(loadAsFloat(initial_state, read_idx)); + } + __syncthreads(); + + for (int64_t token_idx = token_begin; token_idx < token_end; ++token_idx) { + ptrdiff_t q_base = static_cast(token_batch) * q_s0 + + static_cast(token_idx) * q_s1 + + static_cast(key_head_idx) * q_s2; + ptrdiff_t k_base = static_cast(token_batch) * k_s0 + + static_cast(token_idx) * k_s1 + + static_cast(key_head_idx) * k_s2; + + Tcompute q_sum = 0; + Tcompute k_sum = 0; + for (int dk = threadIdx.x; dk < static_cast(Dk); dk += blockDim.x) { + Tcompute q_raw = static_cast(loadAsFloat(q, q_base + dk)); + Tcompute k_raw = static_cast(loadAsFloat(k, k_base + dk)); + q_vec[dk] = q_raw; + k_vec[dk] = k_raw; + q_sum += q_raw * q_raw; + k_sum += k_raw * k_raw; + } + + q_sum = blockReduceSum(q_sum); + k_sum = blockReduceSum(k_sum); + + if (threadIdx.x == 0) { + ptrdiff_t gate_offset = static_cast(token_batch) * g_s0 + + static_cast(token_idx) * g_s1 + + static_cast(value_head_idx) * g_s2; + ptrdiff_t beta_offset = static_cast(token_batch) * beta_s0 + + static_cast(token_idx) * beta_s1 + + static_cast(value_head_idx) * beta_s2; + + scalar_buf[0] = expf(static_cast(loadAsFloat(g, gate_offset))); + scalar_buf[1] = static_cast(loadAsFloat(beta, beta_offset)); + scalar_buf[2] = use_qk_l2norm + ? rsqrtf(q_sum + static_cast(1e-6)) * scale + : scale; + } + __syncthreads(); + + const Tcompute decay = scalar_buf[0]; + const Tcompute beta_t = scalar_buf[1]; + const Tcompute q_scale = scalar_buf[2]; + const Tcompute k_scale = use_qk_l2norm + ? rsqrtf(k_sum + static_cast(1e-6)) + : static_cast(1); + + for (int dk = threadIdx.x; dk < static_cast(Dk); dk += blockDim.x) { + q_vec[dk] *= q_scale; + k_vec[dk] *= k_scale; + } + __syncthreads(); + + for (int dv = threadIdx.x; dv < static_cast(Dv); dv += blockDim.x) { + Tcompute projected = 0; + for (int dk = 0; dk < static_cast(Dk); ++dk) { + projected += k_vec[dk] * state_local[dk * Dv + dv]; + } + + ptrdiff_t v_base = static_cast(token_batch) * v_s0 + + static_cast(token_idx) * v_s1 + + static_cast(value_head_idx) * v_s2; + Tcompute v_raw = static_cast(loadAsFloat(v, v_base + dv)); + v_new[dv] = beta_t * (v_raw - decay * projected); + } + __syncthreads(); + + CGDR_FOR(x, Dk * Dv) { + int dk = x / Dv; + int dv = x % Dv; + state_local[x] = decay * state_local[x] + k_vec[dk] * v_new[dv]; + } + __syncthreads(); + + for (int dv = threadIdx.x; dv < static_cast(Dv); dv += blockDim.x) { + Tcompute out_val = 0; + for (int dk = 0; dk < static_cast(Dk); ++dk) { + out_val += q_vec[dk] * state_local[dk * Dv + dv]; + } + + 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 + dv] = static_cast(out_val); + } + __syncthreads(); + } + + CGDR_FOR(i, Dk * Dv) { + int dk = i / Dv; + int dv = i % Dv; + + const ptrdiff_t s2 = final_state_indices != nullptr ? initial_s2 : final_s2; + const ptrdiff_t s3 = final_state_indices != nullptr ? initial_s3 : final_s3; + ptrdiff_t write_idx = final_base + static_cast(dv) * s2 + + static_cast(dk) * s3; + + final_state_target[write_idx] = static_cast(state_local[i]); + } +} +#undef CGDR_FOR +} // namespace infini::ops::chunk_native_detail 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 new file mode 100644 index 000000000..000fa36ca --- /dev/null +++ b/src/native/cuda/nvidia/ops/chunk_gated_delta_rule/kernel.h @@ -0,0 +1,180 @@ +#pragma once +#include +#include +#include + +#include + +#include "base/chunk_gated_delta_rule.h" +#include "chunk_gated_delta_rule_legacy_kernel.cuh" +#include "native/cuda/nvidia/caster.cuh" +#include "native/cuda/nvidia/runtime_.h" +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; +}; +template +__global__ void kernel(float* workspace, const Data* q, const Data* k, + const Data* v, const Gate* g, const Gate* beta, + 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, + 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 <> +class Operator + : public ChunkGatedDeltaRule { + public: + Operator(const Tensor q, const Tensor k, const Tensor v, const Tensor g, + const Tensor beta, Tensor initial_state, + const std::optional cu_seqlens, + const std::optional initial_state_indices, + const std::optional final_state_indices, Tensor out, + std::optional final_state, bool use_qk_l2norm, + size_t chunk_size) + : ChunkGatedDeltaRule(q, k, v, g, beta, initial_state, cu_seqlens, + initial_state_indices, final_state_indices, out, + 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 || + property.maxThreadsPerBlock < int(D_)) + throw std::runtime_error("chunk delta CUDA block size unsupported"); + if (B_ >= static_cast(INT_MAX) || Hv_ > 65535 || + Hk_ > static_cast(INT_MAX)) + throw std::invalid_argument("chunk delta launch extent unsupported"); + const auto maximum = std::numeric_limits::max(); + const size_t per_head = D_ * D_ * sizeof(float); + if (B_ > maximum / Hv_ || B_ * Hv_ > maximum / per_head) + throw std::overflow_error("chunk delta workspace size overflows"); + bytes_ = B_ * Hv_ * per_head; + if (bytes_ / sizeof(float) > + static_cast(std::numeric_limits::max())) + throw std::overflow_error( + "chunk delta workspace element offset exceeds ptrdiff"); + p_.pool = pool_; + p_.Hk = Hk_; + p_.ratio = Hv_ / Hk_; + p_.norm = norm_; + p_.T = T_; + 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) + 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, + const Tensor g, const Tensor beta, Tensor initial_state, + const std::optional cu_seqlens, + const std::optional initial_state_indices, + const std::optional final_state_indices, Tensor out, + std::optional final_state, bool use_qk_l2norm, + size_t chunk_size) const override { + const Views views{q, + k, + v, + g, + beta, + initial_state, + cu_seqlens, + initial_state_indices, + final_state_indices, + out, + final_state}; + CheckInvocation(views, use_qk_l2norm, chunk_size); + CheckDevice(); + if (!B_) return; + if (!workspace_ || workspace_size_in_bytes_ < bytes_) + throw std::invalid_argument("chunk delta workspace missing or too small"); + switch (data_type_) { + case DataType::kFloat16: + gate(views); + break; + case DataType::kBFloat16: + gate<__nv_bfloat16>(views); + break; + case DataType::kFloat32: + gate(views); + break; + default: + throw std::invalid_argument("chunk delta data dtype unsupported"); + } + if (cudaPeekAtLastError() != cudaSuccess) + throw std::runtime_error("chunk delta CUDA launch failed"); + } + + private: + size_t bytes_{0}; + chunk_native_detail::Layout p_{}; + void CheckDevice() const { + int device = -1; + if (cudaGetDevice(&device) != cudaSuccess) + throw std::runtime_error("chunk delta device query failed"); + if (device != device_.index()) + throw std::invalid_argument("chunk delta owner device mismatch"); + } + template + void gate(const Views& v) const { + switch (gate_type_) { + case DataType::kFloat16: + dimension(v); + break; + case DataType::kBFloat16: + dimension(v); + break; + case DataType::kFloat32: + dimension(v); + break; + default: + throw std::invalid_argument("chunk delta gate dtype unsupported"); + } + } + template + void dimension(const Views& v) const { + if (D_ == 64) + launch(v); + else + launch(v); + } + template + void launch(const Views& v) const { + auto ptr = [&](size_t i) -> void* { + return v[i] ? const_cast(v[i]->data()) : nullptr; + }; + 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_); + } +}; +} // namespace infini::ops 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/copy/kernel.h b/src/native/cuda/nvidia/ops/copy/kernel.h index ebc3f441f..38304282a 100644 --- a/src/native/cuda/nvidia/ops/copy/kernel.h +++ b/src/native/cuda/nvidia/ops/copy/kernel.h @@ -14,6 +14,8 @@ class Operator : public CudaCopy> { public: using CudaCopy>::CudaCopy; + // All dispatched types use kernels and constructor-owned metadata on replay. + bool capture_safe() const override { return true; } }; } // namespace infini::ops diff --git a/src/native/cuda/nvidia/ops/deepseek_moe/kernel.cu b/src/native/cuda/nvidia/ops/deepseek_moe/kernel.cu new file mode 100644 index 000000000..a76794471 --- /dev/null +++ b/src/native/cuda/nvidia/ops/deepseek_moe/kernel.cu @@ -0,0 +1,1091 @@ +// 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 +#include +#include +#include +#include +#include +#include +#include +namespace infini::ops { namespace { + +enum MoeStatus { MOE_SUCCESS=0, MOE_BAD_PARAM, MOE_BAD_DTYPE, MOE_WORKSPACE, MOE_INTERNAL, MOE_DEVICE }; +#define MOE_CHECK_STATUS(expr) do { auto s=(expr); if(s!=MOE_SUCCESS) return s; } while(0) +#define MOE_CHECK_CUDA(expr) do { auto s=(expr); if(s!=cudaSuccess) throw std::runtime_error(cudaGetErrorString(s)); } while(0) +#define MOE_CHECK_CUBLAS(expr) do { auto s=(expr); if(s!=CUBLAS_STATUS_SUCCESS) throw std::runtime_error("MoE cuBLAS failure: "+std::to_string(int(s))); } while(0) +inline void Check(MoeStatus s) { if(s!=MOE_SUCCESS) throw std::runtime_error("MoE native failure: "+std::to_string(int(s))); } +struct DeepseekMoeInfo { DataType dtype; size_t ntokens, hidden_size, topk, intermediate_size, num_experts; }; + +// Native by-value kernel arguments are copied by the CUDA launch. No temporary +// pageable/pinned host pointer array survives submission; the old FP16/BF16 weights are not repacked. +constexpr size_t kPointerChunk=128; +struct WeightPointerChunk { const void* gate[kPointerChunk]; const void* up[kPointerChunk]; const void* down[kPointerChunk]; }; +static_assert(sizeof(WeightPointerChunk)==3072,"Fits pre-Volta CUDA's 4096-byte argument limit with outputs/scalars"); +__global__ void StageWeightPointers(WeightPointerChunk pointers,const void** gate,const void** up,const void** down,size_t begin,size_t count) { + const auto j=size_t(threadIdx.x); if(j(align_up(reinterpret_cast(ptr), alignment)); + char *out = ptr; + ptr += bytes; + return out; +} + +template +size_t route_batched_workspace_size(const DeepseekMoeInfo &info, size_t base_offset) { + const size_t routes = info.ntokens * info.topk; + size_t bytes = base_offset; + bytes = align_up(bytes, alignof(void *)); + bytes += routes * sizeof(void *) * 9; + bytes = align_up(bytes, 256); + bytes += routes * info.intermediate_size * sizeof(T) * 3; + bytes = align_up(bytes, 256); + bytes += routes * info.hidden_size * sizeof(T); + return align_up(bytes, 256); +} + +template +size_t expert_grouped_workspace_size(const DeepseekMoeInfo &info, size_t base_offset) { + const size_t routes = info.ntokens * info.topk; + size_t bytes = base_offset; + for (int i = 0; i < 3; ++i) { + bytes = align_up(bytes, 256); + bytes += (info.num_experts + 1) * sizeof(int); + } + for (int i = 0; i < 2; ++i) { + bytes = align_up(bytes, 256); + bytes += routes * sizeof(int); + } + bytes = align_up(bytes, 256); + bytes += routes * info.hidden_size * sizeof(T); + for (int i = 0; i < 2; ++i) { + bytes = align_up(bytes, 256); + bytes += routes * info.intermediate_size * sizeof(T); + } + return align_up(bytes, 256); +} + +template +__device__ float to_float(T value) { + return static_cast(value); +} + +template <> +__device__ float to_float(half value) { + return __half2float(value); +} + +template <> +__device__ float to_float<__nv_bfloat16>(__nv_bfloat16 value) { + return __bfloat162float(value); +} + +template +__device__ T from_float(float value) { + return static_cast(value); +} + +template <> +__device__ half from_float(float value) { + return __float2half_rn(value); +} + +template <> +__device__ __nv_bfloat16 from_float<__nv_bfloat16>(float value) { + return __float2bfloat16_rn(value); +} + +__device__ __forceinline__ float dot_bf16x8(const uint4 &a, const uint4 &b) { + const auto *a2 = reinterpret_cast(&a); + const auto *b2 = reinterpret_cast(&b); + float sum = 0.0f; +#pragma unroll + for (int pair = 0; pair < 4; ++pair) { + const float2 av = __bfloat1622float2(a2[pair]); + const float2 bv = __bfloat1622float2(b2[pair]); + sum += av.x * bv.x + av.y * bv.y; + } + return sum; +} + +template +__global__ void gate_up_kernel( + T *intermediate, + const T *hidden, + const int *topk_indices, + const float *topk_weights, + const void *const *gate_weights, + const void *const *up_weights, + size_t ntokens, + size_t hidden_size, + size_t topk, + size_t intermediate_size, + size_t num_experts) { + + const size_t route = blockIdx.x / intermediate_size; + const size_t j = blockIdx.x - route * intermediate_size; + if (route >= ntokens * topk) { + return; + } + const int expert = topk_indices[route]; + if (expert < 0 || static_cast(expert) >= num_experts) { + return; + } + const size_t token = route / topk; + const T *x = hidden + token * hidden_size; + const T *gate = reinterpret_cast(gate_weights[expert]) + j * hidden_size; + const T *up = reinterpret_cast(up_weights[expert]) + j * hidden_size; + + float gate_sum = 0.0f; + float up_sum = 0.0f; + if constexpr (std::is_same_v) { + if ((hidden_size & 7) == 0) { + const auto *x8 = reinterpret_cast(x); + const auto *gate8 = reinterpret_cast(gate); + const auto *up8 = reinterpret_cast(up); + const size_t hidden_chunks = hidden_size / 8; + for (size_t h8 = threadIdx.x; h8 < hidden_chunks; h8 += blockDim.x) { + const uint4 xv = x8[h8]; + gate_sum += dot_bf16x8(xv, gate8[h8]); + up_sum += dot_bf16x8(xv, up8[h8]); + } + } else if ((hidden_size & 1) == 0) { + const auto *x2 = reinterpret_cast(x); + const auto *gate2 = reinterpret_cast(gate); + const auto *up2 = reinterpret_cast(up); + const size_t hidden_pairs = hidden_size / 2; + for (size_t h2 = threadIdx.x; h2 < hidden_pairs; h2 += blockDim.x) { + const float2 xv = __bfloat1622float2(x2[h2]); + const float2 gv = __bfloat1622float2(gate2[h2]); + const float2 uv = __bfloat1622float2(up2[h2]); + gate_sum += xv.x * gv.x + xv.y * gv.y; + up_sum += xv.x * uv.x + xv.y * uv.y; + } + } else { + for (size_t h = threadIdx.x; h < hidden_size; h += blockDim.x) { + const float xv = to_float(x[h]); + gate_sum += xv * to_float(gate[h]); + up_sum += xv * to_float(up[h]); + } + } + } else { + for (size_t h = threadIdx.x; h < hidden_size; h += blockDim.x) { + const float xv = to_float(x[h]); + gate_sum += xv * to_float(gate[h]); + up_sum += xv * to_float(up[h]); + } + } + + __shared__ float gate_shared[FUSED_KERNEL_THREADS]; + __shared__ float up_shared[FUSED_KERNEL_THREADS]; + gate_shared[threadIdx.x] = gate_sum; + up_shared[threadIdx.x] = up_sum; + __syncthreads(); + + for (unsigned int stride = blockDim.x / 2; stride > 0; stride >>= 1) { + if (threadIdx.x < stride) { + gate_shared[threadIdx.x] += gate_shared[threadIdx.x + stride]; + up_shared[threadIdx.x] += up_shared[threadIdx.x + stride]; + } + __syncthreads(); + } + + if (threadIdx.x == 0) { + const float g = gate_shared[0]; + const float silu = g / (1.0f + __expf(-g)); + intermediate[route * intermediate_size + j] = from_float(silu * up_shared[0] * topk_weights[route]); + } +} + +template +__global__ void down_kernel( + T *out, + const T *intermediate, + const int *topk_indices, + const void *const *down_weights, + size_t ntokens, + size_t hidden_size, + size_t topk, + size_t intermediate_size, + size_t num_experts) { + + const size_t linear = blockIdx.x; + const size_t token = linear / hidden_size; + const size_t h = linear - token * hidden_size; + if (token >= ntokens) { + return; + } + + float acc = 0.0f; + const size_t route_base = token * topk; + const size_t count = topk * intermediate_size; + if constexpr (std::is_same_v) { + if ((intermediate_size & 7) == 0) { + const size_t chunks_per_expert = intermediate_size / 8; + const size_t chunk_count = topk * chunks_per_expert; + for (size_t chunk_idx = threadIdx.x; chunk_idx < chunk_count; chunk_idx += blockDim.x) { + const size_t k = chunk_idx / chunks_per_expert; + const size_t j8 = chunk_idx - k * chunks_per_expert; + const size_t route = route_base + k; + const int expert = topk_indices[route]; + if (expert >= 0 && static_cast(expert) < num_experts) { + const auto *intermediate8 = reinterpret_cast( + intermediate + route * intermediate_size); + const auto *down8 = reinterpret_cast( + reinterpret_cast(down_weights[expert]) + h * intermediate_size); + acc += dot_bf16x8(intermediate8[j8], down8[j8]); + } + } + } else if ((intermediate_size & 1) == 0) { + const size_t pairs_per_expert = intermediate_size / 2; + const size_t pair_count = topk * pairs_per_expert; + for (size_t pair_idx = threadIdx.x; pair_idx < pair_count; pair_idx += blockDim.x) { + const size_t k = pair_idx / pairs_per_expert; + const size_t j2 = pair_idx - k * pairs_per_expert; + const size_t route = route_base + k; + const int expert = topk_indices[route]; + if (expert >= 0 && static_cast(expert) < num_experts) { + const auto *intermediate2 = reinterpret_cast( + intermediate + route * intermediate_size); + const auto *down2 = reinterpret_cast( + reinterpret_cast(down_weights[expert]) + h * intermediate_size); + const float2 iv = __bfloat1622float2(intermediate2[j2]); + const float2 dv = __bfloat1622float2(down2[j2]); + acc += iv.x * dv.x + iv.y * dv.y; + } + } + } else { + for (size_t idx = threadIdx.x; idx < count; idx += blockDim.x) { + const size_t k = idx / intermediate_size; + const size_t j = idx - k * intermediate_size; + const size_t route = route_base + k; + const int expert = topk_indices[route]; + if (expert >= 0 && static_cast(expert) < num_experts) { + const T *down = reinterpret_cast(down_weights[expert]) + h * intermediate_size; + acc += to_float(intermediate[route * intermediate_size + j]) * to_float(down[j]); + } + } + } + } else { + for (size_t idx = threadIdx.x; idx < count; idx += blockDim.x) { + const size_t k = idx / intermediate_size; + const size_t j = idx - k * intermediate_size; + const size_t route = route_base + k; + const int expert = topk_indices[route]; + if (expert >= 0 && static_cast(expert) < num_experts) { + const T *down = reinterpret_cast(down_weights[expert]) + h * intermediate_size; + acc += to_float(intermediate[route * intermediate_size + j]) * to_float(down[j]); + } + } + } + + __shared__ float shared[FUSED_KERNEL_THREADS]; + shared[threadIdx.x] = acc; + __syncthreads(); + for (unsigned int stride = blockDim.x / 2; stride > 0; stride >>= 1) { + if (threadIdx.x < stride) { + shared[threadIdx.x] += shared[threadIdx.x + stride]; + } + __syncthreads(); + } + if (threadIdx.x == 0) { + out[token * hidden_size + h] = from_float(shared[0]); + } +} + +__global__ void count_experts_kernel( + const int *topk_indices, + int *counts, + int routes, + int num_experts) { + int route = blockIdx.x * blockDim.x + threadIdx.x; + if (route >= routes) { + return; + } + int expert = topk_indices[route]; + if (expert >= 0 && expert < num_experts) { + atomicAdd(counts + expert, 1); + } +} + +__global__ void prefix_counts_kernel( + const int *counts, + int *offsets, + int num_experts) { + extern __shared__ int scan[]; + int tid = threadIdx.x; + if (tid < num_experts) { + scan[tid] = counts[tid]; + } else { + scan[tid] = 0; + } + __syncthreads(); + + for (int stride = 1; stride < blockDim.x; stride <<= 1) { + int value = 0; + if (tid >= stride) { + value = scan[tid - stride]; + } + __syncthreads(); + scan[tid] += value; + __syncthreads(); + } + + if (tid == 0) { + offsets[0] = 0; + } + if (tid < num_experts) { + offsets[tid + 1] = scan[tid]; + } +} + +template +__global__ void pack_grouped_hidden_kernel( + const T *hidden, + const int *topk_indices, + const int *offsets, + int *positions, + int *row_to_route, + int *route_to_row, + T *packed_hidden, + int routes, + int hidden_size, + int topk, + int num_experts) { + int route = blockIdx.x; + int tid = threadIdx.x; + if (route >= routes) { + return; + } + int expert = topk_indices[route]; + if (expert < 0 || expert >= num_experts) { + return; + } + __shared__ int row; + if (tid == 0) { + int local = atomicAdd(positions + expert, 1); + row = offsets[expert] + local; + row_to_route[row] = route; + route_to_row[route] = row; + } + __syncthreads(); + int token = route / topk; + for (int h = tid; h < hidden_size; h += blockDim.x) { + packed_hidden[static_cast(row) * hidden_size + h] = hidden[static_cast(token) * hidden_size + h]; + } +} + +template +__global__ void swiglu_weight_grouped_kernel( + T *activated, + const T *gate, + const T *up, + const int *row_to_route, + const float *topk_weights, + int routes, + int intermediate_size) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + int total = routes * intermediate_size; + if (idx >= total) { + return; + } + int row = idx / intermediate_size; + int route = row_to_route[row]; + float g = to_float(gate[idx]); + float u = to_float(up[idx]); + float silu = g / (1.0f + __expf(-g)); + activated[idx] = from_float(silu * u * topk_weights[route]); +} + +template +__global__ void sum_grouped_out_kernel( + T *out, + const T *packed_out, + const int *route_to_row, + int ntokens, + int hidden_size, + int topk) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + int total = ntokens * hidden_size; + if (idx >= total) { + return; + } + int token = idx / hidden_size; + int h = idx - token * hidden_size; + float acc = 0.0f; + for (int k = 0; k < topk; ++k) { + int route = token * topk + k; + int row = route_to_row[route]; + acc += to_float(packed_out[static_cast(row) * hidden_size + h]); + } + out[idx] = from_float(acc); +} + +template +__global__ void setup_route_batched_ptrs_kernel( + const void **a_array, + const void **b_array, + void **c_array, + const T *hidden, + const int *topk_indices, + const void *const *weights, + T *gemm_out, + int routes, + int hidden_size, + int topk, + int out_size, + int num_experts) { + int route = blockIdx.x * blockDim.x + threadIdx.x; + if (route >= routes) { + return; + } + int expert = topk_indices[route]; + if (expert < 0 || expert >= num_experts) { + expert = 0; + } + int token = route / topk; + a_array[route] = weights[expert]; + b_array[route] = hidden + static_cast(token) * hidden_size; + c_array[route] = gemm_out + static_cast(route) * out_size; +} + +template +__global__ void setup_down_batched_ptrs_kernel( + const void **a_array, + const void **b_array, + void **c_array, + const T *activated, + const int *topk_indices, + const void *const *down_weights, + T *route_out, + int routes, + int hidden_size, + int topk, + int intermediate_size, + int num_experts) { + int route = blockIdx.x * blockDim.x + threadIdx.x; + if (route >= routes) { + return; + } + int expert = topk_indices[route]; + if (expert < 0 || expert >= num_experts) { + expert = 0; + } + a_array[route] = down_weights[expert]; + b_array[route] = activated + static_cast(route) * intermediate_size; + c_array[route] = route_out + static_cast(route) * hidden_size; +} + +template +__global__ void setup_all_batched_ptrs_kernel( + const void **gate_a_array, + const void **gate_b_array, + void **gate_c_array, + const void **up_a_array, + const void **up_b_array, + void **up_c_array, + const void **down_a_array, + const void **down_b_array, + void **down_c_array, + const T *hidden, + const int *topk_indices, + const void *const *gate_weights, + const void *const *up_weights, + const void *const *down_weights, + T *gate_buf, + T *up_buf, + T *activated, + T *route_out, + int routes, + int hidden_size, + int topk, + int intermediate_size, + int num_experts) { + int route = blockIdx.x * blockDim.x + threadIdx.x; + if (route >= routes) { + return; + } + int expert = topk_indices[route]; + if (expert < 0 || expert >= num_experts) { + expert = 0; + } + int token = route / topk; + const T *token_hidden = hidden + static_cast(token) * hidden_size; + + gate_a_array[route] = gate_weights[expert]; + gate_b_array[route] = token_hidden; + gate_c_array[route] = gate_buf + static_cast(route) * intermediate_size; + + up_a_array[route] = up_weights[expert]; + up_b_array[route] = token_hidden; + up_c_array[route] = up_buf + static_cast(route) * intermediate_size; + + down_a_array[route] = down_weights[expert]; + down_b_array[route] = activated + static_cast(route) * intermediate_size; + down_c_array[route] = route_out + static_cast(route) * hidden_size; +} + +template +__global__ void swiglu_weight_kernel( + T *activated, + const T *gate, + const T *up, + const float *topk_weights, + int routes, + int intermediate_size) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + int total = routes * intermediate_size; + if (idx >= total) { + return; + } + int route = idx / intermediate_size; + float g = to_float(gate[idx]); + float u = to_float(up[idx]); + float silu = g / (1.0f + __expf(-g)); + activated[idx] = from_float(silu * u * topk_weights[route]); +} + +template +__global__ void sum_route_out_kernel( + T *out, + const T *route_out, + int ntokens, + int hidden_size, + int topk) { + int linear = blockIdx.x * blockDim.x + threadIdx.x; + int total = ntokens * hidden_size; + if (linear >= total) { + return; + } + int token = linear / hidden_size; + int h = linear - token * hidden_size; + float acc = 0.0f; + for (int k = 0; k < topk; ++k) { + int route = token * topk + k; + acc += to_float(route_out[static_cast(route) * hidden_size + h]); + } + out[linear] = from_float(acc); +} + +template +MoeStatus run_batched_gemm( + const NvidiaGemmHandle &internal, + cudaStream_t stream, + const void **a_array, + const void **b_array, + void **c_array, + int m, + int n, + int k, + int lda, + int ldb, + int ldc, + int batch_count, + cudaDataType dtype) { + float alpha = 1.0f; + float beta = 0.0f; + const auto handle=internal.get(); + MOE_CHECK_CUBLAS(cublasSetStream(handle,stream)); + MOE_CHECK_CUBLAS(cublasSetPointerMode(handle,CUBLAS_POINTER_MODE_HOST)); + cublasComputeType_t compute_type = CUBLAS_COMPUTE_32F; + MOE_CHECK_CUBLAS(cublasGemmBatchedEx( + handle, + CUBLAS_OP_T, + CUBLAS_OP_N, + m, + n, + k, + &alpha, + reinterpret_cast(a_array), + dtype, + lda, + reinterpret_cast(b_array), + dtype, + ldb, + &beta, + reinterpret_cast(c_array), + dtype, + ldc, + batch_count, + compute_type, + CUBLAS_GEMM_DEFAULT_TENSOR_OP)); + return MOE_SUCCESS; +} + +template +MoeStatus run_gemm( + const NvidiaGemmHandle &internal, + cudaStream_t stream, + const void *a, + const void *b, + void *c, + int m, + int n, + int k, + int lda, + int ldb, + int ldc, + cudaDataType dtype) { + float alpha = 1.0f; + float beta = 0.0f; + const auto handle=internal.get(); + MOE_CHECK_CUBLAS(cublasSetStream(handle,stream)); + MOE_CHECK_CUBLAS(cublasSetPointerMode(handle,CUBLAS_POINTER_MODE_HOST)); + cublasComputeType_t compute_type = CUBLAS_COMPUTE_32F; + MOE_CHECK_CUBLAS(cublasGemmEx( + handle, + CUBLAS_OP_T, + CUBLAS_OP_N, + m, + n, + k, + &alpha, + a, + dtype, + lda, + b, + dtype, + ldb, + &beta, + c, + dtype, + ldc, + compute_type, + CUBLAS_GEMM_DEFAULT_TENSOR_OP)); + return MOE_SUCCESS; +} + +template +MoeStatus run_chunked_grouped_gemm( + char *base, + size_t workspace_size, + size_t intermediate_offset, + const DeepseekMoeInfo &info, + T *out, + const T *hidden, + const int *topk_indices, + const float *topk_weights, + const void *const *gate_weights, + const void *const *up_weights, + const void *const *down_weights, + cudaStream_t stream, + const NvidiaGemmHandle &internal) { + + const int hidden_size = static_cast(info.hidden_size); + const int topk = static_cast(info.topk); + const int intermediate_size = static_cast(info.intermediate_size); + const int num_experts = static_cast(info.num_experts); + const cudaDataType blas_dtype = info.dtype == DataType::kFloat16 ? CUDA_R_16F : CUDA_R_16BF; + + std::vector host_counts(info.num_experts + 1); + std::vector host_offsets(info.num_experts + 1); + std::vector host_gate(info.num_experts); + std::vector host_up(info.num_experts); + std::vector host_down(info.num_experts); + MOE_CHECK_CUDA(cudaMemcpyAsync(host_gate.data(), gate_weights, info.num_experts * sizeof(void *), cudaMemcpyDeviceToHost, stream)); + MOE_CHECK_CUDA(cudaMemcpyAsync(host_up.data(), up_weights, info.num_experts * sizeof(void *), cudaMemcpyDeviceToHost, stream)); + MOE_CHECK_CUDA(cudaMemcpyAsync(host_down.data(), down_weights, info.num_experts * sizeof(void *), cudaMemcpyDeviceToHost, stream)); + MOE_CHECK_CUDA(cudaStreamSynchronize(stream)); + + for (size_t token_begin = 0; token_begin < info.ntokens; token_begin += CHUNKED_GROUPED_GEMM_TOKENS) { + const int chunk_tokens = static_cast(std::min( + CHUNKED_GROUPED_GEMM_TOKENS, info.ntokens - token_begin)); + const int routes = chunk_tokens * topk; + + char *gptr = base + intermediate_offset; + auto *counts = reinterpret_cast(advance_workspace(gptr, (info.num_experts + 1) * sizeof(int))); + auto *offsets = reinterpret_cast(advance_workspace(gptr, (info.num_experts + 1) * sizeof(int))); + auto *positions = reinterpret_cast(advance_workspace(gptr, (info.num_experts + 1) * sizeof(int))); + auto *row_to_route = reinterpret_cast(advance_workspace(gptr, static_cast(routes) * sizeof(int))); + auto *route_to_row = reinterpret_cast(advance_workspace(gptr, static_cast(routes) * sizeof(int))); + auto *packed_hidden = reinterpret_cast(advance_workspace(gptr, static_cast(routes) * info.hidden_size * sizeof(T))); + auto *gate_buf = reinterpret_cast(advance_workspace(gptr, static_cast(routes) * info.intermediate_size * sizeof(T))); + auto *up_buf = reinterpret_cast(advance_workspace(gptr, static_cast(routes) * info.intermediate_size * sizeof(T))); + if (static_cast(gptr - base) > workspace_size) { + return MOE_WORKSPACE; + } + + const T *hidden_chunk = hidden + token_begin * info.hidden_size; + const int *indices_chunk = topk_indices + token_begin * info.topk; + const float *weights_chunk = topk_weights + token_begin * info.topk; + T *out_chunk = out + token_begin * info.hidden_size; + + MOE_CHECK_CUDA(cudaMemsetAsync(counts, 0, (info.num_experts + 1) * sizeof(int), stream)); + MOE_CHECK_CUDA(cudaMemsetAsync(positions, 0, (info.num_experts + 1) * sizeof(int), stream)); + count_experts_kernel<<<(routes + 255) / 256, 256, 0, stream>>>( + indices_chunk, counts, routes, num_experts); + prefix_counts_kernel<<<1, 256, 256 * sizeof(int), stream>>>( + counts, offsets, num_experts); + pack_grouped_hidden_kernel<<>>( + hidden_chunk, indices_chunk, offsets, positions, + row_to_route, route_to_row, packed_hidden, + routes, hidden_size, topk, num_experts); + + MOE_CHECK_CUDA(cudaMemcpyAsync(host_counts.data(), counts, (info.num_experts + 1) * sizeof(int), cudaMemcpyDeviceToHost, stream)); + MOE_CHECK_CUDA(cudaMemcpyAsync(host_offsets.data(), offsets, (info.num_experts + 1) * sizeof(int), cudaMemcpyDeviceToHost, stream)); + MOE_CHECK_CUDA(cudaStreamSynchronize(stream)); + + for (int e = 0; e < num_experts; ++e) { + const int count = host_counts[e]; + if (count == 0) { + continue; + } + const int offset = host_offsets[e]; + MOE_CHECK_STATUS(run_gemm( + internal, stream, host_gate[e], packed_hidden + static_cast(offset) * hidden_size, + gate_buf + static_cast(offset) * intermediate_size, + intermediate_size, count, hidden_size, + hidden_size, hidden_size, intermediate_size, blas_dtype)); + MOE_CHECK_STATUS(run_gemm( + internal, stream, host_up[e], packed_hidden + static_cast(offset) * hidden_size, + up_buf + static_cast(offset) * intermediate_size, + intermediate_size, count, hidden_size, + hidden_size, hidden_size, intermediate_size, blas_dtype)); + } + + swiglu_weight_grouped_kernel<<<(static_cast(routes) * info.intermediate_size + 255) / 256, 256, 0, stream>>>( + gate_buf, gate_buf, up_buf, row_to_route, weights_chunk, routes, intermediate_size); + + for (int e = 0; e < num_experts; ++e) { + const int count = host_counts[e]; + if (count == 0) { + continue; + } + const int offset = host_offsets[e]; + MOE_CHECK_STATUS(run_gemm( + internal, stream, host_down[e], gate_buf + static_cast(offset) * intermediate_size, + packed_hidden + static_cast(offset) * hidden_size, + hidden_size, count, intermediate_size, + intermediate_size, intermediate_size, hidden_size, blas_dtype)); + } + + sum_grouped_out_kernel<<<(static_cast(chunk_tokens) * info.hidden_size + 255) / 256, 256, 0, stream>>>( + out_chunk, packed_hidden, route_to_row, chunk_tokens, hidden_size, topk); + } + + // Keep the large-prefill path memory-bounded across transformer layers. + // All earlier chunks are synchronized by the next chunk's routing copy; + // this final sync drains the last chunk before its workspace can be reused. + MOE_CHECK_CUDA(cudaStreamSynchronize(stream)); + + return MOE_SUCCESS; +} + +template +MoeStatus launch_typed( + void *workspace, + size_t workspace_size, + const DeepseekMoeInfo &info, + void *out, + const void *hidden, + const void *topk_indices, + const void *topk_weights, + const void *const *gate_weights, + const void *const *up_weights, + const void *const *down_weights, + cudaStream_t stream, + const NvidiaGemmHandle &internal) { + + const size_t ptr_bytes = align_up(info.num_experts * sizeof(void *), 256); + const size_t ptr_workspace = ptr_bytes * 3; + const size_t intermediate_offset = align_up(ptr_workspace, 256); + size_t intermediate_tokens = info.ntokens; + if (info.ntokens > ROUTE_BATCHED_GEMM_MAX_TOKENS && info.num_experts <= 256) { + intermediate_tokens = std::min(info.ntokens, CHUNKED_GROUPED_GEMM_TOKENS); + } + const size_t intermediate_bytes = intermediate_tokens * info.topk * info.intermediate_size * sizeof(T); + if (workspace_size < intermediate_offset + intermediate_bytes) { + return MOE_WORKSPACE; + } + + auto *base = reinterpret_cast(workspace); + const void *const *gate_ptrs = gate_weights; + const void *const *up_ptrs = up_weights; + const void *const *down_ptrs = down_weights; + auto *intermediate = reinterpret_cast(base + intermediate_offset); + + if (info.ntokens > ROUTE_BATCHED_GEMM_MAX_TOKENS && info.num_experts <= 256) { + return run_chunked_grouped_gemm( + base, workspace_size, intermediate_offset, info, + reinterpret_cast(out), reinterpret_cast(hidden), + reinterpret_cast(topk_indices), reinterpret_cast(topk_weights), + gate_weights, up_weights, down_weights, + stream, internal); + } + + if (info.ntokens > FUSED_DECODE_MAX_TOKENS + && info.ntokens <= ROUTE_BATCHED_GEMM_MAX_TOKENS) { + const int routes = static_cast(info.ntokens * info.topk); + const int hidden_size_i = static_cast(info.hidden_size); + const int topk_i = static_cast(info.topk); + const int intermediate_size_i = static_cast(info.intermediate_size); + const int num_experts_i = static_cast(info.num_experts); + const cudaDataType blas_dtype = info.dtype == DataType::kFloat16 ? CUDA_R_16F : CUDA_R_16BF; + + if (info.ntokens >= EXPERT_GROUPED_GEMM_MIN_TOKENS && info.num_experts <= 256) { + char *gptr = base + intermediate_offset; + auto *counts = reinterpret_cast(advance_workspace(gptr, (info.num_experts + 1) * sizeof(int))); + auto *offsets = reinterpret_cast(advance_workspace(gptr, (info.num_experts + 1) * sizeof(int))); + auto *positions = reinterpret_cast(advance_workspace(gptr, (info.num_experts + 1) * sizeof(int))); + auto *row_to_route = reinterpret_cast(advance_workspace(gptr, static_cast(routes) * sizeof(int))); + auto *route_to_row = reinterpret_cast(advance_workspace(gptr, static_cast(routes) * sizeof(int))); + auto *packed_hidden = reinterpret_cast(advance_workspace(gptr, static_cast(routes) * info.hidden_size * sizeof(T))); + auto *gate_buf_g = reinterpret_cast(advance_workspace(gptr, static_cast(routes) * info.intermediate_size * sizeof(T))); + auto *up_buf_g = reinterpret_cast(advance_workspace(gptr, static_cast(routes) * info.intermediate_size * sizeof(T))); + if (static_cast(gptr - base) <= workspace_size) { + MOE_CHECK_CUDA(cudaMemsetAsync(counts, 0, (info.num_experts + 1) * sizeof(int), stream)); + MOE_CHECK_CUDA(cudaMemsetAsync(positions, 0, (info.num_experts + 1) * sizeof(int), stream)); + count_experts_kernel<<<(routes + 255) / 256, 256, 0, stream>>>( + reinterpret_cast(topk_indices), counts, routes, num_experts_i); + prefix_counts_kernel<<<1, 256, 256 * sizeof(int), stream>>>( + counts, offsets, num_experts_i); + pack_grouped_hidden_kernel<<>>( + reinterpret_cast(hidden), reinterpret_cast(topk_indices), + offsets, positions, row_to_route, route_to_row, packed_hidden, + routes, hidden_size_i, topk_i, num_experts_i); + + std::vector host_counts(info.num_experts + 1); + std::vector host_offsets(info.num_experts + 1); + std::vector host_gate(info.num_experts); + std::vector host_up(info.num_experts); + std::vector host_down(info.num_experts); + MOE_CHECK_CUDA(cudaMemcpyAsync(host_counts.data(), counts, (info.num_experts + 1) * sizeof(int), cudaMemcpyDeviceToHost, stream)); + MOE_CHECK_CUDA(cudaMemcpyAsync(host_offsets.data(), offsets, (info.num_experts + 1) * sizeof(int), cudaMemcpyDeviceToHost, stream)); + MOE_CHECK_CUDA(cudaMemcpyAsync(host_gate.data(), gate_weights, info.num_experts * sizeof(void *), cudaMemcpyDeviceToHost, stream)); + MOE_CHECK_CUDA(cudaMemcpyAsync(host_up.data(), up_weights, info.num_experts * sizeof(void *), cudaMemcpyDeviceToHost, stream)); + MOE_CHECK_CUDA(cudaMemcpyAsync(host_down.data(), down_weights, info.num_experts * sizeof(void *), cudaMemcpyDeviceToHost, stream)); + MOE_CHECK_CUDA(cudaStreamSynchronize(stream)); + + for (int e = 0; e < num_experts_i; ++e) { + int count = host_counts[e]; + if (count == 0) { + continue; + } + int offset = host_offsets[e]; + MOE_CHECK_STATUS(run_gemm( + internal, stream, host_gate[e], packed_hidden + static_cast(offset) * hidden_size_i, + gate_buf_g + static_cast(offset) * intermediate_size_i, + intermediate_size_i, count, hidden_size_i, + hidden_size_i, hidden_size_i, intermediate_size_i, blas_dtype)); + MOE_CHECK_STATUS(run_gemm( + internal, stream, host_up[e], packed_hidden + static_cast(offset) * hidden_size_i, + up_buf_g + static_cast(offset) * intermediate_size_i, + intermediate_size_i, count, hidden_size_i, + hidden_size_i, hidden_size_i, intermediate_size_i, blas_dtype)); + } + + swiglu_weight_grouped_kernel<<<(static_cast(routes) * info.intermediate_size + 255) / 256, 256, 0, stream>>>( + gate_buf_g, gate_buf_g, up_buf_g, row_to_route, reinterpret_cast(topk_weights), routes, intermediate_size_i); + + for (int e = 0; e < num_experts_i; ++e) { + int count = host_counts[e]; + if (count == 0) { + continue; + } + int offset = host_offsets[e]; + MOE_CHECK_STATUS(run_gemm( + internal, stream, host_down[e], gate_buf_g + static_cast(offset) * intermediate_size_i, + packed_hidden + static_cast(offset) * hidden_size_i, + hidden_size_i, count, intermediate_size_i, + intermediate_size_i, intermediate_size_i, hidden_size_i, blas_dtype)); + } + + sum_grouped_out_kernel<<<(info.ntokens * info.hidden_size + 255) / 256, 256, 0, stream>>>( + reinterpret_cast(out), packed_hidden, route_to_row, static_cast(info.ntokens), hidden_size_i, topk_i); + return MOE_SUCCESS; + } + } + + char *bptr = base + intermediate_offset; + auto **gate_a_array = reinterpret_cast(advance_workspace(bptr, static_cast(routes) * sizeof(void *), alignof(void *))); + auto **gate_b_array = reinterpret_cast(advance_workspace(bptr, static_cast(routes) * sizeof(void *), alignof(void *))); + auto **gate_c_array = reinterpret_cast(advance_workspace(bptr, static_cast(routes) * sizeof(void *), alignof(void *))); + auto **up_a_array = reinterpret_cast(advance_workspace(bptr, static_cast(routes) * sizeof(void *), alignof(void *))); + auto **up_b_array = reinterpret_cast(advance_workspace(bptr, static_cast(routes) * sizeof(void *), alignof(void *))); + auto **up_c_array = reinterpret_cast(advance_workspace(bptr, static_cast(routes) * sizeof(void *), alignof(void *))); + auto **down_a_array = reinterpret_cast(advance_workspace(bptr, static_cast(routes) * sizeof(void *), alignof(void *))); + auto **down_b_array = reinterpret_cast(advance_workspace(bptr, static_cast(routes) * sizeof(void *), alignof(void *))); + auto **down_c_array = reinterpret_cast(advance_workspace(bptr, static_cast(routes) * sizeof(void *), alignof(void *))); + auto *gate_buf = reinterpret_cast(advance_workspace(bptr, static_cast(routes) * info.intermediate_size * sizeof(T))); + auto *up_buf = reinterpret_cast(advance_workspace(bptr, static_cast(routes) * info.intermediate_size * sizeof(T))); + auto *activated = reinterpret_cast(advance_workspace(bptr, static_cast(routes) * info.intermediate_size * sizeof(T))); + auto *route_out = reinterpret_cast(advance_workspace(bptr, static_cast(routes) * info.hidden_size * sizeof(T))); + if (static_cast(bptr - base) <= workspace_size) { + const int blocks = (routes + 255) / 256; + if (info.ntokens >= EXPERT_GROUPED_GEMM_MIN_TOKENS) { + setup_all_batched_ptrs_kernel<<>>( + gate_a_array, gate_b_array, gate_c_array, + up_a_array, up_b_array, up_c_array, + down_a_array, down_b_array, down_c_array, + reinterpret_cast(hidden), reinterpret_cast(topk_indices), + gate_ptrs, up_ptrs, down_ptrs, gate_buf, up_buf, activated, route_out, + routes, hidden_size_i, topk_i, intermediate_size_i, num_experts_i); + MOE_CHECK_STATUS(run_batched_gemm( + internal, stream, gate_a_array, gate_b_array, gate_c_array, + intermediate_size_i, 1, hidden_size_i, hidden_size_i, hidden_size_i, intermediate_size_i, + routes, blas_dtype)); + + MOE_CHECK_STATUS(run_batched_gemm( + internal, stream, up_a_array, up_b_array, up_c_array, + intermediate_size_i, 1, hidden_size_i, hidden_size_i, hidden_size_i, intermediate_size_i, + routes, blas_dtype)); + + swiglu_weight_kernel<<<(static_cast(routes) * info.intermediate_size + 255) / 256, 256, 0, stream>>>( + activated, gate_buf, up_buf, reinterpret_cast(topk_weights), routes, intermediate_size_i); + + MOE_CHECK_STATUS(run_batched_gemm( + internal, stream, down_a_array, down_b_array, down_c_array, + hidden_size_i, 1, intermediate_size_i, intermediate_size_i, intermediate_size_i, hidden_size_i, + routes, blas_dtype)); + } else { + setup_route_batched_ptrs_kernel<<>>( + gate_a_array, gate_b_array, gate_c_array, + reinterpret_cast(hidden), reinterpret_cast(topk_indices), gate_ptrs, + gate_buf, routes, hidden_size_i, topk_i, intermediate_size_i, num_experts_i); + MOE_CHECK_STATUS(run_batched_gemm( + internal, stream, gate_a_array, gate_b_array, gate_c_array, + intermediate_size_i, 1, hidden_size_i, hidden_size_i, hidden_size_i, intermediate_size_i, + routes, blas_dtype)); + + setup_route_batched_ptrs_kernel<<>>( + gate_a_array, gate_b_array, gate_c_array, + reinterpret_cast(hidden), reinterpret_cast(topk_indices), up_ptrs, + up_buf, routes, hidden_size_i, topk_i, intermediate_size_i, num_experts_i); + MOE_CHECK_STATUS(run_batched_gemm( + internal, stream, gate_a_array, gate_b_array, gate_c_array, + intermediate_size_i, 1, hidden_size_i, hidden_size_i, hidden_size_i, intermediate_size_i, + routes, blas_dtype)); + + swiglu_weight_kernel<<<(static_cast(routes) * info.intermediate_size + 255) / 256, 256, 0, stream>>>( + activated, gate_buf, up_buf, reinterpret_cast(topk_weights), routes, intermediate_size_i); + + setup_down_batched_ptrs_kernel<<>>( + gate_a_array, gate_b_array, gate_c_array, + activated, reinterpret_cast(topk_indices), down_ptrs, + route_out, routes, hidden_size_i, topk_i, intermediate_size_i, num_experts_i); + MOE_CHECK_STATUS(run_batched_gemm( + internal, stream, gate_a_array, gate_b_array, gate_c_array, + hidden_size_i, 1, intermediate_size_i, intermediate_size_i, intermediate_size_i, hidden_size_i, + routes, blas_dtype)); + } + + sum_route_out_kernel<<<(info.ntokens * info.hidden_size + 255) / 256, 256, 0, stream>>>( + reinterpret_cast(out), route_out, static_cast(info.ntokens), hidden_size_i, topk_i); + return MOE_SUCCESS; + } + } + + constexpr int threads = FUSED_KERNEL_THREADS; + const dim3 gate_blocks(static_cast(info.ntokens * info.topk * info.intermediate_size)); + gate_up_kernel<<>>( + intermediate, + reinterpret_cast(hidden), + reinterpret_cast(topk_indices), + reinterpret_cast(topk_weights), + gate_ptrs, + up_ptrs, + info.ntokens, + info.hidden_size, + info.topk, + info.intermediate_size, + info.num_experts); + + const dim3 down_blocks(static_cast(info.ntokens * info.hidden_size)); + down_kernel<<>>( + reinterpret_cast(out), + intermediate, + reinterpret_cast(topk_indices), + down_ptrs, + info.ntokens, + info.hidden_size, + info.topk, + info.intermediate_size, + info.num_experts); + + return MOE_SUCCESS; +} + +DeepseekMoeInfo Info(const std::vector& v,size_t m,size_t e) { return {v[1].dtype(),v[1].size(0),v[1].size(1),v[2].size(1),m,e}; } +size_t WorkspaceBytes(const DeepseekMoeInfo& info) { + const size_t dtype_size = info.dtype == DataType::kFloat16 ? sizeof(half) : sizeof(__nv_bfloat16); + const size_t ptr_bytes = align_up(info.num_experts * sizeof(void *), 256); + const size_t intermediate_offset = align_up(ptr_bytes * 3, 256); + size_t workspace_tokens = info.ntokens; + if (info.ntokens > ROUTE_BATCHED_GEMM_MAX_TOKENS && info.num_experts <= 256) { + workspace_tokens = std::min(info.ntokens, CHUNKED_GROUPED_GEMM_TOKENS); + } + const size_t intermediate_bytes = workspace_tokens * info.topk * info.intermediate_size * dtype_size; + const size_t old_workspace_size = intermediate_offset + intermediate_bytes; + size_t batched_workspace_size = old_workspace_size; + size_t grouped_workspace_size = old_workspace_size; + if (info.ntokens > ROUTE_BATCHED_GEMM_MAX_TOKENS && info.num_experts <= 256) { + auto chunk_info = info; + chunk_info.ntokens = workspace_tokens; + if (info.dtype == DataType::kFloat16) { + grouped_workspace_size = expert_grouped_workspace_size(chunk_info, intermediate_offset); + } else { + grouped_workspace_size = expert_grouped_workspace_size<__nv_bfloat16>(chunk_info, intermediate_offset); + } + } else if (info.ntokens > FUSED_DECODE_MAX_TOKENS + && info.ntokens <= ROUTE_BATCHED_GEMM_MAX_TOKENS) { + if (info.dtype == DataType::kFloat16) { + batched_workspace_size = route_batched_workspace_size(info, intermediate_offset); + grouped_workspace_size = expert_grouped_workspace_size(info, intermediate_offset); + } else { + batched_workspace_size = route_batched_workspace_size<__nv_bfloat16>(info, intermediate_offset); + grouped_workspace_size = expert_grouped_workspace_size<__nv_bfloat16>(info, intermediate_offset); + } + } + const size_t workspace_size = std::max(old_workspace_size, std::max(batched_workspace_size, grouped_workspace_size)); + +return workspace_size + 4096; // per-region alignment headroom in the old pointer/buffer layout +} + +} // namespace +Operator::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) : DeepseekMoe(output, hidden, ids, weights, gate, up, down, intermediate, num_experts), blas_(NvidiaGemmHandle::Acquire(output.device().index())) { moe_detail::Require(SupportsLayout(hidden,gate,up,down,intermediate)); } +size_t Operator::workspace_size_in_bytes() const { return WorkspaceBytes(Info(metadata_,intermediate_,num_experts_)); } +void Operator::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 { + MoeDeviceGuard device_guard{metadata_[0].device().index()}; + Invocation(output, hidden, ids, weights, gate, up, down, intermediate, num_experts); if(hidden.size(0)==0) return; + moe_detail::Require(SupportsLayout(hidden,gate,up,down,intermediate)); + const auto info=Info(metadata_,intermediate,num_experts); + moe_detail::Workspace(workspace_,workspace_size_in_bytes_,WorkspaceBytes(info),Collect(output,hidden,ids,weights,gate,up,down)); + const auto stream=static_cast(stream_); + const size_t ptr_bytes=align_up(num_experts*sizeof(void*),256); + auto base=static_cast(workspace_); + auto gp=reinterpret_cast(base); + auto uptr=reinterpret_cast(base+ptr_bytes); + auto dp=reinterpret_cast(base+2*ptr_bytes); + for(size_t begin=0;begin>>(pointers,gp,uptr,dp,begin,count); + } + if(hidden.dtype()==DataType::kFloat16) Check(launch_typed(workspace_,workspace_size_in_bytes_,info,output.data(),hidden.data(),ids.data(),weights.data(),gp,uptr,dp,stream,*blas_)); + else Check(launch_typed<__nv_bfloat16>(workspace_,workspace_size_in_bytes_,info,output.data(),hidden.data(),ids.data(),weights.data(),gp,uptr,dp,stream,*blas_)); + MOE_CHECK_CUDA(cudaGetLastError()); + +} + +#undef MOE_CHECK_STATUS +#undef MOE_CHECK_CUDA +#undef MOE_CHECK_CUBLAS +} // 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 new file mode 100644 index 000000000..571b7c688 --- /dev/null +++ b/src/native/cuda/nvidia/ops/deepseek_moe/kernel.h @@ -0,0 +1,52 @@ +#pragma once +#include "base/deepseek_moe.h" +#include "native/cuda/nvidia/ops/gemm/handle_owner.h" +namespace infini::ops { +template <> +class Operator : public DeepseekMoe { + public: + 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); + size_t workspace_size_in_bytes() const override; + + private: + std::shared_ptr blas_; + + public: + 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) { + return DeepseekMoe::supports(output, hidden, ids, weights, gate, up, down, + intermediate, num_experts) && + SupportsLayout(hidden, gate, up, down, intermediate); + } + 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 override; + + private: + static bool SupportsLayout(const Tensor hidden, + const std::vector& gate, + const std::vector& up, + const std::vector& down, + size_t intermediate) { + if (hidden.dtype() != DataType::kBFloat16) return true; + const auto align_h = + hidden.size(1) % 8 == 0 ? 16u : (hidden.size(1) % 2 == 0 ? 4u : 2u); + const auto align_m = + intermediate % 8 == 0 ? 16u : (intermediate % 2 == 0 ? 4u : 2u); + if (reinterpret_cast(hidden.data()) % align_h) return false; + for (size_t e = 0; e < gate.size(); ++e) + if (reinterpret_cast(gate[e].data()) % align_h || + reinterpret_cast(up[e].data()) % align_h || + reinterpret_cast(down[e].data()) % align_m) + return false; + return true; + } +}; +} // 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..7486c1e58 --- /dev/null +++ b/src/native/cuda/nvidia/ops/flash_attn_varlen_func/kernel.cc @@ -0,0 +1,251 @@ +#include "kernel.h" + +#include +#include +#include +#include +#include + +#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 { + +using native_flash::DeviceScope; +using native_flash::plain_layout; +using native_flash::vector_aligned; + +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; + // These native kernels require SM80+. The attribute query allocates no + // scratch. + 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) + native_flash::launch_splitkv(p, stream); + else + native_flash::launch_varlen(p, stream); + 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..8bb148462 --- /dev/null +++ b/src/native/cuda/nvidia/ops/flash_attn_varlen_func/kernel.h @@ -0,0 +1,65 @@ +#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 rounded widths and vector layouts 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/kernel.cc b/src/native/cuda/nvidia/ops/flash_attn_with_kvcache/kernel.cc new file mode 100644 index 000000000..c278c79d0 --- /dev/null +++ b/src/native/cuda/nvidia/ops/flash_attn_with_kvcache/kernel.cc @@ -0,0 +1,216 @@ +#include "kernel.h" + +#include +#include +#include +#include + +#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 { +namespace { +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; + 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) || !native_flash::head_dim_bin(q.size(3)) || + q.device().type() != Device::Type::kNvidia) + 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 (!vector_aligned(tensor)) 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; + int major = 0; + return cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, + q.device().index()) == cudaSuccess && + major >= 8; +} + +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 = native_flash::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; + native_flash::launch_splitkv(p, stream); + 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..8d6ac332d --- /dev/null +++ b/src/native/cuda/nvidia/ops/flash_attn_with_kvcache/kernel.h @@ -0,0 +1,190 @@ +#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: rounded kernel widths, read-only cache, global or +// finite left-window/right-zero attention, no ALiBi/softcap. +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); + } + } + + 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. + bool capture_safe() const override { return true; } + 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_common.cc b/src/native/cuda/nvidia/ops/flash_common.cc new file mode 100644 index 000000000..d1921a423 --- /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/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/flash_head_dim.h b/src/native/cuda/nvidia/ops/flash_head_dim.h new file mode 100644 index 000000000..99990830a --- /dev/null +++ b/src/native/cuda/nvidia/ops/flash_head_dim.h @@ -0,0 +1,17 @@ +#pragma once + +namespace infini::ops::native_flash { + +// 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) { + 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) { + 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..170a0f4cd --- /dev/null +++ b/src/native/cuda/nvidia/ops/flash_workspace.h @@ -0,0 +1,109 @@ +#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/fused_gated_delta_net_gating/kernel.h b/src/native/cuda/nvidia/ops/fused_gated_delta_net_gating/kernel.h new file mode 100644 index 000000000..41699f41f --- /dev/null +++ b/src/native/cuda/nvidia/ops/fused_gated_delta_net_gating/kernel.h @@ -0,0 +1,120 @@ +#pragma once +#include +#include +#include + +#include +#include +#include + +#include "base/fused_gated_delta_net_gating.h" +#include "native/cuda/nvidia/caster.cuh" +#include "native/cuda/nvidia/runtime_.h" + +namespace infini::ops { +namespace gating_detail { +struct Layout { + ptrdiff_t strides[6][3]{}; + size_t total, seq, hidden; + float beta, threshold; +}; +template +__global__ void kernel(const T* A, const T* a, const T* b, const T* dt, + float* g, float* betao, Layout p) { + const size_t linear = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (linear >= p.total) return; + const size_t h = linear % p.hidden, s = (linear / p.hidden) % p.seq; + const size_t batch = linear / p.hidden / p.seq; + ptrdiff_t offsets[6]; + for (int i = 0; i < 6; ++i) + offsets[i] = static_cast(batch) * p.strides[i][0] + + static_cast(s) * p.strides[i][1] + + static_cast(h) * p.strides[i][2]; + const float x = static_cast(a[offsets[1]]) + + static_cast(dt[h * p.strides[3][0]]); + const float bx = p.beta * x; + const float softplus = bx <= p.threshold ? log1pf(expf(bx)) / p.beta : x; + g[offsets[4]] = -expf(static_cast(A[h * p.strides[0][0]])) * softplus; + const float bv = static_cast(b[offsets[2]]); + const float z = expf(bv >= 0 ? -bv : bv); + betao[offsets[5]] = bv >= 0 ? 1.0f / (1.0f + z) : z / (1.0f + z); +} +} // namespace gating_detail + +template <> +class Operator + : public FusedGatedDeltaNetGating { + public: + Operator(Tensor A, Tensor a, Tensor b, Tensor dt, float beta, float threshold, + Tensor g, Tensor betao) + : FusedGatedDeltaNetGating(A, a, b, dt, beta, threshold, g, betao) { + if (device_.type() != Device::Type::kNvidia) + throw std::invalid_argument("gating requires NVIDIA device"); + check_device(); + if (total_ / 256 + (total_ % 256 != 0) > + std::numeric_limits::max()) + throw std::overflow_error("gating grid exceeds CUDA limit"); + p_.total = total_; + p_.seq = seq_; + p_.hidden = hidden_; + p_.beta = beta; + p_.threshold = threshold; + for (size_t i = 0; i < 6; ++i) + 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 + // is capturable. Make owns no device resources and requires zero workspace. + bool capture_safe() const override { return true; } + void operator()(Tensor A, Tensor a, Tensor b, Tensor dt, float beta, + float threshold, Tensor g, Tensor betao) const override { + check_device(); + std::array tensors{A, a, b, dt, g, betao}; + for (size_t i = 0; i < 6; ++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))) + throw std::invalid_argument("gating attributes changed after Make"); + if (!total_) return; + const dim3 grid( + static_cast(total_ / 256 + (total_ % 256 != 0))); + const auto stream = static_cast(stream_ ? stream_ : nullptr); + switch (metadata_[1].dtype()) { + case DataType::kFloat32: + launch(A, a, b, dt, g, betao, grid, stream); + break; + case DataType::kFloat16: + launch(A, a, b, dt, g, betao, grid, stream); + break; + case DataType::kBFloat16: + launch<__nv_bfloat16>(A, a, b, dt, g, betao, grid, stream); + break; + default: + throw std::invalid_argument("gating dtype unsupported"); + } + if (cudaPeekAtLastError() != cudaSuccess) + throw std::runtime_error("gating CUDA launch failed"); + } + + private: + gating_detail::Layout p_{}; + void check_device() const { + int current = -1; + if (cudaGetDevice(¤t) != cudaSuccess) + throw std::runtime_error("gating device query failed"); + if (current != device_.index()) + throw std::invalid_argument("gating owner device mismatch"); + } + template + void launch(Tensor A, Tensor a, Tensor b, Tensor dt, Tensor g, Tensor betao, + dim3 grid, cudaStream_t stream) const { + gating_detail::kernel<<>>( + static_cast(A.data()), static_cast(a.data()), + static_cast(b.data()), static_cast(dt.data()), + static_cast(g.data()), static_cast(betao.data()), p_); + } +}; +} // 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/gemm/cublas.h b/src/native/cuda/nvidia/ops/gemm/cublas.h index eefe0b7af..795b71e9c 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,42 @@ template <> class Operator : public BlasGemm> { public: - using BlasGemm>::BlasGemm; + 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) + : BlasGemm>{a, b, c, alpha, + beta, trans_a, trans_b, y}, + handle_{NvidiaGemmHandle::Acquire( + y.device().index(), + static_cast( + CUBLAS_DEFAULT_MATH | + CUBLAS_MATH_DISALLOW_REDUCED_PRECISION_REDUCTION))} {} + + 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..540a1beb4 100644 --- a/src/native/cuda/nvidia/ops/gemm/cublaslt.h +++ b/src/native/cuda/nvidia/ops/gemm/cublaslt.h @@ -18,6 +18,14 @@ namespace infini::ops { template <> class Operator : public Gemm { public: + 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) @@ -80,7 +88,9 @@ 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..8f2d58e10 --- /dev/null +++ b/src/native/cuda/nvidia/ops/gemm/handle_owner.h @@ -0,0 +1,138 @@ +#ifndef INFINI_OPS_NVIDIA_GEMM_HANDLE_OWNER_H_ +#define INFINI_OPS_NVIDIA_GEMM_HANDLE_OWNER_H_ + +#include +#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, cublasMath_t math_mode = CUBLAS_DEFAULT_MATH) { + 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"); + } + const auto key = std::make_pair(context, math_mode); + static thread_local std::map, + std::weak_ptr> + owners; + for (auto it = owners.begin(); it != owners.end();) { + if (it->second.expired()) + it = owners.erase(it); + else + ++it; + } + if (auto it = owners.find(key); it != owners.end()) { + if (auto owner = it->second.lock()) return owner; + } + auto owner = std::shared_ptr( + new NvidiaGemmHandle(context, math_mode)); + owners[key] = 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, cublasMath_t math_mode) + : 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))); + } + if (math_mode != CUBLAS_DEFAULT_MATH) { + const auto math_status = cublasSetMathMode(handle_, math_mode); + if (math_status != CUBLAS_STATUS_SUCCESS) { + const auto cleanup = cublasDestroy(handle_); + handle_ = nullptr; + if (cleanup != CUBLAS_STATUS_SUCCESS) { + ReportCleanupError("cleaning failed math mode setup", + static_cast(cleanup)); + } + throw std::runtime_error( + "setting GEMM BLAS reduction math mode failed: " + + std::to_string(static_cast(math_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/i8_gemm/kernel.cu b/src/native/cuda/nvidia/ops/i8_gemm/kernel.cu new file mode 100644 index 000000000..ebdbcb038 --- /dev/null +++ b/src/native/cuda/nvidia/ops/i8_gemm/kernel.cu @@ -0,0 +1,28 @@ +#include "native/cuda/nvidia/ops/i8_gemm/kernel.h" +#include +#include +#include +namespace infini::ops::nvidia_i8_gemm_detail { +template __global__ void Post(T* c,const std::int32_t* accum, + const float* as,const float* bs,const T* bias,Layout layout) { + for(std::size_t i=std::size_t(blockIdx.x)*blockDim.x+threadIdx.x; + i((layout.m*layout.n-1)/256+1,65535)); + switch(dtype) { + case DataType::kFloat16:Post<<>>(static_cast<__half*>(c),accum,as,bs,static_cast(bias),layout);break; + case DataType::kBFloat16:Post<<>>(static_cast<__nv_bfloat16*>(c),accum,as,bs,static_cast(bias),layout);break; + case DataType::kFloat32:Post<<>>(static_cast(c),accum,as,bs,static_cast(bias),layout);break; + default:throw std::invalid_argument("unsupported I8Gemm output dtype"); + } + const auto error=cudaGetLastError();if(error!=cudaSuccess) throw std::runtime_error(cudaGetErrorString(error)); +} +} diff --git a/src/native/cuda/nvidia/ops/i8_gemm/kernel.h b/src/native/cuda/nvidia/ops/i8_gemm/kernel.h new file mode 100644 index 000000000..ecb1a0c63 --- /dev/null +++ b/src/native/cuda/nvidia/ops/i8_gemm/kernel.h @@ -0,0 +1,160 @@ +#pragma once +#include +#include +#include + +#include "base/i8_gemm.h" +#include "native/cuda/nvidia/ops/binary_view_capture.h" +#include "native/cuda/nvidia/ops/gemm/handle_owner.h" +namespace infini::ops::nvidia_i8_gemm_detail { +struct Layout { + std::size_t m, n; + std::int64_t cr, cc, ar, br, sar, sbr, biasr, biasc; +}; +void Epilogue(void*, const std::int32_t*, const float*, const float*, + const void*, DataType, Layout, cudaStream_t); +inline void Check(cublasStatus_t status) { + if (status != CUBLAS_STATUS_SUCCESS) + throw std::runtime_error("I8Gemm cublas error " + + std::to_string(int(status))); +} +} // namespace infini::ops::nvidia_i8_gemm_detail +namespace infini::ops { +template <> +class Operator : public I8Gemm { + public: + static constexpr std::size_t kBlasWorkspace = 4 * 1024 * 1024; + static bool supports(Tensor c, const Tensor a, const Tensor as, + const Tensor b, const Tensor bs, + const std::optional bias) { + using namespace nvidia_binary_view; + if (c.device().type() != Device::Type::kNvidia || c.ndim() != 2 || + a.ndim() != 2 || b.ndim() != 2 || a.dtype() != DataType::kInt8 || + b.dtype() != DataType::kInt8 || + (c.dtype() != DataType::kFloat16 && c.dtype() != DataType::kBFloat16 && + c.dtype() != DataType::kFloat32) || + as.dtype() != DataType::kFloat32 || bs.dtype() != DataType::kFloat32 || + c.device() != a.device() || c.device() != b.device() || + c.device() != as.device() || c.device() != bs.device()) + return false; + const auto m = c.shape()[0], n = c.shape()[1], k = a.shape()[1]; + if (!m || !n || !k || m > INT_MAX || n > INT_MAX || k > INT_MAX || n % 4 || + k % 4 || a.shape()[0] != m || b.shape() != Tensor::Shape{k, n} || + as.shape() != Tensor::Shape{m, 1} || + bs.shape() != Tensor::Shape{n, 1} || a.strides().size() != 2 || + b.strides().size() != 2 || as.strides().size() != 2 || + bs.strides().size() != 2 || a.stride(1) != 1 || b.stride(0) != 1 || + a.stride(0) < std::int64_t(k) || b.stride(1) < std::int64_t(k) || + a.stride(0) > INT_MAX || b.stride(1) > INT_MAX || a.stride(0) % 4 || + b.stride(1) % 4 || !NonOverlappingOutput(c) || as.stride(0) < 1 || + bs.stride(0) < 1) + return false; + // Worst-case signed INT8 dot cannot overflow int32 under this bound. + if (k > std::size_t(INT_MAX) / 16384 || + m > std::numeric_limits::max() / n / sizeof(std::int32_t)) + return false; + int device = -1, major = 0, minor = 0; + if (cudaGetDevice(&device) != cudaSuccess || device != c.device().index() || + cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, + device) != cudaSuccess || + cudaDeviceGetAttribute(&minor, cudaDevAttrComputeCapabilityMinor, + device) != cudaSuccess || + major * 10 + minor < 75 || + (c.dtype() == DataType::kBFloat16 && major < 8)) + return false; + if (bias && (bias->device() != c.device() || bias->dtype() != c.dtype() || + bias->shape() != c.shape())) + return false; + std::size_t cb = 0; + if (!SpanBytes(c, &cb)) return false; + const auto clear = [&](Tensor input) { + std::size_t bytes = 0; + return SpanBytes(input, &bytes) && !Overlaps(c, cb, input, bytes) && + (!bytes || reinterpret_cast(input.data()) % + input.element_size() == + 0); + }; + if (!clear(a) || !clear(b) || !clear(as) || !clear(bs) || + (bias && !clear(*bias))) + return false; + return reinterpret_cast(a.data()) % 4 == 0 && + reinterpret_cast(b.data()) % 4 == 0 && + reinterpret_cast(c.data()) % c.element_size() == 0; + } + Operator(Tensor c, const Tensor a, const Tensor as, const Tensor b, + const Tensor bs, const std::optional bias) + : I8Gemm(c, a, as, b, bs, bias), dtype_(c.dtype()) { + const auto m = c.shape()[0], n = c.shape()[1]; + const auto bytes = m * n * sizeof(std::int32_t); + if (bytes > std::numeric_limits::max() - 255 - kBlasWorkspace) + throw std::invalid_argument("I8Gemm workspace size overflow"); + scratch_offset_ = (bytes + 255) & ~std::size_t(255); + owner_ = NvidiaGemmHandle::Acquire(c.device().index()); + layout_ = {m, + n, + c.stride(0), + c.stride(1), + a.stride(0), + b.stride(1), + as.stride(0), + bs.stride(0), + bias ? bias->stride(0) : 0, + bias ? bias->stride(1) : 0}; + } + std::size_t workspace_size_in_bytes() const override { + return scratch_offset_ + kBlasWorkspace; + } + bool capture_safe() const override { return true; } + void operator()(Tensor c, const Tensor a, const Tensor as, const Tensor b, + const Tensor bs, + const std::optional bias) const override { + if (!supports(c, a, as, b, bs, bias)) + throw std::invalid_argument("unsupported I8Gemm operands"); + const std::array tensors{c, a, as, b, bs}; + if (!workspace_ || workspace_size_in_bytes_ < workspace_size_in_bytes() || + reinterpret_cast(workspace_) % 256) + throw std::invalid_argument("I8Gemm requires aligned caller workspace"); + for (const auto& t : tensors) + if (WorkspaceOverlaps(t)) + throw std::invalid_argument("I8Gemm workspace aliases tensor"); + if (bias && WorkspaceOverlaps(*bias)) + throw std::invalid_argument("I8Gemm workspace aliases bias"); + const auto handle = owner_->get(); + const auto stream = static_cast(stream_ ? stream_ : 0); + using nvidia_i8_gemm_detail::Check; + Check(cublasSetStream(handle, stream)); + Check(cublasSetWorkspace(handle, + static_cast(workspace_) + scratch_offset_, + kBlasWorkspace)); + Check(cublasSetPointerMode(handle, CUBLAS_POINTER_MODE_HOST)); + Check(cublasSetMathMode(handle, CUBLAS_DEFAULT_MATH)); + const std::int32_t alpha = 1, beta = 0; + // Column-major C^T = B^T * A^T. B's [K,N] view stores each + // original [N,K] weight row contiguously, hence OP_T for B. + Check(cublasGemmEx( + handle, CUBLAS_OP_T, CUBLAS_OP_N, int(layout_.n), int(layout_.m), + int(a.shape()[1]), &alpha, b.data(), CUDA_R_8I, int(layout_.br), + a.data(), CUDA_R_8I, int(layout_.ar), &beta, workspace_, CUDA_R_32I, + int(layout_.n), CUBLAS_COMPUTE_32I, CUBLAS_GEMM_DEFAULT)); + nvidia_i8_gemm_detail::Epilogue( + c.data(), static_cast(workspace_), + static_cast(as.data()), + static_cast(bs.data()), bias ? bias->data() : nullptr, + dtype_, layout_, stream); + } + + private: + bool WorkspaceOverlaps(Tensor t) const { + std::size_t bytes = 0; + if (!nvidia_binary_view::SpanBytes(t, &bytes)) return true; + const auto a = reinterpret_cast(workspace_), + b = reinterpret_cast(t.data()); + return bytes && + (a <= b ? b - a < workspace_size_in_bytes() : a - b < bytes); + } + DataType dtype_; + std::size_t scratch_offset_; + nvidia_i8_gemm_detail::Layout layout_{}; + std::shared_ptr owner_; +}; +} // namespace infini::ops diff --git a/src/native/cuda/nvidia/ops/infini_flash_cuda.h b/src/native/cuda/nvidia/ops/infini_flash_cuda.h new file mode 100644 index 000000000..c72ab625b --- /dev/null +++ b/src/native/cuda/nvidia/ops/infini_flash_cuda.h @@ -0,0 +1,23 @@ +#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 +#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()) + +// 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/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..e8e7bd21e --- /dev/null +++ b/src/native/cuda/nvidia/ops/layer_norm/kernel.h @@ -0,0 +1,188 @@ +#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) || + !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) { + 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_meta_.strides() == out_meta_.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; } + + void operator()(const Tensor input, const Tensor weight, + const std::optional bias, float eps, + Tensor out) const override { + 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"); + + 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(), + input_meta_.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)); + 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/mamba_selective_scan/kernel.h b/src/native/cuda/nvidia/ops/mamba_selective_scan/kernel.h new file mode 100644 index 000000000..ddb2b2c40 --- /dev/null +++ b/src/native/cuda/nvidia/ops/mamba_selective_scan/kernel.h @@ -0,0 +1,18 @@ +#pragma once +#include "native/cuda/nvidia/caster.cuh" +#include "native/cuda/nvidia/runtime_.h" +#include "native/cuda/ops/mamba_selective_scan/kernel.h" +namespace infini::ops { +struct MambaScanNvidiaBackend : Runtime { + static void CheckLaunch() { + if (cudaGetLastError() != cudaSuccess) + throw std::runtime_error("MambaSelectiveScan CUDA launch failed"); + } +}; +template <> +class Operator + : public CudaMambaSelectiveScan { + public: + using CudaMambaSelectiveScan::CudaMambaSelectiveScan; +}; +} // namespace infini::ops diff --git a/src/native/cuda/nvidia/ops/moe_align_block_size/kernel.cu b/src/native/cuda/nvidia/ops/moe_align_block_size/kernel.cu index f1897a5eb..9b9c91baf 100644 --- a/src/native/cuda/nvidia/ops/moe_align_block_size/kernel.cu +++ b/src/native/cuda/nvidia/ops/moe_align_block_size/kernel.cu @@ -70,7 +70,8 @@ void Operator::Run( reinterpret_cast(experts_ids.data()), reinterpret_cast(num_tokens_post_pad.data()), numel_, static_cast(num_experts_), static_cast(block_size_), - sorted_token_ids_size_, experts_ids_size_); + sorted_token_ids_size_, experts_ids_size_, + maybe_expert_map ? static_cast(maybe_expert_map->numel()) : 0); const auto status = cudaGetLastError(); assert(status == cudaSuccess && diff --git a/src/native/cuda/nvidia/ops/moe_align_block_size/kernel.cuh b/src/native/cuda/nvidia/ops/moe_align_block_size/kernel.cuh index 467383221..3e90713e4 100644 --- a/src/native/cuda/nvidia/ops/moe_align_block_size/kernel.cuh +++ b/src/native/cuda/nvidia/ops/moe_align_block_size/kernel.cuh @@ -9,8 +9,9 @@ namespace moe_align_block_size_detail { __device__ __forceinline__ int32_t MapExpertId(int32_t expert_id, const int32_t* expert_map, - int32_t num_experts) { - if (expert_id < 0 || expert_id >= num_experts) { + int32_t num_experts, + int32_t map_size) { + if (expert_id < 0 || expert_id >= (expert_map ? map_size : num_experts)) { return -1; } if (expert_map != nullptr) { @@ -26,7 +27,7 @@ __global__ void MoeAlignBlockSizeKernel( int32_t* __restrict__ sorted_token_ids, int32_t* __restrict__ experts_ids, int32_t* __restrict__ num_tokens_post_pad, size_t numel, int32_t num_experts, int32_t block_size, size_t sorted_token_ids_size, - size_t experts_ids_size) { + size_t experts_ids_size, int32_t map_size) { extern __shared__ int32_t expert_offsets[]; for (size_t i = threadIdx.x; i < sorted_token_ids_size; i += blockDim.x) { @@ -44,7 +45,8 @@ __global__ void MoeAlignBlockSizeKernel( __syncthreads(); for (size_t i = threadIdx.x; i < numel; i += blockDim.x) { - const auto expert_id = MapExpertId(topk_ids[i], expert_map, num_experts); + const auto expert_id = + MapExpertId(topk_ids[i], expert_map, num_experts, map_size); if (expert_id >= 0) { atomicAdd(&expert_offsets[expert_id], 1); } @@ -69,7 +71,8 @@ __global__ void MoeAlignBlockSizeKernel( if (threadIdx.x == 0) { for (size_t i = 0; i < numel; ++i) { - const auto expert_id = MapExpertId(topk_ids[i], expert_map, num_experts); + const auto expert_id = + MapExpertId(topk_ids[i], expert_map, num_experts, map_size); if (expert_id >= 0) { sorted_token_ids[expert_offsets[expert_id]++] = static_cast(i); } diff --git a/src/native/cuda/nvidia/ops/moe_align_block_size/kernel.h b/src/native/cuda/nvidia/ops/moe_align_block_size/kernel.h index 95e5a272c..79283640b 100644 --- a/src/native/cuda/nvidia/ops/moe_align_block_size/kernel.h +++ b/src/native/cuda/nvidia/ops/moe_align_block_size/kernel.h @@ -18,6 +18,8 @@ class Operator using MoeAlignBlockSize::operator(); + bool capture_safe() const override { return true; } + protected: void Run(const Tensor topk_ids, const std::optional maybe_expert_map, const int64_t num_experts, const int64_t block_size, diff --git a/src/native/cuda/nvidia/ops/moe_fused_dense/kernel.cu b/src/native/cuda/nvidia/ops/moe_fused_dense/kernel.cu new file mode 100644 index 000000000..f80fb5ce8 --- /dev/null +++ b/src/native/cuda/nvidia/ops/moe_fused_dense/kernel.cu @@ -0,0 +1,699 @@ +// 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 +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace infini::ops { namespace { + +enum MoeStatus { MOE_SUCCESS=0, MOE_BAD_PARAM, MOE_BAD_DTYPE, MOE_WORKSPACE, MOE_INTERNAL, MOE_DEVICE }; +#define MOE_CHECK_STATUS(expr) do { auto s=(expr); if(s!=MOE_SUCCESS) return s; } while(0) +#define MOE_CHECK_CUDA(expr) do { auto s=(expr); if(s!=cudaSuccess) throw std::runtime_error(cudaGetErrorString(s)); } while(0) +inline void Check(MoeStatus s) { if(s!=MOE_SUCCESS) throw std::runtime_error("MoE native failure: "+std::to_string(int(s))); } +struct MoeFusedDenseInfo { size_t num_tokens, hidden_size, num_experts, intermediate_size, topk, max_num_tokens_padded, max_num_blocks, block_size; DataType dtype; }; + + +size_t align_up(size_t value, size_t alignment) { + return (value + alignment - 1) / alignment * alignment; +} + +void *advance_workspace(uint8_t *&ptr, size_t &remaining, size_t bytes, size_t alignment = 16) { + auto address = reinterpret_cast(ptr); + auto aligned = align_up(address, alignment); + auto padding = aligned - address; + if (padding + bytes > remaining) { + return nullptr; + } + ptr += padding; + remaining -= padding; + void *out = ptr; + ptr += bytes; + remaining -= bytes; + return out; +} + +size_t workspace_size(const MoeFusedDenseInfo &info) { + const size_t pairs = info.num_tokens * info.topk; + const size_t rows = info.max_num_tokens_padded; + const size_t dtype_size = info.dtype == DataType::kFloat16 ? sizeof(half) : sizeof(__nv_bfloat16); + size_t bytes = 0; + bytes += align_up((info.num_experts + 1) * sizeof(int), 16); + bytes += align_up((info.num_experts + 1) * sizeof(int), 16); + bytes += align_up(pairs * sizeof(int), 16); + bytes += align_up(rows * info.hidden_size * dtype_size, 16); + bytes += align_up(rows * info.intermediate_size * 2 * dtype_size, 16); + bytes += align_up(rows * info.intermediate_size * dtype_size, 16); + bytes += align_up(rows * info.hidden_size * dtype_size, 16); + bytes += align_up(info.num_experts * sizeof(cutlass::gemm::GemmCoord), 16); + bytes += align_up(info.num_experts * sizeof(void *), 16) * 4; + bytes += align_up(info.num_experts * sizeof(int64_t), 16) * 4; + bytes += align_up(sizeof(int), 16); + return bytes + 256; +} + +template +__device__ float to_float(T v); + +template <> +__device__ float to_float(half v) { + return __half2float(v); +} + +template <> +__device__ float to_float<__nv_bfloat16>(__nv_bfloat16 v) { + return __bfloat162float(v); +} + +template +__device__ T from_float(float v); + +template <> +__device__ half from_float(float v) { + return __float2half(v); +} + +template <> +__device__ __nv_bfloat16 from_float<__nv_bfloat16>(float v) { + return __float2bfloat16(v); +} + +__global__ void exclusive_prefix_counts_kernel(const int *counts, int *offsets, int num_experts) { + extern __shared__ int scan[]; + int tid = threadIdx.x; + if (tid < num_experts) { + scan[tid] = counts[tid]; + } + if (tid >= num_experts && tid < blockDim.x) { + scan[tid] = 0; + } + __syncthreads(); + + for (int stride = 1; stride < blockDim.x; stride <<= 1) { + int value = 0; + if (tid >= stride) { + value = scan[tid - stride]; + } + __syncthreads(); + scan[tid] += value; + __syncthreads(); + } + + if (tid == 0) { + offsets[0] = 0; + } + if (tid < num_experts) { + offsets[tid + 1] = scan[tid]; + } +} + +__global__ void count_aligned_experts_kernel(const int *expert_ids, + const int *num_tokens_post_padded, + int *counts, + int num_experts, + int block_size) { + int block = blockIdx.x * blockDim.x + threadIdx.x; + int num_blocks = (*num_tokens_post_padded + block_size - 1) / block_size; + if (block >= num_blocks) { + return; + } + int expert = expert_ids[block]; + if (expert >= 0 && expert < num_experts) { + atomicAdd(counts + expert, block_size); + } +} + +template +__global__ void pack_hidden_aligned_kernel(const T *hidden, + const int *sorted_token_ids, + int *output_permutation, + T *packed_hidden, + int pairs, + int topk, + int hidden_size, + int max_num_tokens_padded) { + int row = blockIdx.x; + int tid = threadIdx.x; + if (row >= max_num_tokens_padded) { + return; + } + int pair = sorted_token_ids[row]; + if (pair >= 0 && pair < pairs) { + if (tid == 0) { + output_permutation[pair] = row; + } + int token = pair / topk; + for (int h = tid; h < hidden_size; h += blockDim.x) { + packed_hidden[static_cast(row) * hidden_size + h] = hidden[static_cast(token) * hidden_size + h]; + } + } else { + for (int h = tid; h < hidden_size; h += blockDim.x) { + packed_hidden[static_cast(row) * hidden_size + h] = from_float(0.0f); + } + } +} + +template +__global__ void fill_prefill_bias_kernel(T *gate_up, T *expert_out, + const T *bias1, const T *bias2, + const int *counts, const int *offsets, + int intermediate_size, int hidden_size) { + const int expert = blockIdx.x; + const int rows = counts[expert]; + const int offset = offsets[expert]; + const int gate_columns = intermediate_size * 2; + if (bias1) { + for (size_t i = threadIdx.x; i < static_cast(rows) * gate_columns; i += blockDim.x) { + gate_up[static_cast(offset) * gate_columns + i] = + bias1[static_cast(expert) * gate_columns + i % gate_columns]; + } + } + if (bias2) { + for (size_t i = threadIdx.x; i < static_cast(rows) * hidden_size; i += blockDim.x) { + expert_out[static_cast(offset) * hidden_size + i] = + bias2[static_cast(expert) * hidden_size + i % hidden_size]; + } + } +} + +template +__global__ void fill_decode_bias_kernel(T *gate_up, T *expert_out, + const T *bias1, const T *bias2, + const int *sorted_token_ids, const int *expert_ids, + const int *num_tokens_post_padded, const int *output_permutation, + int pairs, int num_experts, int intermediate_size, + int hidden_size, int block_size) { + const int row = blockIdx.x; + if (row >= *num_tokens_post_padded) return; + const int pair = sorted_token_ids[row]; + if (pair < 0 || pair >= pairs) return; + const int expert = expert_ids[row / block_size]; + if (expert < 0 || expert >= num_experts) return; + const int idx = output_permutation[pair]; + if (idx < 0 || idx >= pairs) return; + const int gate_columns = intermediate_size * 2; + if (bias1) { + for (size_t col = threadIdx.x; col < static_cast(gate_columns); col += blockDim.x) { + gate_up[static_cast(idx) * gate_columns + col] = + bias1[static_cast(expert) * gate_columns + col]; + } + } + if (bias2) { + for (size_t col = threadIdx.x; col < static_cast(hidden_size); col += blockDim.x) { + expert_out[static_cast(idx) * hidden_size + col] = + bias2[static_cast(expert) * hidden_size + col]; + } + } +} + +template +__global__ void swiglu_kernel(const T *gate_up, T *activated, int rows, int intermediate_size) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + int total = rows * intermediate_size; + if (idx >= total) { + return; + } + int row = idx / intermediate_size; + int col = idx - row * intermediate_size; + const T *base = gate_up + static_cast(row) * intermediate_size * 2; + float gate = to_float(base[col]); + float up = to_float(base[intermediate_size + col]); + float silu = gate / (1.0f + expf(-gate)); + activated[idx] = from_float(up * silu); +} + +template +__global__ void apply_shuffle_mul_sum_kernel(const T *__restrict__ expert_out, + T *__restrict__ out, + const int *__restrict__ output_permutation, + const float *__restrict__ topk_weights, + int num_tokens, + int topk, + int hidden_size) { + int token = blockIdx.x; + if (token >= num_tokens) { + return; + } + + for (int h = threadIdx.x; h < hidden_size; h += blockDim.x) { + float sum = 0.0f; + for (int k = 0; k < topk; ++k) { + int pair = token * topk + k; + int src_row = output_permutation[pair]; + if (src_row >= 0) { + sum += to_float(expert_out[static_cast(src_row) * hidden_size + h]) * topk_weights[pair]; + } + } + out[static_cast(token) * hidden_size + h] = from_float(sum); + } +} + +template +__global__ void setup_prefill_gemm1_kernel(cutlass::gemm::GemmCoord *problems, + void **ptr_a, + void **ptr_b, + void **ptr_c, + void **ptr_d, + int64_t *lda, + int64_t *ldb, + int64_t *ldc, + int64_t *ldd, + const int *counts, + const int *offsets, + const T *packed_hidden, + const T *w13, + T *gate_up, + int num_experts, + int hidden_size, + int intermediate_size) { + int expert = blockIdx.x * blockDim.x + threadIdx.x; + if (expert >= num_experts) { + return; + } + int m = counts[expert]; + int off = offsets[expert]; + problems[expert] = cutlass::gemm::GemmCoord(m, intermediate_size * 2, hidden_size); + ptr_a[expert] = const_cast(packed_hidden + static_cast(off) * hidden_size); + ptr_b[expert] = const_cast(w13 + static_cast(expert) * intermediate_size * 2 * hidden_size); + ptr_c[expert] = gate_up + static_cast(off) * intermediate_size * 2; + ptr_d[expert] = gate_up + static_cast(off) * intermediate_size * 2; + lda[expert] = hidden_size; + ldb[expert] = hidden_size; + ldc[expert] = intermediate_size * 2; + ldd[expert] = intermediate_size * 2; +} + +template +__global__ void setup_prefill_gemm2_kernel(cutlass::gemm::GemmCoord *problems, + void **ptr_a, + void **ptr_b, + void **ptr_c, + void **ptr_d, + int64_t *lda, + int64_t *ldb, + int64_t *ldc, + int64_t *ldd, + const int *counts, + const int *offsets, + const T *activated, + const T *w2, + T *expert_out, + int num_experts, + int hidden_size, + int intermediate_size) { + int expert = blockIdx.x * blockDim.x + threadIdx.x; + if (expert >= num_experts) { + return; + } + int m = counts[expert]; + int off = offsets[expert]; + problems[expert] = cutlass::gemm::GemmCoord(m, hidden_size, intermediate_size); + ptr_a[expert] = const_cast(activated + static_cast(off) * intermediate_size); + ptr_b[expert] = const_cast(w2 + static_cast(expert) * hidden_size * intermediate_size); + ptr_c[expert] = expert_out + static_cast(off) * hidden_size; + ptr_d[expert] = expert_out + static_cast(off) * hidden_size; + lda[expert] = intermediate_size; + ldb[expert] = intermediate_size; + ldc[expert] = hidden_size; + ldd[expert] = hidden_size; +} + +template +__global__ void setup_decode_gemm1_aligned_compact_kernel(cutlass::gemm::GemmCoord *problems, + void **ptr_a, + void **ptr_b, + void **ptr_c, + void **ptr_d, + int64_t *lda, + int64_t *ldb, + int64_t *ldc, + int64_t *ldd, + int *active_count, + int *output_permutation, + const T *hidden, + const T *w13, + T *gate_up, + const int *sorted_token_ids, + const int *expert_ids, + const int *num_tokens_post_padded, + int pairs, + int topk, + int num_experts, + int hidden_size, + int intermediate_size, + int block_size) { + int row = blockIdx.x * blockDim.x + threadIdx.x; + int total_rows = *num_tokens_post_padded; + if (row >= total_rows) { + return; + } + int pair = sorted_token_ids[row]; + if (pair < 0 || pair >= pairs) { + return; + } + int expert = expert_ids[row / block_size]; + if (expert < 0 || expert >= num_experts) { + return; + } + + int idx = atomicAdd(active_count, 1); + int token = pair / topk; + output_permutation[pair] = idx; + problems[idx] = cutlass::gemm::GemmCoord(1, intermediate_size * 2, hidden_size); + ptr_a[idx] = const_cast(hidden + static_cast(token) * hidden_size); + ptr_b[idx] = const_cast(w13 + static_cast(expert) * intermediate_size * 2 * hidden_size); + ptr_c[idx] = gate_up + static_cast(idx) * intermediate_size * 2; + ptr_d[idx] = gate_up + static_cast(idx) * intermediate_size * 2; + lda[idx] = hidden_size; + ldb[idx] = hidden_size; + ldc[idx] = intermediate_size * 2; + ldd[idx] = intermediate_size * 2; +} + +template +__global__ void setup_decode_gemm2_aligned_compact_kernel(cutlass::gemm::GemmCoord *problems, + void **ptr_a, + void **ptr_b, + void **ptr_c, + void **ptr_d, + int64_t *lda, + int64_t *ldb, + int64_t *ldc, + int64_t *ldd, + const T *activated, + const T *w2, + T *expert_out, + const int *sorted_token_ids, + const int *expert_ids, + const int *num_tokens_post_padded, + const int *output_permutation, + int pairs, + int num_experts, + int hidden_size, + int intermediate_size, + int block_size) { + int row = blockIdx.x * blockDim.x + threadIdx.x; + int total_rows = *num_tokens_post_padded; + if (row >= total_rows) { + return; + } + int pair = sorted_token_ids[row]; + if (pair < 0 || pair >= pairs) { + return; + } + int expert = expert_ids[row / block_size]; + if (expert < 0 || expert >= num_experts) { + return; + } + int idx = output_permutation[pair]; + if (idx < 0) { + return; + } + problems[idx] = cutlass::gemm::GemmCoord(1, hidden_size, intermediate_size); + ptr_a[idx] = const_cast(activated + static_cast(idx) * intermediate_size); + ptr_b[idx] = const_cast(w2 + static_cast(expert) * hidden_size * intermediate_size); + ptr_c[idx] = expert_out + static_cast(idx) * hidden_size; + ptr_d[idx] = expert_out + static_cast(idx) * hidden_size; + lda[idx] = intermediate_size; + ldb[idx] = intermediate_size; + ldc[idx] = hidden_size; + ldd[idx] = hidden_size; +} + +template +MoeStatus launch_cutlass_gemm_grouped_device_meta(int problem_count, + cutlass::gemm::GemmCoord *d_problems, + void **d_ptr_a, + void **d_ptr_b, + void **d_ptr_c, + void **d_ptr_d, + int64_t *d_lda, + int64_t *d_ldb, + int64_t *d_ldc, + int64_t *d_ldd, + bool full_occupancy_wave, + bool use_bias, + cudaStream_t stream) { + if (problem_count == 0) { + return MOE_SUCCESS; + } + + using Element = CutlassT; + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::RowMajor; + using OutputOp = cutlass::epilogue::thread::LinearCombination; + using GemmKernel = typename cutlass::gemm::kernel::DefaultGemmGrouped< + Element, + LayoutA, + cutlass::ComplexTransform::kNone, + 8, + Element, + LayoutB, + cutlass::ComplexTransform::kNone, + 8, + Element, + LayoutC, + float, + cutlass::arch::OpClassTensorOp, + cutlass::arch::Sm80, + cutlass::gemm::GemmShape<128, 128, 32>, + cutlass::gemm::GemmShape<64, 64, 32>, + cutlass::gemm::GemmShape<16, 8, 16>, + OutputOp, + cutlass::gemm::threadblock::GemmBatchedIdentityThreadblockSwizzle, + 4, + cutlass::gemm::kernel::GroupScheduleMode::kDeviceOnly>::GemmKernel; + using Gemm = cutlass::gemm::device::GemmGrouped; + + const int occupancy_wave = Gemm::sufficient(); + // Decode routes to only top-k experts, but each expert still contains many + // N-dimension tiles. A full persistent wave lets CUTLASS distribute those + // tiles across the GPU instead of leaving most SMs idle. Prefill retains + // the previous problem-count cap to avoid scheduler overhead from excess + // threadblocks when most experts receive no tokens. + const int threadblock_count = full_occupancy_wave + ? occupancy_wave + : std::min(occupancy_wave, problem_count); + if (threadblock_count <= 0) { + return MOE_DEVICE; + } + + typename Gemm::EpilogueOutputOp::Params epilogue_op(1.0f, use_bias ? 1.0f : 0.0f); + typename Gemm::Arguments args( + d_problems, + problem_count, + threadblock_count, + epilogue_op, + reinterpret_cast(d_ptr_a), + reinterpret_cast(d_ptr_b), + reinterpret_cast(d_ptr_c), + reinterpret_cast(d_ptr_d), + d_lda, + d_ldb, + d_ldc, + d_ldd, + nullptr); + + Gemm gemm; + auto status = gemm(args, nullptr, stream); + return status == cutlass::Status::kSuccess ? MOE_SUCCESS : MOE_INTERNAL; +} + +template +MoeStatus calculate_typed(const MoeFusedDenseInfo &info, + void *workspace, + size_t workspace_size, + void *output, + const void *hidden_states, + const void *w13, + const void *w2, + const void *topk_weights, + const void *topk_ids, + const void *sorted_token_ids, + const void *expert_ids, + const void *num_tokens_post_padded, + const void *bias1, + const void *bias2, + cudaStream_t stream) { + const int num_tokens = static_cast(info.num_tokens); + const int hidden_size = static_cast(info.hidden_size); + const int num_experts = static_cast(info.num_experts); + const int intermediate_size = static_cast(info.intermediate_size); + const int topk = static_cast(info.topk); + const int pairs = num_tokens * topk; + const int max_num_tokens_padded = static_cast(info.max_num_tokens_padded); + const int block_size = static_cast(info.block_size); + + uint8_t *ptr = reinterpret_cast(workspace); + size_t remaining = workspace_size; + auto counts = reinterpret_cast(advance_workspace(ptr, remaining, (num_experts + 1) * sizeof(int))); + auto offsets = reinterpret_cast(advance_workspace(ptr, remaining, (num_experts + 1) * sizeof(int))); + auto output_permutation = reinterpret_cast(advance_workspace(ptr, remaining, pairs * sizeof(int))); + auto packed_hidden = reinterpret_cast(advance_workspace(ptr, remaining, static_cast(max_num_tokens_padded) * hidden_size * sizeof(T))); + auto gate_up = reinterpret_cast(advance_workspace(ptr, remaining, static_cast(max_num_tokens_padded) * intermediate_size * 2 * sizeof(T))); + auto activated = reinterpret_cast(advance_workspace(ptr, remaining, static_cast(max_num_tokens_padded) * intermediate_size * sizeof(T))); + auto expert_out = reinterpret_cast(advance_workspace(ptr, remaining, static_cast(max_num_tokens_padded) * hidden_size * sizeof(T))); + auto grouped_problems = reinterpret_cast( + advance_workspace(ptr, remaining, static_cast(num_experts) * sizeof(cutlass::gemm::GemmCoord), alignof(cutlass::gemm::GemmCoord))); + auto grouped_ptr_a = reinterpret_cast(advance_workspace(ptr, remaining, static_cast(num_experts) * sizeof(void *), alignof(void *))); + auto grouped_ptr_b = reinterpret_cast(advance_workspace(ptr, remaining, static_cast(num_experts) * sizeof(void *), alignof(void *))); + auto grouped_ptr_c = reinterpret_cast(advance_workspace(ptr, remaining, static_cast(num_experts) * sizeof(void *), alignof(void *))); + auto grouped_ptr_d = reinterpret_cast(advance_workspace(ptr, remaining, static_cast(num_experts) * sizeof(void *), alignof(void *))); + auto grouped_lda = reinterpret_cast(advance_workspace(ptr, remaining, static_cast(num_experts) * sizeof(int64_t), alignof(int64_t))); + auto grouped_ldb = reinterpret_cast(advance_workspace(ptr, remaining, static_cast(num_experts) * sizeof(int64_t), alignof(int64_t))); + auto grouped_ldc = reinterpret_cast(advance_workspace(ptr, remaining, static_cast(num_experts) * sizeof(int64_t), alignof(int64_t))); + auto grouped_ldd = reinterpret_cast(advance_workspace(ptr, remaining, static_cast(num_experts) * sizeof(int64_t), alignof(int64_t))); + auto active_count = reinterpret_cast(advance_workspace(ptr, remaining, sizeof(int), alignof(int))); + + if (!counts || !offsets || !output_permutation || !packed_hidden || !gate_up || !activated || !expert_out || !grouped_problems || !grouped_ptr_a || !grouped_ptr_b || !grouped_ptr_c || !grouped_ptr_d || !grouped_lda || !grouped_ldb || !grouped_ldc || !grouped_ldd || !active_count) { + return MOE_WORKSPACE; + } + + auto w13_t = reinterpret_cast(w13); + auto w2_t = reinterpret_cast(w2); + if (num_tokens == 1 && topk <= num_experts) { + cudaMemsetAsync(output_permutation, 0xff, pairs * sizeof(int), stream); + cudaMemsetAsync(grouped_problems, 0, topk * sizeof(cutlass::gemm::GemmCoord), stream); + cudaMemsetAsync(active_count, 0, sizeof(int), stream); + setup_decode_gemm1_aligned_compact_kernel<<<(max_num_tokens_padded + 255) / 256, 256, 0, stream>>>( + grouped_problems, grouped_ptr_a, grouped_ptr_b, grouped_ptr_c, grouped_ptr_d, + grouped_lda, grouped_ldb, grouped_ldc, grouped_ldd, active_count, output_permutation, + reinterpret_cast(hidden_states), w13_t, gate_up, + reinterpret_cast(sorted_token_ids), + reinterpret_cast(expert_ids), + reinterpret_cast(num_tokens_post_padded), + pairs, topk, num_experts, hidden_size, intermediate_size, block_size); + + if (bias1 || bias2) { + fill_decode_bias_kernel<<>>( + gate_up, expert_out, reinterpret_cast(bias1), reinterpret_cast(bias2), + reinterpret_cast(sorted_token_ids), reinterpret_cast(expert_ids), + reinterpret_cast(num_tokens_post_padded), output_permutation, + pairs, num_experts, intermediate_size, hidden_size, block_size); + MOE_CHECK_CUDA(cudaGetLastError()); + } + + // Decode has one token and exactly one valid aligned row for each of + // its top-k routes, so the compact problem count is always topk. + // A D2H count copy and stream sync cannot be captured for replay. + MOE_CHECK_STATUS(launch_cutlass_gemm_grouped_device_meta( + topk, grouped_problems, grouped_ptr_a, grouped_ptr_b, grouped_ptr_c, grouped_ptr_d, + grouped_lda, grouped_ldb, grouped_ldc, grouped_ldd, true, bias1 != nullptr, stream)); + + swiglu_kernel<<<(topk * intermediate_size + 255) / 256, 256, 0, stream>>>(gate_up, activated, topk, intermediate_size); + + setup_decode_gemm2_aligned_compact_kernel<<<(max_num_tokens_padded + 255) / 256, 256, 0, stream>>>( + grouped_problems, grouped_ptr_a, grouped_ptr_b, grouped_ptr_c, grouped_ptr_d, + grouped_lda, grouped_ldb, grouped_ldc, grouped_ldd, activated, w2_t, expert_out, + reinterpret_cast(sorted_token_ids), + reinterpret_cast(expert_ids), + reinterpret_cast(num_tokens_post_padded), + output_permutation, pairs, num_experts, hidden_size, intermediate_size, block_size); + MOE_CHECK_STATUS(launch_cutlass_gemm_grouped_device_meta( + topk, grouped_problems, grouped_ptr_a, grouped_ptr_b, grouped_ptr_c, grouped_ptr_d, + grouped_lda, grouped_ldb, grouped_ldc, grouped_ldd, true, bias2 != nullptr, stream)); + + apply_shuffle_mul_sum_kernel<<<1, std::min(hidden_size, 1024), 0, stream>>>( + expert_out, reinterpret_cast(output), output_permutation, + reinterpret_cast(topk_weights), 1, topk, hidden_size); + return MOE_SUCCESS; + } + + cudaMemsetAsync(output_permutation, 0xff, pairs * sizeof(int), stream); + cudaMemsetAsync(counts, 0, (num_experts + 1) * sizeof(int), stream); + count_aligned_experts_kernel<<<(info.max_num_blocks + 255) / 256, 256, 0, stream>>>( + reinterpret_cast(expert_ids), + reinterpret_cast(num_tokens_post_padded), + counts, + num_experts, + block_size); + + int scan_threads = 1; + while (scan_threads < num_experts) { + scan_threads <<= 1; + } + scan_threads = std::max(32, scan_threads); + exclusive_prefix_counts_kernel<<<1, scan_threads, scan_threads * sizeof(int), stream>>>(counts, offsets, num_experts); + + pack_hidden_aligned_kernel<<>>( + reinterpret_cast(hidden_states), + reinterpret_cast(sorted_token_ids), + output_permutation, + packed_hidden, + pairs, + topk, + hidden_size, + max_num_tokens_padded); + + setup_prefill_gemm1_kernel<<<(num_experts + 255) / 256, 256, 0, stream>>>( + grouped_problems, grouped_ptr_a, grouped_ptr_b, grouped_ptr_c, grouped_ptr_d, + grouped_lda, grouped_ldb, grouped_ldc, grouped_ldd, counts, offsets, packed_hidden, + w13_t, gate_up, num_experts, hidden_size, intermediate_size); + + if (bias1 || bias2) { + fill_prefill_bias_kernel<<>>( + gate_up, expert_out, reinterpret_cast(bias1), reinterpret_cast(bias2), + counts, offsets, intermediate_size, hidden_size); + MOE_CHECK_CUDA(cudaGetLastError()); + } + + // Keep one grouped problem per expert. Experts with no routed tokens have + // m == 0 and are skipped by CUTLASS's device-only grouped scheduler. This + // avoids making the runtime problem count depend on device routing data. + MOE_CHECK_STATUS(launch_cutlass_gemm_grouped_device_meta( + num_experts, grouped_problems, grouped_ptr_a, grouped_ptr_b, grouped_ptr_c, grouped_ptr_d, + grouped_lda, grouped_ldb, grouped_ldc, grouped_ldd, false, bias1 != nullptr, stream)); + + swiglu_kernel<<<(max_num_tokens_padded * intermediate_size + 255) / 256, 256, 0, stream>>>( + gate_up, activated, max_num_tokens_padded, intermediate_size); + + setup_prefill_gemm2_kernel<<<(num_experts + 255) / 256, 256, 0, stream>>>( + grouped_problems, grouped_ptr_a, grouped_ptr_b, grouped_ptr_c, grouped_ptr_d, + grouped_lda, grouped_ldb, grouped_ldc, grouped_ldd, counts, offsets, activated, + w2_t, expert_out, num_experts, hidden_size, intermediate_size); + MOE_CHECK_STATUS(launch_cutlass_gemm_grouped_device_meta( + num_experts, grouped_problems, grouped_ptr_a, grouped_ptr_b, grouped_ptr_c, grouped_ptr_d, + grouped_lda, grouped_ldb, grouped_ldc, grouped_ldd, false, bias2 != nullptr, stream)); + + apply_shuffle_mul_sum_kernel<<>>( + expert_out, reinterpret_cast(output), output_permutation, + reinterpret_cast(topk_weights), num_tokens, topk, hidden_size); + return MOE_SUCCESS; +} + +} // namespace +Operator::Operator(Tensor output, const Tensor hidden, const Tensor w13, const Tensor w2, const Tensor weights, const Tensor ids, const Tensor sorted, const Tensor experts, const Tensor total, size_t block_size, const std::optional bias1, const std::optional bias2) : MoeFusedDense(output, hidden, w13, w2, weights, ids, sorted, experts, total, block_size, bias1, bias2) { moe_detail::Require(SupportsLayout(hidden,w13,w2)); } +size_t Operator::workspace_size_in_bytes() const { + const auto& v=metadata_; MoeFusedDenseInfo info{v[1].size(0),v[1].size(1),v[2].size(0),v[3].size(2),v[5].size(1),v[6].numel(),v[7].numel(),block_size_,v[1].dtype()}; + return workspace_size(info); +} +void Operator::operator()(Tensor output, const Tensor hidden, const Tensor w13, const Tensor w2, const Tensor weights, const Tensor ids, const Tensor sorted, const Tensor experts, const Tensor total, size_t block_size, const std::optional bias1, const std::optional bias2) const { + MoeDeviceGuard device_guard{metadata_[0].device().index()}; + Invocation(output, hidden, w13, w2, weights, ids, sorted, experts, total, block_size, bias1, bias2); MoeFusedDenseInfo info{hidden.size(0),hidden.size(1),w13.size(0),w2.size(2),ids.size(1),sorted.numel(),experts.numel(),block_size,hidden.dtype()}; + moe_detail::Require(SupportsLayout(hidden,w13,w2)); + std::vector views{output, hidden, w13, w2, weights, ids, sorted, experts, total}; + if (bias1) views.push_back(*bias1); + if (bias2) views.push_back(*bias2); + moe_detail::Workspace(workspace_,workspace_size_in_bytes_,workspace_size(info),views); + auto stream=static_cast(stream_); + if(hidden.dtype()==DataType::kFloat16) Check(calculate_typed(info,workspace_,workspace_size_in_bytes_,output.data(),hidden.data(),w13.data(),w2.data(),weights.data(),ids.data(),sorted.data(),experts.data(),total.data(),bias1 ? bias1->data() : nullptr,bias2 ? bias2->data() : nullptr,stream)); + else Check(calculate_typed<__nv_bfloat16,cutlass::bfloat16_t>(info,workspace_,workspace_size_in_bytes_,output.data(),hidden.data(),w13.data(),w2.data(),weights.data(),ids.data(),sorted.data(),experts.data(),total.data(),bias1 ? bias1->data() : nullptr,bias2 ? bias2->data() : nullptr,stream)); + MOE_CHECK_CUDA(cudaGetLastError()); +} + +#undef MOE_CHECK_STATUS +#undef MOE_CHECK_CUDA +} // namespace infini::ops diff --git a/src/native/cuda/nvidia/ops/moe_fused_dense/kernel.h b/src/native/cuda/nvidia/ops/moe_fused_dense/kernel.h new file mode 100644 index 000000000..76bad5c2f --- /dev/null +++ b/src/native/cuda/nvidia/ops/moe_fused_dense/kernel.h @@ -0,0 +1,42 @@ +#pragma once +#include "base/moe_fused_dense.h" +namespace infini::ops { +template <> +class Operator : public MoeFusedDense { + public: + Operator(Tensor output, const Tensor hidden, const Tensor w13, + const Tensor w2, const Tensor weights, const Tensor ids, + const Tensor sorted, const Tensor experts, const Tensor total, + size_t block_size, const std::optional bias1 = std::nullopt, + const std::optional bias2 = std::nullopt); + size_t workspace_size_in_bytes() const override; + 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, + const Tensor total, size_t block_size, + const std::optional bias1 = std::nullopt, + const std::optional bias2 = std::nullopt) { + return MoeFusedDense::supports(output, hidden, w13, w2, weights, ids, + sorted, experts, total, block_size, bias1, + bias2) && + SupportsLayout(hidden, w13, w2); + } + void operator()( + Tensor output, const Tensor hidden, const Tensor w13, const Tensor w2, + const Tensor weights, const Tensor ids, const Tensor sorted, + const Tensor experts, const Tensor total, size_t block_size, + const std::optional bias1 = std::nullopt, + const std::optional bias2 = std::nullopt) const override; + + private: + static bool SupportsLayout(const Tensor hidden, const Tensor w13, + const Tensor w2) { + return w13.size(0) <= 1024 && hidden.size(1) % 8 == 0 && + w2.size(2) % 8 == 0 && + reinterpret_cast(w13.data()) % 16 == 0 && + reinterpret_cast(w2.data()) % 16 == 0 && + (hidden.size(0) != 1 || + reinterpret_cast(hidden.data()) % 16 == 0); + } +}; +} // namespace infini::ops diff --git a/src/native/cuda/nvidia/ops/moe_fused_gate/device_guard.cuh b/src/native/cuda/nvidia/ops/moe_fused_gate/device_guard.cuh new file mode 100644 index 000000000..6d78efd5e --- /dev/null +++ b/src/native/cuda/nvidia/ops/moe_fused_gate/device_guard.cuh @@ -0,0 +1,31 @@ +#pragma once +#include + +#include +#include +namespace infini::ops { +class MoeDeviceGuard { + public: + explicit MoeDeviceGuard(int device) { + auto s = cudaGetDevice(&previous_); + if (s != cudaSuccess) throw std::runtime_error(cudaGetErrorString(s)); + if (previous_ != device) { + s = cudaSetDevice(device); + if (s != cudaSuccess) throw std::runtime_error(cudaGetErrorString(s)); + restore_ = true; + } + } + ~MoeDeviceGuard() { + if (restore_) { + const auto s = cudaSetDevice(previous_); + if (s != cudaSuccess) + std::fprintf(stderr, "MoE restore device failed: %s\n", + cudaGetErrorString(s)); + } + } + + private: + int previous_ = 0; + bool restore_ = false; +}; +} // namespace infini::ops diff --git a/src/native/cuda/nvidia/ops/moe_fused_gate/kernel.cu b/src/native/cuda/nvidia/ops/moe_fused_gate/kernel.cu new file mode 100644 index 000000000..bfd364dcf --- /dev/null +++ b/src/native/cuda/nvidia/ops/moe_fused_gate/kernel.cu @@ -0,0 +1,239 @@ +// 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 + * + * Copyright 2025 SGLang Team. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0. + */ + +#include "native/cuda/nvidia/ops/moe_fused_gate/kernel.h" +#include "native/cuda/nvidia/ops/moe_fused_gate/device_guard.cuh" +#include +#include +#include +#include +#include +namespace infini::ops { namespace { + +#define MOE_CHECK_CUDA(expr) do { auto s=(expr); if(s!=cudaSuccess) throw std::runtime_error(cudaGetErrorString(s)); } while(0) + +constexpr int MAX_VPT = 32; + +template +__device__ float convert_to_float(T x) { + if constexpr (std::is_same_v) { + return __half2float(x); + } else if constexpr (std::is_same_v) { + return __bfloat162float(x); + } else if constexpr (std::is_same_v) { + return x; + } else { + return static_cast(x); + } +} + +struct KernelParamsDynamic { + int VPT; + int NUM_EXPERTS; + int THREADS_PER_ROW; + int ROWS_PER_WARP; + int ROWS_PER_CTA; +}; + +template +__device__ void moe_fused_gate_impl( + const T *__restrict__ input, + const float *__restrict__ bias, + float *__restrict__ output, + int32_t *__restrict__ indices, + int64_t num_rows, + int64_t topk_group, + int64_t topk, + int64_t num_fused_shared_experts, + float routed_scaling_factor, + bool apply_routed_scaling_factor_on_output, + Params params) { + const int tidx = threadIdx.x; + const int64_t thread_row = blockIdx.x * params.ROWS_PER_CTA + threadIdx.y * params.ROWS_PER_WARP + tidx / params.THREADS_PER_ROW; + if (thread_row >= num_rows) { + return; + } + + const int64_t routed_topk = topk - num_fused_shared_experts; + const T *thread_row_ptr = input + thread_row * params.NUM_EXPERTS; + const int thread_group_idx = tidx % params.THREADS_PER_ROW; + const int first_elt_read_by_thread = thread_group_idx * params.VPT; + + float row_chunk[MAX_VPT]; + float bias_chunk[MAX_VPT]; +#pragma unroll + for (int ii = 0; ii < MAX_VPT; ++ii) { + if (ii < params.VPT) { + const int expert = first_elt_read_by_thread + ii; + const float sigmoid = 1.0f / (1.0f + expf(-convert_to_float(thread_row_ptr[expert]))); + row_chunk[ii] = sigmoid; + bias_chunk[ii] = sigmoid + bias[expert]; + } + } + + __syncthreads(); + + for (int k_idx = 0; k_idx < params.THREADS_PER_ROW - topk_group; ++k_idx) { + int expert = first_elt_read_by_thread; + float max_val = -FLT_MAX; + float max_val_second = -FLT_MAX; +#pragma unroll + for (int ii = 0; ii < MAX_VPT; ++ii) { + if (ii < params.VPT) { + const float val = bias_chunk[ii]; + if (val > max_val) { + max_val_second = max_val; + max_val = val; + } else if (val > max_val_second) { + max_val_second = val; + } + } + } + float max_sum = max_val + max_val_second; + +#pragma unroll + for (int mask = params.THREADS_PER_ROW / 2; mask > 0; mask /= 2) { + const float other_max_sum = __shfl_xor_sync(__activemask(), max_sum, mask, params.THREADS_PER_ROW); + const int other_expert = __shfl_xor_sync(__activemask(), expert, mask, params.THREADS_PER_ROW); + if (max_sum > other_max_sum || (other_max_sum == max_sum && other_expert > expert)) { + max_sum = other_max_sum; + expert = other_expert; + } + } + + const int thread_to_clear_in_group = expert / params.VPT; + if (thread_group_idx == thread_to_clear_in_group) { +#pragma unroll + for (int ii = 0; ii < MAX_VPT; ++ii) { + if (ii < params.VPT) { + bias_chunk[ii] = FLT_MAX; + } + } + } + } + + __syncthreads(); + + float output_sum = 0.0f; + for (int k_idx = 0; k_idx < routed_topk; ++k_idx) { + float max_val = bias_chunk[0]; + int expert = first_elt_read_by_thread; + if (max_val != FLT_MAX) { +#pragma unroll + for (int ii = 1; ii < MAX_VPT; ++ii) { + if (ii < params.VPT) { + const float val = bias_chunk[ii]; + if (val > max_val) { + max_val = val; + expert = first_elt_read_by_thread + ii; + } + } + } + } else { + max_val = -FLT_MAX; + } + +#pragma unroll + for (int mask = params.THREADS_PER_ROW / 2; mask > 0; mask /= 2) { + const float other_max = __shfl_xor_sync(__activemask(), max_val, mask, params.THREADS_PER_ROW); + const int other_expert = __shfl_xor_sync(__activemask(), expert, mask, params.THREADS_PER_ROW); + if (other_max > max_val || (other_max == max_val && other_expert < expert)) { + max_val = other_max; + expert = other_expert; + } + } + + const int thread_to_clear_in_group = expert / params.VPT; + const int64_t idx = topk * thread_row + k_idx; + if (thread_group_idx == thread_to_clear_in_group) { + const int expert_to_clear_in_thread = expert % params.VPT; + bias_chunk[expert_to_clear_in_thread] = -FLT_MAX; + output[idx] = row_chunk[expert_to_clear_in_thread]; + indices[idx] = static_cast(expert); + } + + __syncthreads(); + if (thread_group_idx == 0) { + output_sum += output[idx]; + } + __syncthreads(); + } + + if (thread_group_idx == 0 && num_fused_shared_experts > 0) { + int64_t idx = topk * thread_row + routed_topk; + for (int i = 0; i < num_fused_shared_experts; ++i) { + indices[idx + i] = static_cast(params.NUM_EXPERTS + i); + output[idx + i] = output_sum / routed_scaling_factor; + } + } + __syncthreads(); + + if (thread_group_idx == 0) { +#pragma unroll + for (int ii = 0; ii < topk; ++ii) { + const int64_t idx = topk * thread_row + ii; + output[idx] = output[idx] / output_sum; + if (apply_routed_scaling_factor_on_output) { + output[idx] *= routed_scaling_factor; + } + } + } +} + +template +__global__ void moe_fused_gate_kernel_dynamic( + const T *__restrict__ input, + const float *__restrict__ bias, + float *__restrict__ output, + int32_t *__restrict__ indices, + int64_t num_rows, + int64_t num_experts, + int64_t num_expert_group, + int64_t topk_group, + int64_t topk, + int64_t num_fused_shared_experts, + float routed_scaling_factor, + bool apply_routed_scaling_factor_on_output) { + KernelParamsDynamic params; + params.NUM_EXPERTS = static_cast(num_experts); + params.VPT = static_cast(num_experts / num_expert_group); + params.THREADS_PER_ROW = static_cast(num_expert_group); + params.ROWS_PER_WARP = 1; + params.ROWS_PER_CTA = 1; + moe_fused_gate_impl( + input, bias, output, indices, num_rows, topk_group, topk, num_fused_shared_experts, + routed_scaling_factor, apply_routed_scaling_factor_on_output, params); +} + + +} // namespace +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(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<<>>( + static_cast(input.data()),static_cast(bias.data()), + static_cast(weights.data()),static_cast(indices.data()), + input.size(0),input.size(1),groups,selected_groups,weights.size(1),shared,scale,apply_scale); + }; + switch(input.dtype()) { case DataType::kFloat16: run(half{}); break; + case DataType::kBFloat16: run(__nv_bfloat16{}); break; case DataType::kFloat32: run(float{}); break; + default: throw std::invalid_argument("Gate dtype"); } + MOE_CHECK_CUDA(cudaGetLastError()); +} + +#undef MOE_CHECK_CUDA +} // namespace infini::ops diff --git a/src/native/cuda/nvidia/ops/moe_fused_gate/kernel.h b/src/native/cuda/nvidia/ops/moe_fused_gate/kernel.h new file mode 100644 index 000000000..e4734e793 --- /dev/null +++ b/src/native/cuda/nvidia/ops/moe_fused_gate/kernel.h @@ -0,0 +1,21 @@ +#pragma once +#include "base/moe_fused_gate.h" +namespace infini::ops { +template <> +class Operator : public MoeFusedGate { + public: + using MoeFusedGate::MoeFusedGate; + static bool supports(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 auto e = input.ndim() == 2 ? input.size(1) : 0; + return MoeFusedGate::supports(weights, indices, input, bias, groups, + selected_groups, shared, scale, + apply_scale) && + 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, + size_t shared, float scale, bool apply_scale) const override; +}; +} // namespace infini::ops 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..dfd653692 --- /dev/null +++ b/src/native/cuda/nvidia/ops/mrope/kernel.h @@ -0,0 +1,12 @@ +#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; +}; +} // namespace infini::ops diff --git a/src/native/cuda/nvidia/ops/mul/kernel.h b/src/native/cuda/nvidia/ops/mul/kernel.h index 547dd3270..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,6 +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.cu b/src/native/cuda/nvidia/ops/mul/scalar.cu new file mode 100644 index 000000000..6deaa0b74 --- /dev/null +++ b/src/native/cuda/nvidia/ops/mul/scalar.cu @@ -0,0 +1,98 @@ +#include +#include + +#include +#include + +#include "native/cuda/nvidia/ops/mul/scalar.h" + +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 scalar Mul unsupported dtype"); + } +} +} // namespace infini::ops::nvidia_mul_scalar_detail diff --git a/src/native/cuda/nvidia/ops/mul/scalar.h b/src/native/cuda/nvidia/ops/mul/scalar.h new file mode 100644 index 000000000..c96553aed --- /dev/null +++ b/src/native/cuda/nvidia/ops/mul/scalar.h @@ -0,0 +1,148 @@ +#ifndef INFINI_OPS_NVIDIA_MUL_SCALAR_H_ +#define INFINI_OPS_NVIDIA_MUL_SCALAR_H_ + +#include + +#include +#include +#include +#include +#include + +#include "base/mul.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 { +class NvidiaScalarMul { + 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; + } + + 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(); + 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]); + } + } + + 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 scalar Mul 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_meta_.strides() == out_meta_.strides()))) + throw std::invalid_argument( + "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_, stream); + } + + private: + static Tensor Checked(const Tensor input, const double alpha, Tensor out) { + if (!supports(input, alpha, out)) + throw std::invalid_argument( + "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 + +#endif // INFINI_OPS_NVIDIA_MUL_SCALAR_KERNEL_H_ diff --git a/src/native/cuda/nvidia/ops/pad/kernel.cu b/src/native/cuda/nvidia/ops/pad/kernel.cu new file mode 100644 index 000000000..0a4f99d63 --- /dev/null +++ b/src/native/cuda/nvidia/ops/pad/kernel.cu @@ -0,0 +1,75 @@ +#include +#include + +#include +#include + +#include "native/cuda/nvidia/ops/pad/kernel.h" + +namespace infini::ops::nvidia_pad_detail { +namespace { +template +__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; + 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)); + PadKernel<<>>( + 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 Pad unsupported dtype"); + } + const auto status = cudaPeekAtLastError(); + if (status != cudaSuccess) + throw std::runtime_error(cudaGetErrorString(status)); +} +} // namespace infini::ops::nvidia_pad_detail diff --git a/src/native/cuda/nvidia/ops/pad/kernel.h b/src/native/cuda/nvidia/ops/pad/kernel.h new file mode 100644 index 000000000..ad6c54733 --- /dev/null +++ b/src/native/cuda/nvidia/ops/pad/kernel.h @@ -0,0 +1,161 @@ +#ifndef INFINI_OPS_NVIDIA_PAD_KERNEL_H_ +#define INFINI_OPS_NVIDIA_PAD_KERNEL_H_ + +#include + +#include +#include +#include +#include +#include + +#include "base/pad.h" +#include "native/cuda/nvidia/ops/binary_view_capture.h" +#include "native/cuda/nvidia/runtime_.h" + +namespace infini::ops::nvidia_pad_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_detail + +namespace infini::ops { +template <> +class Operator : public Pad { + public: + 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() || + (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; + 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 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_ = + input.shape() == out.shape() && input.strides() == out.strides(); + } + + bool capture_safe() const override { return true; } + + 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_) || pad != padding_ || + mode != mode_ || std::memcmp(&fill, &value_, sizeof(fill))) + throw std::invalid_argument( + "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 Pad current device, pointer or alias changed after " + "Make"); + if (out.numel() == 0 || + (input.data() == out.data() && exact_inplace_layout_)) + return; + 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& pad, + const std::string& mode, std::optional value, + Tensor out) { + if (!supports(input, pad, mode, value, out)) + throw std::invalid_argument( + "unsupported NVIDIA Pad metadata, mode, 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_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_KERNEL_H_ diff --git a/src/native/cuda/nvidia/ops/per_channel_quant_i8/kernel.cu b/src/native/cuda/nvidia/ops/per_channel_quant_i8/kernel.cu new file mode 100644 index 000000000..daf48dc43 --- /dev/null +++ b/src/native/cuda/nvidia/ops/per_channel_quant_i8/kernel.cu @@ -0,0 +1,43 @@ +#include "native/cuda/nvidia/ops/per_channel_quant_i8/kernel.h" +#include +#include +#include +#include +namespace infini::ops::nvidia_per_channel_i8_detail { +template __global__ void Kernel(const T* x,std::int8_t* p,float* scale, + std::size_t rows,std::size_t k) { + __shared__ float maximum[256]; + for(std::size_t row=blockIdx.x;row=0?int(floorf(fabsf(v)+0.5f)):-int(floorf(fabsf(v)+0.5f)); + q=q>127?127:(q< -127?-127:q); p[row*k+c]=std::int8_t(q); + } + __syncthreads(); + } +} +void Launch(const void* x,std::int8_t* p,float* scale,DataType dtype, + std::size_t rows,std::size_t k,cudaStream_t stream) { + if(!rows) return; + const unsigned blocks=unsigned(std::min(rows,65535)); + switch(dtype) { + case DataType::kFloat16: Kernel<<>>(static_cast(x),p,scale,rows,k);break; + case DataType::kBFloat16: Kernel<<>>(static_cast(x),p,scale,rows,k);break; + case DataType::kFloat32: Kernel<<>>(static_cast(x),p,scale,rows,k);break; + default: throw std::invalid_argument("unsupported per-channel INT8 dtype"); + } + const auto error=cudaGetLastError(); + if(error!=cudaSuccess) throw std::runtime_error(cudaGetErrorString(error)); +} +} diff --git a/src/native/cuda/nvidia/ops/per_channel_quant_i8/kernel.h b/src/native/cuda/nvidia/ops/per_channel_quant_i8/kernel.h new file mode 100644 index 000000000..016a161f3 --- /dev/null +++ b/src/native/cuda/nvidia/ops/per_channel_quant_i8/kernel.h @@ -0,0 +1,60 @@ +#pragma once +#include + +#include "base/per_channel_quant_i8.h" +#include "native/cuda/nvidia/ops/binary_view_capture.h" +#include "native/cuda/nvidia/runtime_.h" +namespace infini::ops::nvidia_per_channel_i8_detail { +void Launch(const void*, std::int8_t*, float*, DataType, std::size_t, + std::size_t, cudaStream_t); +} +namespace infini::ops { +template <> +class Operator + : public PerChannelQuantI8 { + public: + static bool supports(const Tensor x, Tensor packed, Tensor scale) { + using namespace nvidia_binary_view; + if (x.device().type() != Device::Type::kNvidia || + x.device() != packed.device() || x.device() != scale.device() || + x.ndim() != 2 || x.shape() != packed.shape() || !x.IsContiguous() || + !packed.IsContiguous() || !scale.IsContiguous() || + scale.shape() != Tensor::Shape{x.shape()[0], 1} || + packed.dtype() != DataType::kInt8 || + scale.dtype() != DataType::kFloat32 || + (x.dtype() != DataType::kFloat16 && x.dtype() != DataType::kBFloat16 && + x.dtype() != DataType::kFloat32) || + !x.shape()[1]) + return false; + int device = -1; + if (cudaGetDevice(&device) != cudaSuccess || device != x.device().index()) + return false; + if (x.dtype() == DataType::kBFloat16) { + int major = 0; + if (cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, + device) != cudaSuccess || + major < 8) + return false; + } + std::size_t xb = 0, pb = 0, sb = 0; + if (!SpanBytes(x, &xb) || !SpanBytes(packed, &pb) || !SpanBytes(scale, &sb)) + return false; + if ((xb && reinterpret_cast(x.data()) % x.element_size()) || + (sb && reinterpret_cast(scale.data()) % sizeof(float))) + return false; + return !Overlaps(x, xb, packed, pb) && !Overlaps(x, xb, scale, sb) && + !Overlaps(packed, pb, scale, sb); + } + Operator(const Tensor x, Tensor packed, Tensor scale) + : PerChannelQuantI8(x, packed, scale) {} + bool capture_safe() const override { return true; } + void operator()(const Tensor x, Tensor packed, Tensor scale) const override { + if (!supports(x, packed, scale)) + throw std::invalid_argument("unsupported per-channel INT8 operands"); + nvidia_per_channel_i8_detail::Launch( + x.data(), static_cast(packed.data()), + static_cast(scale.data()), x.dtype(), x.shape()[0], + x.shape()[1], static_cast(stream_ ? stream_ : 0)); + } +}; +} // namespace infini::ops diff --git a/src/native/cuda/nvidia/ops/per_tensor_i8/kernel.cu b/src/native/cuda/nvidia/ops/per_tensor_i8/kernel.cu new file mode 100644 index 000000000..3ac684cbc --- /dev/null +++ b/src/native/cuda/nvidia/ops/per_tensor_i8/kernel.cu @@ -0,0 +1,96 @@ +#include "native/cuda/nvidia/ops/per_tensor_i8/kernel.h" +#include +#include +#include +#include +namespace infini::ops::nvidia_per_tensor_i8_detail { +namespace { +__device__ std::int64_t Offset(std::size_t linear, const Layout& l, + bool packed) { + std::int64_t offset = 0; + for (std::int64_t d = l.ndim; d-- > 0;) { + auto coord = linear % l.shape[d]; linear /= l.shape[d]; + offset += coord * (packed ? l.packed_stride[d] : l.x_stride[d]); + } + return offset; +} +__device__ void AtomicMax(float* ptr, float value) { + auto bits = reinterpret_cast(ptr); + int old = *bits, assumed; + do { + assumed = old; + old = atomicCAS(bits, assumed, + __float_as_int(fmaxf(value, __int_as_float(assumed)))); + } while (old != assumed); +} +template __global__ void Absmax(const T* x, float* scale, + Layout l, std::size_t count) { + float maximum = 0; + for (std::size_t i = blockIdx.x * blockDim.x + threadIdx.x; i < count; + i += std::size_t(gridDim.x) * blockDim.x) + maximum = fmaxf(maximum, fabsf(float(x[Offset(i, l, false)]))); + // Preserve the original dynamic contract: atomic max with caller-provided + // scale, without epsilon or clearing scale between recorded executions. + AtomicMax(scale, maximum / 127.0f); +} +template __global__ void Quant(const T* x, std::int8_t* packed, + const float* scale, Layout l, std::size_t count) { + const float inverse = 1.0f / scale[0]; + for (std::size_t i = blockIdx.x * blockDim.x + threadIdx.x; i < count; + i += std::size_t(gridDim.x) * blockDim.x) { + const float value = float(x[Offset(i, l, false)]) * inverse; + const float rounded = floorf(fabsf(value) + 0.5f); + int q = value >= 0 ? int(rounded) : -int(rounded); + q = q > 127 ? 127 : (q < -127 ? -127 : q); + packed[Offset(i, l, true)] = std::int8_t(q); + } +} +template __global__ void Dequant(T* x, const std::int8_t* packed, + const float* scale, Layout l, std::size_t count) { + for (std::size_t i = blockIdx.x * blockDim.x + threadIdx.x; i < count; + i += std::size_t(gridDim.x) * blockDim.x) + x[Offset(i, l, false)] = T(float(packed[Offset(i, l, true)]) * scale[0]); +} +void CheckLaunch() { + const auto error = cudaGetLastError(); + if (error != cudaSuccess) + throw std::runtime_error(cudaGetErrorString(error)); +} +template void SubmitQuant(const void* x, std::int8_t* packed, + float* scale, Layout l, std::size_t count, bool is_static, cudaStream_t stream) { + const unsigned blocks = unsigned(std::min((count-1)/256+1,65535)); + if (!is_static) { + Absmax<<>>(static_cast(x),scale,l,count); + CheckLaunch(); + } + Quant<<>>(static_cast(x),packed,scale,l,count); + CheckLaunch(); +} +template void SubmitDequant(void* x, const std::int8_t* packed, + const float* scale, Layout l, std::size_t count, cudaStream_t stream) { + const unsigned blocks = unsigned(std::min((count-1)/256+1,65535)); + Dequant<<>>(static_cast(x),packed,scale,l,count); + CheckLaunch(); +} +} +void Quantize(const void* x, std::int8_t* packed, float* scale, DataType dtype, + Layout l, std::size_t count, bool is_static, cudaStream_t stream) { + if (!count) return; + switch (dtype) { + case DataType::kFloat16: SubmitQuant<__half>(x,packed,scale,l,count,is_static,stream); break; + case DataType::kBFloat16: SubmitQuant<__nv_bfloat16>(x,packed,scale,l,count,is_static,stream); break; + case DataType::kFloat32: SubmitQuant(x,packed,scale,l,count,is_static,stream); break; + default: throw std::invalid_argument("unsupported per-tensor INT8 dtype"); + } +} +void Dequantize(void* x, const std::int8_t* packed, const float* scale, + DataType dtype, Layout l, std::size_t count, cudaStream_t stream) { + if (!count) return; + switch (dtype) { + case DataType::kFloat16: SubmitDequant<__half>(x,packed,scale,l,count,stream); break; + case DataType::kBFloat16: SubmitDequant<__nv_bfloat16>(x,packed,scale,l,count,stream); break; + case DataType::kFloat32: SubmitDequant(x,packed,scale,l,count,stream); break; + default: throw std::invalid_argument("unsupported per-tensor INT8 dtype"); + } +} +} diff --git a/src/native/cuda/nvidia/ops/per_tensor_i8/kernel.h b/src/native/cuda/nvidia/ops/per_tensor_i8/kernel.h new file mode 100644 index 000000000..40a3437ac --- /dev/null +++ b/src/native/cuda/nvidia/ops/per_tensor_i8/kernel.h @@ -0,0 +1,160 @@ +#pragma once +#include + +#include +#include +#include + +#include "base/per_tensor_dequant_i8.h" +#include "base/per_tensor_quant_i8.h" +#include "native/cuda/nvidia/ops/binary_view_capture.h" +#include "native/cuda/nvidia/runtime_.h" +namespace infini::ops::nvidia_per_tensor_i8_detail { +struct Layout { + std::int64_t ndim = 0, shape[4]{}, x_stride[4]{}, packed_stride[4]{}; +}; +void Quantize(const void*, std::int8_t*, float*, DataType, Layout, std::size_t, + bool, cudaStream_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, bool quant, + bool is_static = true) { + using namespace nvidia_binary_view; + if (zero || + (x.dtype() != DataType::kFloat16 && x.dtype() != DataType::kFloat32 && + x.dtype() != DataType::kBFloat16) || + packed.dtype() != DataType::kInt8 || + 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 || + 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()) + return false; + if (x.dtype() == DataType::kBFloat16) { + int major = 0; + if (cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, + device) != cudaSuccess || + major < 8) + return false; + } + for (std::size_t d = 0; d < x.ndim(); ++d) + if (x.strides()[d] < 0 || packed.strides()[d] < 0 || + x.shape()[d] > static_cast(INT64_MAX)) + return false; + std::size_t xb = 0, pb = 0, sb = 0; + if (!SpanBytes(x, &xb) || !SpanBytes(packed, &pb) || !SpanBytes(scale, &sb) || + xb > INT64_MAX || pb > INT64_MAX || sb > INT64_MAX) + return false; + if ((xb && reinterpret_cast(x.data()) % x.element_size()) || + reinterpret_cast(scale.data()) % sizeof(float)) + return false; + 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; + 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]; + layout.x_stride[d] = x.strides()[d]; + 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 { +template <> +class Operator + : public PerTensorQuantI8 { + public: + static bool supports(const Tensor x, Tensor packed, Tensor scale, + 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_(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) || + !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()), + static_cast(scale.data()), meta_.dtype, meta_.layout, + meta_.count, is_static, + static_cast(stream_ ? stream_ : 0)); + } + + private: + nvidia_per_tensor_i8_detail::Metadata meta_; + bool is_static_; +}; +template <> +class Operator + : public PerTensorDequantI8 { + 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, false); + } + Operator(Tensor x, const Tensor packed, const Tensor scale, + const std::optional zero) + : 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) || !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()), + static_cast(scale.data()), meta_.dtype, meta_.layout, + meta_.count, static_cast(stream_ ? stream_ : 0)); + } + + private: + nvidia_per_tensor_i8_detail::Metadata meta_; +}; +} // namespace infini::ops 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..e5668f654 --- /dev/null +++ b/src/native/cuda/nvidia/ops/quick_gelu/kernel.h @@ -0,0 +1,136 @@ +#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 > 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) { + 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 (!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"); + 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(), 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( + "unsupported NVIDIA QuickGelu metadata, device or alias"); + return input; + } +}; +} // namespace infini::ops +#endif 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 new file mode 100644 index 000000000..ddd5c8c71 --- /dev/null +++ b/src/native/cuda/nvidia/ops/recurrent_gated_delta_rule/kernel.h @@ -0,0 +1,158 @@ +#pragma once +#include +#include +#include + +#include + +#include "base/recurrent_gated_delta_rule.h" +#include "native/cuda/nvidia/caster.cuh" +#include "native/cuda/nvidia/runtime_.h" +#include "recurrent_gated_delta_rule_legacy_kernel.cuh" +namespace infini::ops { +namespace recurrent_native_detail { +struct Layout { + ptrdiff_t s[10][4]{}; + bool initial_i64, final_i64, norm, indexed; + size_t pool, Hk, ratio; +}; +template +__global__ void kernel(const Data* q, const Data* k, const Data* v, + const Gate* g, const Gate* beta, Data* state, + const void* initial_indices, const void* final_indices, + Data* out, Data* final_state, Layout p) { + extern __shared__ float shared[]; + recurrentGatedDeltaRuleIndexedPoolWarpKernel( + out, state, final_state, q, k, v, g, beta, initial_indices, final_indices, + p.initial_i64, p.final_i64, p.norm, p.indexed, p.pool, p.Hk, p.ratio, + p.s[8][0], p.s[8][1], p.s[8][2], p.s[5][0], p.s[5][1], p.s[5][2], + p.s[5][3], p.s[9][0], p.s[9][1], p.s[9][2], p.s[9][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], shared); +} +} // namespace recurrent_native_detail +template <> +class Operator + : public RecurrentGatedDeltaRule { + public: + 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, + const std::optional final_state_indices, Tensor out, + std::optional final_state, bool use_qk_l2norm) + : RecurrentGatedDeltaRule(q, k, v, g, beta, initial_state, + initial_state_indices, final_state_indices, out, + 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 || + property.maxThreadsPerBlock < 256) + throw std::runtime_error( + "recurrent delta requires 256-thread CUDA blocks"); + if (B_ >= static_cast(INT_MAX) || Hv_ > 65535 || + Hk_ > static_cast(INT_MAX)) + throw std::invalid_argument("recurrent delta launch extent unsupported"); + p_.pool = pool_; + p_.Hk = Hk_; + p_.ratio = Hv_ / Hk_; + p_.norm = norm_; + 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) + 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, + const std::optional initial_state_indices, + const std::optional final_state_indices, Tensor out, + std::optional final_state, + bool use_qk_l2norm) const override { + const Views views{q, + k, + v, + g, + beta, + initial_state, + initial_state_indices, + final_state_indices, + out, + final_state}; + CheckInvocation(views, use_qk_l2norm); + CheckDevice(); + if (!B_) return; + switch (data_type_) { + case DataType::kFloat16: + gate(views); + break; + case DataType::kBFloat16: + gate<__nv_bfloat16>(views); + break; + case DataType::kFloat32: + gate(views); + break; + default: + throw std::invalid_argument("recurrent delta data dtype unsupported"); + } + if (cudaPeekAtLastError() != cudaSuccess) + throw std::runtime_error("recurrent delta CUDA launch failed"); + } + + private: + recurrent_native_detail::Layout p_{}; + void CheckDevice() const { + int device = -1; + if (cudaGetDevice(&device) != cudaSuccess) + throw std::runtime_error("recurrent delta device query failed"); + if (device != device_.index()) + throw std::invalid_argument("recurrent delta owner device mismatch"); + } + template + void gate(const Views& v) const { + switch (gate_type_) { + case DataType::kFloat16: + dimension(v); + break; + case DataType::kBFloat16: + dimension(v); + break; + case DataType::kFloat32: + dimension(v); + break; + default: + throw std::invalid_argument("recurrent delta gate dtype unsupported"); + } + } + template + void dimension(const Views& v) const { + if (D_ == 64) + launch(v); + else + launch(v); + } + template + void launch(const Views& v) const { + auto ptr = [&](size_t i) -> void* { + return v[i] ? const_cast(v[i]->data()) : nullptr; + }; + dim3 grid(static_cast(B_), static_cast(Hv_), + D / 8); + const auto stream = static_cast(stream_ ? stream_ : nullptr); + recurrent_native_detail::kernel + <<>>( + 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), static_cast(ptr(8)), + static_cast(ptr(9)), p_); + } +}; +} // namespace infini::ops 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 new file mode 100644 index 000000000..9a715fdc7 --- /dev/null +++ b/src/native/cuda/nvidia/ops/recurrent_gated_delta_rule/recurrent_delta_rule_common.cuh @@ -0,0 +1,265 @@ +// Ported from InfiniCore 1ab85ef1d3c2bce49edf1fc73a5b48becaf7552e: +// src/infiniop/ops/recurrent_gated_delta_rule/cuda/recurrent_delta_rule_common.cuh. +// 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__ + +#include +#include + +// Retain the original NVIDIA warp width. +#define INFINIOP_RECURRENT_DELTA_RULE_WARP_SIZE 32 +#define INFINIOP_RECURRENT_DELTA_RULE_FULL_MASK 0xffffffffU + +namespace op::recurrent_gated_delta_rule::cuda { + +template +__device__ inline float loadAsFloat(const T* ptr, ptrdiff_t offset) { + return static_cast(ptr[offset]); +} + +template <> +__device__ inline float loadAsFloat(const half* ptr, ptrdiff_t offset) { + return __half2float(ptr[offset]); +} + +template <> +__device__ inline float loadAsFloat<__nv_bfloat16>(const __nv_bfloat16* ptr, + ptrdiff_t offset) { + return __bfloat162float(ptr[offset]); +} + +__device__ inline int64_t loadOptionalIndex(const void* indices, bool is_i64, + int index, int fallback) { + if (indices == nullptr) { + return static_cast(fallback); + } + return is_i64 ? static_cast(indices)[index] + : static_cast( + static_cast(indices)[index]); +} + +template +__device__ inline Tcompute warpReduceSum(Tcompute value) { + constexpr int WARP_SIZE = INFINIOP_RECURRENT_DELTA_RULE_WARP_SIZE; +#pragma unroll + for (int offset = WARP_SIZE / 2; offset > 0; offset >>= 1) { + value += __shfl_down_sync(INFINIOP_RECURRENT_DELTA_RULE_FULL_MASK, value, + offset, WARP_SIZE); + } + return __shfl_sync(INFINIOP_RECURRENT_DELTA_RULE_FULL_MASK, value, 0, + WARP_SIZE); +} + +template +__device__ inline Tcompute blockReduceSum(Tcompute value, Tcompute* scratch) { + scratch[threadIdx.x] = value; + __syncthreads(); + for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) { + if (threadIdx.x < stride) { + scratch[threadIdx.x] += scratch[threadIdx.x + stride]; + } + __syncthreads(); + } + const Tcompute result = scratch[0]; + __syncthreads(); + return result; +} + +// 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* 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; + + const int batch_idx = blockIdx.x; + const int value_head_idx = blockIdx.y; + const int key_head_idx = + value_head_idx / static_cast(value_heads_per_key_head); + const int warp_idx = threadIdx.x / WARP_SIZE; + const int lane_idx = threadIdx.x & (WARP_SIZE - 1); + const int value_dim_idx = blockIdx.z * WARPS_PER_BLOCK + warp_idx; + const bool valid_value_dim = value_dim_idx < static_cast(Dv); + + if (key_head_idx >= static_cast(num_key_heads)) { + return; + } + + int64_t read_slot = batch_idx; + int64_t write_slot = batch_idx; + if (indexed_state_pool) { + read_slot = loadOptionalIndex( + initial_state_indices, initial_state_indices_i64, batch_idx, batch_idx); + write_slot = + final_state_indices == nullptr + ? 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) || + (final_state_indices != nullptr && + (write_slot < 0 || write_slot >= static_cast(pool_size)))) { + if (valid_value_dim && lane_idx == 0) { + 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; + } + } + + const ptrdiff_t initial_base = + static_cast(read_slot) * initial_s0 + + static_cast(value_head_idx) * initial_s1 + + static_cast(value_dim_idx) * initial_s2; + + Tdata* final_state_target = + final_state_indices == nullptr ? final_state : initial_state; + const ptrdiff_t final_base = + final_state_indices == nullptr + ? static_cast(batch_idx) * final_s0 + + static_cast(value_head_idx) * final_s1 + + static_cast(value_dim_idx) * final_s2 + : static_cast(write_slot) * initial_s0 + + static_cast(value_head_idx) * initial_s1 + + static_cast(value_dim_idx) * initial_s2; + const ptrdiff_t final_k_stride = + final_state_indices == nullptr ? final_s3 : initial_s3; + + Tcompute state[STATE_VALUES_PER_LANE]; +#pragma unroll + for (int i = 0; i < STATE_VALUES_PER_LANE; ++i) { + const int key_dim_idx = lane_idx + i * WARP_SIZE; + state[i] = valid_value_dim && key_dim_idx < static_cast(Dk) + ? static_cast(loadAsFloat( + initial_state, + initial_base + + static_cast(key_dim_idx) * initial_s3)) + : static_cast(0); + } + + Tcompute* q_local = shared; + Tcompute* k_local = q_local + Dk; + Tcompute* decay_local = k_local + Dk; + Tcompute* reduction_scratch = decay_local + Dk; + Tcompute* beta_shared = reduction_scratch + NUM_THREADS; + + 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; + const ptrdiff_t k_base = static_cast(token_batch) * k_s0 + + static_cast(token_idx) * k_s1 + + static_cast(key_head_idx) * k_s2; + + Tcompute q_sum = 0; + Tcompute k_sum = 0; + for (int key_dim_idx = threadIdx.x; key_dim_idx < static_cast(Dk); + key_dim_idx += NUM_THREADS) { + const Tcompute q_value = + static_cast(loadAsFloat(q, q_base + key_dim_idx)); + const Tcompute k_value = + static_cast(loadAsFloat(k, k_base + key_dim_idx)); + q_local[key_dim_idx] = q_value; + k_local[key_dim_idx] = k_value; + q_sum += q_value * q_value; + k_sum += k_value * k_value; + } + q_sum = blockReduceSum(q_sum, reduction_scratch); + k_sum = blockReduceSum(k_sum, reduction_scratch); + + const Tcompute q_norm = use_qk_l2norm + ? rsqrtf(q_sum + static_cast(1e-6)) + : static_cast(1); + const Tcompute k_norm = use_qk_l2norm + ? rsqrtf(k_sum + static_cast(1e-6)) + : static_cast(1); + for (int key_dim_idx = threadIdx.x; key_dim_idx < static_cast(Dk); + key_dim_idx += NUM_THREADS) { + q_local[key_dim_idx] *= q_norm * query_scale; + k_local[key_dim_idx] *= k_norm; + } + + gate_policy.prepare(token_batch, token_idx, key_head_idx, value_head_idx, + decay_local, beta_shared); + __syncthreads(); + + Tcompute kv_memory = 0; + Tcompute hq_memory = 0; + Tcompute kq_memory = 0; +#pragma unroll + for (int i = 0; i < STATE_VALUES_PER_LANE; ++i) { + const int key_dim_idx = lane_idx + i * WARP_SIZE; + if (valid_value_dim && key_dim_idx < static_cast(Dk)) { + const Tcompute decayed_state = state[i] * decay_local[key_dim_idx]; + const Tcompute k_value = k_local[key_dim_idx]; + const Tcompute q_value = q_local[key_dim_idx]; + kv_memory += decayed_state * k_value; + hq_memory += decayed_state * q_value; + kq_memory += k_value * q_value; + } + } + kv_memory = warpReduceSum(kv_memory); + hq_memory = warpReduceSum(hq_memory); + kq_memory = warpReduceSum(kq_memory); + + Tcompute delta = 0; + if (valid_value_dim && lane_idx == 0) { + const ptrdiff_t v_base = static_cast(token_batch) * v_s0 + + static_cast(token_idx) * v_s1 + + static_cast(value_head_idx) * v_s2; + const Tcompute v_value = + static_cast(loadAsFloat(v, v_base + value_dim_idx)); + delta = (v_value - kv_memory) * beta_shared[0]; + 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(hq_memory + delta * kq_memory); + } + delta = __shfl_sync(INFINIOP_RECURRENT_DELTA_RULE_FULL_MASK, delta, 0, + WARP_SIZE); + +#pragma unroll + for (int i = 0; i < STATE_VALUES_PER_LANE; ++i) { + const int key_dim_idx = lane_idx + i * WARP_SIZE; + if (valid_value_dim && key_dim_idx < static_cast(Dk)) { + state[i] = + state[i] * decay_local[key_dim_idx] + k_local[key_dim_idx] * delta; + } + } + __syncthreads(); + } + +#pragma unroll + for (int i = 0; i < STATE_VALUES_PER_LANE; ++i) { + const int key_dim_idx = lane_idx + i * WARP_SIZE; + if (valid_value_dim && key_dim_idx < static_cast(Dk)) { + final_state_target[final_base + + static_cast(key_dim_idx) * final_k_stride] = + static_cast(state[i]); + } + } +} + +} // namespace op::recurrent_gated_delta_rule::cuda + +#endif // __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 new file mode 100644 index 000000000..d5dd6e14b --- /dev/null +++ b/src/native/cuda/nvidia/ops/recurrent_gated_delta_rule/recurrent_gated_delta_rule_legacy_kernel.cuh @@ -0,0 +1,78 @@ +// Copied from InfiniCore 1ab85ef1d3c2bce49edf1fc73a5b48becaf7552e: +// src/infiniop/ops/recurrent_gated_delta_rule/cuda/kernel.cuh. +// 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__ + +#include "recurrent_delta_rule_common.cuh" + +template +struct ScalarGatedDeltaRulePolicy { + const Tgate* g; + const Tgate* beta; + ptrdiff_t g_s0; + ptrdiff_t g_s1; + ptrdiff_t g_s2; + ptrdiff_t beta_s0; + ptrdiff_t beta_s1; + ptrdiff_t beta_s2; + + __device__ void prepare(int token_batch, int64_t token_idx, int, + int value_head_idx, Tcompute* decay, + Tcompute* beta_out) const { + if (threadIdx.x == 0) { + const ptrdiff_t gate_offset = + static_cast(token_batch) * g_s0 + + static_cast(token_idx) * g_s1 + + static_cast(value_head_idx) * g_s2; + const ptrdiff_t beta_offset = + static_cast(token_batch) * beta_s0 + + static_cast(token_idx) * beta_s1 + + static_cast(value_head_idx) * beta_s2; + decay[0] = expf(static_cast( + op::recurrent_gated_delta_rule::cuda::loadAsFloat(g, gate_offset))); + beta_out[0] = static_cast( + op::recurrent_gated_delta_rule::cuda::loadAsFloat(beta, beta_offset)); + } + __syncthreads(); + const Tcompute scalar_decay = decay[0]; + for (int key_dim_idx = threadIdx.x; key_dim_idx < static_cast(Dk); + key_dim_idx += blockDim.x) { + decay[key_dim_idx] = scalar_decay; + } + } +}; + +template +__device__ void recurrentGatedDeltaRuleIndexedPoolWarpKernel( + Tdata* out, Tdata* initial_state, Tdata* final_state, const Tdata* q, + const Tdata* k, const Tdata* v, const Tgate* g, const Tgate* beta, + 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, 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, ptrdiff_t g_s0, ptrdiff_t g_s1, + ptrdiff_t g_s2, ptrdiff_t beta_s0, ptrdiff_t beta_s1, ptrdiff_t beta_s2, + Tcompute* shared) { + const ScalarGatedDeltaRulePolicy gate_policy{ + g, beta, g_s0, g_s1, g_s2, beta_s0, beta_s1, beta_s2, + }; + op::recurrent_gated_delta_rule::cuda::recurrentDeltaRuleWarpSequence< + Tdata, Tcompute, Dk, Dv, WARPS_PER_BLOCK>( + 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/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..85769299c 100644 --- a/src/native/cuda/nvidia/ops/rms_norm/kernel.h +++ b/src/native/cuda/nvidia/ops/rms_norm/kernel.h @@ -1,8 +1,6 @@ #ifndef INFINI_OPS_NVIDIA_RMS_NORM_KERNEL_H_ #define INFINI_OPS_NVIDIA_RMS_NORM_KERNEL_H_ -#include - #include "native/cuda/nvidia/caster.cuh" #include "native/cuda/nvidia/runtime_.h" #include "native/cuda/ops/rms_norm/kernel.h" @@ -13,7 +11,17 @@ template <> class Operator : public CudaRmsNorm> { public: - using CudaRmsNorm>::CudaRmsNorm; + using Base = CudaRmsNorm>; + using Base::Base; + + // Mixed weight types dispatch to kernels too; replay allocates no metadata. + bool capture_safe() const override { return true; } + + 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); + } }; } // 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/separate_table_rope/kernel.h b/src/native/cuda/nvidia/ops/separate_table_rope/kernel.h new file mode 100644 index 000000000..72808598e --- /dev/null +++ b/src/native/cuda/nvidia/ops/separate_table_rope/kernel.h @@ -0,0 +1,27 @@ +#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< + Runtime>::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/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/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..e0d1132be --- /dev/null +++ b/src/native/cuda/nvidia/ops/sum/kernel.h @@ -0,0 +1,188 @@ +#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) { + 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 && *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 && + input.dtype() != DataType::kBFloat16 && + input.dtype() != DataType::kFloat32) + return false; + if (input.dtype() == DataType::kBFloat16) { + cudaDeviceProp properties{}; + if (cudaGetDeviceProperties(&properties, input.device().index()) != + 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.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 > 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)) + 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) + : 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 (tensor.shape()[i - 1] == 1) continue; + 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, true)) { + 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..26d7bfb28 --- /dev/null +++ b/src/native/cuda/nvidia/ops/take/kernel.h @@ -0,0 +1,168 @@ +#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) { + 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) || + !CheckedContiguous(out) || index.dtype() != DataType::kInt64 || + input.dtype() != out.dtype() || index.shape() != out.shape()) { + 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 (!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 && + 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, out.data(), + output_count * width) && + !Overlaps(index.data(), output_count * sizeof(std::int64_t), + out.data(), output_count * width); + } + + public: + 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 (shape[i - 1] == 1) continue; + 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, true)) + 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/topk_softmax/kernel.cu b/src/native/cuda/nvidia/ops/topk_softmax/kernel.cu index 15d860fc4..c1e625c9d 100644 --- a/src/native/cuda/nvidia/ops/topk_softmax/kernel.cu +++ b/src/native/cuda/nvidia/ops/topk_softmax/kernel.cu @@ -82,7 +82,7 @@ void Operator::operator()( reinterpret_cast(token_expert_indices.data()), static_cast(num_tokens_), static_cast(num_experts_), static_cast(topk_), - renormalize_); + renormalize_, softcap_, min_sum_); }, "Operator::operator()"); diff --git a/src/native/cuda/nvidia/ops/topk_softmax/kernel.h b/src/native/cuda/nvidia/ops/topk_softmax/kernel.h index 8e59a9605..2579d5728 100644 --- a/src/native/cuda/nvidia/ops/topk_softmax/kernel.h +++ b/src/native/cuda/nvidia/ops/topk_softmax/kernel.h @@ -12,6 +12,22 @@ class Operator : public TopkSoftmax { public: using TopkSoftmax::TopkSoftmax; + Operator(const Config& config, const Tensor input, std::optional bias, + std::optional padding, bool renormalize, Tensor weights, + Tensor indices, Tensor scratch) + : TopkSoftmax(input, bias, padding, renormalize, weights, indices, + scratch) { + softcap_ = TopkSoftmaxConfig::SoftcapFrom(config); + min_sum_ = TopkSoftmaxConfig::MinSumFrom(config); + } + + static bool supports(const Config&, const Tensor, std::optional, + std::optional, bool, Tensor, Tensor, Tensor) { + return true; + } + + bool capture_safe() const override { return true; } + void operator()(const Tensor gating_output, std::optional bias, std::optional is_padding, const bool renormalize, Tensor topk_weights, Tensor topk_indices, 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..cd2ff5f54 --- /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, + Tril, false> { + public: + using triangular_detail::CudaTriangular, 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..a6fff2ac3 --- /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, + Triu, true> { + public: + using triangular_detail::CudaTriangular, Triu, + true>::CudaTriangular; +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_NVIDIA_TRIU_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..5c3893f2f --- /dev/null +++ b/src/native/cuda/nvidia/ops/upsample_bilinear2d/kernel.h @@ -0,0 +1,209 @@ +#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) { + 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 || + 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; + if (input.dtype() == DataType::kBFloat16) { + cudaDeviceProp properties{}; + if (cudaGetDeviceProperties(&properties, input.device().index()) != + 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 || + 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)) + 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) + : 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 (tensor.shape()[i - 1] == 1) continue; + 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, + true)) + 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/nvidia/ops/w4_marlin/gemm.cu b/src/native/cuda/nvidia/ops/w4_marlin/gemm.cu new file mode 100644 index 000000000..e9d8b0e12 --- /dev/null +++ b/src/native/cuda/nvidia/ops/w4_marlin/gemm.cu @@ -0,0 +1,93 @@ +#include + +#include "kernel.h" +#include "launch.h" +#include "core/scalar_type.hpp" +#include "marlin.cuh" + +namespace infini::ops::native_w4 { +inline std::size_t align256(std::size_t bytes) { + if (bytes > std::numeric_limits::max()-255) + throw std::overflow_error("W4 workspace alignment overflow"); + return (bytes+255)&~std::size_t(255); +} +inline std::size_t append(std::size_t& total, std::size_t size) { + const auto result = align256(total); + if (size > std::numeric_limits::max()-result) + throw std::overflow_error("W4 workspace sum overflow"); + total = result + size; + return result; +} +template +__global__ void epilogue(T* output, const T* bias, std::size_t count, int n, float alpha) { + const std::size_t i = std::size_t(blockIdx.x)*blockDim.x+threadIdx.x; + if (i < count) { + const float value = float(output[i])*alpha + (bias ? float(bias[i%n]) : 0.f); + output[i] = T(value); + } +} +} // namespace infini::ops::native_w4 + +namespace infini::ops { +Operator::Operator( + const Tensor a, const Tensor b, const Tensor scales, const Tensor zeros, + const Tensor g_idx, const Tensor perm, const Tensor bias, + const int64_t scheme, const int64_t group_size, const bool is_k_full, + 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); + ctmp_ = native_w4::append(bytes_,w4::mul(w4::mul(sms_,max_m),marlin::max_thread_n*sizeof(float))); + atmp_ = native_w4::append(bytes_,perm.numel() ? w4::mul(w4::mul(m_,k_),a.element_size()) : 0); + locks_ = native_w4::append(bytes_,w4::mul(sms_,sizeof(int))); + bytes_ = native_w4::align256(bytes_); +} +void Operator::operator()( + const Tensor a, const Tensor b, const Tensor scales, const Tensor zeros, + const Tensor g_idx, const Tensor perm, const Tensor bias, + const int64_t scheme, const int64_t group_size, const bool is_k_full, + const float alpha, Tensor out) const { + 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"); + for (const auto& t : {a,b,scales,zeros,g_idx,perm,bias,out}) + w4::disjoint(workspace_,bytes_,t.data(),w4::span(t)); + native_w4::DeviceGuard guard(device_); + auto stream = static_cast(stream_); + native_w4::no_capture(stream); + auto* memory = static_cast(workspace_); + // Locks are state, not merely scratch. Every eager/replayed call resets them + // on the same stream before the original reduction kernel consumes them. + native_w4::check(cudaMemsetAsync(memory+locks_,0,sms_*sizeof(int),stream), "W4 lock reset"); + const auto type = a.dtype() == DataType::kFloat16 ? vllm::kFloat16 : vllm::kBFloat16; + const auto btype = scheme == 0 ? vllm::kU4 : vllm::kU4B8; + const bool act = perm.numel() != 0; + const int group = act && !full_ ? 0 : (scales.size(0) == 1 ? -1 : int(group_)); + // The original unified AWQ kernel generator already contains GPTQ U4B8 + // specializations. This avoids the old separate GPTQ TVM-FFI tensor wrapper. + marlin::marlin_mm(a.data(),b.data(),out.data(),memory+ctmp_,nullptr, + nullptr,const_cast(scales.data()),nullptr, + zeros.numel() ? const_cast(zeros.data()) : nullptr, + act ? const_cast(g_idx.data()) : nullptr, + act ? const_cast(perm.data()) : nullptr, + act ? memory+atmp_ : nullptr,m_,n_,k_,a.stride(0),memory+locks_, + type,btype,type,type,false,act,full_,scheme == 0,scales.size(0),group, + device_,stream,-1,-1,sms_,false,true,false); + native_w4::check(cudaGetLastError(), "W4 Marlin GEMM launch"); + if (alpha != 1.f || bias.numel()) { + const auto count = out.numel(); + if (a.dtype() == DataType::kFloat16) + native_w4::epilogue<<<(count+255)/256,256,0,stream>>>(static_cast(out.data()),bias.numel() ? static_cast(bias.data()) : nullptr,count,n_,alpha); + else + native_w4::epilogue<<<(count+255)/256,256,0,stream>>>(static_cast(out.data()),bias.numel() ? static_cast(bias.data()) : nullptr,count,n_,alpha); + native_w4::check(cudaGetLastError(), "W4 alpha/bias epilogue"); + } +} +} // namespace infini::ops diff --git a/src/native/cuda/nvidia/ops/w4_marlin/kernel.h b/src/native/cuda/nvidia/ops/w4_marlin/kernel.h new file mode 100644 index 000000000..9df1d31a7 --- /dev/null +++ b/src/native/cuda/nvidia/ops/w4_marlin/kernel.h @@ -0,0 +1,58 @@ +#pragma once +#include "base/w4_marlin_gemm.h" +#include "base/w4_marlin_parameters.h" +#include "base/w4_marlin_repack.h" + +namespace infini::ops { +template <> +class Operator + : public W4MarlinRepack { + public: + Operator(const Tensor qweight, const Tensor perm, const int64_t scheme, + const int64_t size_k, const int64_t size_n, Tensor out); + using W4MarlinRepack::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; + + private: + int sms_{0}, shared_{0}, device_{0}; +}; +template <> +class Operator + : public W4MarlinParameters { + public: + Operator(const Tensor scales, const Tensor qzeros, 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); + using W4MarlinParameters::operator(); + void operator()(const Tensor scales, const Tensor qzeros, + 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; + + private: + int device_{0}; +}; +template <> +class Operator : public W4MarlinGemm { + public: + Operator(const Tensor a, const Tensor b, const Tensor scales, + const Tensor zeros, const Tensor g_idx, const Tensor perm, + const Tensor bias, const int64_t scheme, const int64_t group_size, + const bool is_k_full, const float alpha, Tensor out); + using W4MarlinGemm::operator(); + void operator()(const Tensor a, const Tensor b, const Tensor scales, + const Tensor zeros, const Tensor g_idx, const Tensor perm, + const Tensor bias, const int64_t scheme, + 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_; } + // Original host selector still sets CUDA function attributes and queries the + // device at invocation. No capture promise without a qualified frozen launch. + + private: + int device_{0}, sms_{0}; + std::size_t ctmp_{0}, atmp_{0}, locks_{0}, bytes_{0}; +}; +} // namespace infini::ops diff --git a/src/native/cuda/nvidia/ops/w4_marlin/launch.h b/src/native/cuda/nvidia/ops/w4_marlin/launch.h new file mode 100644 index 000000000..c0e0e2a8f --- /dev/null +++ b/src/native/cuda/nvidia/ops/w4_marlin/launch.h @@ -0,0 +1,76 @@ +#pragma once +#include + +#include +#include +#include + +namespace infini::ops::native_w4 { +inline void check(cudaError_t e, const char* operation) { + if (e != cudaSuccess) + throw std::runtime_error(std::string(operation) + ": " + + cudaGetErrorString(e)); +} +// Restores the previous device; an instance records a device index, never a +// CUDA context or workspace cache. Runtime already owns the active stream. +class DeviceGuard { + public: + explicit DeviceGuard(int device) : target_{device} { + check(cudaGetDevice(&previous_), "W4 cudaGetDevice"); + if (previous_ != target_) check(cudaSetDevice(target_), "W4 cudaSetDevice"); + } + ~DeviceGuard() { + if (previous_ != target_) cudaSetDevice(previous_); + } + + private: + int previous_{0}, target_{0}; +}; +inline void no_capture(cudaStream_t stream) { + cudaStreamCaptureStatus state; + check(cudaStreamIsCapturing(stream, &state), "W4 capture status"); + if (state != cudaStreamCaptureStatusNone) + throw std::runtime_error( + "W4 Marlin capture=false; use explicit recorded host replay"); +} +inline void info(int device, int& sms, int& shared) { + DeviceGuard guard(device); + int major = 0; + check( + cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, device), + "W4 SM query"); + if (major < 8) + throw std::runtime_error("W4 Marlin builds SM80+ kernels only"); + check(cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, device), + "W4 SM count"); + check(cudaDeviceGetAttribute(&shared, cudaDevAttrMaxSharedMemoryPerBlockOptin, + device), + "W4 shared memory"); + if (sms <= 0 || shared <= 0) + throw std::runtime_error("W4 invalid device limits"); +} +void prepare_awq_repack(int shared); +void prepare_gptq_repack(int shared, bool permutation); +void launch_awq_repack(const uint32_t* input, uint32_t* output, int k, int n, + int sms, int shared, cudaStream_t stream); +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..a4d0536d9 --- /dev/null +++ b/src/native/cuda/nvidia/ops/w4_marlin/marlin_check.h @@ -0,0 +1,29 @@ +#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) + +// Keep upstream check call sites intact inside the private prepared sources. +#define STD_TORCH_CHECK INFINI_MARLIN_CHECK diff --git a/src/native/cuda/nvidia/ops/w4_marlin/parameters.cu b/src/native/cuda/nvidia/ops/w4_marlin/parameters.cu new file mode 100644 index 000000000..c6ce73586 --- /dev/null +++ b/src/native/cuda/nvidia/ops/w4_marlin/parameters.cu @@ -0,0 +1,71 @@ +// 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 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); + const int j[8] = {0,1,8,9,16,17,24,25}; + return 2*(i/8) + j[i%8]; +} +__global__ void scales_reorder(const uint16_t* source, uint16_t* dest, std::size_t count, bool single) { + const std::size_t i = std::size_t(blockIdx.x)*blockDim.x + threadIdx.x; + if (i >= count) return; + const int block = single ? 32 : 64; + dest[i] = source[(i/block)*block + scale_source(i%block,single)]; +} +__global__ void zeros_reorder(const uint32_t* source, uint32_t* dest, std::size_t words, int n) { + const std::size_t i = std::size_t(blockIdx.x)*blockDim.x + threadIdx.x; + if (i >= words) return; + const int undo[8] = {0,4,1,5,2,6,3,7}; + const int interleave[8] = {0,2,4,6,1,3,5,7}; + const std::size_t row = i/(n/8), col = i%(n/8); + uint32_t word = 0; + for (int lane = 0; lane < 8; ++lane) { + const int permuted = int(col*8) + interleave[lane]; + const int logical = (permuted/64)*64 + scale_source(permuted%64,false); + const uint32_t raw = source[row*(n/8) + logical/8]; + word |= ((raw >> (4*undo[logical%8])) & 15) << (4*lane); + } + dest[i] = word; +} +} // namespace infini::ops::native_w4 + +namespace infini::ops { +Operator::Operator( + const Tensor scales, const Tensor qzeros, 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) + : W4MarlinParameters(scales,qzeros,scheme,size_k,size_n,group_size,scales_out,zeros_out), device_{scales.device().index()} { + w4::require(scales.device().type() == Device::Type::kNvidia, "W4 requires one NVIDIA device"); + int sms, shared; native_w4::info(device_,sms,shared); +} +void Operator::operator()( + const Tensor scales, const Tensor qzeros, 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 { + w4::require(scheme == scheme_ && size_k == k_ && size_n == n_ && group_size == group_, "W4 parameter attributes changed"); + for (const auto& t : {scales,qzeros,scales_out,zeros_out}) w4::address(t,4); + for (const auto& t : {scales,qzeros}) { w4::disjoint(scales_out,t); w4::disjoint(zeros_out,t); } + w4::disjoint(scales_out,zeros_out); + native_w4::DeviceGuard guard(device_); + auto stream = static_cast(stream_); + native_w4::no_capture(stream); + const bool single = group_size == -1 || group_size >= size_k; + native_w4::scales_reorder<<<(scales.numel()+255)/256,256,0,stream>>>( + static_cast(scales.data()),static_cast(scales_out.data()),scales.numel(),single); + native_w4::check(cudaGetLastError(), "W4 scale permutation"); + if (scheme == 0) { + native_w4::zeros_reorder<<<(qzeros.numel()+255)/256,256,0,stream>>>( + static_cast(qzeros.data()),static_cast(zeros_out.data()),qzeros.numel(),n_); + native_w4::check(cudaGetLastError(), "W4 AWQ zero permutation"); + } +} +} // namespace infini::ops diff --git a/src/native/cuda/nvidia/ops/w4_marlin/repack.cc b/src/native/cuda/nvidia/ops/w4_marlin/repack.cc new file mode 100644 index 000000000..0786e7da5 --- /dev/null +++ b/src/native/cuda/nvidia/ops/w4_marlin/repack.cc @@ -0,0 +1,43 @@ +#include "kernel.h" +#include "launch.h" + +namespace infini::ops { +Operator::Operator( + const Tensor qweight, const Tensor perm, const int64_t scheme, + const int64_t size_k, const int64_t size_n, Tensor out) + : W4MarlinRepack(qweight, perm, scheme, size_k, size_n, out), + 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) + native_w4::prepare_awq_repack(shared_); + else + native_w4::prepare_gptq_repack(shared_, perm.numel() != 0); +} +void Operator::operator()( + const Tensor qweight, const Tensor perm, const int64_t scheme, + const int64_t size_k, const int64_t size_n, Tensor out) const { + w4::require(scheme == scheme_ && size_k == k_ && size_n == n_, + "W4 repack attributes changed"); + w4::address(qweight, 16); + w4::address(out, 16); + w4::address(perm, scheme == 1 && perm.numel() ? 16 : 4); + w4::disjoint(out, qweight); + w4::disjoint(out, perm); + native_w4::DeviceGuard guard(device_); + auto stream = static_cast(stream_); + native_w4::no_capture(stream); + auto* output = static_cast(out.data()); + auto* input = static_cast(qweight.data()); + if (scheme == 0) + native_w4::launch_awq_repack(input, output, k_, n_, sms_, shared_, stream); + else + native_w4::launch_gptq_repack( + input, + perm.numel() ? static_cast(perm.data()) : nullptr, + output, k_, n_, sms_, shared_, stream); +} +} // namespace infini::ops diff --git a/src/native/cuda/nvidia/ops/w4_marlin/repack_awq.cu b/src/native/cuda/nvidia/ops/w4_marlin/repack_awq.cu new file mode 100644 index 000000000..6bc12c71c --- /dev/null +++ b/src/native/cuda/nvidia/ops/w4_marlin/repack_awq.cu @@ -0,0 +1,13 @@ +#include "launch.h" +#include "awq_marlin_repack.cu" + +namespace infini::ops::native_w4 { +void prepare_awq_repack(int shared) { + 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) { + 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 new file mode 100644 index 000000000..bdfdf99fd --- /dev/null +++ b/src/native/cuda/nvidia/ops/w4_marlin/repack_gptq.cu @@ -0,0 +1,15 @@ +#include "launch.h" +#include "gptq_marlin_repack.cu" + +namespace infini::ops::native_w4 { +void prepare_gptq_repack(int shared, bool permutation) { + 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) 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/ops/add/kernel.h b/src/native/cuda/ops/add/kernel.h index 911d706a2..82941cca9 100644 --- a/src/native/cuda/ops/add/kernel.h +++ b/src/native/cuda/ops/add/kernel.h @@ -4,11 +4,13 @@ #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" @@ -21,12 +23,13 @@ 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 (!metadata_size) return; std::vector metadata(metadata_size); - Backend::Malloc((void**)&d_metadata_, metadata_size); + auto* d_metadata_ = metadata_.Allocate(metadata_size); size_t offset = 0; d_input_shape_ = reinterpret_cast(d_metadata_ + offset); @@ -52,15 +55,14 @@ class CudaAdd : public Add { 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); + 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() { Backend::Free(d_metadata_); } - void operator()(const Tensor input, const Tensor other, const double alpha, Tensor out) const override { if (output_size_ == 0) { @@ -70,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); @@ -80,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()); @@ -98,7 +99,9 @@ class CudaAdd : public Add { } private: - std::byte* d_metadata_{nullptr}; + const int block_size_ = + RuntimeUtils::GetOptimalBlockSize(); + CudaDeviceBuffer metadata_; Tensor::Size* d_input_shape_{nullptr}; 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/copy/kernel.h b/src/native/cuda/ops/copy/kernel.h index 5caaac264..c91e70f82 100644 --- a/src/native/cuda/ops/copy/kernel.h +++ b/src/native/cuda/ops/copy/kernel.h @@ -4,12 +4,14 @@ #include #include #include +#include #include #include "base/copy.h" #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" @@ -21,12 +23,13 @@ class CudaCopy : public Copy { public: CudaCopy(const Tensor input, const bool non_blocking, Tensor out) : Copy{input, non_blocking, out} { + 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); + auto* d_metadata_ = metadata_.Allocate(metadata_size); size_t offset = 0; d_input_shape_ = reinterpret_cast(d_metadata_ + offset); @@ -44,27 +47,23 @@ class CudaCopy : public Copy { 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); + 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"); } - ~CudaCopy() { Backend::Free(d_metadata_); } - 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); - 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); @@ -79,7 +78,9 @@ class CudaCopy : public Copy { } private: - std::byte* d_metadata_{nullptr}; + 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.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..a0ce60977 100644 --- a/src/native/cuda/ops/embedding/kernel.h +++ b/src/native/cuda/ops/embedding/kernel.h @@ -4,12 +4,14 @@ #include #include #include +#include #include #include "base/embedding.h" #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" @@ -31,42 +33,38 @@ 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_); + 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 (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) { - Backend::Malloc(reinterpret_cast(&d_visited_), - vocab_size_ * sizeof(*d_visited_)); + d_visited_ = reinterpret_cast( + visited_.Allocate(vocab_size_ * sizeof(*d_visited_))); } } @@ -83,14 +81,6 @@ 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_); - } - } - void operator()(const Tensor input, const Tensor weight, const std::optional /*padding_idx*/, const std::optional max_norm, const double norm_type, @@ -113,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>>( @@ -143,7 +126,7 @@ class CudaEmbedding : public Embedding { } EmbeddingKernel - <<>>( + <<>>( reinterpret_cast(out.data()), reinterpret_cast(input.data()), reinterpret_cast(weight.data()), num_indices_, @@ -156,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}; @@ -170,7 +155,8 @@ class CudaEmbedding : public Embedding { int* d_visited_{nullptr}; - std::byte* d_metadata_{nullptr}; + 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 2dbd4e3ed..b1a60249b 100644 --- a/src/native/cuda/ops/fill/kernel.h +++ b/src/native/cuda/ops/fill/kernel.h @@ -5,12 +5,14 @@ #include #include #include +#include #include #include "base/fill.h" #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" @@ -33,13 +35,14 @@ class CudaFill : public Fill { "`CudaFill` requires `input` and `out` on the same device."); assert(!out.HasBroadcastDim() && "`CudaFill` output must not have broadcast dimensions."); + 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); + auto* d_metadata_ = metadata_.Allocate(metadata_size); size_t offset = 0; d_out_shape_ = reinterpret_cast(d_metadata_ + offset); @@ -49,12 +52,11 @@ class CudaFill : public Fill { 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); + if (Backend::Memcpy(d_metadata_, metadata.data(), metadata_size, + Backend::kMemcpyHostToDevice) != Backend::kSuccess) + throw std::runtime_error("CudaFill metadata copy failed"); } - ~CudaFill() override { Backend::Free(d_metadata_); } - void operator()(const Tensor input, const double value, Tensor out) const override { (void)input; @@ -64,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); @@ -93,13 +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}; + 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 70a0a6cd8..f5b8ccd89 100644 --- a/src/native/cuda/ops/gelu/kernel.h +++ b/src/native/cuda/ops/gelu/kernel.h @@ -5,12 +5,14 @@ #include #include #include +#include #include #include "base/gelu.h" #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" @@ -31,16 +33,15 @@ class CudaGelu : public Gelu { "`Gelu` input and output devices must match"); assert(!out.HasBroadcastDim() && "`Gelu` output must not have broadcasted dimensions"); + 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); + auto* d_metadata_ = metadata_.Allocate(metadata_size); size_t offset = 0; d_input_shape_ = reinterpret_cast(d_metadata_ + offset); @@ -58,14 +59,9 @@ class CudaGelu : public Gelu { 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); - } - - ~CudaGelu() { - if (d_metadata_ != nullptr) { - Backend::Free(d_metadata_); - } + 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, @@ -78,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); @@ -111,7 +105,9 @@ class CudaGelu : public Gelu { } private: - std::byte* d_metadata_{nullptr}; + const int block_size_ = + std::min(RuntimeUtils::GetOptimalBlockSize(), 1024); + CudaDeviceBuffer metadata_; Tensor::Size* d_input_shape_{nullptr}; diff --git a/src/native/cuda/ops/gemm/blas.h b/src/native/cuda/ops/gemm/blas.h index ec28a0e1e..1f56f32b3 100644 --- a/src/native/cuda/ops/gemm/blas.h +++ b/src/native/cuda/ops/gemm/blas.h @@ -1,6 +1,8 @@ #ifndef INFINI_OPS_CUDA_GEMM_BLAS_H_ #define INFINI_OPS_CUDA_GEMM_BLAS_H_ +#include +#include #include #include "base/gemm.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,24 +53,31 @@ 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_, - 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() - : a.dtype()), - swap_a_and_b_ ? ldb_ : lda_, - swap_a_and_b_ ? batch_stride_b_ : batch_stride_a_, - swap_a_and_b_ ? a.data() : b.data(), - BlasUtils::GetDataType(swap_a_and_b_ ? a.dtype() - : b.dtype()), - swap_a_and_b_ ? lda_ : ldb_, - 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); + 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() : a.dtype()), + swap_a_and_b_ ? ldb_ : lda_, + swap_a_and_b_ ? batch_stride_b_ : batch_stride_a_, + swap_a_and_b_ ? a.data() : b.data(), + BlasUtils::GetDataType( + swap_a_and_b_ ? a.dtype() : b.dtype()), + swap_a_and_b_ ? lda_ : ldb_, + swap_a_and_b_ ? batch_stride_a_ : batch_stride_b_, beta_ptr, + y.data(), BlasUtils::GetDataType(y.dtype()), + ldy_, batch_stride_y_, batch_count_, compute_type, + Backend::BLAS_GEMM_DEFAULT), + "submitting GEMM"); } protected: @@ -97,15 +108,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/mamba_selective_scan/kernel.h b/src/native/cuda/ops/mamba_selective_scan/kernel.h new file mode 100644 index 000000000..dd4e3cf9d --- /dev/null +++ b/src/native/cuda/ops/mamba_selective_scan/kernel.h @@ -0,0 +1,173 @@ +// 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 +#include +#include + +#include "base/mamba_selective_scan.h" +#include "caster.h" +#include "dispatcher.h" + +namespace infini::ops { +namespace mamba_scan_detail { +template +__device__ __forceinline__ std::size_t Offset(const MambaScanLayout& p, int i, + std::size_t batch, std::size_t t, + std::size_t ch) { + if constexpr (Contiguous) { + if (i == 5) return batch * p.states + t; + if (i == 6 || i == 8) return batch; + if (i == 9) return (batch * p.channels + t) * p.states + ch; + return (batch * p.length + t) * + ((i == 3 || i == 4) ? p.states : p.channels) + + ch; + } + return batch * p.strides[i][0] + t * p.strides[i][1] + ch * p.strides[i][2]; +} +template +__global__ void Scan(T* out, const T* x, const T* dt, const T* b, const T* c, + const P* a_log, const P* d, const T* gate, const T* bias, + float* state, MambaScanLayout p, std::size_t first) { + const std::size_t linear = + first + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (linear >= p.batch * p.channels) return; + const auto ch = linear % p.channels, batch = linear / p.channels; + for (std::size_t t = 0; t < p.length; ++t) { + const float xt = Caster::template Cast( + x[Offset(p, 1, batch, t, ch)]); + const float z = Caster::template Cast( + dt[Offset(p, 2, batch, t, ch)]) + + Caster::template Cast( + bias[Offset(p, 8, ch, 0, 0)]); + const float dtv = z > 20.0f ? z : log1pf(expf(z)); + float y = 0.0f; + for (std::size_t n = 0; n < p.states; ++n) { + const float a = -expf(Caster::template Cast( + a_log[Offset(p, 5, ch, n, 0)])); + const float discrete_a = expf(a * dtv); + const float bn = Caster::template Cast( + b[Offset(p, 3, batch, t, n)]); + const float cn = Caster::template Cast( + c[Offset(p, 4, batch, t, n)]); + const auto si = Offset(p, 9, batch, ch, n); + const float s = discrete_a * state[si] + dtv * bn * xt; + state[si] = s; + y += s * cn; + } + y += xt * Caster::template Cast( + d[Offset(p, 6, ch, 0, 0)]); + const float g = Caster::template Cast( + gate[Offset(p, 7, batch, t, ch)]); + y *= g / (1.0f + expf(-g)); + out[Offset(p, 0, batch, t, ch)] = + Caster::template Cast(y); + } +} +} // namespace mamba_scan_detail + +template +class CudaMambaSelectiveScan : public MambaSelectiveScan { + public: + static bool supports(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) { + try { + ValidateMetadata({out, x, dt, b, c, a_log, d, gate, dt_bias, state}); + 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; + } + } + 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>( + dtype_, + [&](auto tag) { + using T = TypeMapType; + auto launch = [&](auto* parameter_type) { + using P = std::remove_pointer_t; + const auto stream = static_cast(stream_); + constexpr std::size_t cap = 65535u * 256u; + const auto total = layout_.batch * layout_.channels; + for (std::size_t first = 0; first < total;) { + const auto count = std::min(cap, total - first); + const dim3 blocks(static_cast((count + 255) / 256)); + if (layout_.contiguous) + mamba_scan_detail::Scan + <<>>( + static_cast(out.data()), + static_cast(x.data()), + static_cast(dt.data()), + static_cast(b.data()), + static_cast(c.data()), + static_cast(a_log.data()), + static_cast(d.data()), + static_cast(gate.data()), + static_cast(dt_bias.data()), + static_cast(state.data()), layout_, first); + else + mamba_scan_detail::Scan + <<>>( + static_cast(out.data()), + static_cast(x.data()), + static_cast(dt.data()), + static_cast(b.data()), + static_cast(c.data()), + static_cast(a_log.data()), + static_cast(d.data()), + static_cast(gate.data()), + static_cast(dt_bias.data()), + static_cast(state.data()), layout_, first); + Backend::CheckLaunch(); + first += count; + } + }; + if (a_log.dtype() == DataType::kFloat32) + launch(static_cast(nullptr)); + else + launch(static_cast(nullptr)); + }, + "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/mrope/kernel.cuh b/src/native/cuda/ops/mrope/kernel.cuh new file mode 100644 index 000000000..cd0f92fd1 --- /dev/null +++ b/src/native/cuda/ops/mrope/kernel.cuh @@ -0,0 +1,60 @@ +#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); +} +} // namespace infini::ops diff --git a/src/native/cuda/ops/mrope/kernel.h b/src/native/cuda/ops/mrope/kernel.h new file mode 100644 index 000000000..c9ca518e5 --- /dev/null +++ b/src/native/cuda/ops/mrope/kernel.h @@ -0,0 +1,40 @@ +#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()"); + } +}; +} // namespace infini::ops diff --git a/src/native/cuda/ops/mul/kernel.h b/src/native/cuda/ops/mul/kernel.h index 216cf8df6..7cb08d0be 100644 --- a/src/native/cuda/ops/mul/kernel.h +++ b/src/native/cuda/ops/mul/kernel.h @@ -4,10 +4,12 @@ #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" @@ -28,7 +30,7 @@ class CudaMul : public Mul { std::vector metadata(metadata_size); - Backend::Malloc((void**)&d_metadata_, metadata_size); + auto* d_metadata_ = metadata_.Allocate(metadata_size); size_t offset = 0; d_input_shape_ = reinterpret_cast(d_metadata_ + offset); @@ -54,25 +56,21 @@ class CudaMul : public Mul { 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); - } - - ~CudaMul() { - if (d_metadata_ != nullptr) { - Backend::Free(d_metadata_); - } + 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, Tensor out) const override { + if (scalar_mode_) + throw std::invalid_argument("Mul prepared for a scalar operand"); if (output_size_ == 0) { 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); @@ -80,7 +78,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()); @@ -97,8 +95,14 @@ class CudaMul : public Mul { "CudaMul::operator()"); } + protected: + CudaMul(const Tensor input, const double alpha, Tensor out) + : Mul{input, alpha, out}, block_size_(0) {} + private: - std::byte* d_metadata_{nullptr}; + const int block_size_ = + RuntimeUtils::GetOptimalBlockSize(); + CudaDeviceBuffer metadata_; Tensor::Size* d_input_shape_{nullptr}; 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..5d0a44e6a --- /dev/null +++ b/src/native/cuda/ops/separate_table_rope/kernel.h @@ -0,0 +1,81 @@ +#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 the public bool neox convention, 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..9e2cfb0a1 --- /dev/null +++ b/src/native/cuda/ops/separate_table_rope/layout.h @@ -0,0 +1,47 @@ +#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..bdea727f0 100644 --- a/src/native/cuda/ops/sigmoid/kernel.h +++ b/src/native/cuda/ops/sigmoid/kernel.h @@ -5,12 +5,14 @@ #include #include #include +#include #include #include "base/sigmoid.h" #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" @@ -30,6 +32,9 @@ class CudaSigmoid : public Sigmoid { assert(!out.HasBroadcastDim() && "`Sigmoid` output must not have broadcasted dimensions"); + 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,7 +43,7 @@ class CudaSigmoid : public Sigmoid { } std::vector metadata(metadata_size); - Backend::Malloc(reinterpret_cast(&d_metadata_), metadata_size); + auto* d_metadata_ = metadata_.Allocate(metadata_size); size_t offset = 0; d_input_shape_ = reinterpret_cast(d_metadata_ + offset); @@ -56,14 +61,9 @@ class CudaSigmoid : public Sigmoid { 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); - } - - ~CudaSigmoid() { - if (d_metadata_ != nullptr) { - Backend::Free(d_metadata_); - } + 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 { @@ -73,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); @@ -96,7 +94,9 @@ class CudaSigmoid : public Sigmoid { } private: - std::byte* d_metadata_{nullptr}; + 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 f883de566..1fa37e69f 100644 --- a/src/native/cuda/ops/silu/kernel.h +++ b/src/native/cuda/ops/silu/kernel.h @@ -4,12 +4,14 @@ #include #include #include +#include #include #include "base/silu.h" #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" @@ -20,12 +22,15 @@ template class CudaSilu : public Silu { public: CudaSilu(const Tensor input, Tensor out) : Silu{input, out} { + 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); + auto* d_metadata_ = metadata_.Allocate(metadata_size); size_t offset = 0; d_input_shape_ = reinterpret_cast(d_metadata_ + offset); @@ -43,20 +48,18 @@ class CudaSilu : public Silu { 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); + if (Backend::Memcpy(d_metadata_, metadata.data(), metadata_size, + Backend::kMemcpyHostToDevice) != Backend::kSuccess) + throw std::runtime_error("CudaSilu metadata copy failed"); } - ~CudaSilu() { Backend::Free(d_metadata_); } - 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}, + {static_cast(out_type_), block_size_}, [&](auto list_tag) { using T = TypeMapType(list_tag)>; constexpr int kBlockSize = ListGet<1>(list_tag); @@ -64,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()); @@ -80,7 +83,9 @@ class CudaSilu : public Silu { } private: - std::byte* d_metadata_{nullptr}; + 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 0e9b5e620..2174c799c 100644 --- a/src/native/cuda/ops/silu_and_mul/kernel.h +++ b/src/native/cuda/ops/silu_and_mul/kernel.h @@ -4,10 +4,12 @@ #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" @@ -17,13 +19,16 @@ template class CudaSiluAndMul : public SiluAndMul { public: CudaSiluAndMul(const Tensor input, Tensor out) : SiluAndMul{input, out} { + 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); + auto* d_metadata_ = metadata_.Allocate(metadata_size); size_t offset = 0; d_input_shape_ = reinterpret_cast(d_metadata_ + offset); @@ -41,16 +46,18 @@ class CudaSiluAndMul : public SiluAndMul { 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); + if (Backend::Memcpy(d_metadata_, metadata.data(), metadata_size, + Backend::kMemcpyHostToDevice) != Backend::kSuccess) + throw std::runtime_error("CudaSiluAndMul metadata copy failed"); } - ~CudaSiluAndMul() { Backend::Free(d_metadata_); } - void operator()(const Tensor input, Tensor out) const override { - int block_size = RuntimeUtils::GetOptimalBlockSize(); + if (output_size_ == 0) { + return; + } + 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); @@ -58,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()); @@ -73,7 +80,9 @@ class CudaSiluAndMul : public SiluAndMul { } private: - std::byte* d_metadata_{nullptr}; + 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 b2807ca5b..9f909ae12 100644 --- a/src/native/cuda/ops/softmax/kernel.h +++ b/src/native/cuda/ops/softmax/kernel.h @@ -5,12 +5,14 @@ #include #include #include +#include #include #include "base/softmax.h" #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" @@ -36,6 +38,10 @@ class CudaSoftmax : public Softmax { assert(!out.HasBroadcastDim() && "`Softmax` output must not have broadcasted dimensions"); + 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,7 +50,7 @@ class CudaSoftmax : public Softmax { } std::vector metadata(metadata_size); - Backend::Malloc(reinterpret_cast(&d_metadata_), metadata_size); + auto* d_metadata_ = metadata_.Allocate(metadata_size); size_t offset = 0; d_shape_ = reinterpret_cast(d_metadata_ + offset); @@ -58,14 +64,9 @@ class CudaSoftmax : public Softmax { 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); - } - - ~CudaSoftmax() { - if (d_metadata_ != nullptr) { - Backend::Free(d_metadata_); - } + 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, @@ -76,11 +77,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); @@ -101,7 +100,8 @@ class CudaSoftmax : public Softmax { } private: - std::byte* d_metadata_{nullptr}; + CudaDeviceBuffer metadata_; + int block_size_{0}; Tensor::Size* d_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..aadb9f893 --- /dev/null +++ b/src/native/cuda/ops/sum/kernel.cuh @@ -0,0 +1,64 @@ +#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/topk_softmax/kernel.cuh b/src/native/cuda/ops/topk_softmax/kernel.cuh index ee3a0b37c..d1f3152a9 100644 --- a/src/native/cuda/ops/topk_softmax/kernel.cuh +++ b/src/native/cuda/ops/topk_softmax/kernel.cuh @@ -90,6 +90,12 @@ __device__ __forceinline__ float ToFloat(Input value) { return Caster::template Cast(value); } +template +__device__ __forceinline__ float Score(Input value, float softcap) { + const float x = ToFloat(value); + return softcap == 0.0f ? x : tanhf(x / softcap) * softcap; +} + template __global__ void TopkSoftmaxKernel(const Input* __restrict__ gating_output, @@ -99,7 +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) { + 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; @@ -107,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, - ToFloat(gating_output[input_offset + expert])); + 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 = ToFloat(gating_output[input_offset + expert]); + const float value = + Score(gating_output[input_offset + expert], softcap); thread_sum += expf(value - row_max); } const float row_sum = BlockSum(thread_sum); @@ -140,7 +149,7 @@ __global__ void TopkSoftmaxKernel(const Input* __restrict__ gating_output, } const float value = - ToFloat(gating_output[input_offset + expert]); + Score(gating_output[input_offset + expert], softcap); float score = expf(value - row_max) / row_sum; if (!isfinite(score)) { score = 0.0f; @@ -155,7 +164,7 @@ __global__ void TopkSoftmaxKernel(const Input* __restrict__ gating_output, BlockBest(thread_best, thread_index); if (threadIdx.x == 0) { const float value = - ToFloat(gating_output[input_offset + thread_index]); + Score(gating_output[input_offset + thread_index], softcap); float weight = expf(value - row_max) / row_sum; if (!isfinite(weight)) { weight = 0.0f; @@ -173,7 +182,11 @@ __global__ void TopkSoftmaxKernel(const Input* __restrict__ gating_output, renormalize && selected_sum > 0.0f ? 1.0f / selected_sum : 1.0f; const bool padding = is_padding && is_padding[token] != 0; for (int32_t rank = 0; rank < topk; ++rank) { - topk_weights[output_offset + rank] *= scale; + if (renormalize && min_sum > 0.0f) { + topk_weights[output_offset + rank] /= fmaxf(selected_sum, min_sum); + } else { + topk_weights[output_offset + rank] *= scale; + } if (padding) { topk_indices[output_offset + rank] = static_cast(-1); } diff --git a/src/native/cuda/ops/topk_softmax/kernel.h b/src/native/cuda/ops/topk_softmax/kernel.h index 9cd2e0008..3256d3359 100644 --- a/src/native/cuda/ops/topk_softmax/kernel.h +++ b/src/native/cuda/ops/topk_softmax/kernel.h @@ -53,7 +53,7 @@ class CudaTopkSoftmax : public TopkSoftmax { reinterpret_cast(token_expert_indices.data()), static_cast(num_tokens_), static_cast(num_experts_), - static_cast(topk_), renormalize_); + static_cast(topk_), renormalize_, 0.0f); }, "CudaTopkSoftmax::operator()"); } diff --git a/src/native/cuda/ops/triangular/kernel.cuh b/src/native/cuda/ops/triangular/kernel.cuh new file mode 100644 index 000000000..f3f9d32e0 --- /dev/null +++ b/src/native/cuda/ops/triangular/kernel.cuh @@ -0,0 +1,71 @@ +#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..d4544f869 --- /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/native/cuda/ops/upsample_bilinear2d/kernel.cuh b/src/native/cuda/ops/upsample_bilinear2d/kernel.cuh new file mode 100644 index 000000000..d8cb65ca0 --- /dev/null +++ b/src/native/cuda/ops/upsample_bilinear2d/kernel.cuh @@ -0,0 +1,95 @@ +#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_ diff --git a/src/operator.h b/src/operator.h index 6e77ac8b8..3ec0c6733 100644 --- a/src/operator.h +++ b/src/operator.h @@ -11,6 +11,7 @@ #include #include #include +#include #include #include #include @@ -28,6 +29,46 @@ 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 +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 +295,10 @@ class OperatorBase { virtual std::size_t workspace_size_in_bytes() const { return 0; } + // 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 +388,28 @@ 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 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) { @@ -411,7 +478,30 @@ 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. Providers validate their own + // workspace requirements before accessing it. + if (handle.workspace() != nullptr || + handle.workspace_size_in_bytes() != 0) { + return (*op)(handle, args...); + } + auto lease = handle.acquire_workspace(required_workspace); + if (!lease.data) { + 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 +517,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{})) { @@ -473,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>( @@ -492,13 +595,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 +612,112 @@ class Operator : public OperatorBase { static std::unique_ptr MakeWithDevice( const Config& config, Device::Type dispatch_device_type, Args&&... args) { std::unique_ptr op_ptr; - auto cache_args = std::forward_as_tuple(args...); + 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; + + 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) { - constexpr std::size_t kImplementationIndex = - decltype(implementation_tag)::value; - if constexpr (std::is_constructible_v< - Operator, - Args...>) { - std::apply( - [&](auto&... cached_args) { - op_ptr = std::make_unique< - Operator>( - cached_args...); - }, - cache_args); - } else { - assert(false && - "operator is not implemented for this device and " - "implementation index"); + 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 kConfigConstructor = + std::is_constructible_v; + constexpr bool kConfigAware = + detail::HasConfigAwareSupports::value; + if constexpr (kConfigConstructor || + std::is_constructible_v) { + bool supported = false; + if constexpr (kConfigAware) { + supported = Provider::supports(config, std::as_const(args)...); + } else if (!config.requires_config_aware_provider()) { + supported = detail::ProviderSupports( + 0, std::as_const(args)...); + } + if (!supported) { + if (automatic) return; + if (strict) { + throw std::invalid_argument( + "requested operator implementation does not support " + "these parameters"); } - }, - "Operator::Make(implementation_index)", - typename ActiveImplementations::type{}); + return; + } + matched = true; + selected(device_tag, implementation_tag, + std::bool_constant{}, 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{}); + 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..091e626ae 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 ( + "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 + 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 aa9b25366..b0254aa35 100644 --- a/tests/test_generate_wrappers.py +++ b/tests/test_generate_wrappers.py @@ -106,6 +106,18 @@ 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 + 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(): module = _load_generator_module() @@ -123,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 ): @@ -371,6 +410,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()))" @@ -427,6 +472,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( @@ -657,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, @@ -793,7 +845,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" @@ -1166,3 +1228,53 @@ 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", + "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" + ) 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