From 66ad19b528026794881f7392eef20fca4a9caf27 Mon Sep 17 00:00:00 2001 From: Ceng23333 <441651826@qq.com> Date: Sat, 10 Oct 2026 09:33:29 +0800 Subject: [PATCH 1/3] feat(qy): add QY ops backend with linked FA and paged KV flash Enable Denglin/QY native and torch paths, linked flash-attn varlen/kvcache providers via *_no_sinks, and reshape_and_cache_flash for paged caching. --- CMakeLists.txt | 103 ++++++++++++++++-- cmake/InfiniOpsConfig.cmake.in | 2 + scripts/dev/platforms.sh | 11 ++ scripts/dlcc_wrapper.sh | 78 +++++++++++++ scripts/generate_torch_ops.py | 2 + src/CMakeLists.txt | 36 +++++- src/device.h | 2 +- src/linked/torch/qy/flash_attn.yaml | 2 + .../ops/flash_attn_varlen_func/flash_attn.cc | 44 ++++++++ .../ops/flash_attn_varlen_func/flash_attn.h | 49 +++++++++ .../flash_attn_varlen_func/flash_attn.yaml | 4 + .../ops/flash_attn_with_kvcache/flash_attn.cc | 50 +++++++++ .../ops/flash_attn_with_kvcache/flash_attn.h | 50 +++++++++ .../flash_attn_with_kvcache/flash_attn.yaml | 4 + src/native/cuda/qy/blas.h | 36 ++++++ src/native/cuda/qy/blas_utils.h | 26 +++++ src/native/cuda/qy/caster.cuh | 77 +++++++++++++ src/native/cuda/qy/data_type_.h | 13 +++ src/native/cuda/qy/device_.h | 6 + src/native/cuda/qy/device_property.h | 42 +++++++ src/native/cuda/qy/ops/add/kernel.h | 19 ++++ src/native/cuda/qy/ops/argmax/kernel.h | 19 ++++ src/native/cuda/qy/ops/convolution/kernel.h | 21 ++++ src/native/cuda/qy/ops/copy/kernel.h | 19 ++++ src/native/cuda/qy/ops/embedding/kernel.h | 19 ++++ src/native/cuda/qy/ops/fill/kernel.h | 21 ++++ .../cuda/qy/ops/fused_add_rms_norm/kernel.h | 19 ++++ src/native/cuda/qy/ops/gelu/kernel.h | 21 ++++ src/native/cuda/qy/ops/gemm/cublas.h | 18 +++ src/native/cuda/qy/ops/mul/kernel.h | 21 ++++ src/native/cuda/qy/ops/relu/kernel.h | 21 ++++ .../qy/ops/reshape_and_cache_flash/kernel.h | 20 ++++ src/native/cuda/qy/ops/rms_norm/kernel.h | 19 ++++ .../cuda/qy/ops/rotary_embedding/kernel.h | 19 ++++ src/native/cuda/qy/ops/sigmoid/kernel.h | 21 ++++ src/native/cuda/qy/ops/silu/kernel.h | 21 ++++ src/native/cuda/qy/ops/silu_and_mul/kernel.h | 19 ++++ src/native/cuda/qy/ops/softmax/kernel.h | 21 ++++ src/native/cuda/qy/ops/topk_softmax/kernel.h | 21 ++++ src/native/cuda/qy/runtime_.h | 9 ++ src/native/cuda/qy/runtime_utils.h | 15 +++ src/pybind11_utils.h | 1 + src/torch/device_.h | 5 + src/torch/qy/c10.h | 28 +++++ tests/conftest.py | 1 + 45 files changed, 1062 insertions(+), 13 deletions(-) create mode 100755 scripts/dlcc_wrapper.sh create mode 100644 src/linked/torch/qy/flash_attn.yaml create mode 100644 src/linked/torch/qy/ops/flash_attn_varlen_func/flash_attn.cc create mode 100644 src/linked/torch/qy/ops/flash_attn_varlen_func/flash_attn.h create mode 100644 src/linked/torch/qy/ops/flash_attn_varlen_func/flash_attn.yaml create mode 100644 src/linked/torch/qy/ops/flash_attn_with_kvcache/flash_attn.cc create mode 100644 src/linked/torch/qy/ops/flash_attn_with_kvcache/flash_attn.h create mode 100644 src/linked/torch/qy/ops/flash_attn_with_kvcache/flash_attn.yaml create mode 100644 src/native/cuda/qy/blas.h create mode 100644 src/native/cuda/qy/blas_utils.h create mode 100644 src/native/cuda/qy/caster.cuh create mode 100644 src/native/cuda/qy/data_type_.h create mode 100644 src/native/cuda/qy/device_.h create mode 100644 src/native/cuda/qy/device_property.h create mode 100644 src/native/cuda/qy/ops/add/kernel.h create mode 100644 src/native/cuda/qy/ops/argmax/kernel.h create mode 100644 src/native/cuda/qy/ops/convolution/kernel.h create mode 100644 src/native/cuda/qy/ops/copy/kernel.h create mode 100644 src/native/cuda/qy/ops/embedding/kernel.h create mode 100644 src/native/cuda/qy/ops/fill/kernel.h create mode 100644 src/native/cuda/qy/ops/fused_add_rms_norm/kernel.h create mode 100644 src/native/cuda/qy/ops/gelu/kernel.h create mode 100644 src/native/cuda/qy/ops/gemm/cublas.h create mode 100644 src/native/cuda/qy/ops/mul/kernel.h create mode 100644 src/native/cuda/qy/ops/relu/kernel.h create mode 100644 src/native/cuda/qy/ops/reshape_and_cache_flash/kernel.h create mode 100644 src/native/cuda/qy/ops/rms_norm/kernel.h create mode 100644 src/native/cuda/qy/ops/rotary_embedding/kernel.h create mode 100644 src/native/cuda/qy/ops/sigmoid/kernel.h create mode 100644 src/native/cuda/qy/ops/silu/kernel.h create mode 100644 src/native/cuda/qy/ops/silu_and_mul/kernel.h create mode 100644 src/native/cuda/qy/ops/softmax/kernel.h create mode 100644 src/native/cuda/qy/ops/topk_softmax/kernel.h create mode 100644 src/native/cuda/qy/runtime_.h create mode 100644 src/native/cuda/qy/runtime_utils.h create mode 100644 src/torch/qy/c10.h diff --git a/CMakeLists.txt b/CMakeLists.txt index 0d8c3e5b3..76e8e8cca 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -22,6 +22,7 @@ option(WITH_MARS "Enable Mars backend" OFF) option(WITH_CAMBRICON "Enable Cambricon backend" OFF) option(WITH_MOORE "Enable Moore backend" OFF) option(WITH_ASCEND "Enable Ascend backend" OFF) +option(WITH_QY "Enable Denglin QY GPU backend" OFF) option(WITH_TORCH "Enable PyTorch C++ backend" OFF) option(WITH_LINKED "Enable linked third-party operator backends" OFF) @@ -173,9 +174,24 @@ if(AUTO_DETECT_DEVICES) message(STATUS "Auto-detected Ascend environment.") endif() + set(_qy_root_detect "") + foreach(_qy_env QY_ROOT QY_HOME QY_PATH) + if(NOT _qy_root_detect AND DEFINED ENV{${_qy_env}} AND + NOT "$ENV{${_qy_env}}" STREQUAL "") + set(_qy_root_detect "$ENV{${_qy_env}}") + endif() + endforeach() + if(NOT _qy_root_detect AND EXISTS "/usr/local/denglin/sdk/bin/dlcc") + set(_qy_root_detect "/usr/local/denglin") + endif() + if(_qy_root_detect AND EXISTS "${_qy_root_detect}/sdk/bin/dlcc") + set(WITH_QY ON) + message(STATUS "Auto-detected Denglin QY environment from ${_qy_root_detect}.") + endif() + if(WITH_NVIDIA) set(_non_nvidia_gpu_detected FALSE) - foreach(_gpu_backend WITH_ILUVATAR WITH_HYGON WITH_THEAD WITH_METAX WITH_MARS WITH_CAMBRICON WITH_MOORE WITH_ASCEND) + foreach(_gpu_backend WITH_ILUVATAR WITH_HYGON WITH_THEAD WITH_METAX WITH_MARS WITH_CAMBRICON WITH_MOORE WITH_ASCEND WITH_QY) if(${_gpu_backend}) set(_non_nvidia_gpu_detected TRUE) endif() @@ -258,29 +274,37 @@ if(WITH_TORCH OR WITH_LINKED) # Query `torch` paths directly instead of using `find_package(Torch)`, # which pulls in Caffe2's CMake config and may fail on platforms with # non-standard CUDA toolchains. + # Denglin/torch may emit extra stdout lines during import; keep only the + # first line of each query so find_library HINTS stay valid. execute_process( - COMMAND ${_TORCH_PYTHON} -c "from torch.utils.cpp_extension import include_paths; print(';'.join(include_paths()))" + COMMAND ${_TORCH_PYTHON} -c "from torch.utils.cpp_extension import include_paths; print(';'.join(include_paths()), flush=True)" OUTPUT_VARIABLE TORCH_INCLUDE_DIRS OUTPUT_STRIP_TRAILING_WHITESPACE RESULT_VARIABLE _torch_result + ERROR_QUIET ) if(NOT _torch_result EQUAL 0) message(FATAL_ERROR "`WITH_TORCH=ON` or `WITH_LINKED=ON` requires an installed `torch` package.") endif() + string(REGEX REPLACE "[\r\n].*" "" TORCH_INCLUDE_DIRS "${TORCH_INCLUDE_DIRS}") execute_process( - COMMAND ${_TORCH_PYTHON} -c "from torch.utils.cpp_extension import library_paths; print(';'.join(library_paths()))" + COMMAND ${_TORCH_PYTHON} -c "from torch.utils.cpp_extension import library_paths; print(';'.join(library_paths()), flush=True)" OUTPUT_VARIABLE _torch_lib_dirs OUTPUT_STRIP_TRAILING_WHITESPACE + ERROR_QUIET ) + string(REGEX REPLACE "[\r\n].*" "" _torch_lib_dirs "${_torch_lib_dirs}") execute_process( - COMMAND ${_TORCH_PYTHON} -c "import pathlib, torch; p = pathlib.Path(torch.__file__).resolve().parent.parent / 'torch.libs'; print(str(p) if p.exists() else '')" + COMMAND ${_TORCH_PYTHON} -c "import pathlib, torch; p = pathlib.Path(torch.__file__).resolve().parent.parent / 'torch.libs'; print(str(p) if p.exists() else '', flush=True)" OUTPUT_VARIABLE _torch_private_lib_dir OUTPUT_STRIP_TRAILING_WHITESPACE + ERROR_QUIET ) + string(REGEX REPLACE "[\r\n].*" "" _torch_private_lib_dir "${_torch_private_lib_dir}") set(TORCH_RUNTIME_DIRS ${_torch_lib_dirs}) if(_torch_private_lib_dir) @@ -292,7 +316,7 @@ if(WITH_TORCH OR WITH_LINKED) find_library(C10_LIB c10 HINTS ${_torch_lib_dirs} REQUIRED) set(TORCH_LIBRARIES ${TORCH_LIB} ${TORCH_CPU_LIB} ${C10_LIB}) - if(WITH_NVIDIA OR WITH_ILUVATAR OR WITH_THEAD OR WITH_MARS) + if(WITH_NVIDIA OR WITH_ILUVATAR OR WITH_THEAD OR WITH_MARS OR WITH_QY) find_library(TORCH_CUDA_LIB torch_cuda HINTS ${_torch_lib_dirs} REQUIRED) find_library(C10_CUDA_LIB c10_cuda HINTS ${_torch_lib_dirs} REQUIRED) # `torch_cuda` registers CUDA kernels through static initializers, so @@ -364,14 +388,14 @@ endif() # Only one CUDA-like GPU backend can be enabled at a time. set(_gpu_backend_count 0) -foreach(_gpu_backend WITH_NVIDIA WITH_ILUVATAR WITH_HYGON WITH_THEAD WITH_METAX WITH_MARS WITH_MOORE WITH_ASCEND) +foreach(_gpu_backend WITH_NVIDIA WITH_ILUVATAR WITH_HYGON WITH_THEAD WITH_METAX WITH_MARS WITH_MOORE WITH_ASCEND WITH_QY) if(${_gpu_backend}) math(EXPR _gpu_backend_count "${_gpu_backend_count} + 1") endif() endforeach() if(_gpu_backend_count GREATER 1) - message(FATAL_ERROR "`WITH_NVIDIA`, `WITH_ILUVATAR`, `WITH_HYGON`, `WITH_THEAD`, `WITH_METAX`, `WITH_MARS`, `WITH_MOORE`, and `WITH_ASCEND` are mutually exclusive. Build one GPU backend at a time.") + message(FATAL_ERROR "`WITH_NVIDIA`, `WITH_ILUVATAR`, `WITH_HYGON`, `WITH_THEAD`, `WITH_METAX`, `WITH_MARS`, `WITH_MOORE`, `WITH_ASCEND`, and `WITH_QY` are mutually exclusive. Build one GPU backend at a time.") endif() if(WITH_NINETOOTHED AND NOT WITH_NVIDIA) @@ -633,12 +657,73 @@ if(WITH_ASCEND) message(STATUS "Using Ascend from `${ASCEND_HOME}`.") endif() +if(WITH_QY) + add_compile_definitions(WITH_QY=1) + + set(QY_ROOT "") + foreach(_qy_env QY_ROOT QY_HOME QY_PATH) + if(NOT QY_ROOT AND DEFINED ENV{${_qy_env}} AND NOT "$ENV{${_qy_env}}" STREQUAL "") + set(QY_ROOT "$ENV{${_qy_env}}") + endif() + endforeach() + if(NOT QY_ROOT AND EXISTS "/usr/local/denglin/sdk/bin/dlcc") + set(QY_ROOT "/usr/local/denglin") + endif() + if(NOT QY_ROOT OR NOT EXISTS "${QY_ROOT}/sdk/bin/dlcc") + message(FATAL_ERROR + "`WITH_QY` is `ON` but Denglin QY SDK was not found. " + "Set `QY_ROOT` (expected `sdk/bin/dlcc`).") + endif() + set(QY_ROOT "${QY_ROOT}" CACHE PATH "Denglin QY SDK root" FORCE) + set(QY_SDK "${QY_ROOT}/sdk") + set(ENV{QY_ROOT} "${QY_ROOT}") + set(ENV{CUDA_HOME} "${QY_SDK}") + set(ENV{CUDA_PATH} "${QY_SDK}") + set(ENV{PATH} "${QY_SDK}/bin:$ENV{PATH}") + + set(CMAKE_CUDA_COMPILER + "${CMAKE_CURRENT_SOURCE_DIR}/scripts/dlcc_wrapper.sh" + CACHE FILEPATH "Denglin QY CUDA compiler (dlcc wrapper)" FORCE) + # Keep CMAKE_CUDA_FLAGS free of --cuda-gpu-arch; the wrapper always injects + # dlgput64 and strips CMake's sm_* detection probes. + set(CMAKE_CUDA_FLAGS "--cuda-path=${QY_SDK} -std=c++17 -fPIC -O2" + CACHE STRING "QY CUDA flags" FORCE) + set(CMAKE_CUDA_SEPARABLE_COMPILATION OFF CACHE BOOL "Disable RDC for QY" FORCE) + # Skip CMake's sm_20/30/52 Clang probes; wrapper maps all arches to dlgput64. + set(CMAKE_CUDA_ARCHITECTURES "80" CACHE STRING "QY placeholder arch" FORCE) + + include_directories("${QY_SDK}/include") + link_directories("${QY_SDK}/lib") + find_library(QY_CURT_LIB NAMES curt HINTS "${QY_SDK}/lib" REQUIRED) + find_library(QY_CUBLAS_LIB NAMES cublas HINTS "${QY_SDK}/lib" REQUIRED) + find_library(QY_CUBLASLT_LIB NAMES cublasLt HINTS "${QY_SDK}/lib" REQUIRED) + find_library(QY_CUDNN_LIB NAMES cudnn HINTS "${QY_SDK}/lib" REQUIRED) + + # Ensure CMake Clang CUDA toolkit discovery finds nvvm/libdevice. + file(MAKE_DIRECTORY "${QY_SDK}/nvvm/libdevice") + set(CMAKE_CUDA_COMPILER_TOOLKIT_ROOT "${QY_SDK}" CACHE PATH + "QY CUDA toolkit root" FORCE) + set(CMAKE_CUDA_COMPILER_LIBRARY_ROOT "${QY_SDK}" CACHE PATH + "QY CUDA library root" FORCE) + set(CUDAToolkit_ROOT "${QY_SDK}" CACHE PATH "QY CUDA toolkit root" FORCE) + + # Final link goes through the host C++ driver, which would otherwise pull + # CMake's implicit -lcudadevrt/-lcudart_static (absent on Denglin). + set(CMAKE_CUDA_RUNTIME_LIBRARY "None" CACHE STRING + "QY links curt explicitly; do not inject CUDA runtime" FORCE) + + message(STATUS "QY: SDK ${QY_SDK}, compiler ${CMAKE_CUDA_COMPILER}") + enable_language(CUDA) + set(CMAKE_CUDA_IMPLICIT_LINK_LIBRARIES "") + set(CMAKE_CUDA_IMPLICIT_LINK_DIRECTORIES "") +endif() + # If all other platforms are not enabled, CPU is enabled by default. -if(NOT WITH_NVIDIA AND NOT WITH_ILUVATAR AND NOT WITH_HYGON AND NOT WITH_THEAD AND NOT WITH_METAX AND NOT WITH_MARS AND NOT WITH_MOORE AND NOT WITH_CAMBRICON AND NOT WITH_ASCEND) +if(NOT WITH_NVIDIA AND NOT WITH_ILUVATAR AND NOT WITH_HYGON AND NOT WITH_THEAD AND NOT WITH_METAX AND NOT WITH_MARS AND NOT WITH_MOORE AND NOT WITH_CAMBRICON AND NOT WITH_ASCEND AND NOT WITH_QY) set(WITH_CPU ON) endif() -if(WITH_TORCH OR WITH_METAX OR WITH_MARS OR WITH_MOORE) +if(WITH_TORCH OR WITH_METAX OR WITH_MARS OR WITH_MOORE OR WITH_QY) set(PYBIND11_ENABLE_EXTRAS OFF) endif() diff --git a/cmake/InfiniOpsConfig.cmake.in b/cmake/InfiniOpsConfig.cmake.in index 943ac6c78..f0254cd43 100644 --- a/cmake/InfiniOpsConfig.cmake.in +++ b/cmake/InfiniOpsConfig.cmake.in @@ -8,5 +8,7 @@ if(@WITH_NVIDIA@ OR @WITH_ILUVATAR@ OR @WITH_HYGON@ OR @WITH_THEAD@) find_dependency(CUDAToolkit) endif() +# WITH_QY links Denglin `curt`/`cublas` directly; no CUDAToolkit dependency. + include("${CMAKE_CURRENT_LIST_DIR}/InfiniOpsTargets.cmake") check_required_components(InfiniOps) diff --git a/scripts/dev/platforms.sh b/scripts/dev/platforms.sh index d5df3ec27..c4c86361e 100644 --- a/scripts/dev/platforms.sh +++ b/scripts/dev/platforms.sh @@ -10,6 +10,7 @@ INFINI_OPS_SUPPORTED_PLATFORMS=( moore cambricon ascend + qy ) infini_ops_supported_platforms_usage() { @@ -87,6 +88,12 @@ _infini_ops_detect_thead() { [[ -x "${cuda_root}/bin/nvcc" ]] } +_infini_ops_detect_qy() { + local qy_root="${QY_ROOT:-${QY_HOME:-${QY_PATH:-/usr/local/denglin}}}" + + [[ -x "${qy_root}/sdk/bin/dlcc" ]] +} + _infini_ops_detect_metax() { if [[ -n "${MACA_PATH:-}" ]]; then return 0 @@ -127,6 +134,10 @@ infini_ops_detect_platforms() { detected+=("thead") fi + if _infini_ops_detect_qy; then + detected+=("qy") + fi + if _infini_ops_detect_metax; then detected+=("metax") fi diff --git a/scripts/dlcc_wrapper.sh b/scripts/dlcc_wrapper.sh new file mode 100755 index 000000000..c1b8aaba3 --- /dev/null +++ b/scripts/dlcc_wrapper.sh @@ -0,0 +1,78 @@ +#!/bin/bash +# Filter CMake/host flags unsupported or harmful when invoking Denglin `dlcc` +# as CMAKE_CUDA_COMPILER. Emulate enough of `nvcc -v` for CMake 3.22 Clang +# CUDA detection (TOP / NVVMIR_LIBRARY_DIR). +set -euo pipefail + +QY_ROOT="${QY_ROOT:-${QY_HOME:-${QY_PATH:-/usr/local/denglin}}}" +DLCC="${QY_ROOT}/sdk/bin/dlcc" +SDK="${QY_ROOT}/sdk" + +if [[ ! -x "${DLCC}" ]]; then + # Container image layout. + if [[ -x /usr/local/dlgpu/sdk/bin/dlcc ]]; then + QY_ROOT=/usr/local/dlgpu + DLCC=/usr/local/dlgpu/sdk/bin/dlcc + SDK=/usr/local/dlgpu/sdk + else + echo "dlcc_wrapper: dlcc not found; set QY_ROOT" >&2 + exit 1 + fi +fi + +# CMakeDetermineCUDACompiler runs `${CMAKE_CUDA_COMPILER} -v __cmake_determine_cuda`. +if [[ "${1:-}" == "-v" ]]; then + echo "clang version 15.0.6" + echo "Cuda compilation tools, release 11.7, V11.7.0" + echo "#\$ TOP=${SDK}" + # Path need only match CMake's `nvvm/libdevice$` regex; Denglin has no nvvm tree. + echo "#\$ NVVMIR_LIBRARY_DIR=${SDK}/nvvm/libdevice" + "${DLCC}" --version || true + exit 0 +fi + +ARGS=() +skip_next=0 +has_cuda_path=0 +replaced_cudart=0 +for arg in "$@"; do + if [[ "${skip_next}" -eq 1 ]]; then + skip_next=0 + continue + fi + case "${arg}" in + -pthread) ;; + -B) skip_next=1 ;; + -B*) ;; + --cuda-path=*) + has_cuda_path=1 + ARGS+=("${arg}") + ;; + # CMake Clang detection probes sm_20/30/52; Denglin only accepts dlgput64. + --cuda-gpu-arch=sm_*) ;; + --cuda-gpu-arch=dlgput64) + ARGS+=("${arg}") + ;; + --cuda-gpu-arch=*) ;; + -Xcompiler=*) ;; + # Match pre-refactor xmake/qy.lua: drop NVIDIA static CUDA libs. + -lcudadevrt|-lcudart_static|-lcudart) + if [[ "${replaced_cudart}" -eq 0 ]]; then + ARGS+=("-lcurt") + replaced_cudart=1 + fi + ;; + *) + ARGS+=("${arg}") + ;; + esac +done + +EXTRA=() +if [[ "${has_cuda_path}" -eq 0 ]]; then + EXTRA+=("--cuda-path=${SDK}") +fi +# Always compile for Denglin GPU; ignore CMake sm_* probes. +EXTRA+=("--cuda-gpu-arch=dlgput64") + +exec "${DLCC}" "${EXTRA[@]}" "${ARGS[@]}" diff --git a/scripts/generate_torch_ops.py b/scripts/generate_torch_ops.py index 9b44c5e2a..481b83cf0 100644 --- a/scripts/generate_torch_ops.py +++ b/scripts/generate_torch_ops.py @@ -77,6 +77,7 @@ "kHygon", "kThead", "kMars", + "kQy", ) _C10_DEVICE_TYPES = ( @@ -88,6 +89,7 @@ "kIluvatar", "kThead", "kMars", + "kQy", ) # YAML scalar-type tokens → C++ types. Reference types (e.g. `const Scalar&`) diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index ce9f75770..747ee72b2 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -315,6 +315,36 @@ if(WITH_THEAD) list(APPEND DEVICE_LIST "thead") endif() +if(WITH_QY) + set(QY_PATTERNS + "native/cuda/qy/*.cc" + "native/cuda/qy/*.cpp" + "native/cuda/qy/*.cu" + ) + + file(GLOB_RECURSE QY_SOURCES CONFIGURE_DEPENDS ${QY_PATTERNS}) + + enable_language(CUDA) + + target_compile_definitions(infiniops PUBLIC WITH_QY=1) + target_sources(infiniops PRIVATE ${QY_SOURCES}) + target_include_directories(infiniops PUBLIC "${QY_SDK}/include") + target_link_libraries(infiniops PUBLIC + ${QY_CURT_LIB} + ${QY_CUBLAS_LIB} + ${QY_CUBLASLT_LIB} + ${QY_CUDNN_LIB}) + + set_target_properties(infiniops PROPERTIES + CUDA_STANDARD 17 + CUDA_STANDARD_REQUIRED ON + CUDA_RUNTIME_LIBRARY None + CUDA_SEPARABLE_COMPILATION OFF + ) + + list(APPEND DEVICE_LIST "qy") +endif() + if(WITH_METAX) set(METAX_PATTERNS "native/cuda/*.cc" @@ -1270,7 +1300,7 @@ if(GENERATE_OPERATOR_CALL_INSTANTIATIONS) file(GLOB_RECURSE OPERATOR_CALL_INSTANTIATION_SOURCES CONFIGURE_DEPENDS "${PROJECT_SOURCE_DIR}/generated/src/operator_call_instantiations_*.cc") - if(WITH_NVIDIA OR WITH_HYGON OR WITH_THEAD) + if(WITH_NVIDIA OR WITH_HYGON OR WITH_THEAD OR WITH_QY) set_source_files_properties(${OPERATOR_CALL_INSTANTIATION_SOURCES} PROPERTIES LANGUAGE CUDA) target_sources(infiniops PRIVATE ${OPERATOR_CALL_INSTANTIATION_SOURCES}) @@ -1346,7 +1376,7 @@ if(GENERATE_OPERATOR_CALL_INSTANTIATIONS) if(OPERATOR_DISPATCH_SOURCES) target_include_directories(infiniops PRIVATE "${PROJECT_SOURCE_DIR}") - if(WITH_NVIDIA OR WITH_HYGON OR WITH_THEAD) + if(WITH_NVIDIA OR WITH_HYGON OR WITH_THEAD OR WITH_QY) set_source_files_properties(${OPERATOR_DISPATCH_SOURCES} PROPERTIES LANGUAGE CUDA) target_sources(infiniops PRIVATE ${OPERATOR_DISPATCH_SOURCES}) @@ -1454,7 +1484,7 @@ if(GENERATE_PYTHON_BINDINGS) endif() list(APPEND PYBIND11_COMPILE_SOURCES ${PYBIND11_DISPATCH_SOURCES}) - if(WITH_NVIDIA OR WITH_HYGON OR WITH_THEAD) + if(WITH_NVIDIA OR WITH_HYGON OR WITH_THEAD OR WITH_QY) set_source_files_properties(${PYBIND11_COMPILE_SOURCES} PROPERTIES LANGUAGE CUDA) elseif(WITH_ILUVATAR) set(_iluvatar_dispatch_include_flags diff --git a/src/device.h b/src/device.h index 6773b77ad..4f3035b9b 100644 --- a/src/device.h +++ b/src/device.h @@ -16,7 +16,7 @@ using AllDeviceTypes = List; + Device::Type::kThead, Device::Type::kQy>; template struct ActiveDevicesImpl { diff --git a/src/linked/torch/qy/flash_attn.yaml b/src/linked/torch/qy/flash_attn.yaml new file mode 100644 index 000000000..c54d82dbd --- /dev/null +++ b/src/linked/torch/qy/flash_attn.yaml @@ -0,0 +1,2 @@ +python_distribution_package: flash-attn +library_glob: flash_attn_2_cuda*.so diff --git a/src/linked/torch/qy/ops/flash_attn_varlen_func/flash_attn.cc b/src/linked/torch/qy/ops/flash_attn_varlen_func/flash_attn.cc new file mode 100644 index 000000000..69ec2ddc9 --- /dev/null +++ b/src/linked/torch/qy/ops/flash_attn_varlen_func/flash_attn.cc @@ -0,0 +1,44 @@ +#include "linked/torch/qy/ops/flash_attn_varlen_func/flash_attn.h" + +namespace flash { + +std::vector mha_varlen_fwd_no_sinks( + at::Tensor& q, const at::Tensor& k, const at::Tensor& v, + std::optional& out, const at::Tensor& cu_seqlens_q, + const at::Tensor& cu_seqlens_k, std::optional& seqused_k, + std::optional& leftpad_k, + std::optional& block_table, + std::optional& alibi_slopes, int max_seqlen_q, int max_seqlen_k, + float dropout_p, float softmax_scale, bool zero_tensors, bool causal, + int window_size_left, int window_size_right, float softcap, + bool return_softmax, std::optional generator); + +} // namespace flash + +namespace infini::ops::linked::torch::qy { + +std::vector FlashAttnVarlen::Call( + at::Tensor& q, const at::Tensor& k, const at::Tensor& v, + std::optional& out, const at::Tensor& cu_seqlens_q, + const at::Tensor& cu_seqlens_k, std::optional& seqused_k, + std::optional& leftpad_k, + std::optional& block_table, + std::optional& alibi_slopes, int max_seqlen_q, int max_seqlen_k, + float dropout_p, float softmax_scale, bool zero_tensors, bool causal, + int window_size_left, int window_size_right, float softcap, + bool return_softmax, std::optional generator) { + return flash::mha_varlen_fwd_no_sinks( + q, k, v, out, cu_seqlens_q, cu_seqlens_k, seqused_k, leftpad_k, + block_table, alibi_slopes, max_seqlen_q, max_seqlen_k, dropout_p, + softmax_scale, zero_tensors, causal, window_size_left, window_size_right, + softcap, return_softmax, generator); +} + +} // namespace infini::ops::linked::torch::qy + +namespace infini::ops::linked::torch { + +template class TorchFlashAttnVarlenFunc< + ::infini::ops::linked::torch::qy::FlashAttnVarlen>; + +} // namespace infini::ops::linked::torch diff --git a/src/linked/torch/qy/ops/flash_attn_varlen_func/flash_attn.h b/src/linked/torch/qy/ops/flash_attn_varlen_func/flash_attn.h new file mode 100644 index 000000000..0a7cbf712 --- /dev/null +++ b/src/linked/torch/qy/ops/flash_attn_varlen_func/flash_attn.h @@ -0,0 +1,49 @@ +#ifndef INFINI_OPS_LINKED_TORCH_QY_OPS_FLASH_ATTN_VARLEN_FUNC_FLASH_ATTN_H_ +#define INFINI_OPS_LINKED_TORCH_QY_OPS_FLASH_ATTN_VARLEN_FUNC_FLASH_ATTN_H_ + +#include + +#include "linked/torch/ops/flash_attn_varlen_func.h" +#include "torch/qy/c10.h" + +namespace infini::ops::linked::torch::qy { + +struct FlashAttnVarlen : C10 { + static std::vector Call( + at::Tensor& q, const at::Tensor& k, const at::Tensor& v, + std::optional& out, const at::Tensor& cu_seqlens_q, + const at::Tensor& cu_seqlens_k, std::optional& seqused_k, + std::optional& leftpad_k, + std::optional& block_table, + std::optional& alibi_slopes, int max_seqlen_q, + int max_seqlen_k, float dropout_p, float softmax_scale, bool zero_tensors, + bool causal, int window_size_left, int window_size_right, float softcap, + bool return_softmax, std::optional generator); +}; + +} // namespace infini::ops::linked::torch::qy + +namespace infini::ops::linked::torch { + +extern template class TorchFlashAttnVarlenFunc< + ::infini::ops::linked::torch::qy::FlashAttnVarlen>; + +} // namespace infini::ops::linked::torch + +namespace infini::ops { + +template <> +class Operator + : public linked::torch::TorchFlashAttnVarlenFunc< + linked::torch::qy::FlashAttnVarlen> { + public: + using linked::torch::TorchFlashAttnVarlenFunc< + linked::torch::qy::FlashAttnVarlen>::TorchFlashAttnVarlenFunc; + + using linked::torch::TorchFlashAttnVarlenFunc< + linked::torch::qy::FlashAttnVarlen>::operator(); +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_LINKED_TORCH_QY_OPS_FLASH_ATTN_VARLEN_FUNC_FLASH_ATTN_H_ diff --git a/src/linked/torch/qy/ops/flash_attn_varlen_func/flash_attn.yaml b/src/linked/torch/qy/ops/flash_attn_varlen_func/flash_attn.yaml new file mode 100644 index 000000000..d1481af1b --- /dev/null +++ b/src/linked/torch/qy/ops/flash_attn_varlen_func/flash_attn.yaml @@ -0,0 +1,4 @@ +library: flash_attn +required_symbols: + - >- + flash::mha_varlen_fwd_no_sinks(at::Tensor&, at::Tensor const&, at::Tensor const&, std::optional&, at::Tensor const&, at::Tensor const&, std::optional&, std::optional&, std::optional&, std::optional&, int, int, float, float, bool, bool, int, int, float, bool, std::optional) diff --git a/src/linked/torch/qy/ops/flash_attn_with_kvcache/flash_attn.cc b/src/linked/torch/qy/ops/flash_attn_with_kvcache/flash_attn.cc new file mode 100644 index 000000000..e2c4acc31 --- /dev/null +++ b/src/linked/torch/qy/ops/flash_attn_with_kvcache/flash_attn.cc @@ -0,0 +1,50 @@ +#include "linked/torch/qy/ops/flash_attn_with_kvcache/flash_attn.h" + +namespace flash { + +std::vector mha_fwd_kvcache_no_sinks( + at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache, + std::optional& k, std::optional& v, + std::optional& cache_seqlens, + std::optional& rotary_cos, + std::optional& rotary_sin, + std::optional& cache_batch_idx, + std::optional& cache_leftpad, + std::optional& block_table, + std::optional& alibi_slopes, std::optional& out, + float softmax_scale, bool causal, int window_size_left, + int window_size_right, float softcap, bool rotary_interleaved, + int num_splits); + +} // namespace flash + +namespace infini::ops::linked::torch::qy { + +std::vector FlashAttnKvcache::Call( + at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache, + std::optional& k, std::optional& v, + std::optional& cache_seqlens, + std::optional& rotary_cos, + std::optional& rotary_sin, + std::optional& cache_batch_idx, + std::optional& cache_leftpad, + std::optional& block_table, + std::optional& alibi_slopes, std::optional& out, + float softmax_scale, bool causal, int window_size_left, + int window_size_right, float softcap, bool rotary_interleaved, + int num_splits) { + return flash::mha_fwd_kvcache_no_sinks( + q, k_cache, v_cache, k, v, cache_seqlens, rotary_cos, rotary_sin, + cache_batch_idx, cache_leftpad, block_table, alibi_slopes, out, + softmax_scale, causal, window_size_left, window_size_right, softcap, + rotary_interleaved, num_splits); +} + +} // namespace infini::ops::linked::torch::qy + +namespace infini::ops::linked::torch { + +template class TorchFlashAttnWithKvcache< + ::infini::ops::linked::torch::qy::FlashAttnKvcache>; + +} // namespace infini::ops::linked::torch diff --git a/src/linked/torch/qy/ops/flash_attn_with_kvcache/flash_attn.h b/src/linked/torch/qy/ops/flash_attn_with_kvcache/flash_attn.h new file mode 100644 index 000000000..11c7e397f --- /dev/null +++ b/src/linked/torch/qy/ops/flash_attn_with_kvcache/flash_attn.h @@ -0,0 +1,50 @@ +#ifndef INFINI_OPS_LINKED_TORCH_QY_OPS_FLASH_ATTN_WITH_KVCACHE_FLASH_ATTN_H_ +#define INFINI_OPS_LINKED_TORCH_QY_OPS_FLASH_ATTN_WITH_KVCACHE_FLASH_ATTN_H_ + +#include "linked/torch/ops/flash_attn_with_kvcache.h" +#include "torch/qy/c10.h" + +namespace infini::ops::linked::torch::qy { + +struct FlashAttnKvcache : C10 { + static std::vector Call( + at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache, + std::optional& k, std::optional& v, + std::optional& cache_seqlens, + std::optional& rotary_cos, + std::optional& rotary_sin, + std::optional& cache_batch_idx, + std::optional& cache_leftpad, + std::optional& block_table, + std::optional& alibi_slopes, std::optional& out, + float softmax_scale, bool causal, int window_size_left, + int window_size_right, float softcap, bool rotary_interleaved, + int num_splits); +}; + +} // namespace infini::ops::linked::torch::qy + +namespace infini::ops::linked::torch { + +extern template class TorchFlashAttnWithKvcache< + ::infini::ops::linked::torch::qy::FlashAttnKvcache>; + +} // namespace infini::ops::linked::torch + +namespace infini::ops { + +template <> +class Operator + : public linked::torch::TorchFlashAttnWithKvcache< + linked::torch::qy::FlashAttnKvcache> { + public: + using linked::torch::TorchFlashAttnWithKvcache< + linked::torch::qy::FlashAttnKvcache>::TorchFlashAttnWithKvcache; + + using linked::torch::TorchFlashAttnWithKvcache< + linked::torch::qy::FlashAttnKvcache>::operator(); +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_LINKED_TORCH_QY_OPS_FLASH_ATTN_WITH_KVCACHE_FLASH_ATTN_H_ diff --git a/src/linked/torch/qy/ops/flash_attn_with_kvcache/flash_attn.yaml b/src/linked/torch/qy/ops/flash_attn_with_kvcache/flash_attn.yaml new file mode 100644 index 000000000..cec65216d --- /dev/null +++ b/src/linked/torch/qy/ops/flash_attn_with_kvcache/flash_attn.yaml @@ -0,0 +1,4 @@ +library: flash_attn +required_symbols: + - >- + flash::mha_fwd_kvcache_no_sinks(at::Tensor&, at::Tensor const&, at::Tensor const&, std::optional&, std::optional&, std::optional&, std::optional&, std::optional&, std::optional&, std::optional&, std::optional&, std::optional&, std::optional&, float, bool, int, int, float, bool, int) diff --git a/src/native/cuda/qy/blas.h b/src/native/cuda/qy/blas.h new file mode 100644 index 000000000..5846a8348 --- /dev/null +++ b/src/native/cuda/qy/blas.h @@ -0,0 +1,36 @@ +#ifndef INFINI_OPS_QY_BLAS_H_ +#define INFINI_OPS_QY_BLAS_H_ + +#include + +// clang-format off +#include +// clang-format on + +#include "data_type.h" +#include "native/cuda/blas.h" +#include "native/cuda/qy/blas_utils.h" +#include "native/cuda/qy/runtime_.h" + +namespace infini::ops { + +template <> +struct Blas : public Runtime { + using BlasHandle = cublasHandle_t; + + static constexpr auto BLAS_OP_N = CUBLAS_OP_N; + static constexpr auto BLAS_OP_T = CUBLAS_OP_T; + static constexpr auto BLAS_GEMM_DEFAULT = CUBLAS_GEMM_DEFAULT; + + static constexpr auto BlasCreate = cublasCreate; + static constexpr auto BlasSetStream = cublasSetStream; + static constexpr auto BlasDestroy = cublasDestroy; + + static constexpr auto BlasGemmStridedBatchedEx = [](auto&&... args) { + return cublasGemmStridedBatchedEx(std::forward(args)...); + }; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/qy/blas_utils.h b/src/native/cuda/qy/blas_utils.h new file mode 100644 index 000000000..fe1f0e718 --- /dev/null +++ b/src/native/cuda/qy/blas_utils.h @@ -0,0 +1,26 @@ +#ifndef INFINI_OPS_QY_BLAS_UTILS_H_ +#define INFINI_OPS_QY_BLAS_UTILS_H_ + +// clang-format off +#include +// clang-format on + +#include "data_type.h" +#include "native/cuda/blas_utils.h" + +namespace infini::ops { + +template <> +struct BlasUtils { + static auto GetDataType(DataType dtype) { + if (dtype == DataType::kFloat16) return CUDA_R_16F; + if (dtype == DataType::kBFloat16) return CUDA_R_16BF; + return CUDA_R_32F; + } + + static auto GetComputeType(DataType) { return CUBLAS_COMPUTE_32F; } +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/qy/caster.cuh b/src/native/cuda/qy/caster.cuh new file mode 100644 index 000000000..26aee2063 --- /dev/null +++ b/src/native/cuda/qy/caster.cuh @@ -0,0 +1,77 @@ +#ifndef INFINI_OPS_QY_CASTER__H_ +#define INFINI_OPS_QY_CASTER__H_ + +#include "native/cuda/caster.cuh" +#include "native/cuda/qy/data_type_.h" + +namespace infini::ops { +namespace detail { + +template <> +struct ToFloat { + __host__ __device__ float operator()(half x) { return __half2float(x); } +}; + +template <> +struct ToFloat { + __host__ __device__ float operator()(__nv_bfloat16 x) { + return __bfloat162float(x); + } +}; + +template <> +struct FromFloat { + __host__ __device__ half operator()(float f) { return __float2half(f); } +}; + +template <> +struct FromFloat { + __host__ __device__ __nv_bfloat16 operator()(float f) { + return __float2bfloat16(f); + } +}; + +template <> +struct HardwareCast { + inline static constexpr bool kSupported = true; + __host__ __device__ __nv_bfloat16 operator()(int x) { + return __int2bfloat16_rn(x); + } +}; + +template <> +struct HardwareCast { + inline static constexpr bool kSupported = true; + __host__ __device__ half operator()(int x) { return __int2half_rn(x); } +}; + +template <> +struct HardwareCast { + inline static constexpr bool kSupported = true; + __host__ __device__ __nv_bfloat16 operator()(double x) { + return __double2bfloat16(x); + } +}; + +template <> +struct HardwareCast { + inline static constexpr bool kSupported = true; + __host__ __device__ half operator()(double x) { return __double2half(x); } +}; + +template <> +struct HardwareCast { + inline static constexpr bool kSupported = true; + __host__ __device__ half operator()(__nv_bfloat16 x) { + return __float2half(__bfloat162float(x)); + } +}; + +} // namespace detail + +template <> +struct Caster : CudaCasterImpl {}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/qy/data_type_.h b/src/native/cuda/qy/data_type_.h new file mode 100644 index 000000000..2861876b4 --- /dev/null +++ b/src/native/cuda/qy/data_type_.h @@ -0,0 +1,13 @@ +#ifndef INFINI_OPS_QY_DATA_TYPE__H_ +#define INFINI_OPS_QY_DATA_TYPE__H_ + +#include + +namespace infini::ops { + +using infini::rt::cuda_bfloat16; +using infini::rt::cuda_bfloat162; + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/qy/device_.h b/src/native/cuda/qy/device_.h new file mode 100644 index 000000000..69a5fc68a --- /dev/null +++ b/src/native/cuda/qy/device_.h @@ -0,0 +1,6 @@ +#ifndef INFINI_OPS_QY_DEVICE__H_ +#define INFINI_OPS_QY_DEVICE__H_ + +#include + +#endif diff --git a/src/native/cuda/qy/device_property.h b/src/native/cuda/qy/device_property.h new file mode 100644 index 000000000..119b0f8c4 --- /dev/null +++ b/src/native/cuda/qy/device_property.h @@ -0,0 +1,42 @@ +#ifndef INFINI_OPS_QY_DEVICE_PROPERTY_H_ +#define INFINI_OPS_QY_DEVICE_PROPERTY_H_ + +#include + +#include +#include + +namespace infini::ops { + +class DevicePropertyCache { + public: + static const cudaDeviceProp& GetCurrentDeviceProps() { + int device_id = 0; + cudaGetDevice(&device_id); + return GetDeviceProps(device_id); + } + + static const cudaDeviceProp& GetDeviceProps(int device_id) { + static std::vector cache = []() { + int count = 0; + cudaGetDeviceCount(&count); + if (count == 0) return std::vector{}; + std::vector props(count); + for (int i = 0; i < count; ++i) { + cudaGetDeviceProperties(&props[i], i); + } + return props; + }(); + + assert(device_id >= 0 && device_id < static_cast(cache.size())); + return cache[device_id]; + } +}; + +inline int QueryMaxThreadsPerBlock() { + return DevicePropertyCache::GetCurrentDeviceProps().maxThreadsPerBlock; +} + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/qy/ops/add/kernel.h b/src/native/cuda/qy/ops/add/kernel.h new file mode 100644 index 000000000..bd679a30e --- /dev/null +++ b/src/native/cuda/qy/ops/add/kernel.h @@ -0,0 +1,19 @@ +#ifndef INFINI_OPS_QY_ADD_KERNEL_H_ +#define INFINI_OPS_QY_ADD_KERNEL_H_ + +#include "native/cuda/ops/add/kernel.h" +#include "native/cuda/qy/caster.cuh" +#include "native/cuda/qy/runtime_.h" + +namespace infini::ops { + +template <> +class Operator + : public CudaAdd> { + public: + using CudaAdd>::CudaAdd; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/qy/ops/argmax/kernel.h b/src/native/cuda/qy/ops/argmax/kernel.h new file mode 100644 index 000000000..121ce6210 --- /dev/null +++ b/src/native/cuda/qy/ops/argmax/kernel.h @@ -0,0 +1,19 @@ +#ifndef INFINI_OPS_QY_ARGMAX_KERNEL_H_ +#define INFINI_OPS_QY_ARGMAX_KERNEL_H_ + +#include "native/cuda/ops/argmax/kernel.h" +#include "native/cuda/qy/caster.cuh" +#include "native/cuda/qy/runtime_.h" + +namespace infini::ops { + +template <> +class Operator + : public CudaArgmax> { + public: + using CudaArgmax>::CudaArgmax; +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_QY_ARGMAX_KERNEL_H_ diff --git a/src/native/cuda/qy/ops/convolution/kernel.h b/src/native/cuda/qy/ops/convolution/kernel.h new file mode 100644 index 000000000..2bba9ec7e --- /dev/null +++ b/src/native/cuda/qy/ops/convolution/kernel.h @@ -0,0 +1,21 @@ +#ifndef INFINI_OPS_QY_CONVOLUTION_KERNEL_H_ +#define INFINI_OPS_QY_CONVOLUTION_KERNEL_H_ + +#include + +#include "native/cuda/ops/convolution/kernel.h" +#include "native/cuda/qy/caster.cuh" +#include "native/cuda/qy/runtime_.h" + +namespace infini::ops { + +template <> +class Operator + : public CudaConv, Convolution> { + public: + using CudaConv, Convolution>::CudaConv; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/qy/ops/copy/kernel.h b/src/native/cuda/qy/ops/copy/kernel.h new file mode 100644 index 000000000..f0977dc55 --- /dev/null +++ b/src/native/cuda/qy/ops/copy/kernel.h @@ -0,0 +1,19 @@ +#ifndef INFINI_OPS_QY_COPY_KERNEL_H_ +#define INFINI_OPS_QY_COPY_KERNEL_H_ + +#include "native/cuda/ops/copy/kernel.h" +#include "native/cuda/qy/caster.cuh" +#include "native/cuda/qy/runtime_.h" + +namespace infini::ops { + +template <> +class Operator + : public CudaCopy> { + public: + using CudaCopy>::CudaCopy; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/qy/ops/embedding/kernel.h b/src/native/cuda/qy/ops/embedding/kernel.h new file mode 100644 index 000000000..3ea2870c3 --- /dev/null +++ b/src/native/cuda/qy/ops/embedding/kernel.h @@ -0,0 +1,19 @@ +#ifndef INFINI_OPS_QY_EMBEDDING_KERNEL_H_ +#define INFINI_OPS_QY_EMBEDDING_KERNEL_H_ + +#include "native/cuda/ops/embedding/kernel.h" +#include "native/cuda/qy/caster.cuh" +#include "native/cuda/qy/runtime_.h" + +namespace infini::ops { + +template <> +class Operator + : public CudaEmbedding> { + public: + using CudaEmbedding>::CudaEmbedding; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/qy/ops/fill/kernel.h b/src/native/cuda/qy/ops/fill/kernel.h new file mode 100644 index 000000000..1c8a03195 --- /dev/null +++ b/src/native/cuda/qy/ops/fill/kernel.h @@ -0,0 +1,21 @@ +#ifndef INFINI_OPS_QY_FILL_KERNEL_H_ +#define INFINI_OPS_QY_FILL_KERNEL_H_ + +#include + +#include "native/cuda/qy/caster.cuh" +#include "native/cuda/qy/runtime_.h" +#include "native/cuda/ops/fill/kernel.h" + +namespace infini::ops { + +template <> +class Operator + : public CudaFill> { + public: + using CudaFill>::CudaFill; +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_QY_FILL_KERNEL_H_ diff --git a/src/native/cuda/qy/ops/fused_add_rms_norm/kernel.h b/src/native/cuda/qy/ops/fused_add_rms_norm/kernel.h new file mode 100644 index 000000000..446e826c4 --- /dev/null +++ b/src/native/cuda/qy/ops/fused_add_rms_norm/kernel.h @@ -0,0 +1,19 @@ +#ifndef INFINI_OPS_QY_FUSED_ADD_RMS_NORM_KERNEL_H_ +#define INFINI_OPS_QY_FUSED_ADD_RMS_NORM_KERNEL_H_ + +#include "native/cuda/ops/fused_add_rms_norm/kernel.h" +#include "native/cuda/qy/caster.cuh" +#include "native/cuda/qy/runtime_.h" + +namespace infini::ops { + +template <> +class Operator + : public CudaFusedAddRmsNorm> { + public: + using CudaFusedAddRmsNorm>::CudaFusedAddRmsNorm; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/qy/ops/gelu/kernel.h b/src/native/cuda/qy/ops/gelu/kernel.h new file mode 100644 index 000000000..6dd6bff58 --- /dev/null +++ b/src/native/cuda/qy/ops/gelu/kernel.h @@ -0,0 +1,21 @@ +#ifndef INFINI_OPS_QY_GELU_KERNEL_H_ +#define INFINI_OPS_QY_GELU_KERNEL_H_ + +#include + +#include "native/cuda/ops/gelu/kernel.h" +#include "native/cuda/qy/caster.cuh" +#include "native/cuda/qy/runtime_.h" + +namespace infini::ops { + +template <> +class Operator + : public CudaGelu> { + public: + using CudaGelu>::CudaGelu; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/qy/ops/gemm/cublas.h b/src/native/cuda/qy/ops/gemm/cublas.h new file mode 100644 index 000000000..269500906 --- /dev/null +++ b/src/native/cuda/qy/ops/gemm/cublas.h @@ -0,0 +1,18 @@ +#ifndef INFINI_OPS_QY_GEMM_CUBLAS_H_ +#define INFINI_OPS_QY_GEMM_CUBLAS_H_ + +#include "native/cuda/ops/gemm/blas.h" +#include "native/cuda/qy/blas.h" + +namespace infini::ops { + +template <> +class Operator + : public BlasGemm> { + public: + using BlasGemm>::BlasGemm; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/qy/ops/mul/kernel.h b/src/native/cuda/qy/ops/mul/kernel.h new file mode 100644 index 000000000..f7a838717 --- /dev/null +++ b/src/native/cuda/qy/ops/mul/kernel.h @@ -0,0 +1,21 @@ +#ifndef INFINI_OPS_QY_MUL_KERNEL_H_ +#define INFINI_OPS_QY_MUL_KERNEL_H_ + +#include + +#include "native/cuda/qy/caster.cuh" +#include "native/cuda/qy/runtime_.h" +#include "native/cuda/ops/mul/kernel.h" + +namespace infini::ops { + +template <> +class Operator + : public CudaMul> { + public: + using CudaMul>::CudaMul; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/qy/ops/relu/kernel.h b/src/native/cuda/qy/ops/relu/kernel.h new file mode 100644 index 000000000..735d82a4b --- /dev/null +++ b/src/native/cuda/qy/ops/relu/kernel.h @@ -0,0 +1,21 @@ +#ifndef INFINI_OPS_QY_RELU_KERNEL_H_ +#define INFINI_OPS_QY_RELU_KERNEL_H_ + +#include + +#include "native/cuda/ops/relu/kernel.h" +#include "native/cuda/qy/caster.cuh" +#include "native/cuda/qy/runtime_.h" + +namespace infini::ops { + +template <> +class Operator + : public CudaRelu> { + public: + using CudaRelu>::CudaRelu; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/qy/ops/reshape_and_cache_flash/kernel.h b/src/native/cuda/qy/ops/reshape_and_cache_flash/kernel.h new file mode 100644 index 000000000..25eefe556 --- /dev/null +++ b/src/native/cuda/qy/ops/reshape_and_cache_flash/kernel.h @@ -0,0 +1,20 @@ +#ifndef INFINI_OPS_QY_RESHAPE_AND_CACHE_FLASH_KERNEL_H_ +#define INFINI_OPS_QY_RESHAPE_AND_CACHE_FLASH_KERNEL_H_ + +#include "native/cuda/ops/reshape_and_cache_flash/kernel.h" +#include "native/cuda/qy/caster.cuh" +#include "native/cuda/qy/runtime_.h" + +namespace infini::ops { + +template <> +class Operator + : public CudaReshapeAndCacheFlash> { + public: + using CudaReshapeAndCacheFlash< + Runtime>::CudaReshapeAndCacheFlash; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/qy/ops/rms_norm/kernel.h b/src/native/cuda/qy/ops/rms_norm/kernel.h new file mode 100644 index 000000000..576326489 --- /dev/null +++ b/src/native/cuda/qy/ops/rms_norm/kernel.h @@ -0,0 +1,19 @@ +#ifndef INFINI_OPS_QY_RMS_NORM_KERNEL_H_ +#define INFINI_OPS_QY_RMS_NORM_KERNEL_H_ + +#include "native/cuda/ops/rms_norm/kernel.h" +#include "native/cuda/qy/caster.cuh" +#include "native/cuda/qy/runtime_.h" + +namespace infini::ops { + +template <> +class Operator + : public CudaRmsNorm> { + public: + using CudaRmsNorm>::CudaRmsNorm; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/qy/ops/rotary_embedding/kernel.h b/src/native/cuda/qy/ops/rotary_embedding/kernel.h new file mode 100644 index 000000000..2bab78145 --- /dev/null +++ b/src/native/cuda/qy/ops/rotary_embedding/kernel.h @@ -0,0 +1,19 @@ +#ifndef INFINI_OPS_QY_ROTARY_EMBEDDING_KERNEL_H_ +#define INFINI_OPS_QY_ROTARY_EMBEDDING_KERNEL_H_ + +#include "native/cuda/ops/rotary_embedding/kernel.h" +#include "native/cuda/qy/caster.cuh" +#include "native/cuda/qy/runtime_.h" + +namespace infini::ops { + +template <> +class Operator + : public CudaRotaryEmbedding> { + public: + using CudaRotaryEmbedding>::CudaRotaryEmbedding; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/qy/ops/sigmoid/kernel.h b/src/native/cuda/qy/ops/sigmoid/kernel.h new file mode 100644 index 000000000..dd7918cf2 --- /dev/null +++ b/src/native/cuda/qy/ops/sigmoid/kernel.h @@ -0,0 +1,21 @@ +#ifndef INFINI_OPS_QY_SIGMOID_KERNEL_H_ +#define INFINI_OPS_QY_SIGMOID_KERNEL_H_ + +#include + +#include "native/cuda/ops/sigmoid/kernel.h" +#include "native/cuda/qy/caster.cuh" +#include "native/cuda/qy/runtime_.h" + +namespace infini::ops { + +template <> +class Operator + : public CudaSigmoid> { + public: + using CudaSigmoid>::CudaSigmoid; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/qy/ops/silu/kernel.h b/src/native/cuda/qy/ops/silu/kernel.h new file mode 100644 index 000000000..310820424 --- /dev/null +++ b/src/native/cuda/qy/ops/silu/kernel.h @@ -0,0 +1,21 @@ +#ifndef INFINI_OPS_QY_SILU_KERNEL_H_ +#define INFINI_OPS_QY_SILU_KERNEL_H_ + +#include + +#include "native/cuda/ops/silu/kernel.h" +#include "native/cuda/qy/caster.cuh" +#include "native/cuda/qy/runtime_.h" + +namespace infini::ops { + +template <> +class Operator + : public CudaSilu> { + public: + using CudaSilu>::CudaSilu; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/qy/ops/silu_and_mul/kernel.h b/src/native/cuda/qy/ops/silu_and_mul/kernel.h new file mode 100644 index 000000000..ba413e1e9 --- /dev/null +++ b/src/native/cuda/qy/ops/silu_and_mul/kernel.h @@ -0,0 +1,19 @@ +#ifndef INFINI_OPS_QY_SILU_AND_MUL_KERNEL_H_ +#define INFINI_OPS_QY_SILU_AND_MUL_KERNEL_H_ + +#include "native/cuda/ops/silu_and_mul/kernel.h" +#include "native/cuda/qy/caster.cuh" +#include "native/cuda/qy/runtime_.h" + +namespace infini::ops { + +template <> +class Operator + : public CudaSiluAndMul> { + public: + using CudaSiluAndMul>::CudaSiluAndMul; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/qy/ops/softmax/kernel.h b/src/native/cuda/qy/ops/softmax/kernel.h new file mode 100644 index 000000000..a00a27183 --- /dev/null +++ b/src/native/cuda/qy/ops/softmax/kernel.h @@ -0,0 +1,21 @@ +#ifndef INFINI_OPS_QY_SOFTMAX_KERNEL_H_ +#define INFINI_OPS_QY_SOFTMAX_KERNEL_H_ + +#include + +#include "native/cuda/ops/softmax/kernel.h" +#include "native/cuda/qy/caster.cuh" +#include "native/cuda/qy/runtime_.h" + +namespace infini::ops { + +template <> +class Operator + : public CudaSoftmax> { + public: + using CudaSoftmax>::CudaSoftmax; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/qy/ops/topk_softmax/kernel.h b/src/native/cuda/qy/ops/topk_softmax/kernel.h new file mode 100644 index 000000000..efe498708 --- /dev/null +++ b/src/native/cuda/qy/ops/topk_softmax/kernel.h @@ -0,0 +1,21 @@ +#ifndef INFINI_OPS_QY_TOPK_SOFTMAX_KERNEL_H_ +#define INFINI_OPS_QY_TOPK_SOFTMAX_KERNEL_H_ + +#include + +#include "native/cuda/ops/topk_softmax/kernel.h" +#include "native/cuda/qy/caster.cuh" +#include "native/cuda/qy/runtime_.h" + +namespace infini::ops { + +template <> +class Operator + : public CudaTopkSoftmax> { + public: + using CudaTopkSoftmax>::CudaTopkSoftmax; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/qy/runtime_.h b/src/native/cuda/qy/runtime_.h new file mode 100644 index 000000000..758379ccb --- /dev/null +++ b/src/native/cuda/qy/runtime_.h @@ -0,0 +1,9 @@ +#ifndef INFINI_OPS_QY_RUNTIME_H_ +#define INFINI_OPS_QY_RUNTIME_H_ + +#include + +#include "native/cuda/qy/runtime_utils.h" +#include "runtime.h" + +#endif diff --git a/src/native/cuda/qy/runtime_utils.h b/src/native/cuda/qy/runtime_utils.h new file mode 100644 index 000000000..4946c74d1 --- /dev/null +++ b/src/native/cuda/qy/runtime_utils.h @@ -0,0 +1,15 @@ +#ifndef INFINI_OPS_QY_RUNTIME_UTILS_H_ +#define INFINI_OPS_QY_RUNTIME_UTILS_H_ + +#include "native/cuda/runtime_utils.h" +#include "native/cuda/qy/device_property.h" + +namespace infini::ops { + +template <> +struct RuntimeUtils + : CudaRuntimeUtils {}; + +} // namespace infini::ops + +#endif diff --git a/src/pybind11_utils.h b/src/pybind11_utils.h index 4acc97d37..6778d7d7b 100644 --- a/src/pybind11_utils.h +++ b/src/pybind11_utils.h @@ -204,6 +204,7 @@ inline std::optional TryDeviceTypeFromString( {"iluvatar", Device::Type::kIluvatar}, {"hygon", Device::Type::kHygon}, {"thead", Device::Type::kThead}, + {"qy", Device::Type::kQy}, }; auto platform_it{ diff --git a/src/torch/device_.h b/src/torch/device_.h index 3d2090e48..bd0dea7ce 100644 --- a/src/torch/device_.h +++ b/src/torch/device_.h @@ -60,6 +60,11 @@ struct TorchDeviceName { static constexpr std::string_view kValue{"cuda"}; }; +template <> +struct TorchDeviceName { + static constexpr std::string_view kValue{"cuda"}; +}; + } // namespace infini::ops::detail #endif diff --git a/src/torch/qy/c10.h b/src/torch/qy/c10.h new file mode 100644 index 000000000..0c27927c6 --- /dev/null +++ b/src/torch/qy/c10.h @@ -0,0 +1,28 @@ +#ifndef INFINI_OPS_TORCH_QY_C10_H_ +#define INFINI_OPS_TORCH_QY_C10_H_ + +#include +#include +#include + +#include "torch/c10.h" + +namespace infini::ops { + +template <> +struct C10 { + static constexpr Device::Type kDeviceType = Device::Type::kQy; + + using StreamGuard = c10::cuda::CUDAStreamGuard; + + static c10::cuda::CUDAStream GetStreamFromExternal(void* stream, + int device_index) { + return c10::cuda::getStreamFromExternal( + reinterpret_cast(stream), + static_cast(device_index)); + } +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_TORCH_QY_C10_H_ diff --git a/tests/conftest.py b/tests/conftest.py index b54f5ed57..406108b16 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -196,6 +196,7 @@ def _set_random_seed(seed): "iluvatar": "cuda", "hygon": "cuda", "thead": "cuda", + "qy": "cuda", "moore": "musa", "cambricon": "mlu", "ascend": "npu", From 65cad16ce36340243bc042dd2da0bb7817795ddb Mon Sep 17 00:00:00 2001 From: Ceng23333 <441651826@qq.com> Date: Tue, 13 Oct 2026 14:28:46 +0800 Subject: [PATCH 2/3] feat(qy): add longctx MoE/RoPE QY kernels for MiniCPM5 Promote mul_scalar, silu_and_mul, moe_sum/align/topk_sigmoid, and rope/rope_pack kernels from the longctx adapting lane onto feat/qy-backend. --- include/moe_align_kernel.h | 26 ++ include/moe_sum_kernel.h | 16 + include/moe_topk_sigmoid_kernel.h | 22 ++ include/mul_scalar_kernel.h | 15 + include/rope_kernel.h | 20 ++ include/rope_pack_kernel.h | 22 ++ include/silu_and_mul_kernel.h | 16 + src/native/cuda/qy/moe_align_kernel.cu | 213 ++++++++++++ src/native/cuda/qy/moe_align_kernel.h | 26 ++ src/native/cuda/qy/moe_sum_kernel.cu | 147 ++++++++ src/native/cuda/qy/moe_sum_kernel.h | 16 + src/native/cuda/qy/moe_topk_sigmoid_kernel.cu | 315 ++++++++++++++++++ src/native/cuda/qy/moe_topk_sigmoid_kernel.h | 22 ++ src/native/cuda/qy/mul_scalar_kernel.cu | 91 +++++ src/native/cuda/qy/mul_scalar_kernel.h | 15 + src/native/cuda/qy/rope_kernel.cu | 171 ++++++++++ src/native/cuda/qy/rope_kernel.h | 20 ++ src/native/cuda/qy/rope_pack_kernel.cu | 148 ++++++++ src/native/cuda/qy/rope_pack_kernel.h | 19 ++ src/native/cuda/qy/silu_and_mul_kernel.cu | 113 +++++++ src/native/cuda/qy/silu_and_mul_kernel.h | 16 + 21 files changed, 1469 insertions(+) create mode 100644 include/moe_align_kernel.h create mode 100644 include/moe_sum_kernel.h create mode 100644 include/moe_topk_sigmoid_kernel.h create mode 100644 include/mul_scalar_kernel.h create mode 100644 include/rope_kernel.h create mode 100644 include/rope_pack_kernel.h create mode 100644 include/silu_and_mul_kernel.h create mode 100644 src/native/cuda/qy/moe_align_kernel.cu create mode 100644 src/native/cuda/qy/moe_align_kernel.h create mode 100644 src/native/cuda/qy/moe_sum_kernel.cu create mode 100644 src/native/cuda/qy/moe_sum_kernel.h create mode 100644 src/native/cuda/qy/moe_topk_sigmoid_kernel.cu create mode 100644 src/native/cuda/qy/moe_topk_sigmoid_kernel.h create mode 100644 src/native/cuda/qy/mul_scalar_kernel.cu create mode 100644 src/native/cuda/qy/mul_scalar_kernel.h create mode 100644 src/native/cuda/qy/rope_kernel.cu create mode 100644 src/native/cuda/qy/rope_kernel.h create mode 100644 src/native/cuda/qy/rope_pack_kernel.cu create mode 100644 src/native/cuda/qy/rope_pack_kernel.h create mode 100644 src/native/cuda/qy/silu_and_mul_kernel.cu create mode 100644 src/native/cuda/qy/silu_and_mul_kernel.h diff --git a/include/moe_align_kernel.h b/include/moe_align_kernel.h new file mode 100644 index 000000000..31c82e47d --- /dev/null +++ b/include/moe_align_kernel.h @@ -0,0 +1,26 @@ +#ifndef INFINI_OPS_MOE_ALIGN_KERNEL_H_ +#define INFINI_OPS_MOE_ALIGN_KERNEL_H_ + +#include +#include + +namespace infini::ops::qy { + +// Workspace bytes for moe_align_launch: 2 * (num_experts + 1) * sizeof(int32). +inline size_t moe_align_workspace_bytes(size_t num_experts) { + return 2 * (num_experts + 1) * sizeof(int32_t); +} + +// Host-free contiguous moe_align (Denglin-safe block/grid caps). +// topk_ids / expert_map: int32; sorted_token_ids / expert_ids / +// num_tokens_post_padded: int32. expert_map may be nullptr. +void moe_align_launch(void *sorted_token_ids, void *expert_ids, + void *num_tokens_post_padded, const void *topk_ids, + const void *expert_map, size_t numel, size_t num_experts, + size_t block_size, size_t sorted_capacity, + size_t expert_ids_capacity, bool pad_sorted_token_ids, + void *workspace, void *stream); + +} // namespace infini::ops::qy + +#endif // INFINI_OPS_MOE_ALIGN_KERNEL_H_ diff --git a/include/moe_sum_kernel.h b/include/moe_sum_kernel.h new file mode 100644 index 000000000..f7051121d --- /dev/null +++ b/include/moe_sum_kernel.h @@ -0,0 +1,16 @@ +#ifndef INFINI_OPS_MOE_SUM_KERNEL_H_ +#define INFINI_OPS_MOE_SUM_KERNEL_H_ + +#include + +namespace infini::ops::qy { + +// Contiguous expert combine: out[t, h] = sum_k in[t, k, h] +// in is [tokens, topk, hidden], out is [tokens, hidden]. +// dtype_tag matches infini::rt::DataType (kFloat16 / kBFloat16 / kFloat32). +void moe_sum_launch(void *out, const void *in, size_t tokens, size_t topk, + size_t hidden, int dtype_tag, void *stream); + +} // namespace infini::ops::qy + +#endif // INFINI_OPS_MOE_SUM_KERNEL_H_ diff --git a/include/moe_topk_sigmoid_kernel.h b/include/moe_topk_sigmoid_kernel.h new file mode 100644 index 000000000..ec932b1f6 --- /dev/null +++ b/include/moe_topk_sigmoid_kernel.h @@ -0,0 +1,22 @@ +#ifndef INFINI_OPS_MOE_TOPK_SIGMOID_KERNEL_H_ +#define INFINI_OPS_MOE_TOPK_SIGMOID_KERNEL_H_ + +#include + +namespace infini::ops::qy { + +// Contiguous MoE top-k after sigmoid (no CUB). +// gating: [tokens, experts] (f16/bf16/f32) +// correction_bias: [experts] (f16/bf16/f32) or null; bias_dtype_tag ignored if null +// topk_weights: [tokens, topk] (f16/bf16/f32); topk_indices: [tokens, topk] (i32/i64) +// dtype tags match infini::rt::DataType. +void moe_topk_sigmoid_launch(void *topk_weights, void *topk_indices, + const void *gating, const void *correction_bias, + size_t tokens, size_t experts, size_t topk, + bool renormalize, int gating_dtype_tag, + int weight_dtype_tag, int index_dtype_tag, + int bias_dtype_tag, void *stream); + +} // namespace infini::ops::qy + +#endif // INFINI_OPS_MOE_TOPK_SIGMOID_KERNEL_H_ diff --git a/include/mul_scalar_kernel.h b/include/mul_scalar_kernel.h new file mode 100644 index 000000000..aff9f36f3 --- /dev/null +++ b/include/mul_scalar_kernel.h @@ -0,0 +1,15 @@ +#ifndef INFINI_OPS_MUL_SCALAR_KERNEL_H_ +#define INFINI_OPS_MUL_SCALAR_KERNEL_H_ + +#include + +namespace infini::ops::qy { + +// Contiguous elementwise out[i] = in[i] * alpha. +// dtype_tag matches infini::rt::DataType (kFloat16 / kBFloat16 / kFloat32 / kFloat64). +void mul_scalar_launch(void *out, const void *in, size_t numel, int dtype_tag, + double alpha, void *stream); + +} // namespace infini::ops::qy + +#endif // INFINI_OPS_MUL_SCALAR_KERNEL_H_ diff --git a/include/rope_kernel.h b/include/rope_kernel.h new file mode 100644 index 000000000..0796c933f --- /dev/null +++ b/include/rope_kernel.h @@ -0,0 +1,20 @@ +#ifndef INFINI_OPS_ROPE_KERNEL_H_ +#define INFINI_OPS_ROPE_KERNEL_H_ + +#include +#include + +namespace infini::ops::qy { + +// Contiguous packed RoPE on [seqlen, nhead, head_dim] (nhead may be 1). +// pos_ids: int64 [seqlen]; sin/cos tables: [max_pos, head_dim/2], same dtype as x. +// algo_tag: 0 = GPT_J (interleaved), 1 = GPT_NEOX (half-split). +// dtype_tag matches infini::rt::DataType (kFloat16 / kBFloat16 / kFloat32). +void rope_launch(void *out, const void *in, const int64_t *pos_ids, + const void *sin_table, const void *cos_table, size_t seqlen, + size_t nhead, size_t head_dim, int algo_tag, int dtype_tag, + void *stream); + +} // namespace infini::ops::qy + +#endif // INFINI_OPS_ROPE_KERNEL_H_ diff --git a/include/rope_pack_kernel.h b/include/rope_pack_kernel.h new file mode 100644 index 000000000..0085e3659 --- /dev/null +++ b/include/rope_pack_kernel.h @@ -0,0 +1,22 @@ +#ifndef INFINI_OPS_ROPE_PACK_KERNEL_H_ +#define INFINI_OPS_ROPE_PACK_KERNEL_H_ + +#include +#include + +namespace infini::ops::qy { + +// Pack strided [seq, heads, dim] (elem_bytes 2 or 4) into contiguous row-major. +// src_stride{0,1,2} are in elements. dst is contiguous [seq, heads, dim]. +void rope_pack_launch(void *dst, const void *src, size_t seq, size_t heads, + size_t dim, int64_t src_stride0, int64_t src_stride1, + int64_t src_stride2, int elem_bytes, void *stream); + +// Unpack contiguous [seq, heads, dim] into strided dst with the given strides. +void rope_unpack_launch(void *dst, const void *src, size_t seq, size_t heads, + size_t dim, int64_t dst_stride0, int64_t dst_stride1, + int64_t dst_stride2, int elem_bytes, void *stream); + +} // namespace infini::ops::qy + +#endif // INFINI_OPS_ROPE_PACK_KERNEL_H_ diff --git a/include/silu_and_mul_kernel.h b/include/silu_and_mul_kernel.h new file mode 100644 index 000000000..863a7ce0b --- /dev/null +++ b/include/silu_and_mul_kernel.h @@ -0,0 +1,16 @@ +#ifndef INFINI_OPS_SILU_AND_MUL_KERNEL_H_ +#define INFINI_OPS_SILU_AND_MUL_KERNEL_H_ + +#include + +namespace infini::ops::qy { + +// Contiguous gated SiLU: out[r, c] = silu(in[r, c]) * in[r, c + hidden] +// in is [rows, 2*hidden], out is [rows, hidden]. +// dtype_tag matches infini::rt::DataType (kFloat16 / kBFloat16 / kFloat32). +void silu_and_mul_launch(void *out, const void *in, size_t rows, size_t hidden, + int dtype_tag, void *stream); + +} // namespace infini::ops::qy + +#endif // INFINI_OPS_SILU_AND_MUL_KERNEL_H_ diff --git a/src/native/cuda/qy/moe_align_kernel.cu b/src/native/cuda/qy/moe_align_kernel.cu new file mode 100644 index 000000000..21d65aa09 --- /dev/null +++ b/src/native/cuda/qy/moe_align_kernel.cu @@ -0,0 +1,213 @@ +#include "native/cuda/qy/moe_align_kernel.h" + +#include +#include +#include +#include + +#include + +namespace infini::ops::qy { +namespace { + +__global__ void moeAlignFillAndZero(int32_t *__restrict__ sorted_token_ids, + int32_t *__restrict__ counts, + int32_t *__restrict__ write_offsets, + size_t sorted_capacity, size_t num_experts, + size_t fill_value) { + // Always seed sorted slots with numel so intra-expert padding is valid; + // unused capacity beyond total is optionally cleared after scan. + size_t n = sorted_capacity > num_experts ? sorted_capacity : num_experts; + size_t idx = blockIdx.x * blockDim.x + threadIdx.x; + size_t stride = blockDim.x * gridDim.x; + for (; idx < n; idx += stride) { + if (idx < sorted_capacity) { + sorted_token_ids[idx] = static_cast(fill_value); + } + if (idx < num_experts) { + counts[idx] = 0; + write_offsets[idx] = 0; + } + } +} + +__global__ void moeAlignZeroTail(int32_t *__restrict__ sorted_token_ids, + const int32_t *__restrict__ total_ptr, + size_t sorted_capacity) { + const size_t total = static_cast(*total_ptr); + size_t idx = blockIdx.x * blockDim.x + threadIdx.x; + size_t stride = blockDim.x * gridDim.x; + for (size_t i = total + idx; i < sorted_capacity; i += stride) { + sorted_token_ids[i] = 0; + } +} + +__global__ void moeAlignCount(const int32_t *__restrict__ topk_ids, + const int32_t *__restrict__ expert_map, + int32_t *__restrict__ counts, size_t numel, + size_t num_experts) { + size_t idx = blockIdx.x * blockDim.x + threadIdx.x; + size_t stride = blockDim.x * gridDim.x; + for (; idx < numel; idx += stride) { + int32_t expert_id = topk_ids[idx]; + if (expert_map != nullptr) { + if (expert_id < 0 || + expert_id >= static_cast(num_experts)) { + continue; + } + expert_id = expert_map[expert_id]; + } + if (expert_id < 0 || expert_id >= static_cast(num_experts)) { + continue; + } + atomicAdd(&counts[expert_id], 1); + } +} + +// Single-block exclusive scan of padded counts + expert_ids + total. +__global__ void moeAlignScan(const int32_t *__restrict__ counts, + int32_t *__restrict__ write_offsets, + int32_t *__restrict__ expert_ids, + int32_t *__restrict__ total_out, int32_t num_experts, + int32_t block_size, int32_t expert_ids_capacity) { + // Serial scan on thread 0 — E is modest (MiniCPM5 typically <= 256). + if (threadIdx.x != 0 || blockIdx.x != 0) { + return; + } + int32_t cursor = 0; + int32_t blocks_written = 0; + for (int32_t e = 0; e < num_experts; ++e) { + write_offsets[e] = cursor; + int32_t c = counts[e]; + if (c <= 0) { + continue; + } + int32_t padded = ((c + block_size - 1) / block_size) * block_size; + int32_t n_blocks = padded / block_size; + for (int32_t b = 0; b < n_blocks; ++b) { + if (blocks_written < expert_ids_capacity) { + expert_ids[blocks_written] = e; + } + ++blocks_written; + } + cursor += padded; + } + *total_out = cursor; + // Zero unused expert_ids slots for determinism. + for (int32_t i = blocks_written; i < expert_ids_capacity; ++i) { + expert_ids[i] = -1; + } +} + +__global__ void moeAlignScatter(const int32_t *__restrict__ topk_ids, + const int32_t *__restrict__ expert_map, + int32_t *__restrict__ sorted_token_ids, + int32_t *__restrict__ write_offsets, + size_t numel, size_t num_experts, + size_t sorted_capacity) { + size_t idx = blockIdx.x * blockDim.x + threadIdx.x; + size_t stride = blockDim.x * gridDim.x; + for (; idx < numel; idx += stride) { + int32_t expert_id = topk_ids[idx]; + if (expert_map != nullptr) { + if (expert_id < 0 || + expert_id >= static_cast(num_experts)) { + continue; + } + expert_id = expert_map[expert_id]; + } + if (expert_id < 0 || expert_id >= static_cast(num_experts)) { + continue; + } + int32_t slot = atomicAdd(&write_offsets[expert_id], 1); + if (slot >= 0 && static_cast(slot) < sorted_capacity) { + sorted_token_ids[slot] = static_cast(idx); + } + } +} + +void checkLaunch(const char *name, cudaError_t err, size_t numel, + size_t num_experts, int grid, int block) { + if (err != cudaSuccess) { + throw std::runtime_error(std::string("QY moe_align_launch ") + name + + " failed: " + cudaGetErrorString(err) + + " (numel=" + std::to_string(numel) + + ", experts=" + std::to_string(num_experts) + + ", grid=" + std::to_string(grid) + + ", block=" + std::to_string(block) + ")"); + } +} + +} // namespace + +void moe_align_launch(void *sorted_token_ids, void *expert_ids, + void *num_tokens_post_padded, const void *topk_ids, + const void *expert_map, size_t numel, size_t num_experts, + size_t block_size, size_t sorted_capacity, + size_t expert_ids_capacity, bool pad_sorted_token_ids, + void *workspace, void *stream) { + if (workspace == nullptr) { + throw std::runtime_error("QY moe_align_launch: workspace is null"); + } + if (block_size == 0) { + throw std::runtime_error("QY moe_align_launch: block_size must be > 0"); + } + + auto *counts = static_cast(workspace); + auto *write_offsets = counts + (num_experts + 1); + auto cuda_stream = reinterpret_cast(stream); + + constexpr int block = 128; + auto grid_for = [](size_t n) { + int g = static_cast((n + block - 1) / block); + return std::min(g, 128); + }; + + (void)cudaGetLastError(); + + { + size_t n = std::max(sorted_capacity, num_experts); + int grid = grid_for(n == 0 ? 1 : n); + moeAlignFillAndZero<<>>( + static_cast(sorted_token_ids), counts, write_offsets, + sorted_capacity, num_experts, numel); + checkLaunch("fill", cudaGetLastError(), numel, num_experts, grid, block); + } + + if (numel > 0 && num_experts > 0) { + int grid = grid_for(numel); + moeAlignCount<<>>( + static_cast(topk_ids), + static_cast(expert_map), counts, numel, num_experts); + checkLaunch("count", cudaGetLastError(), numel, num_experts, grid, block); + } + + moeAlignScan<<<1, 1, 0, cuda_stream>>>( + counts, write_offsets, static_cast(expert_ids), + static_cast(num_tokens_post_padded), + static_cast(num_experts), static_cast(block_size), + static_cast(expert_ids_capacity)); + checkLaunch("scan", cudaGetLastError(), numel, num_experts, 1, 1); + + if (numel > 0 && num_experts > 0) { + int grid = grid_for(numel); + moeAlignScatter<<>>( + static_cast(topk_ids), + static_cast(expert_map), + static_cast(sorted_token_ids), write_offsets, numel, + num_experts, sorted_capacity); + checkLaunch("scatter", cudaGetLastError(), numel, num_experts, grid, block); + } + + if (!pad_sorted_token_ids && sorted_capacity > 0) { + // Match host path: unused capacity beyond total is zero, not numel. + int grid = grid_for(sorted_capacity); + moeAlignZeroTail<<>>( + static_cast(sorted_token_ids), + static_cast(num_tokens_post_padded), sorted_capacity); + checkLaunch("zero_tail", cudaGetLastError(), numel, num_experts, grid, + block); + } +} + +} // namespace infini::ops::qy diff --git a/src/native/cuda/qy/moe_align_kernel.h b/src/native/cuda/qy/moe_align_kernel.h new file mode 100644 index 000000000..6022eda57 --- /dev/null +++ b/src/native/cuda/qy/moe_align_kernel.h @@ -0,0 +1,26 @@ +#ifndef INFINI_OPS_QY_MOE_ALIGN_KERNEL_H_ +#define INFINI_OPS_QY_MOE_ALIGN_KERNEL_H_ + +#include +#include + +namespace infini::ops::qy { + +// Workspace bytes for moe_align_launch: 2 * (num_experts + 1) * sizeof(int32). +inline size_t moe_align_workspace_bytes(size_t num_experts) { + return 2 * (num_experts + 1) * sizeof(int32_t); +} + +// Host-free contiguous moe_align (Denglin-safe block/grid caps). +// topk_ids / expert_map: int32; sorted_token_ids / expert_ids / +// num_tokens_post_padded: int32. expert_map may be nullptr. +void moe_align_launch(void *sorted_token_ids, void *expert_ids, + void *num_tokens_post_padded, const void *topk_ids, + const void *expert_map, size_t numel, size_t num_experts, + size_t block_size, size_t sorted_capacity, + size_t expert_ids_capacity, bool pad_sorted_token_ids, + void *workspace, void *stream); + +} // namespace infini::ops::qy + +#endif // INFINI_OPS_QY_MOE_ALIGN_KERNEL_H_ diff --git a/src/native/cuda/qy/moe_sum_kernel.cu b/src/native/cuda/qy/moe_sum_kernel.cu new file mode 100644 index 000000000..9dd097d58 --- /dev/null +++ b/src/native/cuda/qy/moe_sum_kernel.cu @@ -0,0 +1,147 @@ +#include "native/cuda/qy/moe_sum_kernel.h" + +#include "native/cuda/qy/data_type_.h" + +#include +#include +#include + +#include + +namespace infini::ops::qy { +namespace { + +template +__device__ __forceinline__ float toFloat(T v); + +template <> +__device__ __forceinline__ float toFloat(float v) { + return v; +} + +template <> +__device__ __forceinline__ float toFloat(half v) { + return __half2float(v); +} + +template <> +__device__ __forceinline__ float toFloat(cuda_bfloat16 v) { + return __bfloat162float(v); +} + +template +__device__ __forceinline__ T fromFloat(float v); + +template <> +__device__ __forceinline__ float fromFloat(float v) { + return v; +} + +template <> +__device__ __forceinline__ half fromFloat(float v) { + return __float2half(v); +} + +template <> +__device__ __forceinline__ cuda_bfloat16 fromFloat(float v) { + return __float2bfloat16(v); +} + +// Runtime-topk loop: one block per token, threads stride over hidden. +// Supports arbitrary topk (MiniCPM5 may use topk > 4). +template +__global__ void moeSumKernel(T *__restrict__ output, const T *__restrict__ input, + size_t topk, size_t hidden) { + const size_t token_idx = blockIdx.x; + const T *token_in = input + token_idx * topk * hidden; + T *token_out = output + token_idx * hidden; + for (size_t idx = threadIdx.x; idx < hidden; idx += blockDim.x) { + float acc = 0.0f; + for (size_t k = 0; k < topk; ++k) { + acc += toFloat(token_in[k * hidden + idx]); + } + token_out[idx] = fromFloat(acc); + } +} + +// Flattened grid-stride fallback when tokens exceed modest grid caps. +template +__global__ void moeSumFlatKernel(T *__restrict__ output, + const T *__restrict__ input, size_t tokens, + size_t topk, size_t hidden, size_t n) { + size_t idx = blockIdx.x * blockDim.x + threadIdx.x; + size_t stride = blockDim.x * gridDim.x; + for (; idx < n; idx += stride) { + size_t token = idx / hidden; + size_t h = idx - token * hidden; + const T *token_in = input + token * topk * hidden; + float acc = 0.0f; + for (size_t k = 0; k < topk; ++k) { + acc += toFloat(token_in[k * hidden + h]); + } + output[idx] = fromFloat(acc); + } +} + +template +void launchMoeSum(void *output, const void *input, size_t tokens, size_t topk, + size_t hidden, void *stream) { + if (tokens == 0 || hidden == 0) { + return; + } + if (topk == 0) { + throw std::runtime_error("QY moe_sum_launch: topk must be > 0"); + } + + constexpr int block = 128; + auto cuda_stream = reinterpret_cast(stream); + (void)cudaGetLastError(); + + // Prefer one-block-per-token when token count is modest (Denglin grid cap). + constexpr int kMaxTokenGrid = 128; + if (tokens <= static_cast(kMaxTokenGrid)) { + int threads = static_cast(std::min(hidden, static_cast(block))); + moeSumKernel<<(tokens), threads, 0, cuda_stream>>>( + reinterpret_cast(output), reinterpret_cast(input), topk, + hidden); + } else { + size_t n = tokens * hidden; + int grid = static_cast((n + block - 1) / block); + grid = std::min(grid, 128); + moeSumFlatKernel<<>>( + reinterpret_cast(output), reinterpret_cast(input), + tokens, topk, hidden, n); + } + + auto err = cudaGetLastError(); + if (err != cudaSuccess) { + throw std::runtime_error( + std::string("QY moe_sum_launch failed: ") + cudaGetErrorString(err) + + " (tokens=" + std::to_string(tokens) + ", topk=" + std::to_string(topk) + + ", hidden=" + std::to_string(hidden) + ", block=" + + std::to_string(block) + ")"); + } +} + +} // namespace + +void moe_sum_launch(void *out, const void *in, size_t tokens, size_t topk, + size_t hidden, int dtype_tag, void *stream) { + using infini::rt::DataType; + switch (static_cast(dtype_tag)) { + case DataType::kFloat16: + launchMoeSum(out, in, tokens, topk, hidden, stream); + return; + case DataType::kBFloat16: + launchMoeSum(out, in, tokens, topk, hidden, stream); + return; + case DataType::kFloat32: + launchMoeSum(out, in, tokens, topk, hidden, stream); + return; + default: + throw std::runtime_error("QY moe_sum_launch: unsupported dtype_tag " + + std::to_string(dtype_tag)); + } +} + +} // namespace infini::ops::qy diff --git a/src/native/cuda/qy/moe_sum_kernel.h b/src/native/cuda/qy/moe_sum_kernel.h new file mode 100644 index 000000000..ea0ba35cd --- /dev/null +++ b/src/native/cuda/qy/moe_sum_kernel.h @@ -0,0 +1,16 @@ +#ifndef INFINI_OPS_QY_MOE_SUM_KERNEL_H_ +#define INFINI_OPS_QY_MOE_SUM_KERNEL_H_ + +#include + +namespace infini::ops::qy { + +// Contiguous expert combine: out[t, h] = sum_k in[t, k, h] +// in is [tokens, topk, hidden], out is [tokens, hidden]. +// dtype_tag matches infini::rt::DataType (kFloat16 / kBFloat16 / kFloat32). +void moe_sum_launch(void *out, const void *in, size_t tokens, size_t topk, + size_t hidden, int dtype_tag, void *stream); + +} // namespace infini::ops::qy + +#endif // INFINI_OPS_QY_MOE_SUM_KERNEL_H_ diff --git a/src/native/cuda/qy/moe_topk_sigmoid_kernel.cu b/src/native/cuda/qy/moe_topk_sigmoid_kernel.cu new file mode 100644 index 000000000..db672829a --- /dev/null +++ b/src/native/cuda/qy/moe_topk_sigmoid_kernel.cu @@ -0,0 +1,315 @@ +#include "native/cuda/qy/moe_topk_sigmoid_kernel.h" + +#include "native/cuda/qy/data_type_.h" + +#include +#include +#include +#include +#include + +#include + +namespace infini::ops::qy { +namespace { + +constexpr int kMaxTopk = 64; + +template +__device__ __forceinline__ float toFloat(T v); + +template <> +__device__ __forceinline__ float toFloat(float v) { + return v; +} + +template <> +__device__ __forceinline__ float toFloat(half v) { + return __half2float(v); +} + +template <> +__device__ __forceinline__ float toFloat(cuda_bfloat16 v) { + return __bfloat162float(v); +} + +template +__device__ __forceinline__ T fromFloat(float v); + +template <> +__device__ __forceinline__ float fromFloat(float v) { + return v; +} + +template <> +__device__ __forceinline__ half fromFloat(float v) { + return __float2half(v); +} + +template <> +__device__ __forceinline__ cuda_bfloat16 fromFloat(float v) { + return __float2bfloat16(v); +} + +__device__ __forceinline__ float sigmoidf_fast(float x) { + return 1.0f / (1.0f + expf(-x)); +} + +// Load correction bias element; convert to float on device (no host cast). +__device__ __forceinline__ float loadBias(const void *__restrict__ bias, + int bias_dtype_tag, int expert) { + using infini::rt::DataType; + switch (static_cast(bias_dtype_tag)) { + case DataType::kFloat32: + return reinterpret_cast(bias)[expert]; + case DataType::kFloat16: + return __half2float(reinterpret_cast(bias)[expert]); + case DataType::kBFloat16: + return __bfloat162float( + reinterpret_cast(bias)[expert]); + default: + return 0.0f; + } +} + +// CUB-free per-row top-k: grid-stride over tokens, block=128. +// Compute in float; write WeightT / IndexT at the end (after optional renorm). +template +__global__ void moeTopkSigmoidKernel( + const GatingT *__restrict__ gating, WeightT *__restrict__ topk_weights, + IndexT *__restrict__ topk_indices, const void *__restrict__ correction_bias, + int bias_dtype_tag, int num_tokens, int num_experts, int topk, + int renormalize) { + __shared__ float s_val[128]; + __shared__ int s_idx[128]; + __shared__ int selected[kMaxTopk]; + __shared__ float s_weights[kMaxTopk]; + + for (int token = static_cast(blockIdx.x); token < num_tokens; + token += static_cast(gridDim.x)) { + const GatingT *row = gating + static_cast(token) * num_experts; + float row_sum = 0.0f; + + for (int k_idx = 0; k_idx < topk; ++k_idx) { + float best_val = -FLT_MAX; + int best_idx = 0; + for (int e = static_cast(threadIdx.x); e < num_experts; + e += static_cast(blockDim.x)) { + bool skip = false; + for (int p = 0; p < k_idx; ++p) { + if (selected[p] == e) { + skip = true; + break; + } + } + if (skip) { + continue; + } + float val = sigmoidf_fast(toFloat(row[e])); + if (correction_bias != nullptr) { + val += loadBias(correction_bias, bias_dtype_tag, e); + } + if (val > best_val || (val == best_val && e < best_idx)) { + best_val = val; + best_idx = e; + } + } + + s_val[threadIdx.x] = best_val; + s_idx[threadIdx.x] = best_idx; + __syncthreads(); + + for (int offset = static_cast(blockDim.x) / 2; offset > 0; + offset >>= 1) { + if (static_cast(threadIdx.x) < offset) { + float ov = s_val[threadIdx.x + offset]; + int oi = s_idx[threadIdx.x + offset]; + if (ov > s_val[threadIdx.x] || + (ov == s_val[threadIdx.x] && oi < s_idx[threadIdx.x])) { + s_val[threadIdx.x] = ov; + s_idx[threadIdx.x] = oi; + } + } + __syncthreads(); + } + + if (threadIdx.x == 0) { + int expert = s_idx[0]; + float val = s_val[0]; + if (correction_bias != nullptr) { + val -= loadBias(correction_bias, bias_dtype_tag, expert); + } + s_weights[k_idx] = val; + selected[k_idx] = expert; + row_sum += val; + } + __syncthreads(); + } + + if (threadIdx.x == 0) { + if (renormalize) { + float inv = 1.0f / fmaxf(row_sum, 1e-20f); + for (int k_idx = 0; k_idx < topk; ++k_idx) { + s_weights[k_idx] *= inv; + } + } + for (int k_idx = 0; k_idx < topk; ++k_idx) { + const int out_i = token * topk + k_idx; + topk_weights[out_i] = fromFloat(s_weights[k_idx]); + topk_indices[out_i] = static_cast(selected[k_idx]); + } + } + __syncthreads(); + } +} + +template +void launchMoeTopkSigmoid(void *topk_weights, void *topk_indices, + const void *gating, const void *correction_bias, + int bias_dtype_tag, size_t tokens, size_t experts, + size_t topk, bool renormalize, void *stream) { + if (tokens == 0) { + return; + } + if (experts == 0 || topk == 0) { + throw std::runtime_error( + "QY moe_topk_sigmoid_launch: experts and topk must be > 0"); + } + if (topk > static_cast(kMaxTopk)) { + throw std::runtime_error( + "QY moe_topk_sigmoid_launch: topk exceeds kMaxTopk=" + + std::to_string(kMaxTopk)); + } + if (topk > experts) { + throw std::runtime_error( + "QY moe_topk_sigmoid_launch: topk must be <= experts"); + } + + constexpr int block = 128; + // One block per token when possible; Denglin OORs on huge grids with large + // blocks, but block=128 tolerates a higher grid than mul_scalar's old 256/256. + constexpr size_t kMaxGrid = 1024; + int grid = static_cast(std::min(tokens, kMaxGrid)); + auto cuda_stream = reinterpret_cast(stream); + (void)cudaGetLastError(); + moeTopkSigmoidKernel + <<>>( + reinterpret_cast(gating), + reinterpret_cast(topk_weights), + reinterpret_cast(topk_indices), correction_bias, + bias_dtype_tag, static_cast(tokens), static_cast(experts), + static_cast(topk), renormalize ? 1 : 0); + auto err = cudaGetLastError(); + if (err != cudaSuccess) { + throw std::runtime_error( + std::string("QY moe_topk_sigmoid_launch failed: ") + + cudaGetErrorString(err) + " (tokens=" + std::to_string(tokens) + + ", experts=" + std::to_string(experts) + + ", topk=" + std::to_string(topk) + ", grid=" + std::to_string(grid) + + ", block=" + std::to_string(block) + ")"); + } +} + +template +void launchWithIndex(void *topk_weights, void *topk_indices, const void *gating, + const void *correction_bias, int bias_dtype_tag, + size_t tokens, size_t experts, size_t topk, + bool renormalize, int index_dtype_tag, void *stream) { + using infini::rt::DataType; + switch (static_cast(index_dtype_tag)) { + case DataType::kInt32: + launchMoeTopkSigmoid( + topk_weights, topk_indices, gating, correction_bias, bias_dtype_tag, + tokens, experts, topk, renormalize, stream); + return; + case DataType::kInt64: + launchMoeTopkSigmoid( + topk_weights, topk_indices, gating, correction_bias, bias_dtype_tag, + tokens, experts, topk, renormalize, stream); + return; + default: + throw std::runtime_error( + "QY moe_topk_sigmoid_launch: unsupported index_dtype_tag " + + std::to_string(index_dtype_tag)); + } +} + +template +void launchWithWeight(void *topk_weights, void *topk_indices, + const void *gating, const void *correction_bias, + int bias_dtype_tag, size_t tokens, size_t experts, + size_t topk, bool renormalize, int weight_dtype_tag, + int index_dtype_tag, void *stream) { + using infini::rt::DataType; + switch (static_cast(weight_dtype_tag)) { + case DataType::kFloat16: + launchWithIndex( + topk_weights, topk_indices, gating, correction_bias, bias_dtype_tag, + tokens, experts, topk, renormalize, index_dtype_tag, stream); + return; + case DataType::kBFloat16: + launchWithIndex( + topk_weights, topk_indices, gating, correction_bias, bias_dtype_tag, + tokens, experts, topk, renormalize, index_dtype_tag, stream); + return; + case DataType::kFloat32: + launchWithIndex( + topk_weights, topk_indices, gating, correction_bias, bias_dtype_tag, + tokens, experts, topk, renormalize, index_dtype_tag, stream); + return; + default: + throw std::runtime_error( + "QY moe_topk_sigmoid_launch: unsupported weight_dtype_tag " + + std::to_string(weight_dtype_tag)); + } +} + +} // namespace + +void moe_topk_sigmoid_launch(void *topk_weights, void *topk_indices, + const void *gating, const void *correction_bias, + size_t tokens, size_t experts, size_t topk, + bool renormalize, int gating_dtype_tag, + int weight_dtype_tag, int index_dtype_tag, + int bias_dtype_tag, void *stream) { + using infini::rt::DataType; + if (correction_bias != nullptr) { + switch (static_cast(bias_dtype_tag)) { + case DataType::kFloat16: + case DataType::kBFloat16: + case DataType::kFloat32: + break; + default: + throw std::runtime_error( + "QY moe_topk_sigmoid_launch: unsupported bias_dtype_tag " + + std::to_string(bias_dtype_tag)); + } + } + + switch (static_cast(gating_dtype_tag)) { + case DataType::kFloat16: + launchWithWeight(topk_weights, topk_indices, gating, + correction_bias, bias_dtype_tag, tokens, experts, + topk, renormalize, weight_dtype_tag, + index_dtype_tag, stream); + return; + case DataType::kBFloat16: + launchWithWeight( + topk_weights, topk_indices, gating, correction_bias, bias_dtype_tag, + tokens, experts, topk, renormalize, weight_dtype_tag, index_dtype_tag, + stream); + return; + case DataType::kFloat32: + launchWithWeight(topk_weights, topk_indices, gating, + correction_bias, bias_dtype_tag, tokens, experts, + topk, renormalize, weight_dtype_tag, + index_dtype_tag, stream); + return; + default: + throw std::runtime_error( + "QY moe_topk_sigmoid_launch: unsupported gating_dtype_tag " + + std::to_string(gating_dtype_tag)); + } +} + +} // namespace infini::ops::qy diff --git a/src/native/cuda/qy/moe_topk_sigmoid_kernel.h b/src/native/cuda/qy/moe_topk_sigmoid_kernel.h new file mode 100644 index 000000000..5c28ccc18 --- /dev/null +++ b/src/native/cuda/qy/moe_topk_sigmoid_kernel.h @@ -0,0 +1,22 @@ +#ifndef INFINI_OPS_QY_MOE_TOPK_SIGMOID_KERNEL_H_ +#define INFINI_OPS_QY_MOE_TOPK_SIGMOID_KERNEL_H_ + +#include + +namespace infini::ops::qy { + +// Contiguous MoE top-k after sigmoid (no CUB). +// gating: [tokens, experts] (f16/bf16/f32) +// correction_bias: [experts] (f16/bf16/f32) or null; bias_dtype_tag ignored if null +// topk_weights: [tokens, topk] (f16/bf16/f32); topk_indices: [tokens, topk] (i32/i64) +// dtype tags match infini::rt::DataType. +void moe_topk_sigmoid_launch(void *topk_weights, void *topk_indices, + const void *gating, const void *correction_bias, + size_t tokens, size_t experts, size_t topk, + bool renormalize, int gating_dtype_tag, + int weight_dtype_tag, int index_dtype_tag, + int bias_dtype_tag, void *stream); + +} // namespace infini::ops::qy + +#endif // INFINI_OPS_QY_MOE_TOPK_SIGMOID_KERNEL_H_ diff --git a/src/native/cuda/qy/mul_scalar_kernel.cu b/src/native/cuda/qy/mul_scalar_kernel.cu new file mode 100644 index 000000000..46392d60d --- /dev/null +++ b/src/native/cuda/qy/mul_scalar_kernel.cu @@ -0,0 +1,91 @@ +#include "native/cuda/qy/mul_scalar_kernel.h" + +#include "native/cuda/qy/data_type_.h" + +#include +#include +#include + +#include + +namespace infini::ops::qy { +namespace { + +template +__device__ __forceinline__ T mulScalarValue(T value, double alpha) { + return value * static_cast(alpha); +} + +template <> +__device__ __forceinline__ half mulScalarValue(half value, double alpha) { + return __hmul(value, __float2half(static_cast(alpha))); +} + +template <> +__device__ __forceinline__ cuda_bfloat16 +mulScalarValue(cuda_bfloat16 value, double alpha) { + return __hmul(value, __float2bfloat16_rn(static_cast(alpha))); +} + +template +__global__ void mulScalarKernel(size_t n, T *output, const T *input, + double alpha) { + size_t idx = blockIdx.x * blockDim.x + threadIdx.x; + size_t stride = blockDim.x * gridDim.x; + for (; idx < n; idx += stride) { + output[idx] = mulScalarValue(input[idx], alpha); + } +} + +template +void launchMulScalarKernel(size_t numel, void *output, const void *input, + double alpha, void *stream) { + if (numel == 0) { + return; + } + + constexpr int block = 128; + int grid = static_cast((numel + block - 1) / block); + // Denglin OORs on large grids even for grid-stride kernels; keep modest. + grid = std::min(grid, 128); + auto cuda_stream = reinterpret_cast(stream); + // Clear any prior sticky launch error so we attribute failures to this launch. + (void)cudaGetLastError(); + mulScalarKernel<<>>( + numel, reinterpret_cast(output), + reinterpret_cast(input), alpha); + auto err = cudaGetLastError(); + if (err != cudaSuccess) { + throw std::runtime_error(std::string("QY mul_scalar_launch failed: ") + + cudaGetErrorString(err) + " (numel=" + + std::to_string(numel) + ", grid=" + + std::to_string(grid) + ", block=" + + std::to_string(block) + ")"); + } +} + +} // namespace + +void mul_scalar_launch(void *out, const void *in, size_t numel, int dtype_tag, + double alpha, void *stream) { + using infini::rt::DataType; + switch (static_cast(dtype_tag)) { + case DataType::kFloat16: + launchMulScalarKernel(numel, out, in, alpha, stream); + return; + case DataType::kBFloat16: + launchMulScalarKernel(numel, out, in, alpha, stream); + return; + case DataType::kFloat32: + launchMulScalarKernel(numel, out, in, alpha, stream); + return; + case DataType::kFloat64: + launchMulScalarKernel(numel, out, in, alpha, stream); + return; + default: + throw std::runtime_error("QY mul_scalar_launch: unsupported dtype_tag " + + std::to_string(dtype_tag)); + } +} + +} // namespace infini::ops::qy diff --git a/src/native/cuda/qy/mul_scalar_kernel.h b/src/native/cuda/qy/mul_scalar_kernel.h new file mode 100644 index 000000000..0df5d48bb --- /dev/null +++ b/src/native/cuda/qy/mul_scalar_kernel.h @@ -0,0 +1,15 @@ +#ifndef INFINI_OPS_QY_MUL_SCALAR_KERNEL_H_ +#define INFINI_OPS_QY_MUL_SCALAR_KERNEL_H_ + +#include + +namespace infini::ops::qy { + +// Contiguous elementwise out[i] = in[i] * alpha. +// dtype_tag matches infini::rt::DataType (kFloat16 / kBFloat16 / kFloat32 / kFloat64). +void mul_scalar_launch(void *out, const void *in, size_t numel, int dtype_tag, + double alpha, void *stream); + +} // namespace infini::ops::qy + +#endif // INFINI_OPS_QY_MUL_SCALAR_KERNEL_H_ diff --git a/src/native/cuda/qy/rope_kernel.cu b/src/native/cuda/qy/rope_kernel.cu new file mode 100644 index 000000000..c678446be --- /dev/null +++ b/src/native/cuda/qy/rope_kernel.cu @@ -0,0 +1,171 @@ +#include "native/cuda/qy/rope_kernel.h" + +#include "native/cuda/qy/data_type_.h" + +#include +#include +#include +#include + +#include + +namespace infini::ops::qy { +namespace { + +template +__device__ __forceinline__ float toFloat(T v); + +template <> +__device__ __forceinline__ float toFloat(float v) { + return v; +} + +template <> +__device__ __forceinline__ float toFloat(half v) { + return __half2float(v); +} + +template <> +__device__ __forceinline__ float toFloat(cuda_bfloat16 v) { + return __bfloat162float(v); +} + +template +__device__ __forceinline__ T fromFloat(float v); + +template <> +__device__ __forceinline__ float fromFloat(float v) { + return v; +} + +template <> +__device__ __forceinline__ half fromFloat(float v) { + return __float2half(v); +} + +template <> +__device__ __forceinline__ cuda_bfloat16 fromFloat(float v) { + return __float2bfloat16(v); +} + +// Flattened over (token, head, rotary_pair). Avoids seqlen×nhead 2D/3D grids. +template +__global__ void ropeFlatKernel(T *y, const T *x, const int64_t *__restrict__ pos_ids, + const T *__restrict__ sin_table, + const T *__restrict__ cos_table, size_t seqlen, + size_t nhead, size_t head_dim, size_t table_dim, + size_t n_pairs) { + size_t idx = blockIdx.x * blockDim.x + threadIdx.x; + size_t stride = blockDim.x * gridDim.x; + const size_t head_stride = head_dim; + const size_t token_stride = nhead * head_dim; + + for (; idx < n_pairs; idx += stride) { + size_t pair = idx % table_dim; + size_t th = idx / table_dim; + size_t head = th % nhead; + size_t token = th / nhead; + if (token >= seqlen) { + continue; + } + + const size_t vec_offset = token * token_stride + head * head_stride; + const size_t pos_id = static_cast(pos_ids[token]); + const size_t table_offset = pos_id * table_dim + pair; + const float sin_v = toFloat(sin_table[table_offset]); + const float cos_v = toFloat(cos_table[table_offset]); + + size_t pos0; + size_t pos1; + if constexpr (IsGptJ) { + pos0 = 2 * pair; + pos1 = 2 * pair + 1; + } else { + pos0 = pair; + pos1 = pair + table_dim; + } + + const float x0 = toFloat(x[vec_offset + pos0]); + const float x1 = toFloat(x[vec_offset + pos1]); + y[vec_offset + pos0] = fromFloat(x0 * cos_v - x1 * sin_v); + y[vec_offset + pos1] = fromFloat(x0 * sin_v + x1 * cos_v); + } +} + +template +void launchRope(void *out, const void *in, const int64_t *pos_ids, + const void *sin_table, const void *cos_table, size_t seqlen, + size_t nhead, size_t head_dim, int algo_tag, void *stream) { + if (seqlen == 0 || nhead == 0 || head_dim == 0) { + return; + } + if ((head_dim & 1) != 0) { + throw std::runtime_error("QY rope_launch: head_dim must be even"); + } + + const size_t table_dim = head_dim / 2; + const size_t n_pairs = seqlen * nhead * table_dim; + constexpr int block = 128; + int grid = static_cast((n_pairs + block - 1) / block); + grid = std::min(grid, 1024); + + auto cuda_stream = reinterpret_cast(stream); + (void)cudaGetLastError(); + + if (algo_tag == 0) { + // GPT_J + ropeFlatKernel<<>>( + reinterpret_cast(out), reinterpret_cast(in), pos_ids, + reinterpret_cast(sin_table), + reinterpret_cast(cos_table), seqlen, nhead, head_dim, + table_dim, n_pairs); + } else if (algo_tag == 1) { + // GPT_NEOX + ropeFlatKernel<<>>( + reinterpret_cast(out), reinterpret_cast(in), pos_ids, + reinterpret_cast(sin_table), + reinterpret_cast(cos_table), seqlen, nhead, head_dim, + table_dim, n_pairs); + } else { + throw std::runtime_error("QY rope_launch: unsupported algo_tag " + + std::to_string(algo_tag)); + } + + auto err = cudaGetLastError(); + if (err != cudaSuccess) { + throw std::runtime_error( + std::string("QY rope_launch failed: ") + cudaGetErrorString(err) + + " (seqlen=" + std::to_string(seqlen) + ", nhead=" + + std::to_string(nhead) + ", head_dim=" + std::to_string(head_dim) + + ", grid=" + std::to_string(grid) + ", block=" + std::to_string(block) + + ")"); + } +} + +} // namespace + +void rope_launch(void *out, const void *in, const int64_t *pos_ids, + const void *sin_table, const void *cos_table, size_t seqlen, + size_t nhead, size_t head_dim, int algo_tag, int dtype_tag, + void *stream) { + using infini::rt::DataType; + switch (static_cast(dtype_tag)) { + case DataType::kFloat16: + launchRope(out, in, pos_ids, sin_table, cos_table, seqlen, nhead, + head_dim, algo_tag, stream); + return; + case DataType::kBFloat16: + launchRope(out, in, pos_ids, sin_table, cos_table, seqlen, + nhead, head_dim, algo_tag, stream); + return; + case DataType::kFloat32: + launchRope(out, in, pos_ids, sin_table, cos_table, seqlen, nhead, + head_dim, algo_tag, stream); + return; + default: + throw std::runtime_error("QY rope_launch: unsupported dtype_tag " + + std::to_string(dtype_tag)); + } +} + +} // namespace infini::ops::qy diff --git a/src/native/cuda/qy/rope_kernel.h b/src/native/cuda/qy/rope_kernel.h new file mode 100644 index 000000000..762380bc9 --- /dev/null +++ b/src/native/cuda/qy/rope_kernel.h @@ -0,0 +1,20 @@ +#ifndef INFINI_OPS_QY_ROPE_KERNEL_H_ +#define INFINI_OPS_QY_ROPE_KERNEL_H_ + +#include +#include + +namespace infini::ops::qy { + +// Contiguous packed RoPE on [seqlen, nhead, head_dim] (nhead may be 1). +// pos_ids: int64 [seqlen]; sin/cos tables: [max_pos, head_dim/2], same dtype as x. +// algo_tag: 0 = GPT_J (interleaved), 1 = GPT_NEOX (half-split). +// dtype_tag matches infini::rt::DataType (kFloat16 / kBFloat16 / kFloat32). +void rope_launch(void *out, const void *in, const int64_t *pos_ids, + const void *sin_table, const void *cos_table, size_t seqlen, + size_t nhead, size_t head_dim, int algo_tag, int dtype_tag, + void *stream); + +} // namespace infini::ops::qy + +#endif // INFINI_OPS_QY_ROPE_KERNEL_H_ diff --git a/src/native/cuda/qy/rope_pack_kernel.cu b/src/native/cuda/qy/rope_pack_kernel.cu new file mode 100644 index 000000000..7174898e5 --- /dev/null +++ b/src/native/cuda/qy/rope_pack_kernel.cu @@ -0,0 +1,148 @@ +#include "native/cuda/qy/rope_pack_kernel.h" + +#include +#include +#include +#include + +#include + +namespace infini::ops::qy { +namespace { + +// One thread per (token, head) vector; copy contiguous last-dim of `dim` elems. +template +__global__ void packVecKernel(T *dst, const T *src, size_t n_vecs, size_t dim, + size_t heads, int64_t s0, int64_t s1) { + size_t idx = blockIdx.x * blockDim.x + threadIdx.x; + size_t stride = blockDim.x * gridDim.x; + for (; idx < n_vecs; idx += stride) { + size_t head = idx % heads; + size_t token = idx / heads; + const T *row = src + token * s0 + head * s1; + T *out = dst + idx * dim; + // Vectorize 8-byte chunks when dim and alignment allow (bf16/fp16 dim=128). + size_t d = 0; + if (dim >= 4 && (reinterpret_cast(row) % 8) == 0 && + (reinterpret_cast(out) % 8) == 0) { + for (; d + 4 <= dim; d += 4) { + *reinterpret_cast(out + d) = + *reinterpret_cast(row + d); + } + } + for (; d < dim; ++d) { + out[d] = row[d]; + } + } +} + +template +__global__ void unpackVecKernel(T *dst, const T *src, size_t n_vecs, size_t dim, + size_t heads, int64_t s0, int64_t s1) { + size_t idx = blockIdx.x * blockDim.x + threadIdx.x; + size_t stride = blockDim.x * gridDim.x; + for (; idx < n_vecs; idx += stride) { + size_t head = idx % heads; + size_t token = idx / heads; + T *row = dst + token * s0 + head * s1; + const T *in = src + idx * dim; + size_t d = 0; + if (dim >= 4 && (reinterpret_cast(row) % 8) == 0 && + (reinterpret_cast(in) % 8) == 0) { + for (; d + 4 <= dim; d += 4) { + *reinterpret_cast(row + d) = + *reinterpret_cast(in + d); + } + } + for (; d < dim; ++d) { + row[d] = in[d]; + } + } +} + +template +void launchTyped(bool pack, void *dst, const void *src, size_t seq, size_t heads, + size_t dim, int64_t stride0, int64_t stride1, int64_t stride2, + void *stream) { + if (stride2 != 1) { + throw std::runtime_error( + "QY rope_pack/unpack: only contiguous last-dim (stride2==1) supported"); + } + const size_t n_vecs = seq * heads; + if (n_vecs == 0) { + return; + } + constexpr int block = 128; + int grid = static_cast((n_vecs + block - 1) / block); + // Prefer covering all vectors in one wave when Denglin allows. + grid = std::min(grid, 4096); + auto cuda_stream = reinterpret_cast(stream); + (void)cudaGetLastError(); + if (pack) { + packVecKernel<<>>( + reinterpret_cast(dst), reinterpret_cast(src), n_vecs, + dim, heads, stride0, stride1); + } else { + unpackVecKernel<<>>( + reinterpret_cast(dst), reinterpret_cast(src), n_vecs, + dim, heads, stride0, stride1); + } + auto err = cudaGetLastError(); + if (err != cudaSuccess) { + // Fall back to smaller grid if Denglin rejects large launch. + grid = std::min(grid, 128); + if (pack) { + packVecKernel<<>>( + reinterpret_cast(dst), reinterpret_cast(src), n_vecs, + dim, heads, stride0, stride1); + } else { + unpackVecKernel<<>>( + reinterpret_cast(dst), reinterpret_cast(src), n_vecs, + dim, heads, stride0, stride1); + } + err = cudaGetLastError(); + } + if (err != cudaSuccess) { + throw std::runtime_error( + std::string(pack ? "QY rope_pack_launch failed: " + : "QY rope_unpack_launch failed: ") + + cudaGetErrorString(err) + " (n_vecs=" + std::to_string(n_vecs) + + ", dim=" + std::to_string(dim) + ")"); + } +} + +void launchPackOrUnpack(bool pack, void *dst, const void *src, size_t seq, + size_t heads, size_t dim, int64_t stride0, + int64_t stride1, int64_t stride2, int elem_bytes, + void *stream) { + if (seq == 0 || heads == 0 || dim == 0) { + return; + } + if (elem_bytes == 2) { + launchTyped(pack, dst, src, seq, heads, dim, stride0, stride1, + stride2, stream); + } else if (elem_bytes == 4) { + launchTyped(pack, dst, src, seq, heads, dim, stride0, stride1, + stride2, stream); + } else { + throw std::runtime_error("QY rope_pack: elem_bytes must be 2 or 4"); + } +} + +} // namespace + +void rope_pack_launch(void *dst, const void *src, size_t seq, size_t heads, + size_t dim, int64_t src_stride0, int64_t src_stride1, + int64_t src_stride2, int elem_bytes, void *stream) { + launchPackOrUnpack(true, dst, src, seq, heads, dim, src_stride0, src_stride1, + src_stride2, elem_bytes, stream); +} + +void rope_unpack_launch(void *dst, const void *src, size_t seq, size_t heads, + size_t dim, int64_t dst_stride0, int64_t dst_stride1, + int64_t dst_stride2, int elem_bytes, void *stream) { + launchPackOrUnpack(false, dst, src, seq, heads, dim, dst_stride0, dst_stride1, + dst_stride2, elem_bytes, stream); +} + +} // namespace infini::ops::qy diff --git a/src/native/cuda/qy/rope_pack_kernel.h b/src/native/cuda/qy/rope_pack_kernel.h new file mode 100644 index 000000000..224ecfd15 --- /dev/null +++ b/src/native/cuda/qy/rope_pack_kernel.h @@ -0,0 +1,19 @@ +#ifndef INFINI_OPS_QY_ROPE_PACK_KERNEL_H_ +#define INFINI_OPS_QY_ROPE_PACK_KERNEL_H_ + +#include +#include + +namespace infini::ops::qy { + +void rope_pack_launch(void *dst, const void *src, size_t seq, size_t heads, + size_t dim, int64_t src_stride0, int64_t src_stride1, + int64_t src_stride2, int elem_bytes, void *stream); + +void rope_unpack_launch(void *dst, const void *src, size_t seq, size_t heads, + size_t dim, int64_t dst_stride0, int64_t dst_stride1, + int64_t dst_stride2, int elem_bytes, void *stream); + +} // namespace infini::ops::qy + +#endif // INFINI_OPS_QY_ROPE_PACK_KERNEL_H_ diff --git a/src/native/cuda/qy/silu_and_mul_kernel.cu b/src/native/cuda/qy/silu_and_mul_kernel.cu new file mode 100644 index 000000000..38dc0505b --- /dev/null +++ b/src/native/cuda/qy/silu_and_mul_kernel.cu @@ -0,0 +1,113 @@ +#include "native/cuda/qy/silu_and_mul_kernel.h" + +#include "native/cuda/qy/data_type_.h" + +#include +#include +#include +#include + +#include + +namespace infini::ops::qy { +namespace { + +template +__device__ __forceinline__ float toFloat(T v); + +template <> +__device__ __forceinline__ float toFloat(float v) { + return v; +} + +template <> +__device__ __forceinline__ float toFloat(half v) { + return __half2float(v); +} + +template <> +__device__ __forceinline__ float toFloat(cuda_bfloat16 v) { + return __bfloat162float(v); +} + +template +__device__ __forceinline__ T fromFloat(float v); + +template <> +__device__ __forceinline__ float fromFloat(float v) { + return v; +} + +template <> +__device__ __forceinline__ half fromFloat(float v) { + return __float2half(v); +} + +template <> +__device__ __forceinline__ cuda_bfloat16 fromFloat(float v) { + return __float2bfloat16(v); +} + +template +__global__ void siluAndMulKernel(T *y, const T *x, size_t n, size_t hidden) { + size_t idx = blockIdx.x * blockDim.x + threadIdx.x; + size_t stride = blockDim.x * gridDim.x; + for (; idx < n; idx += stride) { + size_t row = idx / hidden; + size_t col = idx - row * hidden; + const T *row_x = x + row * hidden * 2; + float gate = toFloat(row_x[col]); + float up = toFloat(row_x[col + hidden]); + float silu = gate / (1.0f + expf(-gate)); + y[idx] = fromFloat(silu * up); + } +} + +template +void launchSiluAndMul(void *out, const void *in, size_t rows, size_t hidden, + void *stream) { + size_t n = rows * hidden; + if (n == 0) { + return; + } + + constexpr int block = 128; + int grid = static_cast((n + block - 1) / block); + grid = std::min(grid, 128); + auto cuda_stream = reinterpret_cast(stream); + (void)cudaGetLastError(); + siluAndMulKernel<<>>( + reinterpret_cast(out), reinterpret_cast(in), n, hidden); + auto err = cudaGetLastError(); + if (err != cudaSuccess) { + throw std::runtime_error( + std::string("QY silu_and_mul_launch failed: ") + cudaGetErrorString(err) + + " (n=" + std::to_string(n) + ", rows=" + std::to_string(rows) + + ", hidden=" + std::to_string(hidden) + ", grid=" + std::to_string(grid) + + ", block=" + std::to_string(block) + ")"); + } +} + +} // namespace + +void silu_and_mul_launch(void *out, const void *in, size_t rows, size_t hidden, + int dtype_tag, void *stream) { + using infini::rt::DataType; + switch (static_cast(dtype_tag)) { + case DataType::kFloat16: + launchSiluAndMul(out, in, rows, hidden, stream); + return; + case DataType::kBFloat16: + launchSiluAndMul(out, in, rows, hidden, stream); + return; + case DataType::kFloat32: + launchSiluAndMul(out, in, rows, hidden, stream); + return; + default: + throw std::runtime_error( + "QY silu_and_mul_launch: unsupported dtype_tag " + + std::to_string(dtype_tag)); + } +} + +} // namespace infini::ops::qy diff --git a/src/native/cuda/qy/silu_and_mul_kernel.h b/src/native/cuda/qy/silu_and_mul_kernel.h new file mode 100644 index 000000000..cc4952476 --- /dev/null +++ b/src/native/cuda/qy/silu_and_mul_kernel.h @@ -0,0 +1,16 @@ +#ifndef INFINI_OPS_QY_SILU_AND_MUL_KERNEL_H_ +#define INFINI_OPS_QY_SILU_AND_MUL_KERNEL_H_ + +#include + +namespace infini::ops::qy { + +// Contiguous gated SiLU: out[r, c] = silu(in[r, c]) * in[r, c + hidden] +// in is [rows, 2*hidden], out is [rows, hidden]. +// dtype_tag matches infini::rt::DataType (kFloat16 / kBFloat16 / kFloat32). +void silu_and_mul_launch(void *out, const void *in, size_t rows, size_t hidden, + int dtype_tag, void *stream); + +} // namespace infini::ops::qy + +#endif // INFINI_OPS_QY_SILU_AND_MUL_KERNEL_H_ From ed0b40db75fc5fdb7e638401fb6bba657b372857 Mon Sep 17 00:00:00 2001 From: Ceng23333 <441651826@qq.com> Date: Wed, 14 Oct 2026 10:10:12 +0800 Subject: [PATCH 3/3] style(qy): clang-format QY backend changes Co-authored-by: Cursor --- include/moe_topk_sigmoid_kernel.h | 6 ++--- include/mul_scalar_kernel.h | 3 ++- include/rope_kernel.h | 6 ++--- src/native/cuda/qy/moe_align_kernel.cu | 15 +++++------ src/native/cuda/qy/moe_sum_kernel.cu | 26 ++++++++++--------- src/native/cuda/qy/moe_topk_sigmoid_kernel.cu | 24 +++++++++-------- src/native/cuda/qy/moe_topk_sigmoid_kernel.h | 6 ++--- src/native/cuda/qy/mul_scalar_kernel.cu | 23 ++++++++-------- src/native/cuda/qy/mul_scalar_kernel.h | 3 ++- src/native/cuda/qy/ops/fill/kernel.h | 2 +- src/native/cuda/qy/ops/mul/kernel.h | 2 +- src/native/cuda/qy/rope_kernel.cu | 23 ++++++++-------- src/native/cuda/qy/rope_kernel.h | 6 ++--- src/native/cuda/qy/rope_pack_kernel.cu | 10 +++---- src/native/cuda/qy/runtime_utils.h | 2 +- src/native/cuda/qy/silu_and_mul_kernel.cu | 19 +++++++------- 16 files changed, 91 insertions(+), 85 deletions(-) diff --git a/include/moe_topk_sigmoid_kernel.h b/include/moe_topk_sigmoid_kernel.h index ec932b1f6..57fa74abf 100644 --- a/include/moe_topk_sigmoid_kernel.h +++ b/include/moe_topk_sigmoid_kernel.h @@ -7,9 +7,9 @@ namespace infini::ops::qy { // Contiguous MoE top-k after sigmoid (no CUB). // gating: [tokens, experts] (f16/bf16/f32) -// correction_bias: [experts] (f16/bf16/f32) or null; bias_dtype_tag ignored if null -// topk_weights: [tokens, topk] (f16/bf16/f32); topk_indices: [tokens, topk] (i32/i64) -// dtype tags match infini::rt::DataType. +// correction_bias: [experts] (f16/bf16/f32) or null; bias_dtype_tag ignored if +// null topk_weights: [tokens, topk] (f16/bf16/f32); topk_indices: [tokens, +// topk] (i32/i64) dtype tags match infini::rt::DataType. void moe_topk_sigmoid_launch(void *topk_weights, void *topk_indices, const void *gating, const void *correction_bias, size_t tokens, size_t experts, size_t topk, diff --git a/include/mul_scalar_kernel.h b/include/mul_scalar_kernel.h index aff9f36f3..99da4d3f9 100644 --- a/include/mul_scalar_kernel.h +++ b/include/mul_scalar_kernel.h @@ -6,7 +6,8 @@ namespace infini::ops::qy { // Contiguous elementwise out[i] = in[i] * alpha. -// dtype_tag matches infini::rt::DataType (kFloat16 / kBFloat16 / kFloat32 / kFloat64). +// dtype_tag matches infini::rt::DataType (kFloat16 / kBFloat16 / kFloat32 / +// kFloat64). void mul_scalar_launch(void *out, const void *in, size_t numel, int dtype_tag, double alpha, void *stream); diff --git a/include/rope_kernel.h b/include/rope_kernel.h index 0796c933f..84ca57048 100644 --- a/include/rope_kernel.h +++ b/include/rope_kernel.h @@ -7,9 +7,9 @@ namespace infini::ops::qy { // Contiguous packed RoPE on [seqlen, nhead, head_dim] (nhead may be 1). -// pos_ids: int64 [seqlen]; sin/cos tables: [max_pos, head_dim/2], same dtype as x. -// algo_tag: 0 = GPT_J (interleaved), 1 = GPT_NEOX (half-split). -// dtype_tag matches infini::rt::DataType (kFloat16 / kBFloat16 / kFloat32). +// pos_ids: int64 [seqlen]; sin/cos tables: [max_pos, head_dim/2], same dtype as +// x. algo_tag: 0 = GPT_J (interleaved), 1 = GPT_NEOX (half-split). dtype_tag +// matches infini::rt::DataType (kFloat16 / kBFloat16 / kFloat32). void rope_launch(void *out, const void *in, const int64_t *pos_ids, const void *sin_table, const void *cos_table, size_t seqlen, size_t nhead, size_t head_dim, int algo_tag, int dtype_tag, diff --git a/src/native/cuda/qy/moe_align_kernel.cu b/src/native/cuda/qy/moe_align_kernel.cu index 21d65aa09..d4a0fdcce 100644 --- a/src/native/cuda/qy/moe_align_kernel.cu +++ b/src/native/cuda/qy/moe_align_kernel.cu @@ -1,11 +1,11 @@ -#include "native/cuda/qy/moe_align_kernel.h" +#include #include #include #include #include -#include +#include "native/cuda/qy/moe_align_kernel.h" namespace infini::ops::qy { namespace { @@ -51,8 +51,7 @@ __global__ void moeAlignCount(const int32_t *__restrict__ topk_ids, for (; idx < numel; idx += stride) { int32_t expert_id = topk_ids[idx]; if (expert_map != nullptr) { - if (expert_id < 0 || - expert_id >= static_cast(num_experts)) { + if (expert_id < 0 || expert_id >= static_cast(num_experts)) { continue; } expert_id = expert_map[expert_id]; @@ -68,8 +67,9 @@ __global__ void moeAlignCount(const int32_t *__restrict__ topk_ids, __global__ void moeAlignScan(const int32_t *__restrict__ counts, int32_t *__restrict__ write_offsets, int32_t *__restrict__ expert_ids, - int32_t *__restrict__ total_out, int32_t num_experts, - int32_t block_size, int32_t expert_ids_capacity) { + int32_t *__restrict__ total_out, + int32_t num_experts, int32_t block_size, + int32_t expert_ids_capacity) { // Serial scan on thread 0 — E is modest (MiniCPM5 typically <= 256). if (threadIdx.x != 0 || blockIdx.x != 0) { return; @@ -110,8 +110,7 @@ __global__ void moeAlignScatter(const int32_t *__restrict__ topk_ids, for (; idx < numel; idx += stride) { int32_t expert_id = topk_ids[idx]; if (expert_map != nullptr) { - if (expert_id < 0 || - expert_id >= static_cast(num_experts)) { + if (expert_id < 0 || expert_id >= static_cast(num_experts)) { continue; } expert_id = expert_map[expert_id]; diff --git a/src/native/cuda/qy/moe_sum_kernel.cu b/src/native/cuda/qy/moe_sum_kernel.cu index 9dd097d58..3599791ee 100644 --- a/src/native/cuda/qy/moe_sum_kernel.cu +++ b/src/native/cuda/qy/moe_sum_kernel.cu @@ -1,12 +1,11 @@ -#include "native/cuda/qy/moe_sum_kernel.h" - -#include "native/cuda/qy/data_type_.h" +#include #include #include #include -#include +#include "native/cuda/qy/data_type_.h" +#include "native/cuda/qy/moe_sum_kernel.h" namespace infini::ops::qy { namespace { @@ -50,8 +49,9 @@ __device__ __forceinline__ cuda_bfloat16 fromFloat(float v) { // Runtime-topk loop: one block per token, threads stride over hidden. // Supports arbitrary topk (MiniCPM5 may use topk > 4). template -__global__ void moeSumKernel(T *__restrict__ output, const T *__restrict__ input, - size_t topk, size_t hidden) { +__global__ void moeSumKernel(T *__restrict__ output, + const T *__restrict__ input, size_t topk, + size_t hidden) { const size_t token_idx = blockIdx.x; const T *token_in = input + token_idx * topk * hidden; T *token_out = output + token_idx * hidden; @@ -100,7 +100,8 @@ void launchMoeSum(void *output, const void *input, size_t tokens, size_t topk, // Prefer one-block-per-token when token count is modest (Denglin grid cap). constexpr int kMaxTokenGrid = 128; if (tokens <= static_cast(kMaxTokenGrid)) { - int threads = static_cast(std::min(hidden, static_cast(block))); + int threads = + static_cast(std::min(hidden, static_cast(block))); moeSumKernel<<(tokens), threads, 0, cuda_stream>>>( reinterpret_cast(output), reinterpret_cast(input), topk, hidden); @@ -115,11 +116,12 @@ void launchMoeSum(void *output, const void *input, size_t tokens, size_t topk, auto err = cudaGetLastError(); if (err != cudaSuccess) { - throw std::runtime_error( - std::string("QY moe_sum_launch failed: ") + cudaGetErrorString(err) + - " (tokens=" + std::to_string(tokens) + ", topk=" + std::to_string(topk) + - ", hidden=" + std::to_string(hidden) + ", block=" + - std::to_string(block) + ")"); + throw std::runtime_error(std::string("QY moe_sum_launch failed: ") + + cudaGetErrorString(err) + + " (tokens=" + std::to_string(tokens) + + ", topk=" + std::to_string(topk) + + ", hidden=" + std::to_string(hidden) + + ", block=" + std::to_string(block) + ")"); } } diff --git a/src/native/cuda/qy/moe_topk_sigmoid_kernel.cu b/src/native/cuda/qy/moe_topk_sigmoid_kernel.cu index db672829a..ee95bc5ab 100644 --- a/src/native/cuda/qy/moe_topk_sigmoid_kernel.cu +++ b/src/native/cuda/qy/moe_topk_sigmoid_kernel.cu @@ -1,6 +1,4 @@ -#include "native/cuda/qy/moe_topk_sigmoid_kernel.h" - -#include "native/cuda/qy/data_type_.h" +#include #include #include @@ -8,7 +6,8 @@ #include #include -#include +#include "native/cuda/qy/data_type_.h" +#include "native/cuda/qy/moe_topk_sigmoid_kernel.h" namespace infini::ops::qy { namespace { @@ -57,7 +56,7 @@ __device__ __forceinline__ float sigmoidf_fast(float x) { // Load correction bias element; convert to float on device (no host cast). __device__ __forceinline__ float loadBias(const void *__restrict__ bias, - int bias_dtype_tag, int expert) { + int bias_dtype_tag, int expert) { using infini::rt::DataType; switch (static_cast(bias_dtype_tag)) { case DataType::kFloat32: @@ -75,11 +74,13 @@ __device__ __forceinline__ float loadBias(const void *__restrict__ bias, // CUB-free per-row top-k: grid-stride over tokens, block=128. // Compute in float; write WeightT / IndexT at the end (after optional renorm). template -__global__ void moeTopkSigmoidKernel( - const GatingT *__restrict__ gating, WeightT *__restrict__ topk_weights, - IndexT *__restrict__ topk_indices, const void *__restrict__ correction_bias, - int bias_dtype_tag, int num_tokens, int num_experts, int topk, - int renormalize) { +__global__ void moeTopkSigmoidKernel(const GatingT *__restrict__ gating, + WeightT *__restrict__ topk_weights, + IndexT *__restrict__ topk_indices, + const void *__restrict__ correction_bias, + int bias_dtype_tag, int num_tokens, + int num_experts, int topk, + int renormalize) { __shared__ float s_val[128]; __shared__ int s_idx[128]; __shared__ int selected[kMaxTopk]; @@ -187,7 +188,8 @@ void launchMoeTopkSigmoid(void *topk_weights, void *topk_indices, constexpr int block = 128; // One block per token when possible; Denglin OORs on huge grids with large - // blocks, but block=128 tolerates a higher grid than mul_scalar's old 256/256. + // blocks, but block=128 tolerates a higher grid than mul_scalar's old + // 256/256. constexpr size_t kMaxGrid = 1024; int grid = static_cast(std::min(tokens, kMaxGrid)); auto cuda_stream = reinterpret_cast(stream); diff --git a/src/native/cuda/qy/moe_topk_sigmoid_kernel.h b/src/native/cuda/qy/moe_topk_sigmoid_kernel.h index 5c28ccc18..a68cdce7b 100644 --- a/src/native/cuda/qy/moe_topk_sigmoid_kernel.h +++ b/src/native/cuda/qy/moe_topk_sigmoid_kernel.h @@ -7,9 +7,9 @@ namespace infini::ops::qy { // Contiguous MoE top-k after sigmoid (no CUB). // gating: [tokens, experts] (f16/bf16/f32) -// correction_bias: [experts] (f16/bf16/f32) or null; bias_dtype_tag ignored if null -// topk_weights: [tokens, topk] (f16/bf16/f32); topk_indices: [tokens, topk] (i32/i64) -// dtype tags match infini::rt::DataType. +// correction_bias: [experts] (f16/bf16/f32) or null; bias_dtype_tag ignored if +// null topk_weights: [tokens, topk] (f16/bf16/f32); topk_indices: [tokens, +// topk] (i32/i64) dtype tags match infini::rt::DataType. void moe_topk_sigmoid_launch(void *topk_weights, void *topk_indices, const void *gating, const void *correction_bias, size_t tokens, size_t experts, size_t topk, diff --git a/src/native/cuda/qy/mul_scalar_kernel.cu b/src/native/cuda/qy/mul_scalar_kernel.cu index 46392d60d..c3f51d1e5 100644 --- a/src/native/cuda/qy/mul_scalar_kernel.cu +++ b/src/native/cuda/qy/mul_scalar_kernel.cu @@ -1,12 +1,11 @@ -#include "native/cuda/qy/mul_scalar_kernel.h" - -#include "native/cuda/qy/data_type_.h" +#include #include #include #include -#include +#include "native/cuda/qy/data_type_.h" +#include "native/cuda/qy/mul_scalar_kernel.h" namespace infini::ops::qy { namespace { @@ -49,18 +48,18 @@ void launchMulScalarKernel(size_t numel, void *output, const void *input, // Denglin OORs on large grids even for grid-stride kernels; keep modest. grid = std::min(grid, 128); auto cuda_stream = reinterpret_cast(stream); - // Clear any prior sticky launch error so we attribute failures to this launch. + // Clear any prior sticky launch error so we attribute failures to this + // launch. (void)cudaGetLastError(); mulScalarKernel<<>>( - numel, reinterpret_cast(output), - reinterpret_cast(input), alpha); + numel, reinterpret_cast(output), reinterpret_cast(input), + alpha); auto err = cudaGetLastError(); if (err != cudaSuccess) { - throw std::runtime_error(std::string("QY mul_scalar_launch failed: ") + - cudaGetErrorString(err) + " (numel=" + - std::to_string(numel) + ", grid=" + - std::to_string(grid) + ", block=" + - std::to_string(block) + ")"); + throw std::runtime_error( + std::string("QY mul_scalar_launch failed: ") + cudaGetErrorString(err) + + " (numel=" + std::to_string(numel) + ", grid=" + std::to_string(grid) + + ", block=" + std::to_string(block) + ")"); } } diff --git a/src/native/cuda/qy/mul_scalar_kernel.h b/src/native/cuda/qy/mul_scalar_kernel.h index 0df5d48bb..3ed51e59b 100644 --- a/src/native/cuda/qy/mul_scalar_kernel.h +++ b/src/native/cuda/qy/mul_scalar_kernel.h @@ -6,7 +6,8 @@ namespace infini::ops::qy { // Contiguous elementwise out[i] = in[i] * alpha. -// dtype_tag matches infini::rt::DataType (kFloat16 / kBFloat16 / kFloat32 / kFloat64). +// dtype_tag matches infini::rt::DataType (kFloat16 / kBFloat16 / kFloat32 / +// kFloat64). void mul_scalar_launch(void *out, const void *in, size_t numel, int dtype_tag, double alpha, void *stream); diff --git a/src/native/cuda/qy/ops/fill/kernel.h b/src/native/cuda/qy/ops/fill/kernel.h index 1c8a03195..5171d8fe1 100644 --- a/src/native/cuda/qy/ops/fill/kernel.h +++ b/src/native/cuda/qy/ops/fill/kernel.h @@ -3,9 +3,9 @@ #include +#include "native/cuda/ops/fill/kernel.h" #include "native/cuda/qy/caster.cuh" #include "native/cuda/qy/runtime_.h" -#include "native/cuda/ops/fill/kernel.h" namespace infini::ops { diff --git a/src/native/cuda/qy/ops/mul/kernel.h b/src/native/cuda/qy/ops/mul/kernel.h index f7a838717..b28d02752 100644 --- a/src/native/cuda/qy/ops/mul/kernel.h +++ b/src/native/cuda/qy/ops/mul/kernel.h @@ -3,9 +3,9 @@ #include +#include "native/cuda/ops/mul/kernel.h" #include "native/cuda/qy/caster.cuh" #include "native/cuda/qy/runtime_.h" -#include "native/cuda/ops/mul/kernel.h" namespace infini::ops { diff --git a/src/native/cuda/qy/rope_kernel.cu b/src/native/cuda/qy/rope_kernel.cu index c678446be..408ea3ced 100644 --- a/src/native/cuda/qy/rope_kernel.cu +++ b/src/native/cuda/qy/rope_kernel.cu @@ -1,13 +1,12 @@ -#include "native/cuda/qy/rope_kernel.h" - -#include "native/cuda/qy/data_type_.h" +#include #include #include #include #include -#include +#include "native/cuda/qy/data_type_.h" +#include "native/cuda/qy/rope_kernel.h" namespace infini::ops::qy { namespace { @@ -50,7 +49,8 @@ __device__ __forceinline__ cuda_bfloat16 fromFloat(float v) { // Flattened over (token, head, rotary_pair). Avoids seqlen×nhead 2D/3D grids. template -__global__ void ropeFlatKernel(T *y, const T *x, const int64_t *__restrict__ pos_ids, +__global__ void ropeFlatKernel(T *y, const T *x, + const int64_t *__restrict__ pos_ids, const T *__restrict__ sin_table, const T *__restrict__ cos_table, size_t seqlen, size_t nhead, size_t head_dim, size_t table_dim, @@ -133,12 +133,13 @@ void launchRope(void *out, const void *in, const int64_t *pos_ids, auto err = cudaGetLastError(); if (err != cudaSuccess) { - throw std::runtime_error( - std::string("QY rope_launch failed: ") + cudaGetErrorString(err) + - " (seqlen=" + std::to_string(seqlen) + ", nhead=" + - std::to_string(nhead) + ", head_dim=" + std::to_string(head_dim) + - ", grid=" + std::to_string(grid) + ", block=" + std::to_string(block) + - ")"); + throw std::runtime_error(std::string("QY rope_launch failed: ") + + cudaGetErrorString(err) + + " (seqlen=" + std::to_string(seqlen) + + ", nhead=" + std::to_string(nhead) + + ", head_dim=" + std::to_string(head_dim) + + ", grid=" + std::to_string(grid) + + ", block=" + std::to_string(block) + ")"); } } diff --git a/src/native/cuda/qy/rope_kernel.h b/src/native/cuda/qy/rope_kernel.h index 762380bc9..1ed894250 100644 --- a/src/native/cuda/qy/rope_kernel.h +++ b/src/native/cuda/qy/rope_kernel.h @@ -7,9 +7,9 @@ namespace infini::ops::qy { // Contiguous packed RoPE on [seqlen, nhead, head_dim] (nhead may be 1). -// pos_ids: int64 [seqlen]; sin/cos tables: [max_pos, head_dim/2], same dtype as x. -// algo_tag: 0 = GPT_J (interleaved), 1 = GPT_NEOX (half-split). -// dtype_tag matches infini::rt::DataType (kFloat16 / kBFloat16 / kFloat32). +// pos_ids: int64 [seqlen]; sin/cos tables: [max_pos, head_dim/2], same dtype as +// x. algo_tag: 0 = GPT_J (interleaved), 1 = GPT_NEOX (half-split). dtype_tag +// matches infini::rt::DataType (kFloat16 / kBFloat16 / kFloat32). void rope_launch(void *out, const void *in, const int64_t *pos_ids, const void *sin_table, const void *cos_table, size_t seqlen, size_t nhead, size_t head_dim, int algo_tag, int dtype_tag, diff --git a/src/native/cuda/qy/rope_pack_kernel.cu b/src/native/cuda/qy/rope_pack_kernel.cu index 7174898e5..6e087156a 100644 --- a/src/native/cuda/qy/rope_pack_kernel.cu +++ b/src/native/cuda/qy/rope_pack_kernel.cu @@ -1,11 +1,11 @@ -#include "native/cuda/qy/rope_pack_kernel.h" +#include #include #include #include #include -#include +#include "native/cuda/qy/rope_pack_kernel.h" namespace infini::ops::qy { namespace { @@ -61,9 +61,9 @@ __global__ void unpackVecKernel(T *dst, const T *src, size_t n_vecs, size_t dim, } template -void launchTyped(bool pack, void *dst, const void *src, size_t seq, size_t heads, - size_t dim, int64_t stride0, int64_t stride1, int64_t stride2, - void *stream) { +void launchTyped(bool pack, void *dst, const void *src, size_t seq, + size_t heads, size_t dim, int64_t stride0, int64_t stride1, + int64_t stride2, void *stream) { if (stride2 != 1) { throw std::runtime_error( "QY rope_pack/unpack: only contiguous last-dim (stride2==1) supported"); diff --git a/src/native/cuda/qy/runtime_utils.h b/src/native/cuda/qy/runtime_utils.h index 4946c74d1..d63656fc6 100644 --- a/src/native/cuda/qy/runtime_utils.h +++ b/src/native/cuda/qy/runtime_utils.h @@ -1,8 +1,8 @@ #ifndef INFINI_OPS_QY_RUNTIME_UTILS_H_ #define INFINI_OPS_QY_RUNTIME_UTILS_H_ -#include "native/cuda/runtime_utils.h" #include "native/cuda/qy/device_property.h" +#include "native/cuda/runtime_utils.h" namespace infini::ops { diff --git a/src/native/cuda/qy/silu_and_mul_kernel.cu b/src/native/cuda/qy/silu_and_mul_kernel.cu index 38dc0505b..8cbb349a1 100644 --- a/src/native/cuda/qy/silu_and_mul_kernel.cu +++ b/src/native/cuda/qy/silu_and_mul_kernel.cu @@ -1,13 +1,12 @@ -#include "native/cuda/qy/silu_and_mul_kernel.h" - -#include "native/cuda/qy/data_type_.h" +#include #include #include #include #include -#include +#include "native/cuda/qy/data_type_.h" +#include "native/cuda/qy/silu_and_mul_kernel.h" namespace infini::ops::qy { namespace { @@ -80,11 +79,13 @@ void launchSiluAndMul(void *out, const void *in, size_t rows, size_t hidden, reinterpret_cast(out), reinterpret_cast(in), n, hidden); auto err = cudaGetLastError(); if (err != cudaSuccess) { - throw std::runtime_error( - std::string("QY silu_and_mul_launch failed: ") + cudaGetErrorString(err) + - " (n=" + std::to_string(n) + ", rows=" + std::to_string(rows) + - ", hidden=" + std::to_string(hidden) + ", grid=" + std::to_string(grid) + - ", block=" + std::to_string(block) + ")"); + throw std::runtime_error(std::string("QY silu_and_mul_launch failed: ") + + cudaGetErrorString(err) + + " (n=" + std::to_string(n) + + ", rows=" + std::to_string(rows) + + ", hidden=" + std::to_string(hidden) + + ", grid=" + std::to_string(grid) + + ", block=" + std::to_string(block) + ")"); } }